lazy load models

gradio_demo.py

@@ -58,25 +58,28 @@ c2w = torch.tensor([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 1], [0, 0, 0, 1]], dty
 
 
 def initialize_models():
-    """Initialize all models (called once at startup)."""
+    """Initialize all models (called lazily on first use)."""
     global sheap_model, flame, fa_model
+    
+    if sheap_model is not None:
+        return  # Already initialized
 
-    print("Loading SHeaP model...")
+    print("Loading SHeaP model...", flush=True)
     sheap_model = load_sheap_model(model_type="expressive").to(device)
     sheap_model.eval()
 
-    print("Loading FLAME model...")
+    print("Loading FLAME model...", flush=True)
     flame_dir = Path("FLAME2020/")
     flame = TinyFlame(flame_dir / "generic_model.pt", eyelids_ckpt=flame_dir / "eyelids.pt").to(
         device
     )
 
-    print("Loading face alignment model...")
+    print("Loading face alignment model...", flush=True)
     fa_model = face_alignment.FaceAlignment(
         face_alignment.LandmarksType.TWO_D, device=str(device), flip_input=False
     )
 
-    print("Models loaded successfully!")
+    print("Models loaded successfully!", flush=True)
 
 
 @spaces.GPU
@@ -90,6 +93,9 @@ def process_image(image: np.ndarray) -> Image.Image:
     Returns:
         PIL Image with three views side-by-side (original, mesh, blended)
     """
+    # Initialize models on first use (lazy loading for @spaces.GPU)
+    initialize_models()
+    
     # Convert to torch tensor for face detection (C, H, W) format with values in [0, 1]
     image_tensor = torch.from_numpy(image).permute(2, 0, 1).float() / 255.0
 
@@ -148,6 +154,9 @@ def process_video(video_path: str, progress=gr.Progress()) -> str:
     """
     Process a video and return path to the rendered output video using background threads.
     """
+    # Initialize models on first use (lazy loading for @spaces.GPU)
+    initialize_models()
+    
     temp_dir = Path(tempfile.mkdtemp())
     render_size = 512
     try:
@@ -239,8 +248,9 @@ def process_input(image: Optional[np.ndarray], video: Optional[str]):
         raise ValueError("Please provide either an image or video!")
 
 
-# Initialize models on startup
-initialize_models()
+# Don't initialize models at startup when using @spaces.GPU
+# They will be loaded lazily on first use
+# initialize_models()
 
 # Create Gradio interface
 with gr.Blocks(title="SHeaP Demo") as demo: