1--加载方式
加载全参数(.ckpt)
加载LoRA(.safetensors)
2--简单实例
python
import sys
sys.path.append("/mnt/dolphinfs/hdd_pool/docker/user/hadoop-waimai-aigc/liujinfu/Codes/v3d-vgen-motion")
import torch
from peft import LoraConfig
from safetensors import safe_open
from svd.models.i2v_svd_unet import UNetSpatioTemporalConditionModel
from svd.utils.util import zero_rank_print
if __name__ == "__main__":
pretrained_model_path = "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-waimai-aigc/liujinfu/Codes/svd_models/models/stable-video-diffusion-img2vid-xt"
unet = UNetSpatioTemporalConditionModel.from_pretrained(pretrained_model_path, subfolder = "unet")
# resume_checkpoint_path = "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-waimai-aigc/liujinfu/Codes/v3d-vgen-motion/results/outputs_motionlora_realRota_0603_1024_stride5/test-0-2024-06-03T14-31-30/checkpoints/checkpoint-500.safetensors"
resume_checkpoint_path = "/mnt/dolphinfs/hdd_pool/docker/user/hadoop-waimai-aigc/liujinfu/Codes/v3d-vgen-motion/results/outputs_motionFull_realRota_0529_stride5/test-0-2024-05-29T10-04-34/checkpoints/checkpoint-step-5000.ckpt"
# Load pretrained unet weights
if resume_checkpoint_path.endswith(".ckpt"):
zero_rank_print(f"resume from checkpoint: {resume_checkpoint_path}")
resume_checkpoint = torch.load(resume_checkpoint_path, map_location="cpu")
# resume dit parameters
print(f'resume_checkpoint keys: {resume_checkpoint.keys()}')
state_dict = resume_checkpoint["state_dict"]
m, u = unet.load_state_dict(state_dict, strict=False)
zero_rank_print(f"dit missing keys: {len(m)}, unexpected keys: {len(u)}")
assert len(u) == 0
# resume global step
resume_global_step = False
if "global_step" in resume_checkpoint and resume_global_step:
zero_rank_print(f"resume global_step: {resume_checkpoint['global_step']}")
global_step = resume_checkpoint['global_step']
elif resume_checkpoint_path.endswith(".safetensors"):
unet_lora_config = LoraConfig(
r = 64,
lora_alpha = 64, # scaling = lora_alpha / r
init_lora_weights = "gaussian",
target_modules = ["to_q","to_k","to_v","to_out.0"],
lora_dropout = 0.1
)
unet.add_adapter(unet_lora_config)
zero_rank_print(f"resume from safetensors: {resume_checkpoint_path}")
state_dict = {}
with safe_open(resume_checkpoint_path, framework="pt", device="cpu") as f:
for key in f.keys():
key_ = key.replace('unet.', '').replace('.weight', '')
state_dict[key_] = f.get_tensor(key)
u = 0
try:
unet.get_submodule(key_+'.default').state_dict()['weight'].data.copy_(state_dict[key_])
except:
u += 1
assert u == 0, "resume unet params failed"
print("All Done!")