Add skeleton template selector (articulation/mixamo/vroid)

demo.py

@@ -199,6 +199,8 @@ MODEL_CKPTS = [
     "experiments/articulation_xl_quantization_256_token_4/grpo_1400.ckpt",
 ]
 
+SKELETON_CLASSES = ["articulation", "mixamo", "vroid"]
+
 HF_PATHS = [
     "None",
 ]
@@ -328,6 +330,7 @@ def run_rig(
     output_paths: List[Path],
     model_ckpt: str,
     hf_path: Optional[str],
+    skeleton_cls: str = "articulation",
 ):
     assert len(filepaths) == len(output_paths)
     ensure_bpy_server_started()
@@ -336,7 +339,7 @@ def run_rig(
     datapath = {
         "data_name": None,
         "loader": "bpy_server",
-        "filepaths": {"articulation": [str(p) for p in filepaths]},
+        "filepaths": {skeleton_cls: [str(p) for p in filepaths]},
     }
 
     dataset_config = DatasetConfig.parse(
@@ -355,7 +358,7 @@ def run_rig(
         process_fn=model._process_fn,
     )
 
-    dataloader = module.predict_dataloader()["articulation"]
+    dataloader = module.predict_dataloader()[skeleton_cls]
 
     results_out = []
     infer_device = model.device if model is not None else "cuda"
@@ -375,7 +378,7 @@ def run_rig(
                 # Hung bpy worker: restart it and retry this file once.
                 print("[run_rig] load timed out — restarting bpy worker and retrying once", flush=True)
                 ensure_bpy_server_started(force=True)
-                dataloader_iter = iter(module.predict_dataloader()["articulation"])
+                dataloader_iter = iter(module.predict_dataloader()[skeleton_cls])
                 for _ in range(i):
                     next(dataloader_iter, None)
                 batch = next(dataloader_iter)
@@ -525,6 +528,7 @@ def run_cli(args):
         outputs,
         args.model_ckpt,
         args.hf_path,
+        args.skeleton,
     )
 
 
@@ -546,6 +550,7 @@ def _gpu_duration(
     use_postprocess,
     model_ckpt,
     hf_path,
+    skeleton_cls="articulation",
 ):
     # Cold workers spend ~30–60 s importing bpy + loading the model before
     # any GPU work. Give every request a generous 240 s floor.
@@ -566,6 +571,7 @@ def run_gradio(
     use_postprocess,
     model_ckpt,
     hf_path,
+    skeleton_cls="articulation",
 ):
     try:
         if not files:
@@ -592,6 +598,7 @@ def run_gradio(
             outputs,
             model_ckpt,
             hf_path,
+            skeleton_cls,
         )
 
         # Collect successfully exported files
@@ -703,6 +710,13 @@ def build_gradio_app():
                         info="Beam-search width. Larger = higher quality but slower; 1 disables beam search.",
                     )
 
+                    skeleton_cls = gr.Dropdown(
+                        choices=SKELETON_CLASSES,
+                        value="articulation",
+                        label="Skeleton template",
+                        info="'mixamo' emits Mixamo bone names (mixamorig:*) so Mixamo animation packs apply directly; 'vroid' emits VRoid/VRM-style humanoid names; 'articulation' is the generic bone_N skeleton.",
+                    )
+
                     gr.Markdown("**Pipeline toggles**")
                     use_skeleton = gr.Checkbox(
                         False,
@@ -751,6 +765,7 @@ def build_gradio_app():
                 use_postprocess,
                 model_ckpt,
                 hf_path,
+                skeleton_cls,
             ],
             outputs=[log, output],
         )
@@ -790,6 +805,8 @@ if __name__ == "__main__":
 
     parser.add_argument("--model_ckpt", default=MODEL_CKPTS[0] if MODEL_CKPTS else "")
     parser.add_argument("--hf_path", default=None)
+    parser.add_argument("--skeleton", default="articulation", choices=SKELETON_CLASSES,
+                        help="Skeleton template: mixamo gives mixamorig:* bone names")
 
     parser.add_argument("--gradio", action="store_true")