lora fix

Author: arcanite24

scripts/cldm_inference.py

@@ -64,29 +64,6 @@ def load_controlnet(
     )
 
     model_patcher = out[0]
-    
-    # model_config = model_config_from_unet(
-    #     state_dict, "model.diffusion_model.", unet_dtype
-    # )
-
-    # # Set the weights
-    # sd_model = model_config.get_model(
-    #     state_dict,
-    #     "model.diffusion_model.",
-    #     device=device,
-    # )
-    # sd_model.load_model_weights(state_dict, "model.diffusion_model.")
-
-    # # Create the comfy model
-    # model_patcher = ModelPatcher(
-    #     sd_model,
-    #     load_device=device,
-    #     current_device=device,
-    #     offload_device=torch.device("cpu"),
-    # )
-
-    # # Move model to GPU
-    # load_model_gpu(model_patcher)
 
     # Apply loras
     lora_model_patcher = model_patcher

scripts/image_server.py

@@ -3114,7 +3114,7 @@ async def server(websocket):
                             autocaption = True
 
                             # Net models, images, and weights in order
-                            controlnets = [{"model_file": "./models/controllora/Tile.safetensors", "image": image_blur, "weight": 1.0}, {"model_file": "./models/controllora/Composition.safetensors", "image": image_blur, "weight": 0.4}]
+                            controlnets = [{"model_file": "./models/controllora/Composition.safetensors", "image": image_blur, "weight": 0.4}]
 
                             for result in neural_img2img(
                                 modelData["file"],

scripts/ldm/lora.py

@@ -11,7 +11,7 @@ def load_lora_for_models(model, clip, lora, strength_model, strength_clip):
     loaded = load_lora(lora, key_map)
     if model is not None:
         new_modelpatcher = model.clone()
-        k = new_modelpatcher.add_patches(loaded, strength_patch=strength_model)
+        k = new_modelpatcher.add_patches(loaded, strength_model)
     else:
         k = ()
         new_modelpatcher = None
@@ -74,7 +74,7 @@ def load_lora(lora, to_load):
             if mid_name is not None and mid_name in lora.keys():
                 mid = lora[mid_name]
                 loaded_keys.add(mid_name)
-            patch_dict[to_load[x]] = (lora[A_name], lora[B_name], alpha, mid)
+            patch_dict[to_load[x]] = ("lora", (lora[A_name], lora[B_name], alpha, mid))
             loaded_keys.add(A_name)
             loaded_keys.add(B_name)
 
@@ -95,13 +95,16 @@ def load_lora(lora, to_load):
                 loaded_keys.add(hada_t2_name)
 
             patch_dict[to_load[x]] = (
-                lora[hada_w1_a_name],
-                lora[hada_w1_b_name],
-                alpha,
-                lora[hada_w2_a_name],
-                lora[hada_w2_b_name],
-                hada_t1,
-                hada_t2,
+                "loha",
+                (
+                    lora[hada_w1_a_name],
+                    lora[hada_w1_b_name],
+                    alpha,
+                    lora[hada_w2_a_name],
+                    lora[hada_w2_b_name],
+                    hada_t1,
+                    hada_t2,
+                ),
             )
             loaded_keys.add(hada_w1_a_name)
             loaded_keys.add(hada_w1_b_name)
@@ -159,43 +162,67 @@ def load_lora(lora, to_load):
             or (lokr_w2_a is not None)
         ):
             patch_dict[to_load[x]] = (
-                lokr_w1,
-                lokr_w2,
-                alpha,
-                lokr_w1_a,
-                lokr_w1_b,
-                lokr_w2_a,
-                lokr_w2_b,
-                lokr_t2,
+                "lokr",
+                (
+                    lokr_w1,
+                    lokr_w2,
+                    alpha,
+                    lokr_w1_a,
+                    lokr_w1_b,
+                    lokr_w2_a,
+                    lokr_w2_b,
+                    lokr_t2,
+                ),
             )
 
+        # glora
+        a1_name = "{}.a1.weight".format(x)
+        a2_name = "{}.a2.weight".format(x)
+        b1_name = "{}.b1.weight".format(x)
+        b2_name = "{}.b2.weight".format(x)
+        if a1_name in lora:
+            patch_dict[to_load[x]] = (
+                "glora",
+                (lora[a1_name], lora[a2_name], lora[b1_name], lora[b2_name], alpha),
+            )
+            loaded_keys.add(a1_name)
+            loaded_keys.add(a2_name)
+            loaded_keys.add(b1_name)
+            loaded_keys.add(b2_name)
+
         w_norm_name = "{}.w_norm".format(x)
         b_norm_name = "{}.b_norm".format(x)
         w_norm = lora.get(w_norm_name, None)
         b_norm = lora.get(b_norm_name, None)
 
         if w_norm is not None:
             loaded_keys.add(w_norm_name)
-            patch_dict[to_load[x]] = (w_norm,)
+            patch_dict[to_load[x]] = ("diff", (w_norm,))
             if b_norm is not None:
                 loaded_keys.add(b_norm_name)
-                patch_dict["{}.bias".format(to_load[x][: -len(".weight")])] = (b_norm,)
+                patch_dict["{}.bias".format(to_load[x][: -len(".weight")])] = (
+                    "diff",
+                    (b_norm,),
+                )
 
         diff_name = "{}.diff".format(x)
         diff_weight = lora.get(diff_name, None)
         if diff_weight is not None:
-            patch_dict[to_load[x]] = (diff_weight,)
+            patch_dict[to_load[x]] = ("diff", (diff_weight,))
             loaded_keys.add(diff_name)
 
         diff_bias_name = "{}.diff_b".format(x)
         diff_bias = lora.get(diff_bias_name, None)
         if diff_bias is not None:
-            patch_dict["{}.bias".format(to_load[x][: -len(".weight")])] = (diff_bias,)
+            patch_dict["{}.bias".format(to_load[x][: -len(".weight")])] = (
+                "diff",
+                (diff_bias,),
+            )
             loaded_keys.add(diff_bias_name)
 
-    for x in lora.keys():
-        if x not in loaded_keys:
-            pass
+    # for x in lora.keys():
+    #     if x not in loaded_keys:
+    #         print("lora key not loaded", x)
     return patch_dict
 
 

scripts/ldm/model_management.py

@@ -86,6 +86,14 @@ def get_torch_device():
             return torch.device(torch.cuda.current_device())
 
 
+def module_size(module):
+    module_mem = 0
+    sd = module.state_dict()
+    for k in sd:
+        t = sd[k]
+        module_mem += t.nelement() * t.element_size()
+    return module_mem
+
 def get_total_memory(dev=None, torch_total_too=False):
     global directml_enabled
     if dev is None: