Refactor app.py to comment out unused imports and disable stage2-only generation, while updating the inference output structure to return only the combined result. Adjust UI elements for clarity and maintainability.

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 stage2-only inference, then combined LoRAs: Lightning + Stage1 + Stage2.
+    Run single inference with 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: (stage2_only_image, result_image, seed_used)
+        tuple: (result_image, seed_used)
     """
 
     # Hardcode the negative prompt
@@ -117,8 +117,8 @@ def infer(
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    def make_generator():
-        return torch.Generator(device=device).manual_seed(seed)
+    # Set up the generator for reproducibility
+    generator = torch.Generator(device=device).manual_seed(seed)
 
     # Load input image into PIL Image
     pil_image = None
@@ -131,27 +131,27 @@ 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
+    # --- 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
+    #
+    # --- Combined generation ---
     print(f"Generating with combined LoRAs...")
     print(f"Prompt: '{STAGE1_PROMPT}'")
     print(f"Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}, Size: {width}x{height}")
@@ -167,7 +167,7 @@ def infer(
         width=width,
         negative_prompt=negative_prompt,
         num_inference_steps=num_inference_steps,
-        generator=make_generator(),
+        generator=generator,
         true_cfg_scale=true_guidance_scale,
         num_images_per_prompt=1,
     ).images
@@ -179,10 +179,12 @@ 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 stage2_only_image, blended_image, seed
+        return blended_image, seed
 
     # Return first result image and seed
-    return stage2_only_image, result_images[0] if result_images else None, seed
+    # return stage2_only_image, result_images[0] if result_images else None, seed
+    return result_images[0] if result_images else None, seed
 
 # --- Examples and UI Layout ---
 examples = []
@@ -201,7 +203,7 @@ with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
         gr.HTML("""
         <div id="logo-title">
-            <h1>🎨✨ QIE-Image2GuideBody</h1>
+            <h1>🎨✨ Qwen Image Edit 2509 - Visualize Body Structure Lines</h1>
             <h3 style="color: #5b47d1;">Anime Character Converter with Combined LoRAs</h3>
         </div>
         """)
@@ -264,14 +266,13 @@ with gr.Blocks(css=css) as demo:
                 </script>
                 """)
 
-            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("### 🧪 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)
+            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")
 
@@ -351,7 +352,8 @@ with gr.Blocks(css=css) as demo:
             stage1_weight,
             stage2_weight,
         ],
-        outputs=[stage2_result, result, seed],
+        # outputs=[stage2_result, result, seed],
+        outputs=[result, seed],
     )
 
 if __name__ == "__main__":