update

app.py

@@ -42,7 +42,16 @@ def prepare_env():
 
     return
 
+
+def server_up():
+    server_up_flag = Path(f'{REPO_ROOT}/server_up')
+    if server_up_flag.exists():
+        while True:
+            time.sleep(600)
+    else:
+        os.system(f'python {REPO_ROOT}/app/entry.py')
+
+
 if __name__ == '__main__':
     prepare_env()
-    while True:
-        time.sleep(600)
+    server_up()

app/demo.py

@@ -84,7 +84,7 @@ def parse_args_to_cfg():
 
 
 @torch.no_grad()
-def run_preprocess(cfg):
+def run_preprocess(cfg, progress):
     Log.info(f"[Preprocess] Start!")
     tic = Log.time()
     video_path = cfg.video_path
@@ -93,6 +93,7 @@ def run_preprocess(cfg):
     verbose = cfg.verbose
 
     # Get bbx tracking result
+    progress(0, '[Preprocess] YoloV8 Tracking')
     if not Path(paths.bbx).exists():
         tracker = Tracker()
         bbx_xyxy = tracker.get_one_track(video_path).float()  # (L, 4)
@@ -109,6 +110,7 @@ def run_preprocess(cfg):
         save_video(video_overlay, cfg.paths.bbx_xyxy_video_overlay)
 
     # Get VitPose
+    progress(1/4, '[Preprocess] ViTPose')
     if not Path(paths.vitpose).exists():
         vitpose_extractor = VitPoseExtractor()
         vitpose = vitpose_extractor.extract(video_path, bbx_xys)
@@ -123,6 +125,7 @@ def run_preprocess(cfg):
         save_video(video_overlay, paths.vitpose_video_overlay)
 
     # Get vit features
+    progress(2/4, '[Preprocess] HMR2 Feature')
     if not Path(paths.vit_features).exists():
         extractor = Extractor()
         vit_features = extractor.extract_video_features(video_path, bbx_xys)
@@ -132,6 +135,7 @@ def run_preprocess(cfg):
         Log.info(f"[Preprocess] vit_features from {paths.vit_features}")
 
     # Get DPVO results
+    progress(3/4, '[Preprocess] DPVO')
     if not static_cam:  # use slam to get cam rotation
         if not Path(paths.slam).exists():
             length, width, height = get_video_lwh(cfg.video_path)
@@ -176,16 +180,16 @@ def load_data_dict(cfg):
     return data
 
 
-def render_incam(cfg):
+def render_incam(cfg, pred, smpl_utils):
     incam_video_path = Path(cfg.paths.incam_video)
     if incam_video_path.exists():
         Log.info(f"[Render Incam] Video already exists at {incam_video_path}")
         return
 
-    pred = torch.load(cfg.paths.hmr4d_results)
-    smplx = make_smplx("supermotion").cuda()
-    smplx2smpl = torch.load("hmr4d/utils/body_model/smplx2smpl_sparse.pt").cuda()
-    faces_smpl = make_smplx("smpl").faces
+    # pred = torch.load(cfg.paths.hmr4d_results)
+    smplx = smpl_utils['smplx']
+    smplx2smpl = smpl_utils['smplx2smpl']
+    faces_smpl = smpl_utils['faces_smpl']
 
     # smpl
     smplx_out = smplx(**to_cuda(pred["smpl_params_incam"]))
@@ -218,18 +222,18 @@ def render_incam(cfg):
     reader.close()
 
 
-def render_global(cfg):
+def render_global(cfg, pred, smpl_utils):
     global_video_path = Path(cfg.paths.global_video)
     if global_video_path.exists():
         Log.info(f"[Render Global] Video already exists at {global_video_path}")
         return
 
     debug_cam = False
-    pred = torch.load(cfg.paths.hmr4d_results)
-    smplx = make_smplx("supermotion").cuda()
-    smplx2smpl = torch.load("hmr4d/utils/body_model/smplx2smpl_sparse.pt").cuda()
-    faces_smpl = make_smplx("smpl").faces
-    J_regressor = torch.load("hmr4d/utils/body_model/smpl_neutral_J_regressor.pt").cuda()
+    # pred = torch.load(cfg.paths.hmr4d_results)
+    smplx = smpl_utils['smplx']
+    smplx2smpl = smpl_utils['smplx2smpl']
+    faces_smpl = smpl_utils['faces_smpl']
+    J_regressor = smpl_utils['J_regressor']
 
     # smpl
     smplx_out = smplx(**to_cuda(pred["smpl_params_global"]))

