Update app.py

app.py

@@ -1,11 +1,11 @@
 """
 UniPic-3 DMD Multi-Image Composition
-Hugging Face Space - ZeroGPU 优化版本
+Hugging Face Space - ZeroGPU 优化版本 V2
 
-架构说明：
-1. 模型在全局作用域加载（ZeroGPU 会拦截 CUDA 调用）
-2. 只有实际推理时才使用 @spaces.GPU 装饰器
-3. 这样避免了每次请求都重新加载模型
+关键修复：使用延迟加载 (Lazy Loading) 方案
+- 模型组件在全局 CPU 上加载
+- Pipeline 在 @spaces.GPU 函数内首次调用时才创建并移动到 GPU
+- 这确保了所有张量都在真实的 GPU 环境中初始化
 """
 
 import gradio as gr
@@ -18,8 +18,10 @@ import sys
 try:
     import spaces
     HF_SPACES = True
+    print("✅ Running in Hugging Face Spaces with ZeroGPU")
 except ImportError:
     HF_SPACES = False
+    print("⚠️ Running locally (no ZeroGPU)")
     # 本地开发时的 mock
     class spaces:
         @staticmethod
@@ -32,60 +34,94 @@ except ImportError:
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 
 # Model configuration
-MODEL_NAME = os.environ.get("MODEL_NAME", "Skywork/Unipic3-DMD")
-TRANSFORMER_PATH = os.environ.get("TRANSFORMER_PATH", "Skywork/Unipic3-DMD/ema_transformer")
+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")
 
 # ============================================================
-# 全局加载模型（ZeroGPU 会拦截 CUDA 调用）
+# 全局变量 - Pipeline 延迟初始化
 # ============================================================
 
-print("🚀 Loading models...")
+# 只在全局加载轻量级组件和 CPU 上的模型权重
+pipe = None  # 延迟初始化
+_models_loaded = False
 
-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
+# 存储 CPU 上的模型组件
+_cpu_components = {}
 
-# 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):
+def load_models_to_cpu():
+    """
+    在 CPU 上加载所有模型组件
+    这一步在全局执行，不需要 GPU
+    """
+    global _cpu_components, _models_loaded
+    
+    if _models_loaded:
+        return
+    
+    print("🚀 Loading models to CPU...")
+    
+    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 = torch.bfloat16
+    
+    # Load scheduler (CPU, 轻量级)
+    print("  Loading scheduler...")
+    _cpu_components['scheduler'] = FlowMatchEulerDiscreteScheduler.from_pretrained(
+        MODEL_NAME, subfolder='scheduler'
+    )
+    
+    # Load tokenizer & processor (CPU, 轻量级)
+    print("  Loading tokenizer & processor...")
+    _cpu_components['tokenizer'] = AutoTokenizer.from_pretrained(MODEL_NAME, subfolder='tokenizer')
+    _cpu_components['processor'] = Qwen2VLProcessor.from_pretrained(MODEL_NAME, subfolder='processor')
+    
+    # Load text encoder to CPU
+    print("  Loading text_encoder to CPU...")
+    _cpu_components['text_encoder'] = AutoModel.from_pretrained(
+        MODEL_NAME,
+        subfolder='text_encoder',
+        torch_dtype=dtype,
+    ).eval()
+    
+    # Load transformer to CPU
+    print("  Loading transformer to CPU...")
+    def load_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):
+                    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:
+            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(
-                    TRANSFORMER_PATH,
+                    repo_id,
+                    subfolder=subfolder,
                     torch_dtype=dtype,
                     use_safetensors=False
                 ).eval()
@@ -96,101 +132,30 @@ def load_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!")
-
-
-def fix_rope_buffers(model):
-    """
-    修复 RoPE (Rotary Position Embedding) 中的 buffer 张量
-    ZeroGPU 环境下，register_buffer 注册的张量可能不会被正确移动到 GPU
     
-    这个函数会遍历模型的所有子模块，检查并修复以下 buffer：
-    - inv_freq: RoPE 的核心频率 buffer
-    - cos_cached / sin_cached: 某些实现会缓存的 cos/sin 值
-    - 其他所有未在 CUDA 上的 buffer
-    """
-    device = 'cuda'
-    fixed_count = 0
+    _cpu_components['transformer'] = load_transformer()
     
