update app

app.py

@@ -363,44 +363,48 @@ def update_selection(evt: gr.SelectData, width, height):
 
 @spaces.GPU
 def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_index, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
-    if selected_index is None:
-        raise gr.Error("You must select a LoRA before proceeding.🧨")
-    
-    selected_lora = loras[selected_index]
-    lora_path = selected_lora["repo"]
-    trigger_word = selected_lora["trigger_word"]
+    # Clean up previous LoRAs in both cases
+    with calculateDuration("Unloading LoRA"):
+        pipe.unload_lora_weights()
     
-    if(trigger_word):
-        if "trigger_position" in selected_lora:
-            if selected_lora["trigger_position"] == "prepend":
-                prompt_mash = f"{trigger_word} {prompt}"
+    # Check if a LoRA is selected
+    if selected_index is not None and selected_index < len(loras):
+        selected_lora = loras[selected_index]
+        lora_path = selected_lora["repo"]
+        trigger_word = selected_lora["trigger_word"]
+        
+        # Prepare Prompt with Trigger Word
+        if(trigger_word):
+            if "trigger_position" in selected_lora:
+                if selected_lora["trigger_position"] == "prepend":
+                    prompt_mash = f"{trigger_word} {prompt}"
+                else:
+                    prompt_mash = f"{prompt} {trigger_word}"
             else:
-                prompt_mash = f"{prompt} {trigger_word}"
+                prompt_mash = f"{trigger_word} {prompt}"
         else:
-            prompt_mash = f"{trigger_word} {prompt}"
+            prompt_mash = prompt
+
+        # Load LoRA
+        with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
+            weight_name = selected_lora.get("weights", None)
+            try:
+                pipe.load_lora_weights(
+                    lora_path, 
+                    weight_name=weight_name, 
+                    adapter_name="default",
+                    low_cpu_mem_usage=True
+                )
+                # Set adapter scale
+                pipe.set_adapters(["default"], adapter_weights=[lora_scale])
+            except Exception as e:
+                print(f"Error loading LoRA: {e}")
+                gr.Warning("Failed to load LoRA weights. Generating with base model.")
     else:
+        # Base Model Case
+        print("No LoRA selected. Running with Base Model.")
         prompt_mash = prompt
-
-    # Unload previous LoRAs to start fresh
-    with calculateDuration("Unloading LoRA"):
-        pipe.unload_lora_weights()
         
-    # LoRA weights flow
-    with calculateDuration(f"Loading LoRA weights for {selected_lora['title']}"):
-        weight_name = selected_lora.get("weights", None)
-        try:
-            pipe.load_lora_weights(
-                lora_path, 
-                weight_name=weight_name, 
-                adapter_name="default",
-                low_cpu_mem_usage=True
-            )
-            # Set adapter scale
-            pipe.set_adapters(["default"], adapter_weights=[lora_scale])
-        except Exception as e:
-            print(f"Error loading LoRA: {e}")
-            gr.Warning("Failed to load LoRA weights. Generating with base model.")
-            
     with calculateDuration("Randomizing seed"):
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
@@ -412,10 +416,6 @@ def run_lora(prompt, image_input, image_strength, cfg_scale, steps, selected_ind
     
     with calculateDuration("Generating image"):
         # For Turbo models, guidance_scale is typically 0.0
-        # The user interface passes cfg_scale, but we override or warn if needed.
-        # However, for flexibility, if the user explicitly sets it, we might check, 
-        # but the reference strongly suggests 0.0 for Turbo.
-        
         forced_guidance = 0.0 # Turbo mode
         
         final_image = pipe(
@@ -536,12 +536,12 @@ with gr.Blocks(delete_cache=(60, 60)) as demo:
     selected_index = gr.State(None)
     with gr.Row():
         with gr.Column(scale=3):
-            prompt = gr.Textbox(label="Prompt", lines=1, placeholder="✦︎ Choose the LoRA and type the prompt ")
+            prompt = gr.Textbox(label="Prompt", lines=1, placeholder="✦︎ Choose the LoRA and type the prompt (or leave LoRA unselected for Base Model)")
         with gr.Column(scale=1, elem_id="gen_column"):
             generate_button = gr.Button("Generate", variant="primary", elem_id="gen_btn")
     with gr.Row():
         with gr.Column():
-            selected_info = gr.Markdown("")
+            selected_info = gr.Markdown("### No LoRA Selected (Base Model)")
             gallery = gr.Gallery(
                 [(item["image"], item["title"]) for item in loras],
                 label="Z-Image LoRAs",
@@ -556,7 +556,7 @@ with gr.Blocks(delete_cache=(60, 60)) as demo:
             custom_lora_button = gr.Button("Remove custom LoRA", visible=False)
         with gr.Column():
             progress_bar = gr.Markdown(elem_id="progress",visible=False)
-            result = gr.Image(label="Generated Image", format="png", height=600)
+            result = gr.Image(label="Generated Image", format="png")
 
     with gr.Row():
         with gr.Accordion("Advanced Settings", open=False):