Update app.py

app.py

@@ -1,179 +1,272 @@
 import gradio as gr
+import numpy as np
+import random
 import torch
-import os
+import spaces
 from PIL import Image
-from spaces import GPU
-from diffusers import QwenImageEditPipeline
-from diffusers.utils import load_image
+from diffusers import FlowMatchEulerDiscreteScheduler
+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
 
-# Model configuration
-MODEL_ID = "Qwen/Qwen-Image-Edit"
+# --- Model Loading & Optimization ---
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Global pipeline variable
-pipeline = None
+# Scheduler configuration for Lightning
+scheduler_config = {
+    "base_image_seq_len": 256,
+    "base_shift": math.log(3),
+    "invert_sigmas": False,
+    "max_image_seq_len": 8192,
+    "max_shift": math.log(3),
+    "num_train_timesteps": 1000,
+    "shift": 1.0,
+    "shift_terminal": None,
+    "stochastic_sampling": False,
+    "time_shift_type": "exponential",
+    "use_beta_sigmas": False,
+    "use_dynamic_shifting": True,
+    "use_exponential_sigmas": False,
+    "use_karras_sigmas": False,
+}
 
-def load_pipeline():
-    """Load the Qwen Image Edit pipeline"""
-    global pipeline
-    
-    try:
-        pipeline = QwenImageEditPipeline.from_pretrained(
-            MODEL_ID,
-            torch_dtype=torch.bfloat16,
-            device_map="auto",
-            trust_remote_code=True
-        )
-        print("✓ Pipeline loaded successfully")
-    except Exception as e:
-        print(f"Error loading pipeline: {e}")
-        raise
-
-@GPU
-def edit_image(image, prompt, negative_prompt="", num_steps=8, guidance_scale=2.5, seed=0):
-    """Edit image based on text prompt using ZeroGPU"""
+# Initialize scheduler with Lightning config
+scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)
+
+# Load the model pipeline
+pipe = QwenImageEditPlusPipeline.from_pretrained(
+    "Qwen/Qwen-Image-Edit-2509", 
+    scheduler=scheduler,
+    torch_dtype=dtype
+).to(device)
+
+pipe.load_lora_weights(
+    "lightx2v/Qwen-Image-Lightning", 
+    weight_name="Qwen-Image-Edit-2509/Qwen-Image-Edit-2509-Lightning-8steps-V1.0-bf16.safetensors"
+)
+pipe.fuse_lora()
+
+# Apply optimizations
+pipe.transformer.__class__ = QwenImageTransformer2DModel
+pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
+
+# Enable memory optimizations
+pipe.enable_attention_slicing()
+
+# Ahead-of-time compilation for faster subsequent runs
+optimize_pipeline_(
+    pipe, 
+    image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], 
+    prompt="remove acne marks and blemishes from the face"
+)
+
+# --- UI Constants ---
+MAX_SEED = np.iinfo(np.int32).max
+
+# Hardcoded prompt for acne removal
+HARDCODED_PROMPT = "remove acne marks and blemishes from the face"
+NEGATIVE_PROMPT = " "
+
+# --- Main Inference Function (Optimized for Speed) ---
+@spaces.GPU()
+def infer(
+    images,
+    seed=42,
+    randomize_seed=False,
+    true_guidance_scale=1.0,
+    num_inference_steps=8,
+    height=1024,
+    width=1024,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """
+    Optimized inference for acne removal with hardcoded prompt.
+    Removes prompt rewriting to save inference time.
+    """
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+
+    # Set up generator for reproducibility
+    generator = torch.Generator(device=device).manual_seed(seed)
     
-    global pipeline
+    # Load and preprocess input images
+    pil_images = []
+    if images is not None:
+        for item in images:
+            try:
+                if isinstance(item[0], Image.Image):
+                    img = item[0].convert("RGB")
+                    # Resize to optimal inference size for speed
+                    img.thumbnail((1024, 1024), Image.Resampling.LANCZOS)
+                    pil_images.append(img)
+                elif isinstance(item[0], str):
+                    img = Image.open(item[0]).convert("RGB")
+                    img.thumbnail((1024, 1024), Image.Resampling.LANCZOS)
+                    pil_images.append(img)
+                elif hasattr(item, "name"):
+                    img = Image.open(item.name).convert("RGB")
+                    img.thumbnail((1024, 1024), Image.Resampling.LANCZOS)
+                    pil_images.append(img)
+            except Exception as e:
+                print(f"Error loading image: {e}")
+                continue
+
+    print(f"Using hardcoded prompt: '{HARDCODED_PROMPT}'")
+    print(f"Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}")
     
-    try:
-        # Load pipeline on first call
-        if pipeline is None:
-            print("Loading Qwen Image Edit pipeline...")
-            load_pipeline()
-        
-        # Ensure image is RGB and proper size
-        if isinstance(image, str):
-            img = load_image(image).convert("RGB")
-        else:
-            img = image.convert("RGB")
-        
-        # Resize to optimal size (model supports up to 1024x1024)
-        max_size = 768
-        img.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
-        
-        # Set random seed for reproducibility
-        generator = torch.manual_seed(seed)
-        
-        # Prepare inputs
-        inputs = {
-            "image": img,
-            "prompt": prompt,
-            "generator": generator,
-            "guidance_scale": guidance_scale,
-            "num_inference_steps": num_steps,
-        }
-        
-        # Add negative prompt if provided
-        if negative_prompt.strip():
-            inputs["negative_prompt"] = negative_prompt
+    # Generate the image with optimized settings
+    with torch.inference_mode():
+        output = pipe(
+            image=pil_images if len(pil_images) > 0 else None,
+            prompt=HARDCODED_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
+
+    return output, seed, gr.update(visible=True)
+
+
+def use_output_as_input(output_images):
+    """Convert output images to input format for the gallery"""
+    if output_images is None or len(output_images) == 0:
+        return []
+    return output_images
+
+
+# --- CSS Styling ---
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 1024px;
+}
+#logo-title {
+    text-align: center;
+}
+#logo-title img {
+    width: 400px;
+}
+#edit_text {
+    margin-top: -62px !important;
+}
+"""
+
+# --- UI Layout ---
+with gr.Blocks(css=css) as demo:
+    with gr.Column(elem_id="col-container"):
+        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: 96px">[Acne Remover] Fast 8-step Lightning LoRA</h2>
+        </div>
+        """)
+        gr.Markdown("""
+        **Remove acne marks and blemishes** from facial images using Qwen-Image-Edit with Lightning LoRA optimization.
         
-        # Run inference
-        with torch.no_grad():
-            output = pipeline(**inputs)
-            output_image = output.images[0]
+        This demo uses [Qwen-Image-Edit-2509](https://huggingface.co/Qwen/Qwen-Image-Edit-2509) with 
+        [Qwen-Image-Lightning](https://huggingface.co/lightx2v/Qwen-Image-Lightning) + FA3 for ultra-fast inference.
         
-        return output_image
-    
-    except Exception as e:
-        print(f"Error during inference: {e}")
-        raise
-
-# Create Gradio interface
-def create_interface():
-    with gr.Blocks(title="Qwen Image Edit - ZeroGPU", theme=gr.themes.Soft()) as demo:
-        gr.Markdown("# 🎨 Qwen Image Edit with ZeroGPU")
-        gr.Markdown(
-            "Edit images using natural language prompts powered by Qwen Image Edit on ZeroGPU."
-        )
+        [Learn more](https://github.com/QwenLM/Qwen-Image) | [Download model](https://huggingface.co/Qwen/Qwen-Image-Edit-2509)
+        """)
         
         with gr.Row():
             with gr.Column():
