Fix custom LoRA loading: remap ai-toolkit diffusion_model.* keys to diffusers transformer.*

app.py

@@ -179,6 +179,32 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scal
     return image
 
 @spaces.GPU(duration=70)
+
+def _load_style_lora(pipe, lora_path, weight_name=None):
+    """Load a style LoRA, converting ai-toolkit / ComfyUI key names to diffusers names."""
+    import os
+    from safetensors.torch import load_file
+    from huggingface_hub import hf_hub_download, list_repo_files
+    try:
+        if os.path.isfile(str(lora_path)):
+            local = str(lora_path)
+        else:
+            wn = weight_name
+            if not wn:
+                cands = [f for f in list_repo_files(lora_path) if f.endswith(".safetensors")]
+                if not cands:
+                    raise FileNotFoundError("no .safetensors in " + str(lora_path))
+                wn = cands[0]
+            local = hf_hub_download(lora_path, wn)
+        sd = load_file(local)
+        if any(k.startswith("diffusion_model.") for k in sd):
+            sd = {("transformer." + k[len("diffusion_model."):]): v for k, v in sd.items()}
+            print("[patch] remapped diffusion_model.* -> transformer.* for", lora_path)
+        pipe.load_lora_weights(sd, adapter_name="style")
+    except Exception as e:
+        print("[patch] remap loader failed (" + repr(e) + "), falling back to default loader")
+        pipe.load_lora_weights(lora_path, weight_name=weight_name, low_cpu_mem_usage=True, adapter_name="style")
+
 def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, aspect_ratio, lora_scale, speed_mode, progress=gr.Progress(track_tqdm=True)):
     if selected_index is None:
         raise gr.Error("You must select a LoRA before proceeding.")
@@ -215,12 +241,7 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, asp
             
             # Load the selected style LoRA
             weight_name = selected_lora.get("weights", None)
-            pipe.load_lora_weights(
-                lora_path, 
-                weight_name=weight_name, 
-                low_cpu_mem_usage=True,
-                adapter_name="style"
-            )
+            _load_style_lora(pipe, lora_path, weight_name)
             
             # Set both adapters active with their weights
             pipe.set_adapters(["lightning", "style"], adapter_weights=[1.0, lora_scale])
@@ -228,12 +249,7 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, asp
         # Quality mode - only load the style LoRA
         with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
             weight_name = selected_lora.get("weights", None)
-            pipe.load_lora_weights(
-                lora_path, 
-                weight_name=weight_name, 
-                low_cpu_mem_usage=True,
-                adapter_name="style"
-            )
+            _load_style_lora(pipe, lora_path, weight_name)
             pipe.set_adapters(["style"], adapter_weights=[lora_scale])
                 
     # Set random seed for reproducibility