Update app.py

app.py

@@ -1,9 +1,11 @@
 """
 UniPic-3 DMD Multi-Image Composition
-Hugging Face Space - ZeroGPU 优化版本 V3
+Hugging Face Space - ZeroGPU 优化版本 V5
 
-关键修复：完全在 @spaces.GPU 内部加载模型
-参考 Qwen 官方的 app.py 实现方式
+关键策略：
+1. 全局只加载不需要 GPU 的组件（scheduler, tokenizer, processor）
+2. 需要 GPU 的模型在 @spaces.GPU 内部加载，显式指定 device='cuda'
+3. 不使用 device_map='auto'，因为它可能在 ZeroGPU 外部被错误地分配
 """
 
 import gradio as gr
@@ -20,7 +22,6 @@ try:
 except ImportError:
     HF_SPACES = False
     print("⚠️ Running locally (no ZeroGPU)")
-    # 本地开发时的 mock
     class spaces:
         @staticmethod
         def GPU(duration=60):
@@ -35,123 +36,47 @@ sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 MODEL_NAME = os.environ.get("MODEL_NAME", "/data_genie/genie/chris/Unipic3-DMD")
 TRANSFORMER_PATH = os.environ.get("TRANSFORMER_PATH", "/data_genie/genie/chris/Unipic3-DMD/ema_transformer")
 
+dtype = torch.bfloat16
+
 # ============================================================
-# 全局变量
+# 全局加载轻量级组件（不需要 GPU）
 # ============================================================
-pipe = None
-dtype = torch.bfloat16
 
+print("🚀 Loading lightweight components (CPU)...")
 
-def load_pipeline():
-    """
-    加载完整的 Pipeline
-    这个函数应该在 @spaces.GPU 装饰的函数内部调用
-    """
-    global pipe
-    
-    if pipe is not None:
-        return pipe
-    
-    print("🚀 Loading pipeline...")
-    
-    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
-    
-    device = 'cuda'
-    
-    # Load scheduler
-    print("  Loading scheduler...")
-    scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
-        MODEL_NAME, subfolder='scheduler'
-    )
-    
-    # Load tokenizer & processor
-    print("  Loading tokenizer & processor...")
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, subfolder='tokenizer')
-    processor = Qwen2VLProcessor.from_pretrained(MODEL_NAME, subfolder='processor')
-    
-    # Load text encoder - 直接加载到 GPU
-    print("  Loading text_encoder...")
-    text_encoder = AutoModel.from_pretrained(
-        MODEL_NAME,
-        subfolder='text_encoder',
-        torch_dtype=dtype,
-    ).to(device).eval()
-    
-    # Load transformer - 直接加载到 GPU
-    print("  Loading transformer...")
-    if os.path.exists(TRANSFORMER_PATH):
-        if os.path.isdir(TRANSFORMER_PATH):
-            config_path = os.path.join(TRANSFORMER_PATH, "config.json")
-            if os.path.exists(config_path):
-                transformer = QwenImageTransformer2DModel.from_pretrained(
-                    TRANSFORMER_PATH,
-                    torch_dtype=dtype,
-                    use_safetensors=False
-                ).to(device).eval()
-            else:
-                transformer = QwenImageTransformer2DModel.from_pretrained(
-                    TRANSFORMER_PATH,
-                    subfolder='transformer',
-                    torch_dtype=dtype,
-                    use_safetensors=False
-                ).to(device).eval()
-    else:
-        path_parts = TRANSFORMER_PATH.split('/')
-        if len(path_parts) >= 3:
-            repo_id = '/'.join(path_parts[:2])
-            subfolder = '/'.join(path_parts[2:])
-            transformer = QwenImageTransformer2DModel.from_pretrained(
-                repo_id,
-                subfolder=subfolder,
-                torch_dtype=dtype,
-                use_safetensors=False
-            ).to(device).eval()
-        else:
-            transformer = QwenImageTransformer2DModel.from_pretrained(
-                TRANSFORMER_PATH,
-                subfolder='transformer',
-                torch_dtype=dtype,
-                use_safetensors=False
-            ).to(device).eval()
-    
-    # Load VAE - 直接加载到 GPU
-    print("  Loading VAE...")
-    vae = AutoencoderKLQwenImage.from_pretrained(
-        MODEL_NAME,
-        subfolder='vae',
-        torch_dtype=dtype,
-    ).to(device).eval()
-    
-    # Create Pipeline
-    print("  Creating pipeline...")
-    pipe = QwenImageEditPipeline(
-        scheduler=scheduler,
-        vae=vae,
-        text_encoder=text_encoder,
-        tokenizer=tokenizer,
-        processor=processor,
-        transformer=transformer
-    )
-    
-    print("✅ Pipeline loaded successfully!")
-    return pipe
+from diffusers import (
+    FlowMatchEulerDiscreteScheduler, 
+    QwenImageTransformer2DModel, 
+    AutoencoderKLQwenImage
+)
+from transformers import AutoModel, AutoTokenizer, Qwen2VLProcessor
+
+try:
+    from pipeline_qwenimage_edit import QwenImageEditPipeline
+except ImportError:
+    from diffusers import QwenImageEditPipeline
+
+# 这些组件不需要 GPU，可以在全局加载
+scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+    MODEL_NAME, subfolder='scheduler'
+)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, subfolder='tokenizer')
+processor = Qwen2VLProcessor.from_pretrained(MODEL_NAME, subfolder='processor')
+
+print("✅ Lightweight components loaded!")
+
+# ============================================================
+# Pipeline 状态
+# ============================================================
+pipe = None
+_models_loaded = False
 
 
 # ============================================================
 # GPU 推理函数 - 模型在这里加载
 # ============================================================
 
