fix rebasin pix2pix+inpainting

s1dlx · Jul 28, 2023 · 23999ff · 23999ff
1 parent 9212347
commit 23999ff
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/sd_meh/merge.py b/sd_meh/merge.py
@@ -414,7 +414,6 @@ def merge_key(
 
         # dealing wiht pix2pix and inpainting models
         if (a_size := merge_args["a"].size()) != (b_size := merge_args["b"].size()):
-            print(key, a_size, b_size)
             if a_size[1] > b_size[1]:
                 merged_key = merge_args["a"]
             else:

diff --git a/sd_meh/rebasin.py b/sd_meh/rebasin.py
@@ -2199,9 +2199,13 @@ def apply_permutation(ps: PermutationSpec, perm, params):
 
 def update_model_a(ps: PermutationSpec, perm, model_a, new_alpha):
     for k in model_a:
-        model_a[k] = model_a[k] * (1 - new_alpha) + new_alpha * get_permuted_param(
-            ps, perm, k, model_a
-        )
+        try:
+            perm_params = get_permuted_param(
+                ps, perm, k, model_a
+            )
+            model_a[k] = model_a[k] * (1 - new_alpha) + new_alpha * perm_params
+        except RuntimeError: # dealing with pix2pix and inpainting models
+            continue
     return model_a