theluyuan/DiffSynth-Studio
1
0
Fork 0
mirror of https://github.com/modelscope/DiffSynth-Studio.git synced 2026-03-18 22:08:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #815 from mi804/lora_checkpoint

Browse Source 
fix bug
This commit is contained in:
Zhongjie Duan
2025-08-19 12:43:04 +08:00
committed by GitHub
parent ce61bef2b0 fff2c89360
commit 3a9621f6da
4 changed files with 5 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
diffsynth/trainers/utils.py
View File

@@ -352,6 +352,8 @@ class DiffusionTrainingModule(torch.nn.Module):
if "lora_A.weight" in key or "lora_B.weight" in key:
new_key = key.replace("lora_A.weight", "lora_A.default.weight").replace("lora_B.weight", "lora_B.default.weight")
new_state_dict[new_key] = value
elif "lora_A.default.weight" in key or "lora_B.default.weight" in key:
new_state_dict[key] = value
return new_state_dict

1
examples/flux/model_training/train.py
View File

@@ -45,6 +45,7 @@ class FluxTrainingModule(DiffusionTrainingModule):
state_dict = load_state_dict(lora_checkpoint)
state_dict = self.mapping_lora_state_dict(state_dict)
load_result = model.load_state_dict(state_dict, strict=False)
print(f"LoRA checkpoint loaded: {lora_checkpoint}, total {len(state_dict)} keys")
if len(load_result[1]) > 0:
print(f"Warning, LoRA key mismatch! Unexpected keys in LoRA checkpoint: {load_result[1]}")
setattr(self.pipe, lora_base_model, model)

1
examples/qwen_image/model_training/train.py
View File

@@ -49,6 +49,7 @@ class QwenImageTrainingModule(DiffusionTrainingModule):
state_dict = load_state_dict(lora_checkpoint)
state_dict = self.mapping_lora_state_dict(state_dict)
load_result = model.load_state_dict(state_dict, strict=False)
print(f"LoRA checkpoint loaded: {lora_checkpoint}, total {len(state_dict)} keys")
if len(load_result[1]) > 0:
print(f"Warning, LoRA key mismatch! Unexpected keys in LoRA checkpoint: {load_result[1]}")
setattr(self.pipe, lora_base_model, model)

1
examples/wanvideo/model_training/train.py
View File

@@ -46,6 +46,7 @@ class WanTrainingModule(DiffusionTrainingModule):
state_dict = load_state_dict(lora_checkpoint)
state_dict = self.mapping_lora_state_dict(state_dict)
load_result = model.load_state_dict(state_dict, strict=False)
print(f"LoRA checkpoint loaded: {lora_checkpoint}, total {len(state_dict)} keys")
if len(load_result[1]) > 0:
print(f"Warning, LoRA key mismatch! Unexpected keys in LoRA checkpoint: {load_result[1]}")
setattr(self.pipe, lora_base_model, model)