fix train_dreambooth_lora_sd3.py loading hook (#9107)

sayakpaul · web-flow · commit 2d753b6fb53a · 2024-08-07T10:09:47.000+05:30
diff --git a/examples/dreambooth/train_dreambooth_lora_sd3.py b/examples/dreambooth/train_dreambooth_lora_sd3.py
@@ -1271,7 +1271,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
`@@ -1271,7 +1271,7 @@ def load_model_hook(models, input_dir):`
`1271`	`1271`	`lora_state_dict = StableDiffusion3Pipeline.lora_state_dict(input_dir)`
`1272`	`1272`
`1273`	`1273`	`transformer_state_dict = {`
`1274`		`- f'{k.replace("transformer.", "")}': v for k, v in lora_state_dict.items() if k.startswith("unet.")`
	`1274`	`+ f'{k.replace("transformer.", "")}': v for k, v in lora_state_dict.items() if k.startswith("transformer.")`
`1275`	`1275`	`}`
`1276`	`1276`	`transformer_state_dict = convert_unet_state_dict_to_peft(transformer_state_dict)`
`1277`	`1277`	`incompatible_keys = set_peft_model_state_dict(transformer_, transformer_state_dict, adapter_name="default")`