Update app.py

app.py

@@ -32,11 +32,9 @@ import importlib, site; site.addsitedir(site.getsitepackages()[0]); importlib.in
 
 from pixel3dmm import env_paths
 
-
 sh("cd src/pixel3dmm/preprocessing/facer && pip install -e . && cd ../../../..")
 sh("cd src/pixel3dmm/preprocessing/PIPNet/FaceBoxesV2/utils && sh make.sh && cd ../../../../../..")
 
-
 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)
@@ -57,12 +55,10 @@ def install_cuda_toolkit():
 install_cuda_toolkit()
 
 from omegaconf import OmegaConf
+from pixel3dmm.network_inference import normals_n_uvs
 
 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 = {}
 
@@ -87,16 +83,16 @@ def reset_all():
         "Awaiting new image upload...",  # status
         {},    # state
         gr.update(interactive=True),   # preprocess_btn
-        gr.update(interactive=False),  # normals_btn
-        gr.update(interactive=False),  # uv_map_btn
-        gr.update(interactive=False)   # track_btn
+        gr.update(interactive=True),  # normals_btn
+        gr.update(interactive=True),  # uv_map_btn
+        gr.update(interactive=True)   # track_btn
     )
 
 # Step 1: Preprocess the input image (Save and Crop)
 @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)
+        return "❌ Please upload an image first.", None, state, gr.update(interactive=True), gr.update(interactive=True)
 
     session_id = str(uuid.uuid4())
     base_dir = os.path.join(os.environ["PIXEL3DMM_PREPROCESSED_DATA"], session_id)
@@ -115,54 +111,63 @@ def preprocess_image(image_array, state):
     except subprocess.CalledProcessError as e:
         err = f"❌ Preprocess failed (exit {e.returncode}).\n\n{e.stdout}\n{e.stderr}"
         shutil.rmtree(base_dir)
-        return err, None, {}, gr.update(interactive=True), gr.update(interactive=False)
+        return err, None, {}, gr.update(interactive=True), gr.update(interactive=True)
 
     crop_dir = os.path.join(base_dir, "cropped")
     image = first_image_from_dir(crop_dir)
-    return "✅ Step 1 complete. Ready for Normals.", image, state, gr.update(interactive=False), gr.update(interactive=True)
+    return "✅ Step 1 complete. Ready for Normals.", image, state, gr.update(interactive=True), gr.update(interactive=True)
 
 # Step 2: Normals inference → normals image
 @spaces.GPU()
 def step2_normals(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), gr.update(interactive=False)
 
-    try:
-        p = subprocess.run([
-            "python", "scripts/network_inference.py", "model.prediction_type=normals", f"video_name={session_id}"
-        ], check=True, capture_output=True, text=True)
-    except subprocess.CalledProcessError as e:
-        err = f"❌ Normal map failed (exit {e.returncode}).\n\n{e.stdout}\n{e.stderr}"
-        return err, None, state, gr.update(interactive=True), gr.update(interactive=False)
+    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   
+
+    session_id = state.get("session_id")
+    base_conf.video_name = f'{session_id}'
+    normals_n_uvs(base_conf, _model_cache["normals_model"])
 
     normals_dir = os.path.join(state["base_dir"], "p3dmm", "normals")
     image = first_image_from_dir(normals_dir)
-    return "✅ Step 2 complete. Ready for UV Map.", image, state, gr.update(interactive=False), gr.update(interactive=True)
+
+    return "✅ Step 2 complete. Ready for UV Map.", image, state, gr.update(interactive=True), gr.update(interactive=True)
 
 # Step 3: UV map inference → uv map image
 @spaces.GPU()
 def step3_uv_map(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), gr.update(interactive=False)
 
-    try:
-        p = subprocess.run([
-            "python", "scripts/network_inference.py", "model.prediction_type=uv_map", f"video_name={session_id}"
-        ], check=True, capture_output=True, text=True)
-    except subprocess.CalledProcessError as e:
-        err = f"❌ UV map failed (exit {e.returncode}).\n\n{e.stdout}\n{e.stderr}"
-        return err, None, state, gr.update(interactive=True), gr.update(interactive=False)
+    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   
+
+    session_id = state.get("session_id")
+    base_conf.video_name = f'{session_id}'
+    base_conf.model.prediction_type = "uv_map"
+    normals_n_uvs(base_conf, _model_cache["uv_model"])
 
     uv_dir = os.path.join(state["base_dir"], "p3dmm", "uv_map")
     image = first_image_from_dir(uv_dir)
-    return "✅ Step 3 complete. Ready for Tracking.", image, state, gr.update(interactive=False), gr.update(interactive=True)
+
+    return "✅ Step 3 complete. Ready for Tracking.", image, state, gr.update(interactive=True), gr.update(interactive=True)
 
 # Step 4: Tracking → final tracking image
 @spaces.GPU()
 def step4_track(state):
 
+    tracking_conf = OmegaConf.load("configs/tracking.yaml")
+
      # Lazy init + caching of FLAME model on GPU
     if "flame_model" not in _model_cache:
         import os
@@ -175,7 +180,7 @@ def step4_track(state):
         from pixel3dmm.tracking.renderer_nvdiffrast import NVDRenderer
         from pixel3dmm.tracking.tracker import Tracker
     
-        flame = FLAME(base_conf)      # CPU instantiation
+        flame = FLAME(tracking_conf)      # CPU instantiation
         flame = flame.to(DEVICE)      # CUDA init happens here
         _model_cache["flame_model"] = flame   
         
@@ -184,24 +189,23 @@ def step4_track(state):
         _obj_faces = load_obj(_mesh_file)[1]
         
         _model_cache["diff_renderer"] = NVDRenderer(
-            image_size=base_conf.size, 
+            image_size=tracking_conf.size, 
             obj_filename=_mesh_file,
             no_sh=False,
             white_bg=True
         ).to(DEVICE)
-
-
+   
     flame_model = _model_cache["flame_model"]
     diff_renderer = _model_cache["diff_renderer"]
     session_id = state.get("session_id")
-    base_conf.video_name = f'{session_id}'
-    tracker = Tracker(base_conf, flame_model, diff_renderer)
+    tracking_conf.video_name = f'{session_id}'
+    tracker = Tracker(tracking_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)
+    return "✅ Pipeline complete!", image, state, gr.update(interactive=True)
 
 # Build Gradio UI
 demo = gr.Blocks()
@@ -212,7 +216,7 @@ with demo:
     with gr.Row():
         with gr.Column():
             image_in = gr.Image(label="Upload Image", type="numpy", height=512)
-            status = gr.Textbox(label="Status", lines=2, interactive=False, value="Upload an image to start.")
+            status = gr.Textbox(label="Status", lines=2, interactive=True, value="Upload an image to start.")
             state = gr.State({})
         with gr.Column():
             with gr.Row():
@@ -224,9 +228,9 @@ with demo:
 
     with gr.Row():
         preprocess_btn = gr.Button("Step 1: Preprocess", interactive=True)
-        normals_btn = gr.Button("Step 2: Normals", interactive=False)
-        uv_map_btn = gr.Button("Step 3: UV Map", interactive=False)
-        track_btn = gr.Button("Step 4: Track", interactive=False)
+        normals_btn = gr.Button("Step 2: Normals", interactive=True)
+        uv_map_btn = gr.Button("Step 3: UV Map", interactive=True)
+        track_btn = gr.Button("Step 4: Track", interactive=True)
 
     # Define component list for reset
     outputs_for_reset = [crop_img, normals_img, uv_img, track_img, status, state, preprocess_btn, normals_btn, uv_map_btn, track_btn]