Update app.py

app.py

@@ -7,7 +7,6 @@ import os
 # Force Dynamo off at import‐time of torch, pytorch3d, etc.
 os.environ["TORCHDYNAMO_DISABLE"] = "1"
 
-from gradio_litmodel3d import LitModel3D
 import subprocess
 import tempfile
 import uuid
@@ -64,6 +63,7 @@ install_cuda_toolkit()
 
 from omegaconf import OmegaConf
 from pixel3dmm.network_inference import normals_n_uvs
+from pixel3dmm.run_facer_segmentation import segment
 
 DEVICE = "cuda"
 
@@ -92,7 +92,6 @@ def reset_all():
         None,  # uv_img
         None,  # track_img
         "Time to Generate!",  # status
-        {},    # state
         gr.update(interactive=True),   # preprocess_btn
         gr.update(interactive=True),  # normals_btn
         gr.update(interactive=True),  # uv_map_btn
@@ -111,15 +110,27 @@ def preprocess_image(image_array, session_id):
     img = Image.fromarray(image_array)
     saved_image_path = os.path.join(os.environ["PIXEL3DMM_PREPROCESSED_DATA"], session_id, f"{session_id}.png")
     img.save(saved_image_path)
-    
-    try:
-        p = subprocess.run([
+
+    import facer
+
+    if "face_detector" not in _model_cache:
+
+        device = 'cuda'
+
+        # This call downloads/loads the RetinaFace Mobilenet weights
+        face_detector = facer.face_detector('retinaface/mobilenet', device=device)
+
+        # This call downloads/loads the FARL parsing model (celeba mask model)
+        face_parser   = facer.face_parser   ('farl/celebm/448',     device=device)
+
+        _model_cache['face_detector'] = face_detector
+        _model_cache['face_parser'] = face_parser
+            
+    subprocess.run([
             "python", "scripts/run_preprocessing.py", "--video_or_images_path", saved_image_path
         ], check=True, capture_output=True, text=True)
-    except subprocess.CalledProcessError as e:
-        err = f"❌ Preprocess failed (exit {e.returncode}).\n\n{e.stdout}\n{e.stderr}"
-        shutil.rmtree(base_dir, ignore_errors=True)
-        return err, None, gr.update(interactive=True), gr.update(interactive=True)
+    
+    segment(f'{session_id}', _model_cache['face_detector'], _model_cache['face_parser'])
 
     crop_dir = os.path.join(os.environ["PIXEL3DMM_PREPROCESSED_DATA"], session_id, "cropped")
     image = first_image_from_dir(crop_dir)
@@ -246,7 +257,7 @@ def generate_results_and_mesh(image, session_id=None):
     # Step 1
     status1, crop_img, _, _ = preprocess_image(image, session_id)
     if "❌" in status1:
-        return status1, None, None, None, None, None
+      return status1, None, None, None, None, None
     # Step 2
     status2, normals_img, _, _ = step2_normals(session_id)
     # Step 3
@@ -348,4 +359,4 @@ with gr.Blocks(css=css) as demo:
 demo.queue(default_concurrency_limit=1,        # ≤ 1 worker per event
            max_size=20)                        # optional: allow 20 waiting jobs
 
-demo.launch()
+demo.launch(share=True)