Update app.py

app.py

@@ -1,11 +1,11 @@
 """
 UniPic-3 DMD Multi-Image Composition
-Hugging Face Space - UI Persistent + GPU On-Demand Architecture
+Hugging Face Space - ZeroGPU 优化版本
 
-核心优化：
-1. UI 常驻 - 页面始终可用，无需等待模型加载
-2. GPU on-demand - 仅在推理时调用 GPU，节省资源
-3. 优化的前端界面 - 现代美观的 UI 设计
+架构说明：
+1. 模型在全局作用域加载（ZeroGPU 会拦截 CUDA 调用）
+2. 只有实际推理时才使用 @spaces.GPU 装饰器
+3. 这样避免了每次请求都重新加载模型
 """
 
 import gradio as gr
@@ -16,14 +16,17 @@ import sys
 
 # Hugging Face Spaces GPU decorator
 try:
-    from spaces import GPU
+    import spaces
     HF_SPACES = True
 except ImportError:
     HF_SPACES = False
-    def GPU(duration=60):
-        def decorator(func):
-            return func
-        return decorator
+    # 本地开发时的 mock
+    class spaces:
+        @staticmethod
+        def GPU(duration=60):
+            def decorator(func):
+                return func
+            return decorator
 
 # Local pipeline import
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
@@ -33,18 +36,119 @@ MODEL_NAME = os.environ.get("MODEL_NAME", "Skywork/Unipic3-DMD")
 TRANSFORMER_PATH = os.environ.get("TRANSFORMER_PATH", "Skywork/Unipic3-DMD/ema_transformer")
 
 # ============================================================
-# GPU On-Demand: Model loading happens inside @GPU decorated function
+# 全局加载模型（ZeroGPU 会拦截 CUDA 调用）
 # ============================================================
 
-def get_device():
-    """Get the appropriate device"""
-    return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print("🚀 Loading models...")
 
