Update app.py

app.py

@@ -1,11 +1,9 @@
 """
 UniPic-3 DMD Multi-Image Composition
-Hugging Face Space - ZeroGPU 优化版本 V2
+Hugging Face Space - ZeroGPU 优化版本 V3
 
-关键修复：使用延迟加载 (Lazy Loading) 方案
-- 模型组件在全局 CPU 上加载
-- Pipeline 在 @spaces.GPU 函数内首次调用时才创建并移动到 GPU
-- 这确保了所有张量都在真实的 GPU 环境中初始化
+关键修复：完全在 @spaces.GPU 内部加载模型
+参考 Qwen 官方的 app.py 实现方式
 """
 
 import gradio as gr
@@ -34,32 +32,27 @@ 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")
 
 # ============================================================
-# 全局变量 - Pipeline 延迟初始化
+# 全局变量
 # ============================================================
+pipe = None
+dtype = torch.bfloat16
 
-# 只在全局加载轻量级组件和 CPU 上的模型权重
-pipe = None  # 延迟初始化
-_models_loaded = False
 
-# 存储 CPU 上的模型组件
-_cpu_components = {}
-
-
-def load_models_to_cpu():
+def load_pipeline():
     """
-    在 CPU 上加载所有模型组件
-    这一步在全局执行，不需要 GPU
+    加载完整的 Pipeline
+    这个函数应该在 @spaces.GPU 装饰的函数内部调用
     """
-    global _cpu_components, _models_loaded
+    global pipe
     
-    if _models_loaded:
-        return
+    if pipe is not None:
+        return pipe
     
-    print("🚀 Loading models to CPU...")
+    print("🚀 Loading pipeline...")
     
     try:
         from pipeline_qwenimage_edit import QwenImageEditPipeline
@@ -73,92 +66,92 @@ def load_models_to_cpu():
     )
     from transformers import AutoModel, AutoTokenizer, Qwen2VLProcessor
     
-    dtype = torch.bfloat16
+    device = 'cuda'
     
-    # Load scheduler (CPU, 轻量级)
+    # Load scheduler
     print("  Loading scheduler...")
-    _cpu_components['scheduler'] = FlowMatchEulerDiscreteScheduler.from_pretrained(
+    scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
         MODEL_NAME, subfolder='scheduler'
     )
     
-    # Load tokenizer & processor (CPU, 轻量级)
+    # Load tokenizer & processor
     print("  Loading tokenizer & processor...")
-    _cpu_components['tokenizer'] = AutoTokenizer.from_pretrained(MODEL_NAME, subfolder='tokenizer')
-    _cpu_components['processor'] = Qwen2VLProcessor.from_pretrained(MODEL_NAME, subfolder='processor')
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, subfolder='tokenizer')
+    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(
+    # Load text encoder - 直接加载到 GPU
+    print("  Loading text_encoder...")
+    text_encoder = AutoModel.from_pretrained(
         MODEL_NAME,
         subfolder='text_encoder',
         torch_dtype=dtype,
-    ).eval()
+    ).to(device).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(
-                    repo_id,
-                    subfolder=subfolder,
+    # 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
-                ).eval()
+                ).to(device).eval()
             else:
-                return QwenImageTransformer2DModel.from_pretrained(
+                transformer = QwenImageTransformer2DModel.from_pretrained(
                     TRANSFORMER_PATH,
                     subfolder='transformer',
                     torch_dtype=dtype,
                     use_safetensors=False
-                ).eval()
-    
-    _cpu_components['transformer'] = load_transformer()
+                ).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 to CPU
-    print("  Loading VAE to CPU...")
-    _cpu_components['vae'] = AutoencoderKLQwenImage.from_pretrained(
+    # Load VAE - 直接加载到 GPU
+    print("  Loading VAE...")
+    vae = AutoencoderKLQwenImage.from_pretrained(
         MODEL_NAME,
         subfolder='vae',
         torch_dtype=dtype,
-    ).eval()
+    ).to(device).eval()
     
-    # 存储 Pipeline 类以便后续使用
-    _cpu_components['pipeline_class'] = QwenImageEditPipeline
+    # Create Pipeline
+    print("  Creating pipeline...")
+    pipe = QwenImageEditPipeline(
+        scheduler=scheduler,
+        vae=vae,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        processor=processor,
+        transformer=transformer
+    )
     
-    _models_loaded = True
-    print("✅ All models loaded to CPU!")
-
-
-# 立即在全局加载到 CPU
-load_models_to_cpu()
+    print("✅ Pipeline loaded successfully!")
+    return pipe
 
 
 # ============================================================
-# GPU 推理函数 - 在这里初始化 Pipeline
+# GPU 推理函数 - 模型在这里加载
 # ============================================================
 
-@spaces.GPU(duration=120)
+@spaces.GPU(duration=180)  # 增加时间以包含首次加载
 def generate_image(
     images: list[Image.Image],
     prompt: str,
@@ -168,7 +161,7 @@ def generate_image(
 ) -> Image.Image:
     """
     GPU 推理函数
-    关键：Pipeline 在这里创建，确保在真实 GPU 环境中初始化
+    关键：Pipeline 完全在这里加载，确保在真实 GPU 环境中初始化
     """
     global pipe
     
@@ -176,33 +169,9 @@ def generate_image(
     print(f"   Prompt: {prompt[:50]}...")
     print(f"   Steps: {num_steps}, CFG: {true_cfg_scale}, Seed: {seed}")
     
-    # 关键修复：在真实 GPU 环境中创建 Pipeline
+    # 在真实 GPU 环境中加载模型（首次调用时）
     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')
+        load_pipeline()
     
     # 验证设备
     print(f"   [DEBUG] text_encoder device: {next(pipe.text_encoder.parameters()).device}")
@@ -527,7 +496,7 @@ def create_demo():
                 
                 status_text = gr.Textbox(
                     label="Status",
-                    value="✨ Ready! Upload images and click Generate.",
+                    value="✨ Ready! Upload images and click Generate. First run will take longer to load the model.",
                     interactive=False,
                 )
                 
@@ -543,7 +512,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 warm-up</li>
+                        <li>First run will take ~60s to load the model</li>
                     </ul>
                 </div>
                 """)