Update app.py

app.py

@@ -69,38 +69,42 @@ def generate_image(
     )
     from transformers import AutoModel, AutoTokenizer, Qwen2VLProcessor
     
-    device = get_device()
-    dtype = get_dtype()
+    # 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
+    # Load scheduler (CPU, no device needed)
     scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
         MODEL_NAME, subfolder='scheduler'
     )
     
-    # Load text encoder - use device_map="cuda" for ZeroGPU compatibility
-    # This ensures all submodules are properly placed on the GPU
+    # 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,
-        device_map="cuda"  # Let transformers handle device placement for ZeroGPU
-    )
-    
-    # Load tokenizer & processor
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, subfolder='tokenizer')
-    processor = Qwen2VLProcessor.from_pretrained(MODEL_NAME, subfolder='processor')
+    ).to(device).eval()
     
-    # Load transformer - also use device_map for consistency
-    transformer = load_transformer(dtype)
+    # 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)
+    ).to(device).eval()
     
     # Create pipeline
     pipe = QwenImageEditPipeline(
@@ -112,7 +116,20 @@ def generate_image(
         transformer=transformer
     )
     
-    print(f"✅ Model loaded! Generating with {len(images)} image(s)...")
+    # 关键修复：手动设置 pipeline 使用的设备
+    # 这确保 _execution_device 返回正确的设备
+    pipe._execution_device = device
+    
+    # 同时确保 processor 也在正确设备上处理
+    # 修改 pipe 的 device 属性（如果存在）
+    if hasattr(pipe, 'device'):
+        pipe.device = device
+    
+    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)...")
     
     # Generate
     with torch.no_grad():
@@ -143,19 +160,18 @@ def generate_image(
     
     # Cleanup to free VRAM
     del pipe, transformer, vae, text_encoder
-    torch.cuda.empty_cache()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
     
     return result
 
 
-def load_transformer(dtype):
+def load_transformer(device, dtype):
     """Load transformer with proper path handling for ZeroGPU"""
     from diffusers import QwenImageTransformer2DModel
     
-    device = get_device()
-    
     if os.path.exists(TRANSFORMER_PATH):
-        # Local path - for ZeroGPU, still use .to(device) for local files
+        # Local path
         if os.path.isdir(TRANSFORMER_PATH):
             config_path = os.path.join(TRANSFORMER_PATH, "config.json")
             if os.path.exists(config_path):
@@ -163,17 +179,17 @@ def load_transformer(dtype):
                     TRANSFORMER_PATH,
                     torch_dtype=dtype,
                     use_safetensors=False
-                ).to(device)
+                ).to(device).eval()
             else:
                 return QwenImageTransformer2DModel.from_pretrained(
                     TRANSFORMER_PATH,
                     subfolder='transformer',
                     torch_dtype=dtype,
                     use_safetensors=False
-                ).to(device)
+                ).to(device).eval()
         raise ValueError(f"Invalid transformer path: {TRANSFORMER_PATH}")
     else:
-        # HuggingFace repo path - use device_map for ZeroGPU
+        # HuggingFace repo path
         path_parts = TRANSFORMER_PATH.split('/')
         if len(path_parts) >= 3:
             repo_id = '/'.join(path_parts[:2])
@@ -182,15 +198,13 @@ def load_transformer(dtype):
                 repo_id,
                 subfolder=subfolder,
                 torch_dtype=dtype,
-                device_map="cuda"
-            )
+            ).to(device).eval()
         else:
             return QwenImageTransformer2DModel.from_pretrained(
                 TRANSFORMER_PATH,
                 subfolder='transformer',
                 torch_dtype=dtype,
-                device_map="cuda"
-            )
+            ).to(device).eval()
 
 
 # ============================================================
@@ -632,4 +646,4 @@ def create_demo():
 demo = create_demo()
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()