Prefetch + cache remapped style LoRA outside the GPU window; raise GPU duration to 120s

app.py

@@ -159,7 +159,7 @@ def handle_speed_mode(speed_mode):
     else: 
         return gr.update(value="Quality mode selected - 45 steps for best quality"), 45, 3.5
 
-@spaces.GPU(duration=70)
+@spaces.GPU(duration=120)
 def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, negative_prompt=""):
     pipe.to("cuda")
     generator = torch.Generator(device="cuda").manual_seed(seed)
@@ -178,31 +178,42 @@ def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scal
         
     return image
 
-@spaces.GPU(duration=70)
+@spaces.GPU(duration=120)
 
-def _load_style_lora(pipe, lora_path, weight_name=None):
-    """Load a style LoRA, converting ai-toolkit / ComfyUI key names to diffusers names."""
+_STYLE_SD_CACHE = {}
+
+def _prefetch_style_lora(lora_path, weight_name=None):
+    """Download + key-remap a style LoRA OUTSIDE the GPU window, cache the state dict."""
     import os
     from safetensors.torch import load_file
     from huggingface_hub import hf_hub_download, list_repo_files
+    key = (str(lora_path), str(weight_name))
+    if key in _STYLE_SD_CACHE:
+        return _STYLE_SD_CACHE[key]
+    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)
+    _STYLE_SD_CACHE[key] = sd
+    print("[patch] prefetched style LoRA", lora_path, "tensors:", len(sd))
+    return sd
+
+def _load_style_lora(pipe, lora_path, weight_name=None):
     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)
+        sd = _prefetch_style_lora(lora_path, weight_name)
         pipe.load_lora_weights(sd, adapter_name="style")
     except Exception as e:
-        print("[patch] remap loader failed (" + repr(e) + "), falling back to default loader")
+        print("[patch] remap loader failed (" + repr(e) + "), falling back")
         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)):
@@ -351,6 +362,10 @@ def add_custom_lora(custom_lora):
     if custom_lora:
         try:
             title, repo, path, trigger_word, image = check_custom_model(custom_lora)
+            try:
+                _prefetch_style_lora(repo, path)
+            except Exception as _e:
+                print("[patch] prefetch failed:", repr(_e))
             print(f"Loaded custom LoRA: {repo}")
             card = f'''
             <div class="custom_lora_card">