Update app.py

app.py

@@ -5,7 +5,12 @@ import torch
 import random
 import os
 
-from diffusers import QwenImageEditInpaintPipeline
+# from diffusers import QwenImageEditInpaintPipeline
+from optimization import optimize_pipeline_
+from qwenimage.pipeline_qwen_image_edit import QwenImageEditInpaintPipeline
+from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
+from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
+
 from PIL import Image
 
 # Set environment variable for parallel loading
@@ -16,8 +21,13 @@ MAX_IMAGE_SIZE = 2048
 
 # Initialize Qwen Image Edit pipeline
 pipe = QwenImageEditInpaintPipeline.from_pretrained("Qwen/Qwen-Image-Edit", torch_dtype=torch.bfloat16).to("cuda")
+pipe.transformer.__class__ = QwenImageTransformer2DModel
+pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
+
+# --- Ahead-of-time compilation ---
+optimize_pipeline_(pipe, image=Image.new("RGB", (1024, 1024)), prompt="prompt")
 
-@spaces.GPU
+@spaces.GPU(duration=120)
 def infer(edit_images, prompt, negative_prompt="", seed=42, randomize_seed=False, strength=1.0, num_inference_steps=35, true_cfg_scale=4.0, progress=gr.Progress(track_tqdm=True)):
     image = edit_images["background"]
     mask = edit_images["layers"][0]
@@ -55,9 +65,18 @@ css="""
 with gr.Blocks(css=css) as demo:
     
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"""# Qwen Image Edit Inpainting
-Advanced image inpainting using Qwen's Image Edit model  
-[[model](https://huggingface.co/Qwen/Qwen-Image-Edit)] [[paper](https://arxiv.org/abs/2412.20710)]
+        gr.HTML("""
+        <div id="logo-title">
+            <img src="https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/qwen_image_edit_logo.png" alt="Qwen-Image Edit Logo" width="400" style="display: block; margin: 0 auto;">
+            <h2 style="font-style: italic;color: #5b47d1;margin-top: -27px !important;margin-left: 133px;">Inapint</h2>
+        </div>
+        """)
+        gr.Markdown("""
+        
+        Inpaint images with Qwen Image Edit. [Learn more](https://github.com/QwenLM/Qwen-Image) about the Qwen-Image series. 
+        
+        This demo uses the [Qwen-Image-Lightning](https://huggingface.co/lightx2v/Qwen-Image-Lightning) LoRA with AoT compilation and FA3 for accelerated 8-step inference.
+        Try on [Qwen Chat](https://chat.qwen.ai/), or [download model](https://huggingface.co/Qwen/Qwen-Image-Edit) to run locally with ComfyUI or diffusers.
         """)
         with gr.Row():
             with gr.Column():