Enhance inference process in app.py to support stage2-only generation, updating the output structure to return both stage2 and combined results. Adjust UI elements for improved clarity and maintainability.

app.py

@@ -93,7 +93,7 @@ def infer(
     progress=gr.Progress(track_tqdm=True),
 ):
     """
-    Run single inference with combined LoRAs: Lightning + Stage1 + Stage2.
+    Run stage2-only inference, then combined LoRAs: Lightning + Stage1 + Stage2.
 
     Parameters:
         image: Input image (PIL Image or path string).
@@ -108,7 +108,7 @@ def infer(
         progress: Gradio progress callback.
 
     Returns:
-        tuple: (result_image, seed_used)
+        tuple: (stage2_only_image, result_image, seed_used)
     """
 
     # Hardcode the negative prompt
@@ -131,26 +131,26 @@ def infer(
     if height==256 and width==256:
         height, width = None, None
 
-    # --- Stage2-only generation (disabled) ---
-    # print("Generating with Stage2 LoRA only...")
-    # print(f"Prompt: '{STAGE2_PROMPT}'")
-    # print(f"Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}, Size: {width}x{height}")
-    # print("LoRA Weights - Stage2: 1.0")
-    #
-    # pipe.set_adapters(["stage2"], adapter_weights=[1.0])
-    # stage2_images = pipe(
-    #     image=[pil_image] if pil_image is not None else None,
-    #     prompt=STAGE2_PROMPT,
-    #     height=height,
-    #     width=width,
-    #     negative_prompt=negative_prompt,
-    #     num_inference_steps=num_inference_steps,
-    #     generator=generator,
-    #     true_cfg_scale=true_guidance_scale,
-    #     num_images_per_prompt=1,
-    # ).images
-    # stage2_only_image = stage2_images[0] if stage2_images else None
-    #
+    # Stage2-only generation
+    print("Generating with Stage2 LoRA only...")
+    print(f"Prompt: '{STAGE2_PROMPT}'")
+    print(f"Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}, Size: {width}x{height}")
+    print("LoRA Weights - Stage2: 1.0")
+
+    pipe.set_adapters(["stage2"], adapter_weights=[1.0])
+    stage2_images = pipe(
+        image=[pil_image] if pil_image is not None else None,
+        prompt=STAGE2_PROMPT,
+        height=height,
+        width=width,
+        negative_prompt=negative_prompt,
+        num_inference_steps=num_inference_steps,
+        generator=generator,
+        true_cfg_scale=true_guidance_scale,
+        num_images_per_prompt=1,
+    ).images
+    stage2_only_image = stage2_images[0] if stage2_images else None
+
     # --- Combined generation ---
     print(f"Generating with combined LoRAs...")
     print(f"Prompt: '{STAGE1_PROMPT}'")
@@ -179,12 +179,10 @@ 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
-        return blended_image, seed
+        return stage2_only_image, blended_image, seed
 
     # Return first result image and seed
-    # return stage2_only_image, result_images[0] if result_images else None, seed
-    return result_images[0] if result_images else None, seed
+    return stage2_only_image, result_images[0] if result_images else None, seed
 
 # --- Examples and UI Layout ---
 examples = []
@@ -266,9 +264,9 @@ with gr.Blocks(css=css) as demo:
                 </script>
                 """)
 
-            # with gr.Column(scale=1):
-            #     gr.Markdown("### 🧪 Result1")
-            #     stage2_result = gr.Image(label="Result1", show_label=False, type="pil", interactive=False, height=350)
+            with gr.Column(scale=1):
+                gr.Markdown("### 🧪 Result1")
+                stage2_result = gr.Image(label="Result1", show_label=False, type="pil", interactive=False, height=350)
 
             with gr.Column(scale=1):
                 gr.Markdown("### 📤 Result2")
@@ -352,8 +350,7 @@ with gr.Blocks(css=css) as demo:
             stage1_weight,
             stage2_weight,
         ],
-        # outputs=[stage2_result, result, seed],
-        outputs=[result, seed],
+        outputs=[stage2_result, result, seed],
     )
 
 if __name__ == "__main__":