fix(api): correctly blend LoRAs into UNet during conversion

ssube · Apr 7, 2023 · 25176fe · 25176fe
1 parent 1cfc538
commit 25176fe
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/api/onnx_web/convert/__main__.py b/api/onnx_web/convert/__main__.py
@@ -347,7 +347,7 @@ def convert_models(ctx: ConversionContext, args, models: Models):
                                 )
 
                             if "unet" not in blend_models:
-                                blend_models["text_encoder"] = load_model(
+                                blend_models["unet"] = load_model(
                                     path.join(dest, "unet", ONNX_MODEL)
                                 )
 
@@ -365,10 +365,15 @@ def convert_models(ctx: ConversionContext, args, models: Models):
                             blend_loras(
                                 ctx,
                                 blend_models["text_encoder"],
-                                [lora_name],
-                                [lora_source],
+                                [(lora_source, lora_weight)],
                                 "text_encoder",
-                                lora_weights=[lora_weight],
+                            )
+
+                            blend_loras(
+                                ctx,
+                                blend_models["unet"],
+                                [(lora_source, lora_weight)],
+                                "unet",
                             )
 
                         if "tokenizer" in blend_models: