This commit is contained in:
mi804
2025-08-19 12:38:33 +08:00
parent ce61bef2b0
commit fff2c89360
4 changed files with 5 additions and 0 deletions

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