Refactor inference process in app.py to support stage2-only generation and update output structure to return both stage2 and combined results. Adjust UI layout for improved result display and enhance generator function for reproducibility.

app.py

@@ -6,10 +6,10 @@ import spaces
 
 from PIL import Image
 from diffusers import FlowMatchEulerDiscreteScheduler, QwenImageEditPlusPipeline
-# from optimization import optimize_pipeline_
-# from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
-# from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
-# from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
+from optimization import optimize_pipeline_
+from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
+from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
+from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
 
 import math
 import os
@@ -68,12 +68,12 @@ pipe.load_lora_weights(STAGE1_LORA_REPO, weight_name=STAGE1_LORA_WEIGHT, adapter
 # Load Stage 2 LoRA
 pipe.load_lora_weights(STAGE2_LORA_REPO, weight_name=STAGE2_LORA_WEIGHT, adapter_name="stage2")
 
-# # Apply the same optimizations from the first version
-# pipe.transformer.__class__ = QwenImageTransformer2DModel
-# pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
+# Apply the same optimizations from the first version
+pipe.transformer.__class__ = QwenImageTransformer2DModel
+pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
 
-# # --- Ahead-of-time compilation ---
-# optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], prompt="prompt")
+# --- Ahead-of-time compilation ---
+optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], prompt="prompt")
 
 # --- UI Constants ---
 MAX_SEED = np.iinfo(np.int32).max
@@ -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
@@ -117,8 +117,8 @@ def infer(
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    # Set up the generator for reproducibility
-    generator = torch.Generator(device=device).manual_seed(seed)
+    def make_generator():
+        return torch.Generator(device=device).manual_seed(seed)
 
     # Load input image into PIL Image
     pil_image = None
@@ -131,6 +131,26 @@ def infer(
     if height==256 and width==256:
         height, width = None, 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=make_generator(),
+        true_cfg_scale=true_guidance_scale,
+        num_images_per_prompt=1,
+    ).images
+    stage2_only_image = stage2_images[0] if stage2_images else None
+
     # Apply all LoRAs: Lightning + Stage1 + Stage2
     print(f"Generating with combined LoRAs...")
     print(f"Prompt: '{STAGE1_PROMPT}'")
@@ -147,7 +167,7 @@ def infer(
         width=width,
         negative_prompt=negative_prompt,
         num_inference_steps=num_inference_steps,
-        generator=generator,
+        generator=make_generator(),
         true_cfg_scale=true_guidance_scale,
         num_images_per_prompt=1,
     ).images
@@ -159,10 +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 blended_image, seed
+        return stage2_only_image, blended_image, seed
 
     # Return first result image and 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 = []
@@ -170,7 +190,7 @@ examples = []
 css = """
 #col-container {
     margin: 0 auto;
-    max-width: 900px;
+    max-width: 1000px;
 }
 #logo-title {
     text-align: center;
@@ -192,7 +212,8 @@ with gr.Blocks(css=css) as demo:
                                        show_label=False,
                                        type="pil",
                                        interactive=True,
-                                       elem_id="input-image")
+                                       elem_id="input-image",
+                                       height=350)
 
                 gr.HTML("""
                 <script>
@@ -243,13 +264,18 @@ with gr.Blocks(css=css) as demo:
                 </script>
                 """)
 
-            with gr.Column(scale=1):
-                gr.Markdown("### 📤 Result")
-                result = gr.Image(label="Result", show_label=False, type="pil", interactive=False)
+            with gr.Row(scale=2):
+                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")
+                    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):
+        with gr.Accordion("Advanced Settings", open=False, visible=False):
             with gr.Row():
                 seed = gr.Slider(
                     label="Seed",
@@ -325,7 +351,7 @@ with gr.Blocks(css=css) as demo:
             stage1_weight,
             stage2_weight,
         ],
-        outputs=[result, seed],
+        outputs=[stage2_result, result, seed],
     )
 
 if __name__ == "__main__":