Update app.py

app.py

@@ -44,7 +44,7 @@ def get_dtype():
     """Get the appropriate dtype"""
     return torch.bfloat16 if torch.cuda.is_available() else torch.float32
 
-@GPU(duration=120)
+@GPU(duration=180)
 def generate_image(
     images: list[Image.Image],
     prompt: str,
@@ -79,25 +79,27 @@ def generate_image(
         MODEL_NAME, subfolder='scheduler'
     )
     
-    # Load text encoder
+    # Load text encoder - use device_map="cuda" for ZeroGPU compatibility
+    # This ensures all submodules are properly placed on the GPU
     text_encoder = AutoModel.from_pretrained(
         MODEL_NAME,
         subfolder='text_encoder',
-        torch_dtype=dtype
-    ).to(device)
+        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')
     
-    # Load transformer
-    transformer = load_transformer(device, dtype)
+    # Load transformer - also use device_map for consistency
+    transformer = load_transformer(dtype)
     
     # Load VAE
     vae = AutoencoderKLQwenImage.from_pretrained(
         MODEL_NAME,
         subfolder='vae',
-        torch_dtype=dtype
+        torch_dtype=dtype,
     ).to(device)
     
     # Create pipeline
@@ -146,12 +148,14 @@ def generate_image(
     return result
 
 
-def load_transformer(device, dtype):
-    """Load transformer with proper path handling"""
+def load_transformer(dtype):
+    """Load transformer with proper path handling for ZeroGPU"""
     from diffusers import QwenImageTransformer2DModel
     
+    device = get_device()
+    
     if os.path.exists(TRANSFORMER_PATH):
-        # Local path
+        # Local path - for ZeroGPU, still use .to(device) for local files
         if os.path.isdir(TRANSFORMER_PATH):
             config_path = os.path.join(TRANSFORMER_PATH, "config.json")
             if os.path.exists(config_path):
@@ -169,7 +173,7 @@ def load_transformer(device, dtype):
                 ).to(device)
         raise ValueError(f"Invalid transformer path: {TRANSFORMER_PATH}")
     else:
-        # HuggingFace repo path
+        # HuggingFace repo path - use device_map for ZeroGPU
         path_parts = TRANSFORMER_PATH.split('/')
         if len(path_parts) >= 3:
             repo_id = '/'.join(path_parts[:2])
@@ -177,14 +181,16 @@ def load_transformer(device, dtype):
             return QwenImageTransformer2DModel.from_pretrained(
                 repo_id,
                 subfolder=subfolder,
-                torch_dtype=dtype
-            ).to(device)
+                torch_dtype=dtype,
+                device_map="cuda"
+            )
         else:
             return QwenImageTransformer2DModel.from_pretrained(
                 TRANSFORMER_PATH,
                 subfolder='transformer',
-                torch_dtype=dtype
-            ).to(device)
+                torch_dtype=dtype,
+                device_map="cuda"
+            )
 
 
 # ============================================================