app/entry.py

@@ -1,3 +1,4 @@
+import os
 import gradio as gr
 
 from app.gui import get_inputs_components, get_outputs_components, get_examples, get_desc
@@ -5,6 +6,8 @@ from app.handler import handler
 
 
 def entry():
+    os.system('touch {REPO_ROOT}/server_up')
+
     demo = gr.Interface(
             fn          = handler,
             inputs      = get_inputs_components(),

app/gui.py

@@ -67,4 +67,6 @@ def get_desc():
         [Ruizhen Hu](https://csse.szu.edu.cn/staff/ruizhenhu/),
         [Xiaowei Zhou](https://xzhou.me/)  
         > SIGGRAPH Asia 2024
+        
+        **Tips: since the GPU quota is limited, we suggest to use short videos (< 5s) for demo.** For longer videos or moving camera videos, please refer to the [colab demo](https://colab.research.google.com/drive/1N9WSchizHv2bfQqkE9Wuiegw_OT7mtGj?usp=sharing).
         '''

app/handler.py

@@ -44,35 +44,44 @@ def prepare_cfg(is_static:bool, video_path:str, demo_id:str):
 def run_demo(cfg, progress, GPU_quota):
     ''' Allow user to adjust GPU quota. '''
 
+    smpl_utils = {
+            'smplx'       : make_smplx("supermotion"),
+            'J_regressor' : torch.load("hmr4d/utils/body_model/smpl_neutral_J_regressor.pt"),
+            'smplx2smpl'  : torch.load("hmr4d/utils/body_model/smplx2smpl_sparse.pt"),
+            'faces_smpl'  : make_smplx("smpl").faces,
+        }
+
     @spaces.GPU(duration=int(GPU_quota))
     def run_GPU_task():
-        paths = cfg.paths
         Log.info(f"[GPU]: {torch.cuda.get_device_name()}")
         Log.info(f'[GPU]: {torch.cuda.get_device_properties("cuda")}')
 
         # ===== Preprocess and save to disk ===== #
-        run_preprocess(cfg)
+        run_preprocess(cfg, progress)
         data = load_data_dict(cfg)
 
         # ===== HMR4D ===== #
-        if not Path(paths.hmr4d_results).exists():
-            Log.info("[HMR4D] Predicting")
-            model: DemoPL = hydra.utils.instantiate(cfg.model, _recursive_=False)
-            model.load_pretrained_model(cfg.ckpt_path)
-            model = model.eval().cuda()
-            tic = Log.sync_time()
-            pred = model.predict(data, static_cam=cfg.static_cam)
-            pred = detach_to_cpu(pred)
-            data_time = data["length"] / 30
-            Log.info(f"[HMR4D] Elapsed: {Log.sync_time() - tic:.2f}s for data-length={data_time:.1f}s")
-            torch.save(pred, paths.hmr4d_results)
+        Log.info("[HMR4D] Predicting")
+        progress(0, '[GVHMR] Initializing pipeline...')
+        model: DemoPL = hydra.utils.instantiate(cfg.model, _recursive_=False)
+        model.load_pretrained_model(cfg.ckpt_path)
+        model = model.eval().cuda()
+        tic = Log.sync_time()
+        progress(1/3, '[GVHMR] Predicting...')
+        pred = model.predict(data, static_cam=cfg.static_cam)
+        pred = detach_to_cpu(pred)
+        data_time = data["length"] / 30
+        Log.info(f"[HMR4D] Elapsed: {Log.sync_time() - tic:.2f}s for data-length={data_time:.1f}s")
+
+        progress(2/3, '[GVHMR] Rendering...')
 
         # ===== Render ===== #
-        render_incam(cfg)
-        render_global(cfg)
-        if not Path(paths.incam_global_horiz_video).exists():
-            Log.info("[Merge Videos]")
-            merge_videos_horizontal([paths.incam_video, paths.global_video], paths.incam_global_horiz_video)
+        smpl_utils['smplx'] = smpl_utils['smplx'].cuda()
+        smpl_utils['J_regressor'] = smpl_utils['J_regressor'].cuda()
+        smpl_utils['smplx2smpl'] = smpl_utils['smplx2smpl'].cuda()
+        render_incam(cfg, pred, smpl_utils)
+        render_global(cfg, pred, smpl_utils)
+        return
 
     run_GPU_task()
     return