Fix GGUF + LoRA + torch.compile

Always something...

gguf/gguf.py

@@ -31,8 +31,6 @@ def _replace_with_gguf_linear(model, compute_dtype, state_dict, prefix="", modul
             and _should_convert_to_gguf(state_dict, module_prefix)
             and name not in modules_to_not_convert
         ):
-            key = "diffusion_model." + module_prefix + "weight"
-            patch = patches.get(key, [])
             in_features = state_dict[module_prefix + "weight"].shape[1]
             out_features = state_dict[module_prefix + "weight"].shape[0]
 
@@ -44,7 +42,7 @@ def _replace_with_gguf_linear(model, compute_dtype, state_dict, prefix="", modul
                     module.bias is not None,
                     compute_dtype=compute_dtype
                 )
-            set_lora_params(model._modules[name], patches, module_prefix)
+            #set_lora_params(model._modules[name], patches, module_prefix)
             model._modules[name].source_cls = type(module)
             # Force requires_grad to False to avoid unexpected errors
             model._modules[name].requires_grad_(False)
@@ -54,14 +52,24 @@ def _replace_with_gguf_linear(model, compute_dtype, state_dict, prefix="", modul
 def set_lora_params(module, patches, module_prefix=""):
     # Recursively set lora_diffs and lora_strengths for all GGUFLinear layers
     for name, child in module.named_children():
-        child_prefix = module_prefix + name + "."
+        child_prefix = (f"{module_prefix}{name}.").replace("_orig_mod.", "")
         set_lora_params(child, patches, child_prefix)
     if isinstance(module, GGUFLinear):
-        key = "diffusion_model." + module_prefix + "weight"
+        key = f"diffusion_model.{module_prefix}weight"
+        print("Setting LoRA params for key:", key)
         patch = patches.get(key, [])
-        lora_diffs = lora_strengths = None
         if len(patch) != 0:
-            lora_diffs = [p[1].weights for p in patch]
+            lora_diffs = []
+            for p in patch:
+                lora_obj = p[1]
+                if "head" in key:
+                    continue  # For now skip LoRA for head layers
+                elif hasattr(lora_obj, "weights"):
+                    lora_diffs.append(lora_obj.weights)
+                elif isinstance(lora_obj, tuple) and lora_obj[0] == "diff":
+                    lora_diffs.append(lora_obj[1])
+                else:
+                    continue
             lora_strengths = [p[0] for p in patch]
         module.lora_diffs = lora_diffs
         module.lora_strengths = lora_strengths