Update app.py

app.py

@@ -105,6 +105,7 @@ pipe = DiffusionPipeline.from_pretrained(
 # Lightning LoRA info (no global state)
 LIGHTNING_LORA_REPO = "lightx2v/Qwen-Image-Lightning"
 LIGHTNING_LORA_WEIGHT = "Qwen-Image-Lightning-4steps-V2.0-bf16.safetensors"
+LIGHTNING8_LORA_WEIGHT = "Qwen-Image-Lightning-8steps-V2.0-bf16.safetensors"
 
 MAX_SEED = np.iinfo(np.int32).max
 
@@ -175,7 +176,9 @@ def update_selection(evt: gr.SelectData, aspect_ratio):
 
 def handle_speed_mode(speed_mode):
     """Update UI based on speed/quality toggle."""
-    if speed_mode == "Speed (8 steps)":
+    if speed_mode == "Speed (4 steps)":
+        return gr.update(value="Speed mode selected - 4 steps with Lightning LoRA"), 4, 1.0
+    elif speed_mode == "Speed (8 steps)":
         return gr.update(value="Speed mode selected - 8 steps with Lightning LoRA"), 8, 1.0
     else: 
         return gr.update(value="Quality mode selected - 45 steps for best quality"), 45, 3.5
@@ -225,7 +228,7 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, asp
         pipe.unload_lora_weights()
 
     # Load LoRAs based on speed mode
-    if speed_mode == "Speed (8 steps)":
+    if speed_mode == "Speed (4 steps)":
         with calculateDuration("Loading Lightning LoRA and style LoRA"):
             # Load Lightning LoRA first
             pipe.load_lora_weights(
@@ -243,6 +246,26 @@ def run_lora(prompt, cfg_scale, steps, selected_index, randomize_seed, seed, asp
                 adapter_name="style"
             )
             
+            # Set both adapters active with their weights
+            pipe.set_adapters(["lightning", "style"], adapter_weights=[1.0, lora_scale])
+    elif 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=LIGHTNING8_LORA_WEIGHT,
+                adapter_name="lightning"
+            )
+            
+            # 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 with their weights
             pipe.set_adapters(["lightning", "style"], adapter_weights=[1.0, lora_scale])
     else:
@@ -448,8 +471,8 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css, delete_cache=(60, 60)) as app:
             with gr.Row():
                 speed_mode = gr.Radio(
                     label="Generation Mode",
-                    choices=["Speed (8 steps)", "Quality (45 steps)"],
-                    value="Speed (8 steps)",
+                    choices=["Speed (4 steps)", "Speed (8 steps)", "Quality (45 steps)"],
+                    value="Speed (4 steps)",
                     info="Speed mode uses Lightning LoRA for faster generation"
                 )
             
@@ -521,7 +544,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css, delete_cache=(60, 60)) as app:
  
     app.load(
         fn=handle_speed_mode,
-        inputs=[gr.State("Speed (8 steps)")],
+        inputs=[gr.State("Speed (4 steps)")],
         outputs=[speed_status, steps, cfg_scale]
     )