fix(api): handle 3x3 kernels in LoHA networks

ssube · Apr 8, 2023 · bd03edc · bd03edc
1 parent 7f7476b
commit bd03edc
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/api/onnx_web/convert/diffusion/lora.py b/api/onnx_web/convert/diffusion/lora.py
@@ -258,9 +258,16 @@ def blend_loras(
             )
 
             if base_weights.shape[-2:] == (1, 1):
-                blended = base_weights.squeeze((3, 2)) + weights.squeeze((3, 2))
+                if weights.shape[-2:] == (1, 1):
+                    blended = base_weights.squeeze((3, 2)) + weights.squeeze((3, 2))
+                else:
+                    blended = base_weights.squeeze((3, 2)) + weights
+
                 blended = np.expand_dims(blended, (2, 3))
             else:
+                if base_weights.shape != weights.shape:
+                    blended = base_weights + weights.reshape(base_weights.shape)
+
                 blended = base_weights + weights
 
             logger.trace("blended weight shape: %s", blended.shape)