lora extract: actually functional for modern models now

Author: mcmonkey4eva

src/BuiltinExtensions/ComfyUIBackend/ComfyUIAPIAbstractBackend.cs

@@ -303,7 +303,7 @@ await socket.SendJson(new JObject()
         float curPercent = 0;
         void yieldProgressUpdate()
         {
-            Logs.Verbose($"Progress [{batchId}]: {nodesDone}/{expectedNodes}, curPercent={curPercent:00.00}");
+            Logs.Verbose($"Progress [{batchId}]: {nodesDone}/{expectedNodes}, curPercent={curPercent * 100:00.0}");
             JObject toSend = new()
             {
                 ["batch_index"] = batchId,

src/BuiltinExtensions/ComfyUIBackend/ExtraNodes/SwarmComfyCommon/SwarmExtractLora.py

@@ -18,7 +18,7 @@ def extract_lora(diff, rank):
         else:
             diff = diff.squeeze()
 
-    U, S, Vh = torch.linalg.svd(diff.float())
+    U, S, Vh = torch.linalg.svd(diff.float(), full_matrices=False)
     U = U[:, :rank]
     S = S[:rank]
     U = U @ torch.diag(S)
@@ -36,65 +36,61 @@ def extract_lora(diff, rank):
     return (U, Vh)
 
 
-def do_lora_handle(base_data, other_data, rank, prefix, require, do_bias, callback):
+def do_lora_handle(base_data, other_data, rank, callback):
     out_data = {}
     device = comfy.model_management.get_torch_device()
     for key in base_data.keys():
         callback()
         if key not in other_data:
             continue
-        base_tensor = base_data[key].float()
-        other_tensor = other_data[key].float()
-        if key.startswith("clip_g"):
-            key = "1." + key[len("clip_g."):]
-        elif key.startswith("clip_l"):
-            key = "0." + key[len("clip_l."):]
-        if require:
-            if not key.startswith(require):
-                print(f"Ignore unmatched key {key} (doesn't match {require})")
-                continue
-            key = key[len(require):]
+        if key.endswith(".weight_scale") or key.endswith(".comfy_quant"):
+            continue
+        base_tensor = base_data[key]
+        other_tensor = other_data[key]
+        if key.endswith(".weight"):
+            fixed_key = key[:-len(".weight")]
+            scale_key = f"{fixed_key}.weight_scale"
+            if scale_key in base_data:
+                scale = base_data[scale_key]
+                base_tensor = base_tensor.to(dtype=torch.bfloat16) * scale
+            if scale_key in other_data:
+                scale = other_data[scale_key]
+                other_tensor = other_tensor.to(dtype=torch.bfloat16) * scale
+        elif key.endswith(".bias") or key.endswith(".scale") or key.endswith(".lin"):
+            fixed_key = key
         if base_tensor.shape != other_tensor.shape:
             continue
         target_dtype = base_tensor.dtype
         if target_dtype == torch.float8_e4m3fn or target_dtype == torch.float8_e5m2:
             target_dtype = torch.bfloat16
         base_tensor = base_tensor.to(dtype=target_dtype)
         other_tensor = other_tensor.to(dtype=target_dtype)
-        diff = other_tensor.to(device) - base_tensor.to(device)
+        diff = other_tensor.to(device, dtype=torch.float32) - base_tensor.to(device, dtype=torch.float32)
         other_tensor = other_tensor.cpu()
         base_tensor = base_tensor.cpu()
         max_diff = float(diff.abs().max())
         if max_diff < 1e-5:
             print(f"discard unaltered key {key} ({max_diff})")
             continue
-        if key.endswith(".weight"):
-            fixed_key = key[:-len(".weight")].replace('.', '_')
-            name = f"lora_{prefix}_{fixed_key}"
-            if len(base_tensor.shape) >= 2:
-                print(f"extract key {name} ({max_diff})")
-                out = extract_lora(diff, rank)
-                up = out[0].contiguous().to(dtype=target_dtype).cpu()
-                down = out[1].contiguous().to(dtype=target_dtype).cpu()
-                if up.isnan().any() or up.isinf().any():
-                    print(f"bad data for {name}.lora_up.weight")
-                    continue
-                if down.isnan().any() or down.isinf().any():
-                    print(f"bad data for {name}.lora_down.weight")
-                    continue
-                out_data[f"{name}.lora_up.weight"] = up
-                out_data[f"{name}.lora_down.weight"] = down
-            else:
-                print(f"ignore valid raw pass-through key {name} ({max_diff})")
-        elif key.endswith(".bias") and do_bias:
-            fixed_key = key[:-len(".bias")].replace('.', '_')
-            name = f"lora_{prefix}_{fixed_key}"
-            print(f"extract bias key {name} ({max_diff})")
-            diff = diff.contiguous().to(dtype=target_dtype).cpu()
-            if diff.isnan().any() or diff.isinf().any():
-                print(f"bad data for {name}.diff_b")
+        if len(base_tensor.shape) >= 2 and base_tensor.numel() < 1024:
+            print(f"extract key {fixed_key} ({max_diff})")
+            out = extract_lora(diff, rank)
+            up = out[0].contiguous().to(dtype=target_dtype).cpu()
+            down = out[1].contiguous().to(dtype=target_dtype).cpu()
+            if up.isnan().any() or up.isinf().any():
+                print(f"bad data for {fixed_key}.lora_up.weight")
+                continue
+            if down.isnan().any() or down.isinf().any():
+                print(f"bad data for {fixed_key}.lora_down.weight")
+                continue
+            out_data[f"{fixed_key}.lora_up.weight"] = up
+            out_data[f"{fixed_key}.lora_down.weight"] = down
+        else:
+            out = diff.contiguous().to(dtype=target_dtype).cpu()
+            if out.isnan().any() or out.isinf().any():
+                print(f"bad data for {fixed_key}")
                 continue
-            out_data[f"{name}.diff_b"] = diff
+            out_data[f"{fixed_key}.diff"] = out
 
 
     return out_data
@@ -108,13 +104,10 @@ def INPUT_TYPES(s):
         return {
             "required": {
                 "base_model": ("MODEL", ),
-                "base_model_clip": ("CLIP", ),
                 "other_model": ("MODEL", ),
-                "other_model_clip": ("CLIP", ),
                 "rank": ("INT", {"default": 16, "min": 1, "max": 320}),
                 "save_rawpath": ("STRING", {"multiline": False}),
                 "save_filename": ("STRING", {"multiline": False}),
-                "save_clip": ("BOOLEAN", {"default": True}),
                 "metadata": ("STRING", {"multiline": True}),
             }
         }
@@ -125,29 +118,18 @@ def INPUT_TYPES(s):
     OUTPUT_NODE = True
     DESCRIPTION = "Internal node, do not use directly - extracts a LoRA from the difference between two models. This is used by SwarmUI Utilities tab."
 
-    def extract_lora(self, base_model, base_model_clip, other_model, other_model_clip, rank, save_rawpath, save_filename, save_clip, metadata):
+    def extract_lora(self, base_model, other_model, rank, save_rawpath, save_filename, metadata):
         base_data = base_model.model_state_dict()
         other_data = other_model.model_state_dict()
         key_count = len(base_data.keys())
-        if save_clip:
-            if base_model_clip is None or other_model_clip is None:
-                print("Warning: save_clip is True but CLIP model(s) are unavailable (model may not have embedded CLIP, e.g. Flux), skipping CLIP extraction")
-                save_clip = False
-            else:
-                key_count += len(base_model_clip.get_sd().keys())
         pbar = comfy.utils.ProgressBar(key_count)
         class Helper:
             steps = 0
             def callback(self):
                 self.steps += 1
                 pbar.update_absolute(self.steps, key_count, None)
         helper = Helper()
-        out_data = do_lora_handle(base_data, other_data, rank, "unet", "diffusion_model.", False, lambda: helper.callback())
-        if save_clip:
-            # TODO: CLIP keys get wonky, this probably doesn't work? Model-arch-dependent.
-            out_clip = do_lora_handle(base_model_clip.get_sd(), other_model_clip.get_sd(), rank, "te_text_model_encoder_layers", "0.transformer.text_model.encoder.layers.", False, lambda: helper.callback())
-            out_clip = do_lora_handle(base_model_clip.get_sd(), other_model_clip.get_sd(), rank, "te2_text_model_encoder_layers", "1.transformer.text_model.encoder.layers.", False, lambda: helper.callback())
-            out_data.update(out_clip)
+        out_data = do_lora_handle(base_data, other_data, rank, lambda: helper.callback())
 
         # Can't easily autodetect all the correct modelspec info, but at least supply some basics
         out_metadata = {