Merge branch 'main' into dev

README.md

@@ -29,7 +29,7 @@ You can have a look at the provided `merge_models.py` cli for an example on how
 
 ## Changelog
 
-### 0.9.0
+### 0.9.1 ... 0.9.3
 - bugfixes
 - support for pix2pix and inpainting models
 

merge_models.py

@@ -5,7 +5,7 @@
 
 from sd_meh.merge import merge_models, save_model
 from sd_meh.presets import BLOCK_WEIGHTS_PRESETS
-from utils import MERGE_METHODS, weights_and_bases
+from sd_meh.utils import MERGE_METHODS, weights_and_bases
 
 
 @click.command()

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sd-meh"
-version = "0.8.0"
+version = "0.9.4"
 description = "stable diffusion merging execution helper"
 authors = ["s1dlx <[email protected]>"]
 license = "MIT"

sd_meh/__init__.py

@@ -1 +1 @@
-__version__ = "0.8.0"
+__version__ = "0.9.4"

sd_meh/merge.py

@@ -152,15 +152,12 @@ def merge_models(
             bases,
             merge_mode,
             precision=precision,
-            weights_clip=False,
+            weights_clip=weights_clip,
             iterations=iterations,
             device=device,
             work_device=work_device,
             threads=threads,
         )
-        # clip only after the last re-basin iteration
-        if weights_clip:
-            merged = clip_weights(thetas, merged)
     else:
         merged = simple_merge(
             thetas,
@@ -298,7 +295,7 @@ def rebasin_merge(
             new_bases,
             merge_mode,
             precision,
-            weights_clip,
+            False,
             device,
             work_device,
             threads,
@@ -344,6 +341,11 @@ def rebasin_merge(
 
         log_vram("model a updated")
 
+    if weights_clip:
+        clip_thetas = thetas.copy()
+        clip_thetas["model_a"] = model_a
+        thetas["model_a"] = clip_weights(thetas, thetas["model_a"])
+
     return thetas["model_a"]
 
 
@@ -412,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:

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