Update app.py

app.py

@@ -67,12 +67,9 @@ pipe = DiffusionPipeline.from_pretrained(
     base_model, scheduler=scheduler, torch_dtype=dtype
 ).to(device)
 
-# Store Lightning LoRA info
-lightning_lora = {
-    "repo": "lightx2v/Qwen-Image-Lightning",
-    "weight_name": "Qwen-Image-Lightning-8steps-V1.0.safetensors",
-    "loaded": False
-}
+# Lightning LoRA info (no global state)
+LIGHTNING_LORA_REPO = "lightx2v/Qwen-Image-Lightning"
+LIGHTNING_LORA_WEIGHT = "Qwen-Image-Lightning-8steps-V1.0.safetensors"
 
 MAX_SEED = np.iinfo(np.int32).max
 
@@ -134,29 +131,11 @@ def update_selection(evt: gr.SelectData, aspect_ratio):
     )
 
 def handle_speed_mode(speed_mode):
-    """Handle the speed/quality toggle for Lightning LoRA."""
-    global lightning_lora
-    
+    """Update UI based on speed/quality toggle."""
     if speed_mode == "Speed (8 steps)":
-        # Load Lightning LoRA if not already loaded
-        if not lightning_lora["loaded"]:
-            with calculateDuration("Loading Lightning LoRA"):
-                pipe.load_lora_weights(
-                    lightning_lora["repo"], 
-                    weight_name=lightning_lora["weight_name"],
-                    adapter_name="lightning"
-                )
-                lightning_lora["loaded"] = True
-            return gr.update(value="Lightning LoRA loaded for fast generation"), 8, 1.0
-        return gr.update(value="Lightning LoRA already loaded"), 8, 1.0
+        return gr.update(value="Speed mode selected - 8 steps with Lightning LoRA"), 8, 1.0
     else:  # Quality mode
-        # Unload Lightning LoRA if loaded
-        if lightning_lora["loaded"]:
-            with calculateDuration("Unloading Lightning LoRA"):
-                pipe.unload_lora_weights()
-                lightning_lora["loaded"] = False
-            return gr.update(value="Lightning LoRA unloaded for quality generation"), 28, 3.5
-        return gr.update(value="Quality mode active"), 28, 3.5
+        return gr.update(value="Quality mode selected - 28 steps for best quality"), 28, 3.5
 
 @spaces.GPU(duration=70)
 def generate_image(prompt_mash, steps, seed, cfg_scale, width, height, lora_scale, negative_prompt=""):
@@ -198,36 +177,35 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, asp
     else:
         prompt_mash = prompt
 
-    # First, handle Lightning LoRA based on speed mode
-    global lightning_lora
-    if speed_mode == "Speed (8 steps)" and not lightning_lora["loaded"]:
-        with calculateDuration("Loading Lightning LoRA"):
+    # Always unload any existing LoRAs first to avoid conflicts
+    with calculateDuration("Unloading existing LoRAs"):
+        pipe.unload_lora_weights()
+
+    # Load LoRAs based on speed mode
+    if speed_mode == "Speed (8 steps)":
+        with calculateDuration("Loading Lightning LoRA and style LoRA"):
+            # Load Lightning LoRA first
             pipe.load_lora_weights(
-                lightning_lora["repo"], 
-                weight_name=lightning_lora["weight_name"],
+                LIGHTNING_LORA_REPO, 
+                weight_name=LIGHTNING_LORA_WEIGHT,
                 adapter_name="lightning"
             )
-            lightning_lora["loaded"] = True
-    elif speed_mode == "Quality (28 steps)" and lightning_lora["loaded"]:
-        with calculateDuration("Unloading Lightning LoRA"):
-            pipe.unload_lora_weights()
-            lightning_lora["loaded"] = False
-
-    # Load the selected style LoRA
-    with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
-        weight_name = selected_lora.get("weights", None)
-        
-        # If Lightning is loaded, we need to handle multiple LoRAs
-        if lightning_lora["loaded"]:
+            
+            # 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"
             )
-            # Set both adapters active
+            
+            # Set both adapters active with their weights
             pipe.set_adapters(["lightning", "style"], adapter_weights=[1.0, lora_scale])
-        else:
+    else:
+        # 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,