fix: loading problem for sdxl lora dreambooth (#7166)

sayakpaul · web-flow · commit 5f150c4cef21 · 2024-03-01T19:30:48.000+05:30
diff --git a/examples/dreambooth/train_dreambooth_lora_sdxl.py b/examples/dreambooth/train_dreambooth_lora_sdxl.py
@@ -1178,7 +1178,7 @@ def load_model_hook(models, input_dir):
             _set_state_dict_into_text_encoder(lora_state_dict, prefix="text_encoder.", text_encoder=text_encoder_one_)
 
             _set_state_dict_into_text_encoder(
-                lora_state_dict, prefix="text_encoder_2.", text_encoder=text_encoder_one_
+                lora_state_dict, prefix="text_encoder_2.", text_encoder=text_encoder_two_
             )
 
         # Make sure the trainable params are in float32. This is again needed since the base models

Original file line number	Diff line number	Diff line change
`@@ -1178,7 +1178,7 @@ def load_model_hook(models, input_dir):`
`1178`	`1178`	`_set_state_dict_into_text_encoder(lora_state_dict, prefix="text_encoder.", text_encoder=text_encoder_one_)`
`1179`	`1179`
`1180`	`1180`	`_set_state_dict_into_text_encoder(`
`1181`		`- lora_state_dict, prefix="text_encoder_2.", text_encoder=text_encoder_one_`
	`1181`	`+ lora_state_dict, prefix="text_encoder_2.", text_encoder=text_encoder_two_`
`1182`	`1182`	`)`
`1183`	`1183`
`1184`	`1184`	`# Make sure the trainable params are in float32. This is again needed since the base models`