-def get_dtype():
-    """Get the appropriate dtype"""
-    return torch.bfloat16 if torch.cuda.is_available() else torch.float32
+try:
+    from pipeline_qwenimage_edit import QwenImageEditPipeline
+except ImportError:
+    from diffusers import QwenImageEditPipeline
+
+from diffusers import (
+    FlowMatchEulerDiscreteScheduler, 
+    QwenImageTransformer2DModel, 
+    AutoencoderKLQwenImage
+)
+from transformers import AutoModel, AutoTokenizer, Qwen2VLProcessor
+
+# 确定 dtype
+dtype = torch.bfloat16
+
+# Load scheduler (CPU)
+print("  Loading scheduler...")
+scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+    MODEL_NAME, subfolder='scheduler'
+)
+
+# Load tokenizer & processor (CPU)
+print("  Loading tokenizer & processor...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, subfolder='tokenizer')
+processor = Qwen2VLProcessor.from_pretrained(MODEL_NAME, subfolder='processor')
+
+# Load text encoder
+print("  Loading text_encoder...")
+text_encoder = AutoModel.from_pretrained(
+    MODEL_NAME,
+    subfolder='text_encoder',
+    torch_dtype=dtype,
+).eval()
+
+# Load transformer
+print("  Loading transformer...")
+def load_transformer():
+    """Load transformer with proper path handling"""
+    if os.path.exists(TRANSFORMER_PATH):
+        # Local path
+        if os.path.isdir(TRANSFORMER_PATH):
+            config_path = os.path.join(TRANSFORMER_PATH, "config.json")
+            if os.path.exists(config_path):
+                return QwenImageTransformer2DModel.from_pretrained(
+                    TRANSFORMER_PATH,
+                    torch_dtype=dtype,
+                    use_safetensors=False
+                ).eval()
+            else:
+                return QwenImageTransformer2DModel.from_pretrained(
+                    TRANSFORMER_PATH,
+                    subfolder='transformer',
+                    torch_dtype=dtype,
+                    use_safetensors=False
+                ).eval()
+        raise ValueError(f"Invalid transformer path: {TRANSFORMER_PATH}")
+    else:
+        # HuggingFace repo path
+        path_parts = TRANSFORMER_PATH.split('/')
+        if len(path_parts) >= 3:
+            repo_id = '/'.join(path_parts[:2])
+            subfolder = '/'.join(path_parts[2:])
+            return QwenImageTransformer2DModel.from_pretrained(
+                repo_id,
+                subfolder=subfolder,
+                torch_dtype=dtype,
+                use_safetensors=False
+            ).eval()
+        else:
+            return QwenImageTransformer2DModel.from_pretrained(
+                TRANSFORMER_PATH,
+                subfolder='transformer',
+                torch_dtype=dtype,
+                use_safetensors=False
+            ).eval()
+
+transformer = load_transformer()
+
+# Load VAE
+print("  Loading VAE...")
+vae = AutoencoderKLQwenImage.from_pretrained(
+    MODEL_NAME,
+    subfolder='vae',
+    torch_dtype=dtype,
+).eval()
+
+# Create pipeline
+print("  Creating pipeline...")
+pipe = QwenImageEditPipeline(
+    scheduler=scheduler,
+    vae=vae,
+    text_encoder=text_encoder,
+    tokenizer=tokenizer,
+    processor=processor,
+    transformer=transformer
+)
+
+# 移动到 CUDA（ZeroGPU 会拦截这个调用）
+pipe.to('cuda')
+
+print("✅ Models loaded successfully!")
 
-@GPU(duration=180)
+
+# ============================================================
+# GPU 推理函数（只包含实际的推理逻辑）
+# ============================================================
+
+@spaces.GPU(duration=120)
 def generate_image(
     images: list[Image.Image],
     prompt: str,
@@ -53,81 +157,16 @@ def generate_image(
     num_steps: int
 ) -> Image.Image:
     """
-    GPU on-demand inference function.
-    Model is loaded fresh each call to work with ZeroGPU.
+    GPU 推理函数 - 只包含实际的推理逻辑
+    模型已在全局加载，这里只执行推理
     """
-    # Import dependencies inside GPU function for ZeroGPU compatibility
-    try:
-        from pipeline_qwenimage_edit import QwenImageEditPipeline
-    except ImportError:
-        from diffusers import QwenImageEditPipeline
-    
-    from diffusers import (
-        FlowMatchEulerDiscreteScheduler, 
-        QwenImageTransformer2DModel, 
-        AutoencoderKLQwenImage
-    )
-    from transformers import AutoModel, AutoTokenizer, Qwen2VLProcessor
-    
-    # ZeroGPU: 必须在 @GPU 函数内部获取设备
-    device = torch.device("cuda:0")  # 明确指定 cuda:0
-    dtype = torch.bfloat16
-    
-    print(f"🚀 Loading model on {device}...")
-    print(f"   CUDA available: {torch.cuda.is_available()}")
-    print(f"   CUDA device count: {torch.cuda.device_count()}")
-    
-    # Load scheduler (CPU, no device needed)
-    scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
-        MODEL_NAME, subfolder='scheduler'
-    )
-    
-    # Load tokenizer & processor (CPU, no device needed)
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, subfolder='tokenizer')
-    processor = Qwen2VLProcessor.from_pretrained(MODEL_NAME, subfolder='processor')
-    
-    # Load text encoder - 直接加载到 CUDA
-    print("  Loading text_encoder...")
-    text_encoder = AutoModel.from_pretrained(
-        MODEL_NAME,
-        subfolder='text_encoder',
-        torch_dtype=dtype,
-    ).to(device).eval()
-    
-    # Load transformer
-    print("  Loading transformer...")
-    transformer = load_transformer(device, dtype)
-    
-    # Load VAE
-    print("  Loading VAE...")
-    vae = AutoencoderKLQwenImage.from_pretrained(
-        MODEL_NAME,
-        subfolder='vae',
-        torch_dtype=dtype,
-    ).to(device).eval()
-    
-    # Create pipeline
-    pipe = QwenImageEditPipeline(
-        scheduler=scheduler,
-        vae=vae,
-        text_encoder=text_encoder,
-        tokenizer=tokenizer,
-        processor=processor,
-        transformer=transformer
-    )
-    
-    # 注意：不需要手动设置 _execution_device
-    # 修复后的 pipeline_qwenimage_edit.py 会直接从 text_encoder 获取设备
-    
-    print(f"✅ Model loaded!")
-    print(f"   text_encoder device: {next(text_encoder.parameters()).device}")
-    print(f"   transformer device: {next(transformer.parameters()).device}")
-    print(f"   vae device: {next(vae.parameters()).device}")
-    print(f"   Generating with {len(images)} image(s)...")
+    print(f"🎨 Generating with {len(images)} image(s)...")
+    print(f"   Prompt: {prompt[:50]}...")
+    print(f"   Steps: {num_steps}, CFG: {true_cfg_scale}, Seed: {seed}")
     
     # Generate
     with torch.no_grad():
