check SDv20

wkpark · Dec 28, 2023 · 73db45e · 73db45e
1 parent f9f39fb
commit 73db45e
Showing 1 changed file with 12 additions and 4 deletions.
diff --git a/scripts/model_mixer.py b/scripts/model_mixer.py
@@ -2951,7 +2951,15 @@ def load_state_dict(checkpoint_info):
             models['model_a'] = sd_models.read_state_dict(checkpoint_info.filename, map_location = "cpu").copy()
             isxl = True
 
-        print("isxl =", isxl)
+        # check SD2
+        isv20 = False
+        if not isxl:
+            for k in ["transformer.resblocks.0.attn.in_proj_weight"]:
+                if f"cond_stage_model.model.{k}" in models['model_a']:
+                    isv20 = True
+                    break
+
+        print("isxl =", isxl, ", sd2 =", isv20)
 
         # get all selected elemental blocks
         elemental_selected = []
@@ -3001,7 +3009,7 @@ def load_state_dict(checkpoint_info):
             BLOCKIDS = BLOCKID if not isxl else BLOCKIDXL
 
             # get all blocks affected by same perm groups by rebasin merge
-            if not isxl and "Rebasin" in mm_calcmodes:
+            if not isxl and not isv20 and "Rebasin" in mm_calcmodes:
                 print("check affected permutation blocks by rebasin merge...")
                 jj = 0
                 while True:
@@ -3294,7 +3302,7 @@ def add_difference(theta0, theta1, base, alpha):
                         print(" - No change blocks detected.")
 
         # check Rebasin mode
-        if not isxl and "Rebasin" in calcmodes:
+        if not isxl and not isv20 and "Rebasin" in calcmodes:
             fullmatching = "fastrebasin" not in calc_settings
             print(" - Dynamic loading rebasin module...")
             load_module(os.path.join(scriptdir, "scripts", "rebasin", "weight_matching.py"))
@@ -3479,7 +3487,7 @@ def cosim(theta_0, theta_1, calcmode):
                                 print(f" +{k}")
                                 theta_0[key] = theta_1[key]
 
-            if not isxl and "Rebasin" in calcmodes[n]:
+            if not isxl and not isv20 and "Rebasin" in calcmodes[n]:
                 print("Rebasin calc...")
                 # rebasin mode
                 # Replace theta_0 with a permutated version using model A and B