Update inference output in app.py to yield stage2-only images, refining the generator's response structure for better integration with the UI.

app.py

@@ -78,9 +78,9 @@ pipe.load_lora_weights(STAGE2_LORA_REPO, weight_name=STAGE2_LORA_WEIGHT, adapter
 # --- UI Constants ---
 MAX_SEED = np.iinfo(np.int32).max
 
-# --- Main Inference Function (Combined LoRA) ---
+# --- Main Inference Function (Split into two stages) ---
 @spaces.GPU()
-def infer(
+def infer_stage2(
     image,
     seed=42,
     randomize_seed=False,
@@ -88,29 +88,14 @@ def infer(
     num_inference_steps=4,
     height=None,
     width=None,
-    stage1_weight=1.0,
-    stage2_weight=1.0,
     progress=gr.Progress(track_tqdm=True),
 ):
     """
-    Run stage2-only inference, then combined LoRAs: Lightning + Stage1 + Stage2.
-
-    Parameters:
-        image: Input image (PIL Image or path string).
-        seed (int): Random seed for reproducibility.
-        randomize_seed (bool): If True, overrides seed with a random value.
-        true_guidance_scale (float): CFG scale used by Qwen-Image.
-        num_inference_steps (int): Number of diffusion steps.
-        height (int | None): Optional output height override.
-        width (int | None): Optional output width override.
-        stage1_weight (float): Weight for Stage1 LoRA.
-        stage2_weight (float): Weight for Stage2 LoRA.
-        progress: Gradio progress callback.
+    Run stage2-only inference.
 
     Returns:
-        generator: yields (stage2_only_image, result_image, seed_used)
+        (stage2_only_image, image, seed, true_guidance_scale, num_inference_steps, height, width)
     """
-
     # Hardcode the negative prompt
     negative_prompt = " "
 
@@ -150,7 +135,43 @@ def infer(
         num_images_per_prompt=1,
     ).images
     stage2_only_image = stage2_images[0] if stage2_images else None
-    yield stage2_only_image, None, seed
+
+    return stage2_only_image, image, seed, true_guidance_scale, num_inference_steps, height, width
+
+@spaces.GPU()
+def infer_combined(
+    image,
+    seed,
+    true_guidance_scale,
+    num_inference_steps,
+    height,
+    width,
+    stage1_weight,
+    stage2_weight,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """
+    Run combined LoRAs inference.
+
+    Returns:
+        result_image
+    """
+    # Hardcode the negative prompt
+    negative_prompt = " "
+
+    # Set up the generator for reproducibility
+    generator = torch.Generator(device=device).manual_seed(seed)
+
+    # Load input image into PIL Image
+    pil_image = None
+    if image is not None:
+        if isinstance(image, Image.Image):
+            pil_image = image.convert("RGB")
+        elif isinstance(image, str):
+            pil_image = Image.open(image).convert("RGB")
+
+    if height==256 and width==256:
+        height, width = None, None
 
     # --- Combined generation ---
     print(f"Generating with combined LoRAs...")
@@ -180,11 +201,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)
-        yield gr.update(), blended_image, seed
-        return
+        return blended_image
 
-    # Return first result image and seed
-    yield gr.update(), result_images[0] if result_images else None, seed
+    # Return first result image
+    return result_images[0] if result_images else None
 
 # --- Examples and UI Layout ---
 examples = []
@@ -339,8 +359,9 @@ with gr.Blocks(css=css) as demo:
                     value=None,
                 )
 
-    run_button.click(
-        fn=infer,
+    # Chain two inference stages using .then()
+    stage2_event = run_button.click(
+        fn=infer_stage2,
         inputs=[
             input_image,
             seed,
@@ -349,10 +370,23 @@ with gr.Blocks(css=css) as demo:
             num_inference_steps,
             height,
             width,
+        ],
+        outputs=[stage2_result, gr.State(), gr.State(), gr.State(), gr.State(), gr.State(), gr.State()],
+    )
+
+    stage2_event.then(
+        fn=infer_combined,
+        inputs=[
+            stage2_event.outputs[1],  # image
+            stage2_event.outputs[2],  # seed
+            stage2_event.outputs[3],  # true_guidance_scale
+            stage2_event.outputs[4],  # num_inference_steps
+            stage2_event.outputs[5],  # height
+            stage2_event.outputs[6],  # width
             stage1_weight,
             stage2_weight,
         ],
-        outputs=[stage2_result, result, seed],
+        outputs=[result],
     )
 
 if __name__ == "__main__":