Update app.py

app.py

@@ -128,12 +128,12 @@ def preprocess_image(image_array, session_id):
 # Step 2: Normals inference → normals image
 @spaces.GPU()
 def step2_normals(session_id):
-
+    from pixel3dmm.lightning.p3dmm_system import system as p3dmm_system
+    
     base_conf = OmegaConf.load("configs/base.yaml")
 
     if "normals_model" not in _model_cache:
-        from pixel3dmm.lightning.p3dmm_system import system as p3dmm_system
-    
+
         model = p3dmm_system.load_from_checkpoint(f"{env_paths.CKPT_N_PRED}", strict=False)
         model =  model.eval().to(DEVICE)
         _model_cache["normals_model"] = model   
@@ -149,12 +149,12 @@ def step2_normals(session_id):
 # Step 3: UV map inference → uv map image
 @spaces.GPU()
 def step3_uv_map(session_id):
-
+    from pixel3dmm.lightning.p3dmm_system import system as p3dmm_system
+    
     base_conf = OmegaConf.load("configs/base.yaml")
 
     if "uv_model" not in _model_cache:
-        from pixel3dmm.lightning.p3dmm_system import system as p3dmm_system
-    
+
         model = p3dmm_system.load_from_checkpoint(f"{env_paths.CKPT_UV_PRED}", strict=False)
         model =  model.eval().to(DEVICE)
         _model_cache["uv_model"] = model   
@@ -171,21 +171,21 @@ def step3_uv_map(session_id):
 # Step 4: Tracking → final tracking image
 @spaces.GPU()
 def step4_track(session_id):
-
+    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.renderer_nvdiffrast import NVDRenderer
+    from pixel3dmm.tracking.tracker import Tracker
+    
     tracking_conf = OmegaConf.load("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.renderer_nvdiffrast import NVDRenderer
-        from pixel3dmm.tracking.tracker import Tracker
-    
+
         flame = FLAME(tracking_conf)      # CPU instantiation
         flame = flame.to(DEVICE)      # CUDA init happens here
         _model_cache["flame_model"] = flame