Update app.py

app.py

@@ -1,144 +1,377 @@
-"""
-Acne-removal demo – Qwen-Image-Edit 4-bit edition  (NO external logo)
-Runs continuously on Hugging-Face Zero-GPU (16 GB)
-"""
 import gradio as gr
-import torch
-import random
 import numpy as np
+import random
+import torch
+import spaces
 from PIL import Image
 import math
 import gc
-import spaces
-from diffusers import (
-    QwenImageEditPlusPipeline,
-    FlowMatchEulerDiscreteScheduler,
-)
-
-# ---------- config ----------
-DTYPE        = torch.float16
-DEVICE       = "cuda" if torch.cuda.is_available() else "cpu"
-MAX_SEED     = np.iinfo(np.int32).max
-PROMPT       = "remove acne marks and blemishes from the face"
-NEG_PROMPT   = " "
-
-# 4-bit model – 4 GB on GPU
-MODEL_ID     = "Qwen/Qwen-Image-Edit-2509-NF4"
-
-scheduler = FlowMatchEulerDiscreteScheduler.from_config({
-    "base_image_seq_len": 256,
-    "base_shift": math.log(3),
-    "max_image_seq_len": 8192,
-    "max_shift": math.log(3),
-    "num_train_timesteps": 1000,
-    "shift": 1.0,
-    "time_shift_type": "exponential",
-    "use_dynamic_shifting": True,
-})
-
-print("🚀 Loading 4-bit NF4 model …")
-pipe = QwenImageEditPlusPipeline.from_pretrained(
-    MODEL_ID,
-    torch_dtype=DTYPE,
-    variant="nf4",
-    use_safetensors=True,
-)
-pipe.scheduler = scheduler
-pipe.enable_attention_slicing(1)
-pipe.enable_vae_tiling()
-pipe.enable_model_cpu_offload()   # keeps only 4-bit weights on GPU
-print("✅ Model ready – <10 GB peak")
-
-# ---------- inference ----------
+import logging
+from typing import List, Optional
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Configuration
+DTYPE = torch.float16
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+MODEL_ID = "Qwen/Qwen-Image-Edit-2509"  # Use standard model [web:44]
+MAX_SEED = np.iinfo(np.int32).max
+HARDCODED_PROMPT = "remove acne marks and blemishes from the face"
+NEGATIVE_PROMPT = " "
+
+# Import pipeline
+try:
+    from diffusers import QwenImageEditPlusPipeline, FlowMatchEulerDiscreteScheduler
+    logger.info("✅ Diffusers imported successfully")
+except ImportError as e:
+    logger.error(f"❌ Import failed: {e}")
+    raise
+
+# Memory management functions
+def cleanup_memory():
+    """Comprehensive memory cleanup"""
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+    gc.collect()
+
+def check_gpu_memory():
+    """Monitor GPU memory usage"""
+    if torch.cuda.is_available():
+        allocated = torch.cuda.memory_allocated() / 1024**3
+        cached = torch.cuda.memory_reserved() / 1024**3
+        logger.info(f"GPU Memory - Allocated: {allocated:.2f}GB, Cached: {cached:.2f}GB")
+
+# Initialize pipeline
+def load_pipeline():
+    """Load and optimize the pipeline"""
+    logger.info(f"🚀 Loading {MODEL_ID}...")
+    
+    # Scheduler configuration [web:39]
+    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,
+    }
+    
+    try:
+        # Create scheduler
+        scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)
+        
+        # Load pipeline [web:38]
+        pipe = QwenImageEditPlusPipeline.from_pretrained(
+            MODEL_ID,
+            scheduler=scheduler,
+            torch_dtype=DTYPE,
+            use_safetensors=True,
+        )
+        
+        # Move to device
+        pipe = pipe.to(DEVICE)
+        
+        # Enable optimizations [web:43]
+        pipe.enable_attention_slicing()  # Memory efficient attention
+        pipe.enable_vae_slicing()        # Sliced VAE decoding
+        pipe.enable_vae_tiling()         # Tiled VAE for large images
+        
+        # Try to load Lightning LoRA for faster inference [web:39]
+        try:
+            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()
+            logger.info("✅ Lightning LoRA loaded (4-step mode)")
+        except Exception as e:
+            logger.warning(f"⚠️ Lightning LoRA skipped: {e}")
+        
+        logger.info("✅ Pipeline loaded and optimized successfully")
+        check_gpu_memory()
+        return pipe
+        
+    except Exception as e:
+        logger.error(f"❌ Pipeline loading failed: {e}")
+        raise
+
+# Load pipeline at startup
+pipe = load_pipeline()
+
 @spaces.GPU()
