Refactor inference function in app.py to yield results instead of returning, enhancing the generator's output structure. Adjust UI layout by increasing input image height and repositioning the generate button for better accessibility.

app.py

@@ -108,7 +108,7 @@ def infer(
         progress: Gradio progress callback.
 
     Returns:
-        tuple: (stage2_only_image, result_image, seed_used)
+        generator: yields (stage2_only_image, result_image, seed_used)
     """
 
     # Hardcode the negative prompt
@@ -150,6 +150,7 @@ def infer(
         num_images_per_prompt=1,
     ).images
     stage2_only_image = stage2_images[0] if stage2_images else None
+    yield stage2_only_image, gr.update(), seed
 
     # --- Combined generation ---
     print(f"Generating with combined LoRAs...")
@@ -179,10 +180,11 @@ def infer(
         if pil_image.size != generated_image.size:
             pil_image = pil_image.resize(generated_image.size, Image.Resampling.LANCZOS)
         blended_image = Image.blend(pil_image, generated_image, alpha=0.75)
-        return stage2_only_image, blended_image, seed
+        yield stage2_only_image, blended_image, seed
+        return
 
     # Return first result image and seed
-    return stage2_only_image, result_images[0] if result_images else None, seed
+    yield stage2_only_image, result_images[0] if result_images else None, seed
 
 # --- Examples and UI Layout ---
 examples = []
@@ -213,7 +215,8 @@ with gr.Blocks(css=css) as demo:
                                        type="pil",
                                        interactive=True,
                                        elem_id="input-image",
-                                       height=350)
+                                       height=380)
+                run_button = gr.Button("🚀 Generate", variant="primary", size="lg")
 
                 gr.HTML("""
                 <script>
@@ -273,8 +276,6 @@ with gr.Blocks(css=css) as demo:
                     gr.Markdown("### 📤 Result2")
                     result = gr.Image(label="Result2", show_label=False, type="pil", interactive=False, height=350)
 
-        run_button = gr.Button("🚀 Generate", variant="primary", size="lg")
-
         with gr.Accordion("Advanced Settings", open=False, visible=False):
             with gr.Row():
                 seed = gr.Slider(
@@ -355,4 +356,4 @@ with gr.Blocks(css=css) as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(mcp_server=True)
+    demo.queue().launch(mcp_server=True)