Fix train_dreambooth_lora_sd3_miniature

hlky · hlky · commit 0ad225c0377a · 2025-01-13T09:23:38.000Z
diff --git a/examples/research_projects/sd3_lora_colab/train_dreambooth_lora_sd3_miniature.py b/examples/research_projects/sd3_lora_colab/train_dreambooth_lora_sd3_miniature.py
@@ -765,7 +765,7 @@ def load_model_hook(models, input_dir):
         lora_state_dict = StableDiffusion3Pipeline.lora_state_dict(input_dir)
 
         transformer_state_dict = {
-            f'{k.replace("transformer.", "")}': v for k, v in lora_state_dict.items() if k.startswith("unet.")
+            f'{k.replace("transformer.", "")}': v for k, v in lora_state_dict.items() if k.startswith("transformer.")
         }
         transformer_state_dict = convert_unet_state_dict_to_peft(transformer_state_dict)
         incompatible_keys = set_peft_model_state_dict(transformer_, transformer_state_dict, adapter_name="default")

Original file line number	Diff line number	Diff line change
`@@ -765,7 +765,7 @@ def load_model_hook(models, input_dir):`
`765`	`765`	`lora_state_dict = StableDiffusion3Pipeline.lora_state_dict(input_dir)`
`766`	`766`
`767`	`767`	`transformer_state_dict = {`
`768`		`- f'{k.replace("transformer.", "")}': v for k, v in lora_state_dict.items() if k.startswith("unet.")`
	`768`	`+ f'{k.replace("transformer.", "")}': v for k, v in lora_state_dict.items() if k.startswith("transformer.")`
`769`	`769`	`}`
`770`	`770`	`transformer_state_dict = convert_unet_state_dict_to_peft(transformer_state_dict)`
`771`	`771`	`incompatible_keys = set_peft_model_state_dict(transformer_, transformer_state_dict, adapter_name="default")`