-def run(
-    gallery,
-    seed=42,
-    randomize_seed=True,
-    guidance=1.0,
-    steps=4,
-    height=512,
-    width=512,
+def infer(
+    images: Optional[List],
+    seed: int = 42,
+    randomize_seed: bool = False,
+    true_guidance_scale: float = 1.0,
+    num_inference_steps: int = 4,
+    height: int = 512,
+    width: int = 512,
     progress=gr.Progress(track_tqdm=True),
 ):
-    torch.cuda.empty_cache()
-    gc.collect()
-
+    """
+    Optimized inference function with proper error handling
+    """
+    # Clean memory before inference
+    cleanup_memory()
+    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-
-    # load / resize images
-    pil_list = []
-    if gallery is not None:
-        for item in gallery:
-            if isinstance(item, Image.Image):
-                img = item.convert("RGB")
-            elif isinstance(item, (list, tuple)):
-                img = item[0].convert("RGB")
-            else:
-                continue
-            img.thumbnail((512, 512), Image.LANCZOS)
-            pil_list.append(img)
-
+    
     generator = torch.Generator(device=DEVICE).manual_seed(seed)
+    
+    # Process input images
+    pil_images = []
+    if images is not None:
+        for item in images:
+            try:
+                # Handle different input types
+                if isinstance(item, tuple) and len(item) > 0:
+                    img_path = item[0]
+                    if isinstance(img_path, Image.Image):
+                        img = img_path.convert("RGB")
+                    elif isinstance(img_path, str):
+                        img = Image.open(img_path).convert("RGB")
+                    else:
+                        continue
+                else:
+                    continue
+                
+                # Resize for memory efficiency [web:38]
+                img.thumbnail((768, 768), Image.Resampling.LANCZOS)
+                pil_images.append(img)
+                
+            except Exception as e:
+                logger.error(f"Error processing image: {e}")
+                continue
+    
+    if not pil_images:
+        raise gr.Error("No valid images provided")
+    
+    logger.info(f"📊 Processing {len(pil_images)} image(s), {height}x{width}, {num_inference_steps} steps")
+    
+    try:
+        # Inference with proper context management [web:27]
+        with torch.inference_mode():
+            with torch.cuda.amp.autocast(enabled=True, dtype=DTYPE):
+                output = pipe(
+                    image=pil_images,
+                    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
+        
+        logger.info("✅ Generation completed successfully")
+        return output, seed, gr.update(visible=True)
+        
+    except torch.cuda.OutOfMemoryError as e:
+        logger.warning("⚠️ CUDA OOM - Trying emergency mode")
+        cleanup_memory()
+        
+        try:
+            # Emergency fallback with reduced settings
+            with torch.inference_mode():
+                with torch.cuda.amp.autocast(enabled=True, dtype=DTYPE):
+                    output = pipe(
+                        image=pil_images,
+                        prompt=HARDCODED_PROMPT,
+                        height=min(height, 384),
+                        width=min(width, 384),
+                        negative_prompt=NEGATIVE_PROMPT,
+                        num_inference_steps=max(2, num_inference_steps // 2),
+                        generator=generator,
+                        true_cfg_scale=1.0,
+                        num_images_per_prompt=1,
+                    ).images
+            
+            logger.info("✅ Emergency mode successful")
+            return output, seed, gr.update(visible=True)
+            
+        except Exception as emergency_e:
+            logger.error(f"❌ Emergency mode failed: {emergency_e}")
+            raise gr.Error(f"GPU memory insufficient. Try smaller images or reduce resolution.")
+    
+    except Exception as e:
+        logger.error(f"❌ Inference failed: {e}")
+        raise gr.Error(f"Generation failed: {str(e)}")
+    
+    finally:
+        # Always clean up after inference [web:32]
+        cleanup_memory()
 
-    # safety shrink
-    if height * width > 512 * 512:
-        height = width = 256
-
-    with torch.inference_mode(), torch.cuda.amp.autocast(dtype=DTYPE):
-        out = pipe(
-            image=pil_list if pil_list else None,
-            prompt=PROMPT,
-            negative_prompt=NEG_PROMPT,
-            height=height,
-            width=width,
-            num_inference_steps=steps,
-            generator=generator,
-            true_cfg_scale=guidance,
-            num_images_per_prompt=1,
-        ).images
-
-    torch.cuda.empty_cache()
-    gc.collect()
-    return out, seed, gr.update(visible=True)
+def use_output_as_input(output_images):
+    """Convert output images to input format"""
+    if output_images is None or len(output_images) == 0:
+        return []
+    return [(img, f"output_{i}.png") for i, img in enumerate(output_images)]
 
-# ---------- UI ----------
+# UI Styles
 css = """
-#col-container{max-width:900px;margin:auto}
+#col-container {
+    margin: 0 auto;
+    max-width: 900px;
+}
+#logo-title {
+    text-align: center;
+}
+#logo-title img {
+    width: 350px;
+}
+.memory-info {
+    font-size: 0.8em;
+    color: #666;
+    margin-top: 5px;
+}
 """
 
-with gr.Blocks(css=css, title="Acne Remover") as demo:
+# Gradio Interface
+with gr.Blocks(css=css, title="Acne Remover - Qwen Image Edit") as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown("## 🚀 Acne Remover – 4-bit edition")
-        gr.Markdown("Upload a facial image and let the model remove acne marks and blemishes.")
-
+        # Header
+        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">
+            <h2 style="font-style: italic;color: #5b47d1;margin-top: -20px">✨ Professional Acne Remover</h2>
+        </div>
+        """)
+        
+        gr.Markdown("""
+        **Remove acne marks and blemishes** using the powerful Qwen-Image-Edit-2509 model.
+        
+        ✅ **State-of-the-art results** with 20B parameter model [web:42]
+        ✅ **Multi-image support** for batch processing [web:45]
+        ✅ **Lightning-fast inference** with 4-step generation [web:39]
+        ✅ **Memory optimized** for stable performance [web:43]
+        """)
+        
         with gr.Row():
             with gr.Column():
-                in_gal = gr.Gallery(label="Upload face", show_label=False, type="pil", interactive=True)
+                input_images = gr.File(
+                    label="📸 Upload facial images", 
+                    file_count="multiple",
+                    file_types=["image"],
+                    height=300
+                )
+                
+                gr.HTML('<div class="memory-info">💡 Tip: Upload multiple images for batch processing</div>')
+
             with gr.Column():
-                out_gal = gr.Gallery(label="Result", show_label=False, type="pil")
-                reuse = gr.Button("↗️ Use as input", size="sm", visible=False)
+                result = gr.Gallery(
+                    label="🎯 Results", 
+                    show_label=True, 
+                    type="pil",
+                    height=300,
+                    columns=2
+                )
+                use_output_btn = gr.Button(
+                    "🔄 Use Results as New Input", 
+                    variant="secondary", 
+                    size="sm", 
+                    visible=False
+                )
+
+        # Main action button
+        run_button = gr.Button(
+            "🚀 Remove Acne & Blemishes!", 
+            variant="primary", 
+            size="lg"
+        )
 
-        run_btn = gr.Button("Remove Acne!", variant="primary", size="lg")
+        # Advanced settings
+        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.Accordion("Advanced", open=False):
-            seed_s = gr.Slider(0, MAX_SEED, step=1, value=42, label="Seed")
-            rand_c = gr.Checkbox(True, label="Randomise seed")
             with gr.Row():
-                guid_s = gr.Slider(1.0, 5.0, step=0.5, value=1.0, label="Guidance")
-                steps_s = gr.Slider(2, 20, step=2, value=4, label="Steps")
+                true_guidance_scale = gr.Slider(
+                    label="📊 Guidance Scale",
+                    minimum=1.0,
+                    maximum=5.0,
+                    step=0.1,
+                    value=1.0,
+                    info="Higher values = stronger prompt adherence"
+                )
+
+                num_inference_steps = gr.Slider(
+                    label="🔄 Inference Steps",
+                    minimum=2,
+                    maximum=20,
+                    step=1,
+                    value=4,
+                    info="More steps = higher quality (slower)"
+                )
+                
             with gr.Row():
-                h_s = gr.Slider(256, 768, step=64, value=512, label="Height")
-                w_s = gr.Slider(256, 768, step=64, value=512, label="Width")
-
-    # events
-    run_btn.click(
-        run,
-        [in_gal, seed_s, rand_c, guid_s, steps_s, h_s, w_s],
-        [out_gal, seed_s, reuse],
+                height = gr.Slider(
+                    label="📏 Height", 
+                    minimum=256, 
+                    maximum=768, 
+                    step=64, 
+                    value=512
+                )
+                width = gr.Slider(
+                    label="📐 Width", 
+                    minimum=256, 
+                    maximum=768, 
+                    step=64, 
+                    value=512
+                )
+
+        # Footer info
+        gr.Markdown("""
+        ---
+        **Model Info**: Qwen-Image-Edit-2509 | **Memory**: Optimized for GPU efficiency | **Speed**: ~4 steps with Lightning LoRA
+        """)
+
+    # Event handlers
+    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],
+        show_progress=True
+    )
+
+    use_output_btn.click(
+        fn=use_output_as_input, 
+        inputs=[result], 
+        outputs=[input_images]
     )
-    reuse.click(lambda x: x, out_gal, in_gal)
 
+# Launch configuration
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        show_error=True,
+        quiet=False
+    )