Update app.py

app.py

@@ -98,16 +98,14 @@ def reset_all():
 
 # Step 1: Preprocess the input image (Save and Crop)
 @spaces.GPU()
-def preprocess_image(image_array, state):
+def preprocess_image(image_array, session_id):
     if image_array is None:
-        return "❌ Please upload an image first.", None, state, gr.update(interactive=True), gr.update(interactive=True)
+        return "❌ Please upload an image first.", None, gr.update(interactive=True), gr.update(interactive=True)
 
-    base_dir = os.path.join(os.environ["PIXEL3DMM_PREPROCESSED_DATA"], session_id)
     os.makedirs(base_dir, exist_ok=True)
-    state.update({"base_dir": base_dir})
 
     img = Image.fromarray(image_array)
-    saved_image_path = os.path.join(base_dir, f"{session_id}.png")
+    saved_image_path = os.path.join(os.environ["PIXEL3DMM_PREPROCESSED_DATA"], session_id, f"{session_id}.png")
     img.save(saved_image_path)
     
     try:
@@ -121,11 +119,11 @@ def preprocess_image(image_array, state):
 
     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=True), gr.update(interactive=True)
+    return "✅ Step 1 complete. Ready for Normals.", image, gr.update(interactive=True), gr.update(interactive=True)
 
 # Step 2: Normals inference → normals image
 @spaces.GPU()
-def step2_normals(state):
+def step2_normals(session_id):
 
     base_conf = OmegaConf.load("configs/base.yaml")
 
@@ -136,18 +134,17 @@ def step2_normals(state):
         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")
+    normals_dir = os.path.join(os.environ["PIXEL3DMM_PREPROCESSED_DATA"], session_id, "p3dmm", "normals")
     image = first_image_from_dir(normals_dir)
 
-    return "✅ Step 2 complete. Ready for UV Map.", image, state, gr.update(interactive=True), gr.update(interactive=True)
+    return "✅ Step 2 complete. Ready for UV Map.", image, gr.update(interactive=True), gr.update(interactive=True)
 
 # Step 3: UV map inference → uv map image
 @spaces.GPU()
-def step3_uv_map(state):
+def step3_uv_map(session_id):
 
     base_conf = OmegaConf.load("configs/base.yaml")
 
@@ -158,19 +155,18 @@ def step3_uv_map(state):
         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")
+    uv_dir = os.path.join(os.environ["PIXEL3DMM_PREPROCESSED_DATA"], session_id, "p3dmm", "uv_map")
     image = first_image_from_dir(uv_dir)
 
-    return "✅ Step 3 complete. Ready for Tracking.", image, state, gr.update(interactive=True), gr.update(interactive=True)
+    return "✅ Step 3 complete. Ready for Tracking.", image, gr.update(interactive=True), gr.update(interactive=True)
 
 # Step 4: Tracking → final tracking image
 @spaces.GPU()
-def step4_track(state):
+def step4_track(session_id):
 
     tracking_conf = OmegaConf.load("configs/tracking.yaml")
 
@@ -203,7 +199,6 @@ def step4_track(state):
    
     flame_model = _model_cache["flame_model"]
     diff_renderer = _model_cache["diff_renderer"]
-    session_id = state.get("session_id")
     tracking_conf.video_name = f'{session_id}'
     tracker = Tracker(tracking_conf, flame_model, diff_renderer)
     tracker.run()
@@ -212,31 +207,30 @@ def step4_track(state):
     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=True)
+    return "✅ Pipeline complete!", image, gr.update(interactive=True)
 
 # New: run all steps sequentially
 @spaces.GPU()
-def run_pipeline(image_array, state, request: gr.Request):
+def run_pipeline(image_array, request: gr.Request):
 
     session_id = request.session_hash
-    state.update({"session_id": session_id, "base_dir": base_dir})
 
     # Step 1
-    status1, crop_img, state, _, _ = preprocess_image(image_array, state)
+    status1, crop_img, _, _ = preprocess_image(image_array)
     if "❌" in status1:
-        return status1, None, None, None, None, None, {}
+        return status1, None, None, None, None, None
     # Step 2
-    status2, normals_img, state, _, _ = step2_normals(state)
+    status2, normals_img, _, _ = step2_normals(session_id)
     # Step 3
-    status3, uv_img, state, _, _ = step3_uv_map(state)
+    status3, uv_img, _, _ = step3_uv_map(session_id)
     # Step 4
-    status4, track_img, state, _ = step4_track(state)
+    status4, track_img, _ = step4_track(session_id)
     # Locate mesh (.ply)
-    mesh_dir = os.path.join(os.environ["PIXEL3DMM_TRACKING_OUTPUT"], state.get("session_id"), "mesh")
+    mesh_dir = os.path.join(os.environ["PIXEL3DMM_TRACKING_OUTPUT"], session_id, "mesh")
     mesh_file = first_file_from_dir(mesh_dir, "ply")
 
     final_status = "\n".join([status1, status2, status3, status4])
-    return final_status, crop_img, normals_img, uv_img, track_img, mesh_file, state
+    return final_status, crop_img, normals_img, uv_img, track_img, mesh_file
 
 # Cleanup on unload
 def cleanup(request: gr.Request):
@@ -275,7 +269,7 @@ with gr.Blocks(css=css) as demo:
             with gr.Column():
                 image_in = gr.Image(label="Upload Image", type="numpy", height=512)
                 status = gr.Textbox(label="Status", lines=6, interactive=True, value="Upload an image to start.")
-                state = gr.State({})
+
             with gr.Column():
                 with gr.Row():
                     crop_img = gr.Image(label="Preprocessed", height=256)
@@ -297,11 +291,11 @@ with gr.Blocks(css=css) as demo:
     
         run_btn.click(
             fn=run_pipeline,
-            inputs=[image_in, state],
-            outputs=[status, crop_img, normals_img, uv_img, track_img, mesh_file, state]
+            inputs=[image_in],
+            outputs=[status, crop_img, normals_img, uv_img, track_img, mesh_file]
         )
-        examples.outputs = [status, crop_img, normals_img, uv_img, track_img, mesh_file, state]
-        image_in.upload(fn=reset_all, inputs=None, outputs=[crop_img, normals_img, uv_img, track_img, mesh_file, status, state, run_btn])
+        examples.outputs = [status, crop_img, normals_img, uv_img, track_img, mesh_file]
+        image_in.upload(fn=reset_all, inputs=None, outputs=[crop_img, normals_img, uv_img, track_img, mesh_file, status, run_btn])
 
     demo.unload(cleanup)