tracker optimised for zero gpu

app.py

@@ -17,11 +17,12 @@ os.environ["PIXEL3DMM_TRACKING_OUTPUT"] = f"{os.getcwd()}/tracking_results"
 
 def sh(cmd): subprocess.check_call(cmd, shell=True)
 
-# only do this once per VM restart
+
 sh("pip install -e .")
 sh("cd src/pixel3dmm/preprocessing/facer && pip install -e .")
 sh("cd src/pixel3dmm/preprocessing/PIPNet/FaceBoxesV2/utils && sh make.sh")
 
+
 def install_cuda_toolkit():
     CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run"
     CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)
@@ -39,9 +40,36 @@ def install_cuda_toolkit():
     os.environ["TORCH_CUDA_ARCH_LIST"] = "9.0"
     print("==> finished installation")
     
-
 install_cuda_toolkit()
 
+import os
+import torch
+import numpy as np
+import trimesh
+from pytorch3d.io import load_obj
+from pixel3dmm.tracking.renderer_nvdiffrast import NVDRenderer
+from pixel3dmm.tracking.flame.FLAME import FLAME
+from pixel3dmm import env_paths
+from omegaconf import OmegaConf
+from pixel3dmm.tracking.tracker import Tracker
+
+DEVICE = "cuda"
+
+base_conf = OmegaConf.load(f'{env_paths.CODE_BASE}/configs/tracking.yaml')
+
+_mesh_file = env_paths.head_template
+flame_model = FLAME(base_conf).to(DEVICE)
+
+_obj_faces = load_obj(_mesh_file)[1]
+
+diff_renderer = NVDRenderer(
+    image_size=base_conf.size, 
+    obj_filename=_mesh_file,
+    no_sh=False,
+    white_bg=True
+).to(DEVICE)
+
+
 # Utility to select first image from a folder
 def first_image_from_dir(directory):
     patterns = ["*.jpg", "*.png", "*.jpeg"]
@@ -68,7 +96,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)
@@ -138,20 +166,13 @@ def step3_uv_map(state):
 @spaces.GPU()
 def step4_track(state):
     session_id = state.get("session_id")
-    if not session_id:
-        return "❌ State lost. Please start from Step 1.", None, state, gr.update(interactive=False)
-
-    script = os.path.join(os.environ["PIXEL3DMM_CODE_BASE"], "scripts", "track.py")
-    try:
-        p = subprocess.run([
-            "python", script, f"video_name={session_id}"
-        ], check=True, capture_output=True, text=True)
-    except subprocess.CalledProcessError as e:
-        err = f"❌ Tracking failed (exit {e.returncode}).\n\n{e.stdout}\n{e.stderr}"
-        return err, None, state, gr.update(interactive=True)
+    base_conf.video_name = f'{session_id}'
+    tracker = Tracker(base_conf, flame_model, diff_renderer)
+    tracker.run()
 
     tracking_dir = os.path.join(os.environ["PIXEL3DMM_TRACKING_OUTPUT"], session_id, "frames")
     image = first_image_from_dir(tracking_dir)
+
     return "✅ Pipeline complete!", image, state, gr.update(interactive=False)
 
 # Build Gradio UI

src/pixel3dmm/tracking/tracker.py

@@ -128,9 +128,13 @@ if COMPILE:
 
 
 class Tracker(object):
-    def __init__(self, config,
+    def __init__(self, config, flame_module, renderer,
                  device='cuda:0',
                  ):
+        self.config = config
+        self.flame = flame_module
+        self.diff_renderer = renderer
+    
         self.config = config
         self.device = device
         self.actor_name = self.config.video_name
@@ -240,7 +244,6 @@ class Tracker(object):
     def setup_renderer(self):
         mesh_file = f'{env_paths.head_template}'
         self.config.image_size = self.get_image_size()
-        self.flame = FLAME(self.config).to(self.device)
         self.flame.vertex_face_mask = self.vertex_face_mask
 
 
@@ -251,11 +254,7 @@ class Tracker(object):
             self.actual_smooth = torch.compile(self.actual_smooth)
 
 
-        self.diff_renderer = NVDRenderer(self.config.size,
-                                         obj_filename=mesh_file,
-                                         no_sh=self.no_sh,
-                                         white_bg= True,
-                                         ).to(self.device)
+        self.renderer = self.diff_renderer  # already global
 
 
         self.faces = load_obj(mesh_file)[1]