Initial commit

app.py

@@ -23,6 +23,7 @@ from easydict import EasyDict as edict
 from einops import rearrange
 from PIL import Image
 from huggingface_hub import snapshot_download
+import spaces
 
 # Install diff-gaussian-rasterization at runtime (requires GPU)
 import subprocess
@@ -98,14 +99,14 @@ class FaceLiftPipeline:
         self.image_size = 512
         self.camera_indices = [2, 1, 0, 5, 4, 3]
         
-        # Load models
+        # Load models (keep on CPU for ZeroGPU compatibility)
         print("Loading models...")
         self.mvdiffusion_pipeline = StableUnCLIPImg2ImgPipeline.from_pretrained(
             str(workspace_dir / "checkpoints/mvdiffusion/pipeckpts"),
             torch_dtype=torch.float16,
         )
-        self.mvdiffusion_pipeline.unet.enable_xformers_memory_efficient_attention()
-        self.mvdiffusion_pipeline.to(self.device)
+        # Don't move to device or enable xformers here - will be done in GPU-decorated function
+        self._models_on_gpu = False
         
         with open(workspace_dir / "configs/gslrm.yaml", "r") as f:
             config = edict(yaml.safe_load(f))
@@ -120,11 +121,11 @@ class FaceLiftPipeline:
             map_location="cpu"
         )
         self.gs_lrm_model.load_state_dict(checkpoint["model"])
-        self.gs_lrm_model.to(self.device)
+        # Keep on CPU initially - will move to GPU in decorated function
         
         self.color_prompt_embedding = torch.load(
             workspace_dir / "mvdiffusion/fixed_prompt_embeds_6view/clr_embeds.pt",
-            map_location=self.device
+            map_location="cpu"
         )
         
         with open(workspace_dir / "utils_folder/opencv_cameras.json", 'r') as f:
@@ -132,6 +133,18 @@ class FaceLiftPipeline:
         
         print("Models loaded successfully!")
     
+    def _move_models_to_gpu(self):
+        """Move models to GPU and enable optimizations. Called within @spaces.GPU context."""
+        if not self._models_on_gpu and torch.cuda.is_available():
+            print("Moving models to GPU...")
+            self.device = torch.device("cuda:0")
+            self.mvdiffusion_pipeline.to(self.device)
+            self.mvdiffusion_pipeline.unet.enable_xformers_memory_efficient_attention()
+            self.gs_lrm_model.to(self.device)
+            self.color_prompt_embedding = self.color_prompt_embedding.to(self.device)
+            self._models_on_gpu = True
+            print("Models on GPU, xformers enabled!")
+    
     def _create_viewer_html(self, splat_path):
         """Create standalone HTML viewer for the gaussian splat."""
         import base64
@@ -246,9 +259,13 @@ class FaceLiftPipeline:
 </html>"""
         return html
     
+    @spaces.GPU
     def generate_3d_head(self, image_path, auto_crop=True, guidance_scale=3.0, 
                          random_seed=4, num_steps=50):
         """Generate 3D head from single image."""
+        # Move models to GPU now that we're in the GPU context
+        self._move_models_to_gpu()
+        
         try:
             # Setup output directory
             timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

requirements.txt

@@ -25,3 +25,4 @@ jaxtyping==0.2.19
 pytorch-msssim==1.0.0
 ffmpeg-python==0.2.0
 tqdm
+spaces