-    for name, module in model.named_modules():
-        # 修复 inv_freq buffer (RoPE 的核心 buffer)
-        if hasattr(module, 'inv_freq') and module.inv_freq is not None:
-            if module.inv_freq.device.type != 'cuda':
-                module.inv_freq = module.inv_freq.to(device)
-                fixed_count += 1
-                print(f"   [FIX] Moved {name}.inv_freq to {device}")
-        
-        # 修复 cos_cached 和 sin_cached (某些 RoPE 实现会缓存这些)
-        if hasattr(module, 'cos_cached') and module.cos_cached is not None:
-            if module.cos_cached.device.type != 'cuda':
-                module.cos_cached = module.cos_cached.to(device)
-                fixed_count += 1
-                print(f"   [FIX] Moved {name}.cos_cached to {device}")
-        
-        if hasattr(module, 'sin_cached') and module.sin_cached is not None:
-            if module.sin_cached.device.type != 'cuda':
-                module.sin_cached = module.sin_cached.to(device)
-                fixed_count += 1
-                print(f"   [FIX] Moved {name}.sin_cached to {device}")
-        
-        # 通用：修复所有 buffer（更全面的修复）
-        for buf_name, buf in module.named_buffers(recurse=False):
-            if buf is not None and buf.device.type != 'cuda':
-                setattr(module, buf_name, buf.to(device))
-                fixed_count += 1
-                print(f"   [FIX] Moved {name}.{buf_name} to {device}")
+    # Load VAE to CPU
+    print("  Loading VAE to CPU...")
+    _cpu_components['vae'] = AutoencoderKLQwenImage.from_pretrained(
+        MODEL_NAME,
+        subfolder='vae',
+        torch_dtype=dtype,
+    ).eval()
+    
+    # 存储 Pipeline 类以便后续使用
+    _cpu_components['pipeline_class'] = QwenImageEditPipeline
     
-    return fixed_count
+    _models_loaded = True
+    print("✅ All models loaded to CPU!")
 
 
+# 立即在全局加载到 CPU
+load_models_to_cpu()
+
 
 # ============================================================
-# GPU 推理函数（只包含实际的推理逻辑）
+# GPU 推理函数 - 在这里初始化 Pipeline
 # ============================================================
 
 @spaces.GPU(duration=120)
@@ -202,31 +167,44 @@ def generate_image(
     num_steps: int
 ) -> Image.Image:
     """
-    GPU 推理函数 - 只包含实际的推理逻辑
-    模型已在全局加载，这里只执行推理
+    GPU 推理函数
+    关键：Pipeline 在这里创建，确保在真实 GPU 环境中初始化
     """
+    global pipe
+    
     print(f"🎨 Generating with {len(images)} image(s)...")
     print(f"   Prompt: {prompt[:50]}...")
     print(f"   Steps: {num_steps}, CFG: {true_cfg_scale}, Seed: {seed}")
     
-   # Step 1: 移动 pipeline 到 CUDA
-    pipe.to('cuda')
-    
-    # Step 2: 关键修复 - 手动修复 RoPE buffer
-    # ZeroGPU 可能没有正确移动 register_buffer 注册的张量
-    print("   [DEBUG] Fixing RoPE buffers...")
-    fixed = 0
-    fixed += fix_rope_buffers(pipe.text_encoder)
-    fixed += fix_rope_buffers(pipe.transformer)
-    fixed += fix_rope_buffers(pipe.vae)
-    print(f"   [DEBUG] Fixed {fixed} buffer(s)")
-    
-    # 调试信息：检查模型设备
-    print(f"   [DEBUG] text_encoder device: {next(pipe.text_encoder.parameters()).device}")
-    print(f"   [DEBUG] transformer device: {next(pipe.transformer.parameters()).device}")
-    print(f"   [DEBUG] vae device: {next(pipe.vae.parameters()).device}")
+    # 关键修复：在真实 GPU 环境中创建 Pipeline
+    if pipe is None:
+        print("   [INIT] Creating pipeline on real GPU...")
+        
+        # 方法：将 CPU 模型移动到 GPU，然后创建 pipeline
+        device = 'cuda'
+        
+        # 移动模型到 GPU
+        text_encoder = _cpu_components['text_encoder'].to(device)
+        transformer = _cpu_components['transformer'].to(device)
+        vae = _cpu_components['vae'].to(device)
+        
+        # 创建 Pipeline
+        PipelineClass = _cpu_components['pipeline_class']
+        pipe = PipelineClass(
+            scheduler=_cpu_components['scheduler'],
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=_cpu_components['tokenizer'],
+            processor=_cpu_components['processor'],
+            transformer=transformer
+        )
+        
+        print("   [INIT] Pipeline created successfully!")
+    else:
+        # Pipeline 已存在，确保在正确的设备上
+        pipe.to('cuda')
     