-        generator = torch.Generator(device=device).manual_seed(int(seed))
+        generator = torch.Generator(device='cuda').manual_seed(int(seed))
         
         if len(images) == 1:
             result = pipe(
@@ -152,60 +191,12 @@ def generate_image(
                 generator=generator
             ).images[0]
     
-    # Cleanup to free VRAM
-    del pipe, transformer, vae, text_encoder
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-    
+    print("✅ Generation complete!")
     return result
 
 
-def load_transformer(device, dtype):
-    """Load transformer with proper path handling for ZeroGPU"""
-    from diffusers import QwenImageTransformer2DModel
-    
-    if os.path.exists(TRANSFORMER_PATH):
-        # Local path
-        if os.path.isdir(TRANSFORMER_PATH):
-            config_path = os.path.join(TRANSFORMER_PATH, "config.json")
-            if os.path.exists(config_path):
-                return QwenImageTransformer2DModel.from_pretrained(
-                    TRANSFORMER_PATH,
-                    torch_dtype=dtype,
-                    use_safetensors=False  # 使用 .bin 文件
-                ).to(device).eval()
-            else:
-                return QwenImageTransformer2DModel.from_pretrained(
-                    TRANSFORMER_PATH,
-                    subfolder='transformer',
-                    torch_dtype=dtype,
-                    use_safetensors=False
-                ).to(device).eval()
-        raise ValueError(f"Invalid transformer path: {TRANSFORMER_PATH}")
-    else:
-        # HuggingFace repo path
-        path_parts = TRANSFORMER_PATH.split('/')
-        if len(path_parts) >= 3:
-            # 路径格式: "Skywork/Unipic3-DMD/ema_transformer"
-            repo_id = '/'.join(path_parts[:2])  # "Skywork/Unipic3-DMD"
-            subfolder = '/'.join(path_parts[2:])  # "ema_transformer"
-            return QwenImageTransformer2DModel.from_pretrained(
-                repo_id,
-                subfolder=subfolder,
-                torch_dtype=dtype,
-                use_safetensors=False  # 使用 .bin 文件
-            ).to(device).eval()
-        else:
-            return QwenImageTransformer2DModel.from_pretrained(
-                TRANSFORMER_PATH,
-                subfolder='transformer',
-                torch_dtype=dtype,
-                use_safetensors=False
-            ).to(device).eval()
-
-
 # ============================================================
-# UI Logic (CPU-only, always available)
+# UI 逻辑（CPU，始终可用）
 # ============================================================
 
 def process_images(
@@ -215,12 +206,12 @@ def process_images(
     seed: int,
     num_steps: int
 ):
-    """Process images - validates input then calls GPU function"""
+    """处理图像 - 验证输入后调用 GPU 函数"""
     
-    # Filter valid images
+    # 过滤有效图像
     images = [img for img in [img1, img2, img3, img4, img5, img6] if img is not None]
     
-    # Validation
+    # 验证
     if len(images) == 0:
         return None, "❌ Please upload at least one image"
     
@@ -231,10 +222,10 @@ def process_images(
         return None, "❌ Please enter an editing instruction"
     
     try:
-        # Convert to RGB
+        # 转换为 RGB
         images = [img.convert("RGB") for img in images]
         
-        # Call GPU function
+        # 调用 GPU 函数
         result = generate_image(
             images=images,
             prompt=prompt,
@@ -252,19 +243,17 @@ def process_images(
 
 
 def update_image_visibility(num):
-    """Update visibility of image upload slots"""
+    """更新图像上传槽的可见性"""
     return [gr.update(visible=(i < num)) for i in range(6)]
 
 
 # ============================================================
-# Custom CSS for Beautiful UI
+# 自定义 CSS
 # ============================================================
 
 CUSTOM_CSS = """
-/* Import distinctive fonts */
 @import url('https://fonts.googleapis.com/css2?family=Outfit:wght@300;400;500;600;700&family=JetBrains+Mono:wght@400;500&display=swap');
 
-/* Root variables */
 :root {
     --primary: #6366f1;
     --primary-dark: #4f46e5;
@@ -278,18 +267,15 @@ CUSTOM_CSS = """
     --success: #10b981;
     --error: #ef4444;
     --gradient-1: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-    --gradient-2: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
     --gradient-hero: linear-gradient(135deg, #0f0f23 0%, #1a1a3e 50%, #252552 100%);
 }
 
-/* Global styles */
 .gradio-container {
     font-family: 'Outfit', sans-serif !important;
     background: var(--gradient-hero) !important;
     min-height: 100vh;
 }
 
-/* Header styling */
 .main-header {
     text-align: center;
     padding: 2rem 1rem;
@@ -316,7 +302,6 @@ CUSTOM_CSS = """
     margin: 0 auto;
 }
 
-/* Feature badges */
 .feature-badges {
     display: flex;
     gap: 1rem;
@@ -338,7 +323,6 @@ CUSTOM_CSS = """
     font-weight: 500;
 }
 
-/* Section headers */
 .section-header {
     display: flex;
     align-items: center;
@@ -355,22 +339,6 @@ CUSTOM_CSS = """
     margin: 0;
 }
 
-/* Card styling */
-.card {
-    background: var(--surface-light) !important;
-    border: 1px solid var(--border) !important;
-    border-radius: 16px !important;
-    padding: 1.5rem !important;
-}
-
-/* Image upload grid */
-.image-grid {
-    display: grid;
-    grid-template-columns: repeat(3, 1fr);
-    gap: 1rem;
-}
-
-/* Button styling */
 .generate-btn {
     background: var(--gradient-1) !important;
     border: none !important;
@@ -389,28 +357,6 @@ CUSTOM_CSS = """
     box-shadow: 0 6px 20px rgba(99, 102, 241, 0.5) !important;
 }
 
-/* Input styling */
-.gr-textbox textarea,
-.gr-textbox input {
-    background: var(--surface) !important;
-    border: 1px solid var(--border) !important;
-    border-radius: 12px !important;
-    color: var(--text) !important;
-    font-family: 'Outfit', sans-serif !important;
-}
-
-.gr-textbox textarea:focus,
-.gr-textbox input:focus {
-    border-color: var(--primary) !important;
-    box-shadow: 0 0 0 3px rgba(99, 102, 241, 0.2) !important;
-}
-
-/* Slider styling */
-.gr-slider input[type="range"] {
-    accent-color: var(--primary) !important;
-}
-
-/* Output image */
 .output-image {
     border-radius: 16px;
     overflow: hidden;
@@ -419,46 +365,10 @@ CUSTOM_CSS = """
                 var(--gradient-1) border-box;
 }
 
-/* Status text */
-.status-success {
-    color: var(--success) !important;
-    font-weight: 500;
-}
-
-.status-error {
-    color: var(--error) !important;
-    font-weight: 500;
-}
-
-/* Accordion */
-.gr-accordion {
-    background: var(--surface-light) !important;
-    border: 1px solid var(--border) !important;
-    border-radius: 12px !important;
-}
-
-/* Labels */
-label {
-    color: var(--text) !important;
-    font-weight: 500 !important;
-}
-
-/* Tooltip / info text */
-.gr-info {
-    color: var(--text-muted) !important;
-    font-size: 0.875rem !important;
-}
-
-/* Responsive adjustments */
 @media (max-width: 768px) {
-    .image-grid {
-        grid-template-columns: repeat(2, 1fr);
-    }
-    
     .main-header h1 {
         font-size: 1.75rem;
     }
-    
     .feature-badges {
         flex-direction: column;
         align-items: center;
@@ -466,8 +376,9 @@ label {
 }
 """
 
+
 # ============================================================
-# Build Gradio Interface
+# 构建 Gradio 界面
 # ============================================================
 
 def create_demo():
@@ -499,7 +410,6 @@ def create_demo():
             # Left Column - Inputs
             with gr.Column(scale=1):
                 
-                # Image Upload Section
                 gr.HTML('<div class="section-header"><span>📸</span><h3>Upload Images</h3></div>')
                 
                 num_images = gr.Slider(
@@ -531,7 +441,6 @@ def create_demo():
                     outputs=image_inputs
                 )
                 
-                # Prompt Section
                 gr.HTML('<div class="section-header"><span>✍️</span><h3>Editing Instruction</h3></div>')
                 
                 prompt_input = gr.Textbox(
@@ -541,7 +450,6 @@ def create_demo():
                     value="Combine the reference images to generate the final result."
                 )
                 
-                # Advanced Settings
                 with gr.Accordion("⚙️ Advanced Settings", open=False):
                     cfg_scale = gr.Slider(
                         minimum=1.0,
@@ -568,7 +476,6 @@ def create_demo():
                             info="8 recommended for DMD"
                         )
                 
-                # Generate Button
                 generate_btn = gr.Button(
                     "🚀 Generate Image",
                     variant="primary",
@@ -594,7 +501,6 @@ def create_demo():
                     show_copy_button=False
                 )
                 
-                # Tips
                 gr.HTML("""
                 <div style="
                     margin-top: 1.5rem;
@@ -607,7 +513,7 @@ def create_demo():
                     <ul style="color: #ffffff; font-size: 0.9rem; margin: 0; padding-left: 1.25rem;">
                         <li>Reference images as "Image1", "Image2", etc. in your prompt</li>
                         <li>Use descriptive prompts for better composition</li>
-                        <li>First run may take longer due to model loading</li>
+                        <li>First run may take longer due to model warm-up</li>
                     </ul>
                 </div>
                 """)
@@ -637,7 +543,7 @@ def create_demo():
 
 
 # ============================================================
-# Launch
+# 启动
 # ============================================================
 
 demo = create_demo()