Update app.py

app.py

@@ -15,19 +15,35 @@ import cv2
 from transformers import pipeline as transformers_pipeline
 from huggingface_hub import hf_hub_download
 import os
+import spaces
 
 # Configuration
 MODEL_REPO = "primerz/pixagram"
-device = "cuda" if torch.cuda.is_available() else "cpu"
-dtype = torch.float16 if device == "cuda" else torch.float32
 
-print(f"Using device: {device}")
-print(f"Loading models from: {MODEL_REPO}")
+# Note: For ZeroGPU, device detection happens dynamically
+# We'll set device inside GPU-decorated functions
+print("Using ZeroGPU - GPU will be allocated on-demand")
 
 class RetroArtConverter:
     def __init__(self):
-        self.device = device
-        self.dtype = dtype
+        self.models_loaded = False
+        self.device = None
+        self.dtype = None
+        self.face_detection_enabled = False
+        print("RetroArtConverter initialized - models will load on first generation")
+    
+    def _initialize_models(self):
+        """Lazy model initialization - called on first generation when GPU is available"""
+        if self.models_loaded:
+            return
+        
+        print("Initializing models...")
+        print(f"Loading models from: {MODEL_REPO}")
+        
+        # Set device (will be cuda when called from GPU-decorated function)
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.dtype = torch.float16 if self.device == "cuda" else torch.float32
+        print(f"Using device: {self.device}")
         
         # Initialize face analysis for InstantID (optional)
         print("Loading face analysis model...")
@@ -78,7 +94,8 @@ class RetroArtConverter:
         print("Loading depth estimator...")
         self.depth_estimator = transformers_pipeline(
             'depth-estimation',
-            model="Intel/dpt-hybrid-midas"
+            model="Intel/dpt-hybrid-midas",
+            device=self.device if self.device == "cuda" else -1
         )
         
         # Load SDXL checkpoint from HuggingFace Hub
@@ -108,7 +125,7 @@ class RetroArtConverter:
                 use_safetensors=True
             ).to(self.device)
         
-        # Load LORA from HuggingFace Hub
+        # Load LORA from HuggingFace Hub (requires PEFT)
         print("Loading LORA (retroart) from HuggingFace Hub...")
         try:
             lora_path = hf_hub_download(
@@ -126,19 +143,25 @@ class RetroArtConverter:
         self.pipe.scheduler = DPMSolverMultistepScheduler.from_config(
             self.pipe.scheduler.config
         )
-        self.pipe.enable_model_cpu_offload()
+        
+        # For ZeroGPU, we don't use model_cpu_offload
+        # self.pipe.enable_model_cpu_offload()
+        
         self.pipe.enable_vae_slicing()
         
         # Enable attention slicing for memory efficiency
         self.pipe.unet.set_attn_processor(AttnProcessor2_0())
         
-        if hasattr(self.pipe, 'enable_xformers_memory_efficient_attention'):
+        # Try to enable xformers if available (only works on GPU)
+        if self.device == "cuda":
             try:
                 self.pipe.enable_xformers_memory_efficient_attention()
+                print("✓ xformers enabled")
             except Exception as e:
-                print(f"xformers not available: {e}")
+                print(f"⚠️ xformers not available: {e}")
         
-        print("Model initialization complete!")
+        self.models_loaded = True
+        print("✓ Model initialization complete!")
     
     def get_depth_map(self, image):
         """Generate depth map from input image"""
@@ -199,6 +222,9 @@ class RetroArtConverter:
     ):
         """Main generation function"""
         
+        # Initialize models on first run (lazy loading for ZeroGPU)
+        self._initialize_models()
+        
         # Resize image maintaining aspect ratio
         original_width, original_height = input_image.size
         target_width, target_height = self.calculate_target_size(original_width, original_height)
@@ -244,7 +270,8 @@ class RetroArtConverter:
 print("Initializing RetroArt Converter...")
 converter = RetroArtConverter()
 
-# Gradio interface
+# Gradio interface with ZeroGPU support
+@spaces.GPU
 def process_image(
     image,
     prompt,