-    # 调试信息：检查模型设备
+    # 验证设备
     print(f"   [DEBUG] text_encoder device: {next(pipe.text_encoder.parameters()).device}")
     print(f"   [DEBUG] transformer device: {next(pipe.transformer.parameters()).device}")
     print(f"   [DEBUG] vae device: {next(pipe.vae.parameters()).device}")
@@ -320,7 +298,6 @@ def update_image_visibility(num):
 
 CUSTOM_CSS = """
 @import url('https://fonts.googleapis.com/css2?family=Outfit:wght@300;400;500;600;700&family=JetBrains+Mono:wght@400;500&display=swap');
-
 :root {
     --primary: #6366f1;
     --primary-dark: #4f46e5;
@@ -336,13 +313,11 @@ CUSTOM_CSS = """
     --gradient-1: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
     --gradient-hero: linear-gradient(135deg, #0f0f23 0%, #1a1a3e 50%, #252552 100%);
 }
-
 .gradio-container {
     font-family: 'Outfit', sans-serif !important;
     background: var(--gradient-hero) !important;
     min-height: 100vh;
 }
-
 .main-header {
     text-align: center;
     padding: 2rem 1rem;
@@ -351,7 +326,6 @@ CUSTOM_CSS = """
     margin-bottom: 2rem;
     border: 1px solid rgba(99, 102, 241, 0.2);
 }
-
 .main-header h1 {
     font-size: 2.5rem;
     font-weight: 700;
@@ -361,14 +335,12 @@ CUSTOM_CSS = """
     background-clip: text;
     margin-bottom: 0.5rem;
 }
-
 .main-header p {
     color: var(--text-muted);
     font-size: 1.1rem;
     max-width: 600px;
     margin: 0 auto;
 }
-
 .feature-badges {
     display: flex;
     gap: 1rem;
@@ -376,7 +348,6 @@ CUSTOM_CSS = """
     flex-wrap: wrap;
     margin-top: 1.5rem;
 }
-
 .badge {
     display: inline-flex;
     align-items: center;
@@ -389,7 +360,6 @@ CUSTOM_CSS = """
     font-size: 0.875rem;
     font-weight: 500;
 }
-
 .section-header {
     display: flex;
     align-items: center;
@@ -398,14 +368,12 @@ CUSTOM_CSS = """
     padding-bottom: 0.75rem;
     border-bottom: 1px solid var(--border);
 }
-
 .section-header h3 {
     font-size: 1.125rem;
     font-weight: 600;
     color: var(--text);
     margin: 0;
 }
-
 .generate-btn {
     background: var(--gradient-1) !important;
     border: none !important;
@@ -418,12 +386,10 @@ CUSTOM_CSS = """
     transition: all 0.3s ease !important;
     box-shadow: 0 4px 15px rgba(99, 102, 241, 0.4) !important;
 }
-
 .generate-btn:hover {
     transform: translateY(-2px) !important;
     box-shadow: 0 6px 20px rgba(99, 102, 241, 0.5) !important;
 }
-
 .output-image {
     border-radius: 16px;
     overflow: hidden;
@@ -431,7 +397,6 @@ CUSTOM_CSS = """
     background: linear-gradient(var(--surface-light), var(--surface-light)) padding-box,
                 var(--gradient-1) border-box;
 }
-
 @media (max-width: 768px) {
     .main-header h1 {
         font-size: 1.75rem;
@@ -558,14 +523,12 @@ def create_demo():
                     type="pil",
                     label="Output",
                     elem_classes=["output-image"],
-                    show_download_button=True
                 )
                 
                 status_text = gr.Textbox(
                     label="Status",
                     value="✨ Ready! Upload images and click Generate.",
                     interactive=False,
-                    show_copy_button=False
                 )
                 
                 gr.HTML("""