-                # Input image
-                input_image = gr.Image(
-                    label="Upload Image",
-                    type="pil",
-                    height=400
+                input_images = gr.Gallery(
+                    label="Upload facial image",
+                    show_label=False, 
+                    type="pil", 
+                    interactive=True
                 )
-                
-                # Edit prompt
-                prompt = gr.Textbox(
-                    label="Edit Instructions",
-                    placeholder="e.g., 'Add a red hat to the person' or 'Change background to sunset'",
-                    lines=2
+
+            with gr.Column():
+                result = gr.Gallery(
+                    label="Acne-removed result", 
+                    show_label=False, 
+                    type="pil"
+                )
+                use_output_btn = gr.Button(
+                    "↗️ Use as input", 
+                    variant="secondary", 
+                    size="sm", 
+                    visible=False
+                )
+
+        with gr.Row():
+            run_button = gr.Button("Remove Acne!", variant="primary", size="lg")
+
+        with gr.Accordion("Advanced Settings", open=False):
+            seed = gr.Slider(
+                label="Seed",
+                minimum=0,
+                maximum=MAX_SEED,
+                step=1,
+                value=0,
+            )
+
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
+            with gr.Row():
+                true_guidance_scale = gr.Slider(
+                    label="Guidance scale",
+                    minimum=1.0,
+                    maximum=10.0,
+                    step=0.1,
+                    value=1.0
+                )
+
+                num_inference_steps = gr.Slider(
+                    label="Inference steps (fewer = faster)",
+                    minimum=1,
+                    maximum=40,
+                    step=1,
+                    value=8,
                 )
                 
-                # Advanced options
-                with gr.Accordion("Advanced Options", open=False):
-                    negative_prompt = gr.Textbox(
-                        label="Negative Prompt (Optional)",
-                        placeholder="What to avoid",
-                        lines=2
-                    )
-                    
-                    num_steps = gr.Slider(
-                        label="Inference Steps",
-                        minimum=4,
-                        maximum=20,
-                        value=8,
-                        step=1,
-                        info="More steps = better quality but slower"
-                    )
-                    
-                    guidance_scale = gr.Slider(
-                        label="Guidance Scale",
-                        minimum=1.0,
-                        maximum=7.5,
-                        value=2.5,
-                        step=0.1,
-                        info="Higher = stronger prompt adherence"
-                    )
-                    
-                    seed = gr.Slider(
-                        label="Seed",
-                        minimum=0,
-                        maximum=2147483647,
-                        value=0,
-                        step=1,
-                        info="For reproducible results"
-                    )
-                
-                # Submit button
-                submit_btn = gr.Button(
-                    "🚀 Edit Image",
-                    variant="primary",
-                    scale=1
+            with gr.Row():
+                height = gr.Slider(
+                    label="Height",
+                    minimum=512,
+                    maximum=1024,
+                    step=64,
+                    value=1024,
                 )
-            
-            with gr.Column():
-                # Output image
-                output_image = gr.Image(
-                    label="Edited Image",
-                    type="pil",
-                    height=400
+                
+                width = gr.Slider(
+                    label="Width",
+                    minimum=512,
+                    maximum=1024,
+                    step=64,
+                    value=1024,
                 )
-        
-        # Connect the function
-        submit_btn.click(
-            fn=edit_image,
-            inputs=[input_image, prompt, negative_prompt, num_steps, guidance_scale, seed],
-            outputs=output_image,
-            show_progress=True
-        )
-        
-        gr.Markdown("""
-        ### 💡 Tips for best results:
-        - Use clear, descriptive prompts
-        - Start with 8 steps (Qwen Image Edit is optimized for low step counts)
-        - Guidance scale 2.0-3.0 for subtle edits, 4.0-6.0 for stronger changes
-        - Image editing takes ~30-60 seconds per inference on ZeroGPU
-        
-        ### ⚙️ Model Info:
-        - **Model**: Qwen/Qwen-Image-Edit (20B parameters)
-        - **Architecture**: Multi-modal Diffusion Transformer
-        - **Input Resolution**: Up to 768x768 (optimized)
-        """)
-    
-    return demo
+
+    # Event handlers
+    gr.on(
+        triggers=[run_button.click],
+        fn=infer,
+        inputs=[
+            input_images,
+            seed,
+            randomize_seed,
+            true_guidance_scale,
+            num_inference_steps,
+            height,
+            width,
+        ],
+        outputs=[result, seed, use_output_btn],
+    )
+
+    use_output_btn.click(
+        fn=use_output_as_input,
+        inputs=[result],
+        outputs=[input_images]
+    )
 
 if __name__ == "__main__":
-    demo = create_interface()
     demo.launch()