-@spaces.GPU(duration=180)  # 增加时间以包含首次加载
+@spaces.GPU(duration=180)
 def generate_image(
     images: list[Image.Image],
     prompt: str,
@@ -161,17 +86,85 @@ def generate_image(
 ) -> Image.Image:
     """
     GPU 推理函数
-    关键：Pipeline 完全在这里加载，确保在真实 GPU 环境中初始化
+    关键：所有需要 GPU 的模型都在这里加载，确保在真实 GPU 环境中
     """
-    global pipe
+    global pipe, _models_loaded
     
     print(f"🎨 Generating with {len(images)} image(s)...")
     print(f"   Prompt: {prompt[:50]}...")
     print(f"   Steps: {num_steps}, CFG: {true_cfg_scale}, Seed: {seed}")
     
     # 在真实 GPU 环境中加载模型（首次调用时）
-    if pipe is None:
-        load_pipeline()
+    if not _models_loaded:
+        print("   [INIT] Loading models on real GPU...")
+        
+        device = 'cuda'
+        
+        # 加载 text_encoder 到 GPU
+        print("   [INIT] Loading text_encoder...")
+        text_encoder = AutoModel.from_pretrained(
+            MODEL_NAME,
+            subfolder='text_encoder',
+            torch_dtype=dtype,
+        ).to(device).eval()
+        
+        # 加载 transformer 到 GPU
+        print("   [INIT] Loading transformer...")
+        if os.path.exists(TRANSFORMER_PATH) and os.path.isdir(TRANSFORMER_PATH):
+            config_path = os.path.join(TRANSFORMER_PATH, "config.json")
+            if os.path.exists(config_path):
+                transformer = QwenImageTransformer2DModel.from_pretrained(
+                    TRANSFORMER_PATH,
+                    torch_dtype=dtype,
+                    use_safetensors=False
+                ).to(device).eval()
+            else:
+                transformer = QwenImageTransformer2DModel.from_pretrained(
+                    TRANSFORMER_PATH,
+                    subfolder='transformer',
+                    torch_dtype=dtype,
+                    use_safetensors=False
+                ).to(device).eval()
+        else:
+            path_parts = TRANSFORMER_PATH.split('/')
+            if len(path_parts) >= 3:
+                repo_id = '/'.join(path_parts[:2])
+                subfolder = '/'.join(path_parts[2:])
+                transformer = QwenImageTransformer2DModel.from_pretrained(
+                    repo_id,
+                    subfolder=subfolder,
+                    torch_dtype=dtype,
+                    use_safetensors=False
+                ).to(device).eval()
+            else:
+                transformer = QwenImageTransformer2DModel.from_pretrained(
+                    TRANSFORMER_PATH,
+                    subfolder='transformer',
+                    torch_dtype=dtype,
+                    use_safetensors=False
+                ).to(device).eval()
+        
+        # 加载 VAE 到 GPU
+        print("   [INIT] Loading VAE...")
+        vae = AutoencoderKLQwenImage.from_pretrained(
+            MODEL_NAME,
+            subfolder='vae',
+            torch_dtype=dtype,
+        ).to(device).eval()
+        
+        # 创建 Pipeline
+        print("   [INIT] Creating pipeline...")
+        pipe = QwenImageEditPipeline(
+            scheduler=scheduler,
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            processor=processor,
+            transformer=transformer
+        )
+        
+        _models_loaded = True
+        print("   [INIT] ✅ Models loaded successfully!")
     
     # 验证设备
     print(f"   [DEBUG] text_encoder device: {next(pipe.text_encoder.parameters()).device}")
@@ -222,10 +215,8 @@ def process_images(
 ):
     """处理图像 - 验证输入后调用 GPU 函数"""
     
-    # 过滤有效图像
     images = [img for img in [img1, img2, img3, img4, img5, img6] if img is not None]
     
-    # 验证
     if len(images) == 0:
         return None, "❌ Please upload at least one image"
     
@@ -236,10 +227,8 @@ def process_images(
         return None, "❌ Please enter an editing instruction"
     
     try:
-        # 转换为 RGB
         images = [img.convert("RGB") for img in images]
         
-        # 调用 GPU 函数
         result = generate_image(
             images=images,
             prompt=prompt,
@@ -257,7 +246,6 @@ def process_images(
 
 
 def update_image_visibility(num):
-    """更新图像上传槽的可见性"""
     return [gr.update(visible=(i < num)) for i in range(6)]
 
 
@@ -367,13 +355,8 @@ CUSTOM_CSS = """
                 var(--gradient-1) border-box;
 }
 @media (max-width: 768px) {
-    .main-header h1 {
-        font-size: 1.75rem;
-    }
-    .feature-badges {
-        flex-direction: column;
-        align-items: center;
-    }
+    .main-header h1 { font-size: 1.75rem; }
+    .feature-badges { flex-direction: column; align-items: center; }
 }
 """
 
@@ -394,7 +377,6 @@ def create_demo():
         css=CUSTOM_CSS
     ) as demo:
         
-        # Header
         gr.HTML("""
         <div class="main-header">
             <h1>🎨 UniPic-3 DMD</h1>
@@ -408,19 +390,11 @@ def create_demo():
         """)
         
         with gr.Row(equal_height=True):
-            # Left Column - Inputs
             with gr.Column(scale=1):
-                
                 gr.HTML('<div class="section-header"><span>📸</span><h3>Upload Images</h3></div>')
                 
-                num_images = gr.Slider(
-                    minimum=1,
-                    maximum=6,
-                    value=2,
-                    step=1,
-                    label="Number of Images",
-                    info="Select how many images to compose"
-                )
+                num_images = gr.Slider(minimum=1, maximum=6, value=2, step=1,
+                    label="Number of Images", info="Select how many images to compose")
                 
                 with gr.Row():
                     img1 = gr.Image(type="pil", label="Image 1", visible=True)
@@ -435,96 +409,57 @@ def create_demo():
                     img6 = gr.Image(type="pil", label="Image 6", visible=False)
                 
                 image_inputs = [img1, img2, img3, img4, img5, img6]
-                
-                num_images.change(
-                    fn=update_image_visibility,
-                    inputs=num_images,
-                    outputs=image_inputs
-                )
+                num_images.change(fn=update_image_visibility, inputs=num_images, outputs=image_inputs)
                 
                 gr.HTML('<div class="section-header"><span>✍️</span><h3>Editing Instruction</h3></div>')
                 
                 prompt_input = gr.Textbox(
                     label="Prompt",
-                    placeholder="e.g., A man from Image1 standing on a surfboard from Image2, riding ocean waves under a bright blue sky.",
+                    placeholder="e.g., A man from Image1 standing on a surfboard from Image2...",
                     lines=3,
                     value="Combine the reference images to generate the final result."
                 )
                 
                 with gr.Accordion("⚙️ Advanced Settings", open=False):
-                    cfg_scale = gr.Slider(
-                        minimum=1.0,
-                        maximum=10.0,
-                        value=4.0,
-                        step=0.5,
-                        label="CFG Scale",
-                        info="Higher = more prompt alignment"
-                    )
+                    cfg_scale = gr.Slider(minimum=1.0, maximum=10.0, value=4.0, step=0.5,
+                        label="CFG Scale", info="Higher = more prompt alignment")
                     
                     with gr.Row():
-                        seed = gr.Number(
-                            value=42,
-                            label="Seed",
-                            info="For reproducibility",
-                            precision=0
-                        )
-                        num_steps = gr.Slider(
-                            minimum=1,
-                            maximum=8,
-                            value=8,
-                            step=1,
-                            label="Steps",
-                            info="8 recommended for DMD"
-                        )
+                        seed = gr.Number(value=42, label="Seed", info="For reproducibility", precision=0)
+                        num_steps = gr.Slider(minimum=1, maximum=8, value=8, step=1,
+                            label="Steps", info="8 recommended for DMD")
                 
-                generate_btn = gr.Button(
-                    "🚀 Generate Image",
-                    variant="primary",
-                    size="lg",
-                    elem_classes=["generate-btn"]
-                )
+                generate_btn = gr.Button("🚀 Generate Image", variant="primary", size="lg",
+                    elem_classes=["generate-btn"])
             
-            # Right Column - Output
             with gr.Column(scale=1):
                 gr.HTML('<div class="section-header"><span>🎨</span><h3>Generated Result</h3></div>')
                 
-                output_image = gr.Image(
-                    type="pil",
-                    label="Output",
-                    elem_classes=["output-image"],
-                )
+                output_image = gr.Image(type="pil", label="Output", elem_classes=["output-image"])
                 
                 status_text = gr.Textbox(
                     label="Status",
-                    value="✨ Ready! Upload images and click Generate. First run will take longer to load the model.",
+                    value="✨ Ready! First run takes ~60s to load models.",
                     interactive=False,
                 )
                 
                 gr.HTML("""
-                <div style="
-                    margin-top: 1.5rem;
-                    padding: 1rem;
-                    background: rgba(99, 102, 241, 0.1);
-                    border-radius: 12px;
-                    border: 1px solid rgba(99, 102, 241, 0.2);
-                ">
+                <div style="margin-top: 1.5rem; padding: 1rem; background: rgba(99, 102, 241, 0.1);
+                    border-radius: 12px; border: 1px solid rgba(99, 102, 241, 0.2);">
                     <p style="color: #ffffff; font-weight: 600; margin-bottom: 0.5rem;">💡 Tips</p>
                     <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 will take ~60s to load the model</li>
+                        <li>Reference images as "Image1", "Image2", etc.</li>
+                        <li>First run loads models (~60s)</li>
                     </ul>
                 </div>
                 """)
         
-        # Connect generate button
         generate_btn.click(
             fn=process_images,
             inputs=[*image_inputs, prompt_input, cfg_scale, seed, num_steps],
             outputs=[output_image, status_text]
         )
         
-        # Examples
         gr.HTML('<div class="section-header" style="margin-top: 2rem;"><span>📚</span><h3>Example Prompts</h3></div>')
         
         gr.Examples(
@@ -532,7 +467,6 @@ def create_demo():
                 ["A person from Image1 wearing the outfit from Image2"],
                 ["Combine Image1 and Image2 into a single cohesive scene"],
                 ["The object from Image1 placed in the environment from Image2"],
-                ["Create a portrait using the face from Image1 and hairstyle from Image2"],
             ],
             inputs=[prompt_input],
             label=""
@@ -541,10 +475,6 @@ def create_demo():
     return demo
 
 
-# ============================================================
-# 启动
-# ============================================================
-
 demo = create_demo()
 
 if __name__ == "__main__":