Update app.py

app.py

@@ -55,7 +55,14 @@ def install_cuda_toolkit():
     print("==> finished installation")
     
 install_cuda_toolkit()
+    
+DEVICE = "cuda"
+
+# 1. Prepare config at import time (no CUDA calls)
+base_conf = OmegaConf.load("configs/tracking.yaml")
 
+# 2. Empty cache for our heavy objects
+_model_cache = {}
 
 
 # Utility to select first image from a folder
@@ -84,7 +91,7 @@ def reset_all():
     )
 
 # Step 1: Preprocess the input image (Save and Crop)
-# @spaces.GPU()
+@spaces.GPU()
 def preprocess_image(image_array, state):
     if image_array is None:
         return "❌ Please upload an image first.", None, state, gr.update(interactive=True), gr.update(interactive=False)
@@ -154,23 +161,23 @@ def step3_uv_map(state):
 @spaces.GPU()
 def step4_track(state):
 
-    import os
-    import torch
-    import numpy as np
-    import trimesh
-    from pytorch3d.io import load_obj
-    
-    from pixel3dmm.tracking.flame.FLAME import FLAME
-    from pixel3dmm.tracking.tracker import Tracker
-    from omegaconf import OmegaConf
-    
-    
-    DEVICE = "cuda"
-    
-    base_conf = OmegaConf.load(f'{env_paths.CODE_BASE}/configs/tracking.yaml')
+     # Lazy init + caching of FLAME model on GPU
+    if "flame_model" not in _model_cache:
+        import os
+        import torch
+        import numpy as np
+        import trimesh
+        from pytorch3d.io import load_obj
+        
+        from pixel3dmm.tracking.flame.FLAME import FLAME
+        from pixel3dmm.tracking.tracker import Tracker
+        from omegaconf import OmegaConf
     
-    flame_model = FLAME(base_conf).to(DEVICE)
+        flame = FLAME(base_conf)      # CPU instantiation
+        flame = flame.to(DEVICE)      # CUDA init happens here
+        _model_cache["flame_model"] = flame   
 
+    flame_model = _model_cache["flame_model"]
     session_id = state.get("session_id")
     base_conf.video_name = f'{session_id}'
     tracker = Tracker(base_conf, flame_model)