Add ZeroGPU duration parameter and improve error handling

app.py

@@ -98,40 +98,46 @@ class FaceLiftPipeline:
         
         # 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,
-        )
-        # 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))
-        
-        module_name, class_name = config.model.class_name.rsplit(".", 1)
-        module = __import__(module_name, fromlist=[class_name])
-        ModelClass = getattr(module, class_name)
-        
-        self.gs_lrm_model = ModelClass(config)
-        checkpoint = torch.load(
-            workspace_dir / "checkpoints/gslrm/ckpt_0000000000021125.pt",
-            map_location="cpu"
-        )
-        # Filter out loss_calculator weights (training-only, not needed for inference)
-        state_dict = {k: v for k, v in checkpoint["model"].items() 
-                      if not k.startswith("loss_calculator.")}
-        self.gs_lrm_model.load_state_dict(state_dict)
-        # 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="cpu"
-        )
-        
-        with open(workspace_dir / "utils_folder/opencv_cameras.json", 'r') as f:
-            self.cameras_data = json.load(f)["frames"]
-        
-        print("Models loaded successfully!")
+        try:
+            self.mvdiffusion_pipeline = StableUnCLIPImg2ImgPipeline.from_pretrained(
+                str(workspace_dir / "checkpoints/mvdiffusion/pipeckpts"),
+                torch_dtype=torch.float16,
+            )
+            # 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))
+            
+            module_name, class_name = config.model.class_name.rsplit(".", 1)
+            module = __import__(module_name, fromlist=[class_name])
+            ModelClass = getattr(module, class_name)
+            
+            self.gs_lrm_model = ModelClass(config)
+            checkpoint = torch.load(
+                workspace_dir / "checkpoints/gslrm/ckpt_0000000000021125.pt",
+                map_location="cpu"
+            )
+            # Filter out loss_calculator weights (training-only, not needed for inference)
+            state_dict = {k: v for k, v in checkpoint["model"].items() 
+                          if not k.startswith("loss_calculator.")}
+            self.gs_lrm_model.load_state_dict(state_dict)
+            # 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="cpu"
+            )
+            
+            with open(workspace_dir / "utils_folder/opencv_cameras.json", 'r') as f:
+                self.cameras_data = json.load(f)["frames"]
+            
+            print("Models loaded successfully!")
+        except Exception as e:
+            print(f"Error loading models: {e}")
+            import traceback
+            traceback.print_exc()
+            raise
     
     def _move_models_to_gpu(self):
         """Move models to GPU and enable optimizations. Called within @spaces.GPU context."""
@@ -145,14 +151,13 @@ class FaceLiftPipeline:
             self._models_on_gpu = True
             print("Models on GPU, xformers enabled!")
     
-    @spaces.GPU
+    @spaces.GPU(duration=120)
     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:
+            # Move models to GPU now that we're in the GPU context
+            self._move_models_to_gpu()
             # Setup output directory
             timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
             output_dir = self.output_dir / timestamp
@@ -260,6 +265,9 @@ class FaceLiftPipeline:
                    str(turntable_path), str(ply_path)
             
         except Exception as e:
+            import traceback
+            error_details = traceback.format_exc()
+            print(f"Error details:\n{error_details}")
             raise gr.Error(f"Generation failed: {str(e)}")