From d24db273a2df5757c73e1cb79489baeaff86fff3 Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Mon, 15 Sep 2025 17:17:13 -0400 Subject: [PATCH 1/3] Convert alphas for embedders for sd-scripts to ai toolkit conversion --- .../loaders/lora_conversion_utils.py | 86 +++++++++---------- 1 file changed, 39 insertions(+), 47 deletions(-) diff --git a/src/diffusers/loaders/lora_conversion_utils.py b/src/diffusers/loaders/lora_conversion_utils.py index 6f584a5f0e0a..89afb6529a50 100644 --- a/src/diffusers/loaders/lora_conversion_utils.py +++ b/src/diffusers/loaders/lora_conversion_utils.py @@ -558,70 +558,62 @@ def assign_remaining_weights(assignments, source): ait_sd[target_key] = value if any("guidance_in" in k for k in sds_sd): - assign_remaining_weights( - [ - ( - "time_text_embed.guidance_embedder.linear_1.{lora_key}.weight", - "lora_unet_guidance_in_in_layer.{orig_lora_key}.weight", - None, - ), - ( - "time_text_embed.guidance_embedder.linear_2.{lora_key}.weight", - "lora_unet_guidance_in_out_layer.{orig_lora_key}.weight", - None, - ), - ], + _convert_to_ai_toolkit( sds_sd, + ait_sd, + "lora_unet_guidance_in_in_layer", + "time_text_embed.guidance_embedder.linear_1", + ) + + _convert_to_ai_toolkit( + sds_sd, + ait_sd, + "lora_unet_guidance_in_out_layer", + "time_text_embed.guidance_embedder.linear_2", ) if any("img_in" in k for k in sds_sd): - assign_remaining_weights( - [ - ("x_embedder.{lora_key}.weight", "lora_unet_img_in.{orig_lora_key}.weight", None), - ], + _convert_to_ai_toolkit( sds_sd, + ait_sd, + "lora_unet_img_in", + "x_embedder", ) if any("txt_in" in k for k in sds_sd): - assign_remaining_weights( - [ - ("context_embedder.{lora_key}.weight", "lora_unet_txt_in.{orig_lora_key}.weight", None), - ], + _convert_to_ai_toolkit( sds_sd, + ait_sd, + "lora_unet_txt_in", + "context_embedder", ) if any("time_in" in k for k in sds_sd): - assign_remaining_weights( - [ - ( - "time_text_embed.timestep_embedder.linear_1.{lora_key}.weight", - "lora_unet_time_in_in_layer.{orig_lora_key}.weight", - None, - ), - ( - "time_text_embed.timestep_embedder.linear_2.{lora_key}.weight", - "lora_unet_time_in_out_layer.{orig_lora_key}.weight", - None, - ), - ], + _convert_to_ai_toolkit( sds_sd, + ait_sd, + "lora_unet_time_in_in_layer", + "time_text_embed.timestep_embedder.linear_1", + ) + _convert_to_ai_toolkit( + sds_sd, + ait_sd, + "lora_unet_time_in_out_layer", + "time_text_embed.timestep_embedder.linear_2", ) if any("vector_in" in k for k in sds_sd): - assign_remaining_weights( - [ - ( - "time_text_embed.text_embedder.linear_1.{lora_key}.weight", - "lora_unet_vector_in_in_layer.{orig_lora_key}.weight", - None, - ), - ( - "time_text_embed.text_embedder.linear_2.{lora_key}.weight", - "lora_unet_vector_in_out_layer.{orig_lora_key}.weight", - None, - ), - ], + _convert_to_ai_toolkit( + sds_sd, + ait_sd, + "lora_unet_vector_in_in_layer", + "time_text_embed.text_embedder.linear_1", + ) + _convert_to_ai_toolkit( sds_sd, + ait_sd, + "lora_unet_vector_in_out_layer", + "time_text_embed.text_embedder.linear_2", ) if any("final_layer" in k for k in sds_sd): From dc53dfeea360fb17a828c30c4e12a664aedaf2c4 Mon Sep 17 00:00:00 2001 From: rockerBOO Date: Mon, 15 Sep 2025 17:34:52 -0400 Subject: [PATCH 2/3] Add kohya embedders conversion test --- tests/lora/test_lora_layers_flux.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/lora/test_lora_layers_flux.py b/tests/lora/test_lora_layers_flux.py index 7d99bcad8087..79f2c5a8dece 100644 --- a/tests/lora/test_lora_layers_flux.py +++ b/tests/lora/test_lora_layers_flux.py @@ -907,6 +907,14 @@ def test_flux_kohya_with_text_encoder(self): assert max_diff < 1e-3 + def test_flux_kohya_embedders_conversion(self): + """Test that embedders load without throwing errors""" + self.pipeline.load_lora_weights("rockerBOO/flux-bpo-po-lora") + self.pipeline.unload_lora_weights() + + assert True + + def test_flux_xlabs(self): self.pipeline.load_lora_weights("XLabs-AI/flux-lora-collection", weight_name="disney_lora.safetensors") self.pipeline.fuse_lora() From 9ab74f6cb3b5aeb1f1f5d3d380112d26044ddfca Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Thu, 18 Sep 2025 06:47:45 +0000 Subject: [PATCH 3/3] Apply style fixes --- tests/lora/test_lora_layers_flux.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/lora/test_lora_layers_flux.py b/tests/lora/test_lora_layers_flux.py index 79f2c5a8dece..e6048f509fd4 100644 --- a/tests/lora/test_lora_layers_flux.py +++ b/tests/lora/test_lora_layers_flux.py @@ -914,7 +914,6 @@ def test_flux_kohya_embedders_conversion(self): assert True - def test_flux_xlabs(self): self.pipeline.load_lora_weights("XLabs-AI/flux-lora-collection", weight_name="disney_lora.safetensors") self.pipeline.fuse_lora()