Deploy Kimodo ZeroGPU motion API (part 2)

.gitattributes

@@ -101,3 +101,4 @@ kimodo/assets/skeletons/g1skel34/meshes/g1/waist_constraint_L.STL filter=lfs dif
 kimodo/assets/skeletons/g1skel34/meshes/g1/waist_constraint_R.STL filter=lfs diff=lfs merge=lfs -text
 kimodo/assets/skeletons/g1skel34/meshes/g1/waist_support_link.STL filter=lfs diff=lfs merge=lfs -text
 kimodo/assets/skeletons/g1skel34/meshes/g1/waist_yaw_link.STL filter=lfs diff=lfs merge=lfs -text
+kimodo/assets/skeletons/g1skel34/meshes/g1/waist_yaw_link_rev_1_0.STL filter=lfs diff=lfs merge=lfs -text

kimodo/assets.py

@@ -0,0 +1,19 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from pathlib import Path
+
+PACKAGE_ROOT = Path(__file__).resolve().parent
+ASSETS_ROOT = PACKAGE_ROOT / "assets"
+DEMO_ASSETS_ROOT = ASSETS_ROOT / "demo"
+DEMO_EXAMPLES_ROOT = DEMO_ASSETS_ROOT / "examples"
+SKELETONS_ROOT = ASSETS_ROOT / "skeletons"
+SOMA_ASSETS_ROOT = ASSETS_ROOT / "SOMA"
+
+
+def skeleton_asset_path(*parts: str) -> Path:
+    return SKELETONS_ROOT.joinpath(*parts)
+
+
+def demo_asset_path(*parts: str) -> Path:
+    return DEMO_ASSETS_ROOT.joinpath(*parts)

kimodo/assets/skeletons/g1skel34/meshes/g1/waist_yaw_link_rev_1_0.STL

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec6db442b11f25eed898b5add07940c85d804f300de24dcbd264ccd8be7d554c
+size 619984

kimodo/assets/skeletons/g1skel34/xml/g1.xml

@@ -0,0 +1,413 @@
+<mujoco model="g1">
+    <compiler angle="radian" meshdir="../meshes/g1"/>
+
+    <default>
+        <default class="g1">
+            <geom contype="0" conaffinity="0"/>
+
+            <joint frictionloss="0.1" solimplimit="0.97 0.995 0.001"/>
+
+            <default class="hip">
+                <default class="hip_pitch">
+                    <joint axis="0 1 0" range="-2.5307 2.8798" actuatorfrcrange="-88 88" armature="0.01017752004"/>
+                </default>
+                <default class="hip_roll">
+                    <joint axis="1 0 0" actuatorfrcrange="-139 139" armature="0.025101925"/>
+                </default>
+                <default class="hip_yaw">
+                    <joint axis="0 0 1" range="-2.7576 2.7576" actuatorfrcrange="-88 88" armature="0.01017752004"/>
+                </default>
+            </default>
+            <default class="knee">
+                <joint axis="0 1 0" range="-0.087267 2.8798" actuatorfrcrange="-139 139" armature="0.025101925"/>
+            </default>
+            <default class="ankle">
+                <default class="ankle_pitch">
+                    <joint axis="0 1 0" range="-0.87267 0.5236" actuatorfrcrange="-50 50" armature="0.00721945"/>
+                </default>
+                <default class="ankle_roll">
+                    <joint axis="1 0 0" range="-0.2618 0.2618" actuatorfrcrange="-50 50" armature="0.00721945"/>
+                </default>
+            </default>
+            <default class="waist_yaw">
+                <joint axis="0 0 1" range="-2.618 2.618" actuatorfrcrange="-88 88" armature="0.01017752004"/>
+            </default>
+            <default class="waist_pitch">
+                <joint axis="0 1 0" range="-0.52 0.52" actuatorfrcrange="-50 50" armature="0.00721945"/>
+            </default>
+            <default class="waist_roll">
+                <joint axis="1 0 0" range="-0.52 0.52" actuatorfrcrange="-50 50" armature="0.00721945"/>
+            </default>
+            <default class="shoulder">
+                <default class="shoulder_pitch">
+                    <joint axis="0 1 0" range="-3.0892 2.6704" actuatorfrcrange="-25 25" armature="0.003609725"/>
+                </default>
+                <default class="shoulder_roll">
+                    <joint axis="1 0 0" actuatorfrcrange="-25 25" armature="0.003609725"/>
+                </default>
+                <default class="shoulder_yaw">
+                    <joint axis="0 0 1" range="-2.618 2.618" actuatorfrcrange="-25 25" armature="0.003609725"/>
+                </default>
+            </default>
+            <default class="elbow">
+                <joint axis="0 1 0" range="-1.0472 2.0944" actuatorfrcrange="-25 25" armature="0.003609725"/>
+            </default>
+            <default class="wrist">
+                <default class="wrist_roll">
+                    <joint axis="1 0 0" range="-1.97222 1.97222" actuatorfrcrange="-25 25" armature="0.003609725"/>
+                </default>
+                <default class="wrist_pitch">
+                    <joint axis="0 1 0" range="-1.61443 1.61443" actuatorfrcrange="-5 5" armature="0.00425"/>
+                </default>
+                <default class="wrist_yaw">
+                    <joint axis="0 0 1" range="-1.61443 1.61443" actuatorfrcrange="-5 5" armature="0.00425"/>
+                </default>
+            </default>
+
+            <default class="visual">
+                <geom group="2" type="mesh" density="0" material="silver"/>
+            </default>
+            <default class="collision">
+                <geom group="3" rgba=".2 .6 .2 .3" contype="1" conaffinity="1"/>
+                <default class="foot">
+                    <geom size="0.085 0.03 0.005"/>
+                </default>
+            </default>
+            <site group="5" rgba="1 0 0 1"/>
+        </default>
+    </default>
+
+    <asset>
+        <material name="silver" rgba="0.7 0.7 0.7 1"/>
+        <material name="black" rgba="0.2 0.2 0.2 1"/>
+
+        <mesh name="pelvis" file="pelvis.STL"/>
+        <mesh name="pelvis_contour_link" file="pelvis_contour_link.STL"/>
+        <mesh name="left_hip_pitch_link" file="left_hip_pitch_link.STL"/>
+        <mesh name="left_hip_roll_link" file="left_hip_roll_link.STL"/>
+        <mesh name="left_hip_yaw_link" file="left_hip_yaw_link.STL"/>
+        <mesh name="left_knee_link" file="left_knee_link.STL"/>
+        <mesh name="left_ankle_pitch_link"
+              file="left_ankle_pitch_link.STL"/>
+        <mesh name="left_ankle_roll_link" file="left_ankle_roll_link.STL"/>
+        <mesh name="right_hip_pitch_link" file="right_hip_pitch_link.STL"/>
+        <mesh name="right_hip_roll_link" file="right_hip_roll_link.STL"/>
+        <mesh name="right_hip_yaw_link" file="right_hip_yaw_link.STL"/>
+        <mesh name="right_knee_link" file="right_knee_link.STL"/>
+        <mesh name="right_ankle_pitch_link"
+              file="right_ankle_pitch_link.STL"/>
+        <mesh name="right_ankle_roll_link"
+              file="right_ankle_roll_link.STL"/>
+        <mesh name="waist_yaw_link" file="waist_yaw_link_rev_1_0.STL"/>
+        <mesh name="waist_roll_link" file="waist_roll_link_rev_1_0.STL"/>
+        <mesh name="torso_link" file="torso_link_rev_1_0.STL"/>
+        <mesh name="logo_link" file="logo_link.STL"/>
+        <mesh name="head_link" file="head_link.STL"/>
+        <mesh name="left_shoulder_pitch_link"
+              file="left_shoulder_pitch_link.STL"/>
+        <mesh name="left_shoulder_roll_link"
+              file="left_shoulder_roll_link.STL"/>
+        <mesh name="left_shoulder_yaw_link"
+              file="left_shoulder_yaw_link.STL"/>
+        <mesh name="left_elbow_link" file="left_elbow_link.STL"/>
+        <mesh name="left_wrist_roll_link" file="left_wrist_roll_link.STL"/>
+        <mesh name="left_wrist_pitch_link"
+              file="left_wrist_pitch_link.STL"/>
+        <mesh name="left_wrist_yaw_link" file="left_wrist_yaw_link.STL"/>
+        <mesh name="left_rubber_hand" file="left_rubber_hand.STL"/>
+        <mesh name="right_shoulder_pitch_link"
+              file="right_shoulder_pitch_link.STL"/>
+        <mesh name="right_shoulder_roll_link"
+              file="right_shoulder_roll_link.STL"/>
+        <mesh name="right_shoulder_yaw_link"
+              file="right_shoulder_yaw_link.STL"/>
+        <mesh name="right_elbow_link" file="right_elbow_link.STL"/>
+        <mesh name="right_wrist_roll_link"
+              file="right_wrist_roll_link.STL"/>
+        <mesh name="right_wrist_pitch_link"
+              file="right_wrist_pitch_link.STL"/>
+        <mesh name="right_wrist_yaw_link" file="right_wrist_yaw_link.STL"/>
+        <mesh name="right_rubber_hand" file="right_rubber_hand.STL"/>
+    </asset>
+
+    <worldbody>
+        <body name="pelvis" pos="0 0 0.793" childclass="g1">
+            <inertial pos="0 0 -0.07605" quat="1 0 -0.000399148 0" mass="3.813" diaginertia="0.010549 0.0093089 0.0079184"/>
+            <freejoint name="floating_base_joint"/>
+            <geom class="visual" material="black" mesh="pelvis"/>
+            <geom class="visual" mesh="pelvis_contour_link"/>
+            <geom mesh="pelvis_contour_link" class="visual"/>
+            <geom name="pelvis_collision" class="collision" type="sphere" size="0.07" pos="0 0 -0.08"/>
+            <site name="imu_in_pelvis" size="0.01" pos="0.04525 0 -0.08339"/>
+            <site name="pelvis" size="0.01" pos="0 0 0"/>
+            <body name="left_hip_pitch_link" pos="0 0.064452 -0.1027">
+                <inertial pos="0.002741 0.047791 -0.02606" quat="0.954862 0.293964 0.0302556 0.030122" mass="1.35"
+                          diaginertia="0.00181517 0.00153422 0.00116212"/>
+                <joint name="left_hip_pitch_joint" class="hip_pitch"/>
+                <geom class="visual" material="black" mesh="left_hip_pitch_link"/>
+                <geom material="black" mesh="left_hip_pitch_link" class="visual"/>
+                <body name="left_hip_roll_link" pos="0 0.052 -0.030465" quat="0.996179 0 -0.0873386 0">
+                    <inertial pos="0.029812 -0.001045 -0.087934" quat="0.977808 -1.97119e-05 0.205576 -0.0403793" mass="1.52"
+                              diaginertia="0.00254986 0.00241169 0.00148755"/>
+                    <joint name="left_hip_roll_joint" class="hip_roll" range="-0.5236 2.9671"/>
+                    <geom class="visual" mesh="left_hip_roll_link"/>
+                    <geom mesh="left_hip_roll_link" class="visual"/>
+                    <geom name="left_thigh" class="collision" type="capsule" size="0.05" fromto="0.02 0 0 0.02 0 -0.2"/>
+                    <body name="left_hip_yaw_link" pos="0.025001 0 -0.12412">
+                        <inertial pos="-0.057709 -0.010981 -0.15078" quat="0.600598 0.15832 0.223482 0.751181" mass="1.702"
+                                  diaginertia="0.00776166 0.00717575 0.00160139"/>
+                        <joint name="left_hip_yaw_joint" class="hip_yaw"/>
+                        <geom class="visual" mesh="left_hip_yaw_link"/>
+                        <body name="left_knee_link" pos="-0.078273 0.0021489 -0.17734" quat="0.996179 0 0.0873386 0">
+                            <inertial pos="0.005457 0.003964 -0.12074" quat="0.923418 -0.0327699 0.0158246 0.382067" mass="1.932"
+                                      diaginertia="0.0113804 0.0112778 0.00146458"/>
+                            <joint name="left_knee_joint" class="knee"/>
+                            <geom class="visual" mesh="left_knee_link"/>
+                            <geom name="left_shin" class="collision" type="capsule" size="0.04" fromto="0.02 0 0 0.02 0 -0.25"/>
+                            <body name="left_ankle_pitch_link" pos="0 -9.4445e-05 -0.30001">
+                                <inertial pos="-0.007269 0 0.011137" quat="0.603053 0.369225 0.369225 0.603053" mass="0.074"
+                                          diaginertia="1.89e-05 1.40805e-05 6.9195e-06"/>
+                                <joint name="left_ankle_pitch_joint" class="ankle_pitch"/>
+                                <geom class="visual" mesh="left_ankle_pitch_link"/>
+                                <body name="left_ankle_roll_link" pos="0 0 -0.017558">
+                                    <site name="left_foot" rgba="1 0 0 1"/>
+                                    <inertial pos="0.026505 0 -0.016425" quat="-0.000481092 0.728482 -0.000618967 0.685065" mass="0.608"
+                                              diaginertia="0.00167218 0.0016161 0.000217621"/>
+                                    <joint name="left_ankle_roll_joint" class="ankle_roll"/>
+                                    <geom class="visual" material="black" mesh="left_ankle_roll_link"/>
+                                    <geom name="left_foot1_collision" type="capsule" size="0.01" class="collision" fromto="0.1 -0.026 -0.025 0.05 -0.027 -0.025"/>
+                                    <geom name="left_foot2_collision" type="capsule" size="0.01" class="collision" fromto="-0.044 -0.018 -0.025 0.123 -0.018 -0.025"/>
+                                    <geom name="left_foot3_collision" type="capsule" size="0.01" class="collision" fromto="-0.052 -0.01 -0.025 0.13 -0.01 -0.025"/>
+                                    <geom name="left_foot4_collision" type="capsule" size="0.01" class="collision" fromto="-0.054 0 -0.025 0.132 0 -0.025"/>
+                                    <geom name="left_foot5_collision" type="capsule" size="0.01" class="collision" fromto="-0.052 0.01 -0.025 0.13 0.01 -0.025"/>
+                                    <geom name="left_foot6_collision" type="capsule" size="0.01" class="collision" fromto="-0.044 0.018 -0.025 0.123 0.018 -0.025"/>
+                                    <geom name="left_foot7_collision" type="capsule" size="0.01" class="collision" fromto="0.1 0.026 -0.025 0.05 0.026 -0.025"/>
+                                </body>
+                            </body>
+                        </body>
+                    </body>
+                </body>
+            </body>
+            <body name="right_hip_pitch_link" pos="0 -0.064452 -0.1027">
+                <inertial pos="0.002741 -0.047791 -0.02606" quat="0.954862 -0.293964 0.0302556 -0.030122" mass="1.35"
+                          diaginertia="0.00181517 0.00153422 0.00116212"/>
+                <joint name="right_hip_pitch_joint" class="hip_pitch"/>
+                <geom class="visual" material="black" mesh="right_hip_pitch_link"/>
+                <body name="right_hip_roll_link" pos="0 -0.052 -0.030465" quat="0.996179 0 -0.0873386 0">
+                    <inertial pos="0.029812 0.001045 -0.087934" quat="0.977808 1.97119e-05 0.205576 0.0403793" mass="1.52"
+                              diaginertia="0.00254986 0.00241169 0.00148755"/>
+                    <joint name="right_hip_roll_joint" class="hip_roll" range="-2.9671 0.5236"/>
+                    <geom class="visual" mesh="right_hip_roll_link"/>
+                    <geom name="right_thigh" class="collision" type="capsule" size="0.05" fromto="0.02 0 0 0.02 0 -0.2"/>
+                    <body name="right_hip_yaw_link" pos="0.025001 0 -0.12412">
+                        <inertial pos="-0.057709 0.010981 -0.15078" quat="0.751181 0.223482 0.15832 0.600598" mass="1.702"
+                                  diaginertia="0.00776166 0.00717575 0.00160139"/>
+                        <joint name="right_hip_yaw_joint" class="hip_yaw"/>
+                        <geom class="visual" mesh="right_hip_yaw_link"/>
+                        <body name="right_knee_link" pos="-0.078273 -0.0021489 -0.17734" quat="0.996179 0 0.0873386 0">
+                            <inertial pos="0.005457 -0.003964 -0.12074" quat="0.923439 0.0345276 0.0116333 -0.382012" mass="1.932"
+                                      diaginertia="0.011374 0.0112843 0.00146452"/>
+                            <joint name="right_knee_joint" class="knee"/>
+                            <geom class="visual" mesh="right_knee_link"/>
+                            <geom name="right_shin" class="collision" type="capsule" size="0.04" fromto="0.02 0 0 0.02 0 -0.25"/>
+                            <body name="right_ankle_pitch_link" pos="0 9.4445e-05 -0.30001">
+                                <inertial pos="-0.007269 0 0.011137" quat="0.603053 0.369225 0.369225 0.603053" mass="0.074"
+                                          diaginertia="1.89e-05 1.40805e-05 6.9195e-06"/>
+                                <joint name="right_ankle_pitch_joint" class="ankle_pitch"/>
+                                <geom class="visual" mesh="right_ankle_pitch_link"/>
+                                <body name="right_ankle_roll_link" pos="0 0 -0.017558">
+                                    <site name="right_foot" pos="0 0 0"/>
+                                    <inertial pos="0.026505 0 -0.016425" quat="0.000481092 0.728482 0.000618967 0.685065" mass="0.608"
+                                              diaginertia="0.00167218 0.0016161 0.000217621"/>
+                                    <joint name="right_ankle_roll_joint" class="ankle_roll"/>
+                                    <geom class="visual" material="black" mesh="right_ankle_roll_link"/>
+                                    <geom name="right_foot1_collision" type="capsule" size="0.01" class="collision" fromto="0.1 -0.026 -0.025 0.05 -0.026 -0.025"/>
+                                    <geom name="right_foot2_collision" type="capsule" size="0.008" class="collision" fromto="-0.044 -0.018 -0.025 0.123 -0.018 -0.025"/>
+                                    <geom name="right_foot3_collision" type="capsule" size="0.01" class="collision" fromto="-0.052 -0.01 -0.025 0.13 -0.01 -0.025"/>
+                                    <geom name="right_foot4_collision" type="capsule" size="0.01" class="collision" fromto="-0.054 0 -0.025 0.132 0 -0.025"/>
+                                    <geom name="right_foot5_collision" type="capsule" size="0.01" class="collision" fromto="-0.052 0.01 -0.025 0.13 0.01 -0.025"/>
+                                    <geom name="right_foot6_collision" type="capsule" size="0.008" class="collision" fromto="-0.044 0.018 -0.025 0.123 0.018 -0.025"/>
+                                    <geom name="right_foot7_collision" type="capsule" size="0.01" class="collision" fromto="0.1 0.026 -0.025 0.05 0.026 -0.025"/>
+                                </body>
+                            </body>
+                        </body>
+                    </body>
+                </body>
+            </body>
+            <body name="waist_yaw_link">
+                <inertial pos="0.003494 0.000233 0.018034" quat="0.289697 0.591001 -0.337795 0.672821" mass="0.214"
+                          diaginertia="0.000163531 0.000107714 0.000102205"/>
+                <joint name="waist_yaw_joint" class="waist_yaw"/>
+                <geom class="visual" mesh="waist_yaw_link"/>
+                <body name="waist_roll_link" pos="-0.0039635 0 0.044">
+                    <inertial pos="0 2.3e-05 0" quat="0.5 0.5 -0.5 0.5" mass="0.086" diaginertia="8.245e-06 7.079e-06 6.339e-06"/>
+                    <joint name="waist_roll_joint" class="waist_roll"/>
+                    <geom class="visual" mesh="waist_roll_link"/>
+                    <body name="torso_link">
+                        <inertial pos="0.00203158 0.000339683 0.184568" quat="0.999803 -6.03319e-05 0.0198256 0.00131986"
+                                  mass="7.818" diaginertia="0.121847 0.109825 0.0273735"/>
+                        <joint name="waist_pitch_joint" class="waist_pitch"/>
+                        <geom class="visual" mesh="torso_link"/>
+                        <geom pos="0.0039635 0 -0.044" quat="1 0 0 0" class="visual" material="black" mesh="logo_link"/>
+                        <geom pos="0.0039635 0 -0.044" class="visual" material="black" mesh="head_link"/>
+
+                        <geom name="torso_collision1" class="collision" type="capsule" size="0.073" fromto="0.005 -0.032 .22 0.005 0.032 .22"/>
+                        <geom name="torso_collision2" class="collision" type="capsule" size="0.07" fromto="0.005 -0.028 .13 0.005 0.028 .13"/>
+                        <geom name="torso_collision3" class="collision" type="capsule" size="0.065" fromto="0.005 -0.02 .06 0.005 0.02 .06"/>
+                        <geom name="head_collision" class="collision" type="capsule" size="0.068" fromto="0.01 0 .41 0.01 0 .42"/>
+
+                        <site name="imu_in_torso" size="0.01" pos="-0.03959 -0.00224 0.14792"/>
+                        <site name="mid360" size="0.01" pos="0.0002835 0.00003 0.41618" quat="0.00000094 -0.99979404 0.00004632 0.02029493"/>
+                        <body name="left_shoulder_pitch_link" pos="0.0039563 0.10022 0.24778"
+                              quat="0.990264 0.139201 1.38722e-05 -9.86868e-05">
+                            <inertial pos="0 0.035892 -0.011628" quat="0.654152 0.0130458 -0.326267 0.68225" mass="0.718"
+                                      diaginertia="0.000465864 0.000432842 0.000406394"/>
+                            <joint name="left_shoulder_pitch_joint" class="shoulder_pitch"/>
+                            <geom class="visual" mesh="left_shoulder_pitch_link"/>
+                            <geom size="0.03 0.025" pos="0 0.04 -0.01" quat="0.707107 0 0.707107 0" type="cylinder" class="visual"/>
+                            <body name="left_shoulder_roll_link" pos="0 0.038 -0.013831" quat="0.990268 -0.139172 0 0">
+                                <inertial pos="-0.000227 0.00727 -0.063243" quat="0.701256 -0.0196223 -0.00710317 0.712604" mass="0.643"
+                                          diaginertia="0.000691311 0.000618011 0.000388977"/>
+                                <joint name="left_shoulder_roll_joint" range="-1.5882 2.2515" class="shoulder_roll"/>
+                                <geom class="visual" mesh="left_shoulder_roll_link"/>
+                                <geom size="0.03 0.015" pos="-0.004 0.006 -0.053" type="cylinder" class="visual"/>
+                                <body name="left_shoulder_yaw_link" pos="0 0.00624 -0.1032">
+                                    <inertial pos="0.010773 -0.002949 -0.072009" quat="0.716879 -0.0964829 -0.0679942 0.687134"
+                                              mass="0.734" diaginertia="0.00106187 0.00103217 0.000400661"/>
+                                    <joint name="left_shoulder_yaw_joint" class="shoulder_yaw"/>
+                                    <geom class="visual" mesh="left_shoulder_yaw_link"/>
+                                    <geom name="left_shoulder_yaw_collision" class="collision" type="capsule" size="0.035" fromto="0 0 -0.08 0 0 0.05"/>
+                                    <body name="left_elbow_link" pos="0.015783 0 -0.080518">
+                                        <inertial pos="0.064956 0.004454 -0.010062" quat="0.541765 0.636132 0.388821 0.388129" mass="0.6"
+                                                  diaginertia="0.000443035 0.000421612 0.000259353"/>
+                                        <joint name="left_elbow_joint" class="elbow"/>
+                                        <geom class="visual" mesh="left_elbow_link"/>
+                                        <geom name="left_elbow_yaw_collision" class="collision" type="capsule" size="0.035" fromto="-0.01 0 -0.01 0.12 0 -0.01"/>
+                                        <body name="left_wrist_roll_link" pos="0.1 0.00188791 -0.01">
+                                            <inertial pos="0.0171394 0.000537591 4.8864e-07" quat="0.575338 0.411667 -0.574906 0.411094"
+                                                      mass="0.085445" diaginertia="5.48211e-05 4.96646e-05 3.57798e-05"/>
+                                            <joint name="left_wrist_roll_joint" class="wrist_roll"/>
+                                            <geom class="visual" mesh="left_wrist_roll_link"/>
+                                            <body name="left_wrist_pitch_link" pos="0.038 0 0">
+                                                <inertial pos="0.0229999 -0.00111685 -0.00111658" quat="0.249998 0.661363 0.293036 0.643608"
+                                                          mass="0.48405" diaginertia="0.000430353 0.000429873 0.000164648"/>
+                                                <joint name="left_wrist_pitch_joint" class="wrist_pitch"/>
+                                                <geom class="visual" mesh="left_wrist_pitch_link"/>
+                                                <body name="left_wrist_yaw_link" pos="0.046 0 0">
+                                                    <inertial pos="0.0708244 0.000191745 0.00161742" quat="0.510571 0.526295 0.468078 0.493188"
+                                                              mass="0.254576" diaginertia="0.000646113 0.000559993 0.000147566"/>
+                                                    <joint name="left_wrist_yaw_joint" class="wrist_yaw"/>
+                                                    <geom class="visual" mesh="left_wrist_yaw_link"/>
+                                                    <geom pos="0.0415 0.003 0" quat="1 0 0 0" class="visual" mesh="left_rubber_hand"/>
+                                                    <site name="left_palm" pos="0.08 0 0" size="0.01"/>
+                                                    <geom name="left_hand_collision" class="collision" type="capsule" size="0.05" fromto="0.05 0 0 0.1 0 0"/>
+                                                </body>
+                                            </body>
+                                        </body>
+                                    </body>
+                                </body>
+                            </body>
+                        </body>
+                        <body name="right_shoulder_pitch_link" pos="0.0039563 -0.10021 0.24778"
+                              quat="0.990264 -0.139201 1.38722e-05 9.86868e-05">
+                            <inertial pos="0 -0.035892 -0.011628" quat="0.68225 -0.326267 0.0130458 0.654152" mass="0.718"
+                                      diaginertia="0.000465864 0.000432842 0.000406394"/>
+                            <joint name="right_shoulder_pitch_joint" class="shoulder_pitch"/>
+                            <geom class="visual" mesh="right_shoulder_pitch_link"/>
+                            <geom size="0.03 0.025" pos="0 -0.04 -0.01" quat="0.707107 0 0.707107 0" type="cylinder" class="visual"/>
+                            <body name="right_shoulder_roll_link" pos="0 -0.038 -0.013831" quat="0.990268 0.139172 0 0">
+                                <inertial pos="-0.000227 -0.00727 -0.063243" quat="0.712604 -0.00710317 -0.0196223 0.701256"
+                                          mass="0.643" diaginertia="0.000691311 0.000618011 0.000388977"/>
+                                <joint name="right_shoulder_roll_joint" range="-2.2515 1.5882" class="shoulder_roll"/>
+                                <geom class="visual" mesh="right_shoulder_roll_link"/>
+                                <geom size="0.03 0.015" pos="-0.004 -0.006 -0.053" type="cylinder" class="visual"/>
+                                <body name="right_shoulder_yaw_link" pos="0 -0.00624 -0.1032">
+                                    <inertial pos="0.010773 0.002949 -0.072009" quat="0.687134 -0.0679942 -0.0964829 0.716879"
+                                              mass="0.734" diaginertia="0.00106187 0.00103217 0.000400661"/>
+                                    <joint name="right_shoulder_yaw_joint" class="shoulder_yaw"/>
+                                    <geom class="visual" mesh="right_shoulder_yaw_link"/>
+                                    <geom name="right_shoulder_yaw_collision" class="collision" type="capsule" size="0.035" fromto="0 0 -0.08 0 0 0.05"/>
+                                    <body name="right_elbow_link" pos="0.015783 0 -0.080518">
+                                        <inertial pos="0.064956 -0.004454 -0.010062" quat="0.388129 0.388821 0.636132 0.541765" mass="0.6"
+                                                  diaginertia="0.000443035 0.000421612 0.000259353"/>
+                                        <joint name="right_elbow_joint" class="elbow"/>
+                                        <geom class="visual" mesh="right_elbow_link"/>
+                                        <geom name="right_elbow_yaw_collision" class="collision" type="capsule" size="0.035" fromto="-0.01 0 -0.01 0.12 0 -0.01"/>
+                                        <body name="right_wrist_roll_link" pos="0.1 -0.00188791 -0.01">
+                                            <inertial pos="0.0171394 -0.000537591 4.8864e-07" quat="0.411667 0.575338 -0.411094 0.574906"
+                                                      mass="0.085445" diaginertia="5.48211e-05 4.96646e-05 3.57798e-05"/>
+                                            <joint name="right_wrist_roll_joint" class="wrist_roll"/>
+                                            <geom class="visual" mesh="right_wrist_roll_link"/>
+                                            <body name="right_wrist_pitch_link" pos="0.038 0 0">
+                                                <inertial pos="0.0229999 0.00111685 -0.00111658" quat="0.643608 0.293036 0.661363 0.249998"
+                                                          mass="0.48405" diaginertia="0.000430353 0.000429873 0.000164648"/>
+                                                <joint name="right_wrist_pitch_joint" class="wrist_pitch"/>
+                                                <geom class="visual" mesh="right_wrist_pitch_link"/>
+                                                <body name="right_wrist_yaw_link" pos="0.046 0 0">
+                                                    <inertial pos="0.0708244 -0.000191745 0.00161742" quat="0.493188 0.468078 0.526295 0.510571"
+                                                              mass="0.254576" diaginertia="0.000646113 0.000559993 0.000147566"/>
+                                                    <joint name="right_wrist_yaw_joint" class="wrist_yaw"/>
+                                                    <geom class="visual" mesh="right_wrist_yaw_link"/>
+                                                    <geom pos="0.0415 -0.003 0" quat="1 0 0 0" class="visual" mesh="right_rubber_hand"/>
+                                                    <site name="right_palm" pos="0.08 0 0" size="0.01"/>
+                                                    <geom name="right_hand_collision" class="collision" type="capsule" size="0.05" fromto="0.05 0 0 0.1 0 0"/>
+                                                </body>
+                                            </body>
+                                        </body>
+                                    </body>
+                                </body>
+                            </body>
+                        </body>
+                    </body>
+                </body>
+            </body>
+        </body>
+    </worldbody>
+
+    <contact>>
+        <!-- left foot - floor -->
+        <pair name="left_foot1_floor" geom1="left_foot1_collision" geom2="floor" solref="0.01 1" friction="0.8 0.8"/>
+        <pair name="left_foot2_floor" geom1="left_foot2_collision" geom2="floor" solref="0.01 1" friction="0.8 0.8"/>
+        <pair name="left_foot3_floor" geom1="left_foot3_collision" geom2="floor" solref="0.01 1" friction="0.8 0.8"/>
+        <pair name="left_foot4_floor" geom1="left_foot4_collision" geom2="floor" solref="0.01 1" friction="0.8 0.8"/>
+        <pair name="left_foot5_floor" geom1="left_foot5_collision" geom2="floor" solref="0.01 1" friction="0.8 0.8"/>
+        <pair name="left_foot6_floor" geom1="left_foot6_collision" geom2="floor" solref="0.01 1" friction="0.8 0.8"/>
+        <pair name="left_foot7_floor" geom1="left_foot7_collision" geom2="floor" solref="0.01 1" friction="0.8 0.8"/>
+        <!-- right foot - floor -->
+        <pair name="right_foot1_floor" geom1="right_foot1_collision" geom2="floor" solref="0.01 1" friction="0.8 0.8"/>
+        <pair name="right_foot2_floor" geom1="right_foot2_collision" geom2="floor" solref="0.01 1" friction="0.8 0.8"/>
+        <pair name="right_foot3_floor" geom1="right_foot3_collision" geom2="floor" solref="0.01 1" friction="0.8 0.8"/>
+        <pair name="right_foot4_floor" geom1="right_foot4_collision" geom2="floor" solref="0.01 1" friction="0.8 0.8"/>
+        <pair name="right_foot5_floor" geom1="right_foot5_collision" geom2="floor" solref="0.01 1" friction="0.8 0.8"/>
+        <pair name="right_foot6_floor" geom1="right_foot6_collision" geom2="floor" solref="0.01 1" friction="0.8 0.8"/>
+        <pair name="right_foot7_floor" geom1="right_foot7_collision" geom2="floor" solref="0.01 1" friction="0.8 0.8"/>
+    </contact>
+
+    <sensor>
+        <framequat name="base_quat" objtype="site" objname="imu_in_pelvis"/>
+        <gyro name="base_gyro" site="imu_in_pelvis"/>
+        <accelerometer name="base_accel" site="imu_in_pelvis"/>
+
+        <framequat name="mid360_quat" objtype="site" objname="mid360"/>
+        <framepos name="mid360_pos" objtype="site" objname="mid360"/>
+    </sensor>
+
+    <visual>
+        <headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0"/>
+        <rgba haze="0.15 0.25 0.35 1"/>
+        <global azimuth="120" elevation="-20"/>
+    </visual>
+
+    <asset>
+        <texture type="skybox" builtin="gradient" rgb1="0.3 0.5 0.7" rgb2="0 0 0" width="512" height="3072"/>
+        <texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3"
+                 markrgb="0.8 0.8 0.8" width="300" height="300"/>
+        <material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance="0.2"/>
+    </asset>
+
+    <worldbody>
+        <light pos="0 0 1.5" dir="0 0 -1" directional="true"/>
+        <geom name="floor" size="0 0 0.01" type="plane" material="groundplane" conaffinity="1"/>
+    </worldbody>
+
+</mujoco>

kimodo/assets/skeletons/smplx22/beta.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:abfc1cc9e819d017580739e4a65b39458ae840808f01819ad54260283575fa11
+size 1328

kimodo/assets/skeletons/smplx22/mean_hands.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eff97ce9d98bfb1495501d8f00a883c08bfdfc05079037e9832f5d64b58d0e39
+size 848

kimodo/assets/skeletons/somaskel30/soma_base_fit_mhr_params.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:61f0b96c8386a19d44823faf4399816642b7e20691ed5b398c7402921090096e
+size 1877

kimodo/assets/skeletons/somaskel77/relaxed_hands_rest_pose.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:64a3828e0d1ef1f1de8228c74eba8040c0810d898169f1892d8997a213b2b64c
+size 2900

kimodo/assets/skeletons/somaskel77/skin_standard.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:90ee2cdf50f168382a7dd7e9c88b118298aba0aca4b72022080c89c4bab0ceb2
+size 531434

kimodo/assets/skeletons/somaskel77/somaskel77_standard_tpose.bvh

@@ -0,0 +1,395 @@
+HIERARCHY
+ROOT Root
+{
+  OFFSET 0.0 0.0 0.0
+  CHANNELS 6 Xposition Yposition Zposition Zrotation Yrotation Xrotation
+  JOINT Hips
+  {
+    OFFSET 0.0 100.0 0.0
+    CHANNELS 6 Xposition Yposition Zposition Zrotation Yrotation Xrotation
+    JOINT Spine1
+    {
+      OFFSET -0.013727 5.003763 -0.053727
+      CHANNELS 3 Zrotation Yrotation Xrotation
+      JOINT Spine2
+      {
+        OFFSET -0.0 7.125301 -0.029825
+        CHANNELS 3 Zrotation Yrotation Xrotation
+        JOINT Chest
+        {
+          OFFSET -1e-06 7.550063 -0.815971
+          CHANNELS 3 Zrotation Yrotation Xrotation
+          JOINT Neck1
+          {
+            OFFSET -0.181677 26.311295 -0.553348
+            CHANNELS 3 Zrotation Yrotation Xrotation
+            JOINT Neck2
+            {
+              OFFSET -3e-06 7.709397 2.302585
+              CHANNELS 3 Zrotation Yrotation Xrotation
+              JOINT Head
+              {
+                OFFSET -5e-06 6.128916 1.953709
+                CHANNELS 3 Zrotation Yrotation Xrotation
+                JOINT HeadEnd
+                {
+                  OFFSET 0.003598 16.065403 -1.835379
+                  CHANNELS 3 Zrotation Yrotation Xrotation
+                }
+                JOINT Jaw
+                {
+                  OFFSET 0.002637 0.475592 3.094941
+                  CHANNELS 3 Zrotation Yrotation Xrotation
+                }
+                JOINT LeftEye
+                {
+                  OFFSET 3.206381 5.380205 7.586883
+                  CHANNELS 3 Zrotation Yrotation Xrotation
+                }
+                JOINT RightEye
+                {
+                  OFFSET -3.22244 5.361869 7.558234
+                  CHANNELS 3 Zrotation Yrotation Xrotation
+                }
+              }
+            }
+          }
+          JOINT LeftShoulder
+          {
+            OFFSET 1.621652 23.237164 5.113413
+            CHANNELS 3 Zrotation Yrotation Xrotation
+            JOINT LeftArm
+            {
+              OFFSET 14.919846 2e-06 -5.502326
+              CHANNELS 3 Zrotation Yrotation Xrotation
+              JOINT LeftForeArm
+              {
+                OFFSET 28.739307 0.0 -0.002588
+                CHANNELS 3 Zrotation Yrotation Xrotation
+                JOINT LeftHand
+                {
+                  OFFSET 27.093981 -1e-06 0.002609
+                  CHANNELS 3 Zrotation Yrotation Xrotation
+                  JOINT LeftHandThumb1
+                  {
+                    OFFSET 2.276482 -1.392045 3.191413
+                    CHANNELS 3 Zrotation Yrotation Xrotation
+                    JOINT LeftHandThumb2
+                    {
+                      OFFSET 4.012836 -1.828127 1.641654
+                      CHANNELS 3 Zrotation Yrotation Xrotation
+                      JOINT LeftHandThumb3
+                      {
+                        OFFSET 2.798515 0.0 -3e-06
+                        CHANNELS 3 Zrotation Yrotation Xrotation
+                        JOINT LeftHandThumbEnd
+                        {
+                          OFFSET 3.180793 -4e-06 4e-06
+                          CHANNELS 3 Zrotation Yrotation Xrotation
+                        }
+                      }
+                    }
+                  }
+                  JOINT LeftHandIndex1
+                  {
+                    OFFSET 3.247555 -0.531998 2.296169
+                    CHANNELS 3 Zrotation Yrotation Xrotation
+                    JOINT LeftHandIndex2
+                    {
+                      OFFSET 6.364578 0.01206 0.1786
+                      CHANNELS 3 Zrotation Yrotation Xrotation
+                      JOINT LeftHandIndex3
+                      {
+                        OFFSET 3.662364 0.0 0.0
+                        CHANNELS 3 Zrotation Yrotation Xrotation
+                        JOINT LeftHandIndex4
+                        {
+                          OFFSET 2.329242 4e-06 4e-06
+                          CHANNELS 3 Zrotation Yrotation Xrotation
+                          JOINT LeftHandIndexEnd
+                          {
+                            OFFSET 2.759615 -0.180537 -0.113024
+                            CHANNELS 3 Zrotation Yrotation Xrotation
+                          }
+                        }
+                      }
+                    }
+                  }
+                  JOINT LeftHandMiddle1
+                  {
+                    OFFSET 3.163495 0.240981 1.000332
+                    CHANNELS 3 Zrotation Yrotation Xrotation
+                    JOINT LeftHandMiddle2
+                    {
+                      OFFSET 6.19078 -0.259278 -1.002548
+                      CHANNELS 3 Zrotation Yrotation Xrotation
+                      JOINT LeftHandMiddle3
+                      {
+                        OFFSET 4.35652 -4e-06 -1e-06
+                        CHANNELS 3 Zrotation Yrotation Xrotation
+                        JOINT LeftHandMiddle4
+                        {
+                          OFFSET 2.996877 -8e-06 0.0
+                          CHANNELS 3 Zrotation Yrotation Xrotation
+                          JOINT LeftHandMiddleEnd
+                          {
+                            OFFSET 2.304287 -0.294569 -0.031741
+                            CHANNELS 3 Zrotation Yrotation Xrotation
+                          }
+                        }
+                      }
+                    }
+                  }
+                  JOINT LeftHandRing1
+                  {
+                    OFFSET 2.882643 -0.053652 -0.322543
+                    CHANNELS 3 Zrotation Yrotation Xrotation
+                    JOINT LeftHandRing2
+                    {
+                      OFFSET 5.854541 -0.486202 -1.373841
+                      CHANNELS 3 Zrotation Yrotation Xrotation
+                      JOINT LeftHandRing3
+                      {
+                        OFFSET 4.350578 0.0 3e-06
+                        CHANNELS 3 Zrotation Yrotation Xrotation
+                        JOINT LeftHandRing4
+                        {
+                          OFFSET 2.651321 7e-06 2e-06
+                          CHANNELS 3 Zrotation Yrotation Xrotation
+                          JOINT LeftHandRingEnd
+                          {
+                            OFFSET 1.936105 0.077687 -7.1e-05
+                            CHANNELS 3 Zrotation Yrotation Xrotation
+                          }
+                        }
+                      }
+                    }
+                  }
+                  JOINT LeftHandPinky1
+                  {
+                    OFFSET 2.8655 -0.310005 -1.600378
+                    CHANNELS 3 Zrotation Yrotation Xrotation
+                    JOINT LeftHandPinky2
+                    {
+                      OFFSET 5.087849 -1.331141 -1.77123
+                      CHANNELS 3 Zrotation Yrotation Xrotation
+                      JOINT LeftHandPinky3
+                      {
+                        OFFSET 3.070974 4e-06 0.0
+                        CHANNELS 3 Zrotation Yrotation Xrotation
+                        JOINT LeftHandPinky4
+                        {
+                          OFFSET 1.549672 0.0 1e-06
+                          CHANNELS 3 Zrotation Yrotation Xrotation
+                          JOINT LeftHandPinkyEnd
+                          {
+                            OFFSET 1.944893 -0.157802 0.057219
+                            CHANNELS 3 Zrotation Yrotation Xrotation
+                          }
+                        }
+                      }
+                    }
+                  }
+                }
+              }
+            }
+          }
+          JOINT RightShoulder
+          {
+            OFFSET -1.380118 23.180309 5.214158
+            CHANNELS 3 Zrotation Yrotation Xrotation
+            JOINT RightArm
+            {
+              OFFSET -15.037196 1.2e-05 -5.545604
+              CHANNELS 3 Zrotation Yrotation Xrotation
+              JOINT RightForeArm
+              {
+                OFFSET -28.736639 2e-06 -0.002597
+                CHANNELS 3 Zrotation Yrotation Xrotation
+                JOINT RightHand
+                {
+                  OFFSET -27.133619 -0.0 0.002613
+                  CHANNELS 3 Zrotation Yrotation Xrotation
+                  JOINT RightHandThumb1
+                  {
+                    OFFSET -2.274032 -1.383988 3.163127
+                    CHANNELS 3 Zrotation Yrotation Xrotation
+                    JOINT RightHandThumb2
+                    {
+                      OFFSET -4.011429 -1.827466 1.640914
+                      CHANNELS 3 Zrotation Yrotation Xrotation
+                      JOINT RightHandThumb3
+                      {
+                        OFFSET -2.794935 -4e-06 -3e-06
+                        CHANNELS 3 Zrotation Yrotation Xrotation
+                        JOINT RightHandThumbEnd
+                        {
+                          OFFSET -3.183852 4e-06 1e-06
+                          CHANNELS 3 Zrotation Yrotation Xrotation
+                        }
+                      }
+                    }
+                  }
+                  JOINT RightHandIndex1
+                  {
+                    OFFSET -3.253266 -0.520057 2.282866
+                    CHANNELS 3 Zrotation Yrotation Xrotation
+                    JOINT RightHandIndex2
+                    {
+                      OFFSET -6.341917 0.012471 0.178266
+                      CHANNELS 3 Zrotation Yrotation Xrotation
+                      JOINT RightHandIndex3
+                      {
+                        OFFSET -3.654871 -8e-06 -0.0
+                        CHANNELS 3 Zrotation Yrotation Xrotation
+                        JOINT RightHandIndex4
+                        {
+                          OFFSET -2.327586 0.0 1e-06
+                          CHANNELS 3 Zrotation Yrotation Xrotation
+                          JOINT RightHandIndexEnd
+                          {
+                            OFFSET -2.76179 -0.180656 -0.113078
+                            CHANNELS 3 Zrotation Yrotation Xrotation
+                          }
+                        }
+                      }
+                    }
+                  }
+                  JOINT RightHandMiddle1
+                  {
+                    OFFSET -3.168106 0.246593 1.00103
+                    CHANNELS 3 Zrotation Yrotation Xrotation
+                    JOINT RightHandMiddle2
+                    {
+                      OFFSET -6.180828 -0.258836 -1.000895
+                      CHANNELS 3 Zrotation Yrotation Xrotation
+                      JOINT RightHandMiddle3
+                      {
+                        OFFSET -4.348901 0.0 -0.0
+                        CHANNELS 3 Zrotation Yrotation Xrotation
+                        JOINT RightHandMiddle4
+                        {
+                          OFFSET -3.00024 -4e-06 -2e-06
+                          CHANNELS 3 Zrotation Yrotation Xrotation
+                          JOINT RightHandMiddleEnd
+                          {
+                            OFFSET -2.30252 -0.29437 -0.031706
+                            CHANNELS 3 Zrotation Yrotation Xrotation
+                          }
+                        }
+                      }
+                    }
+                  }
+                  JOINT RightHandRing1
+                  {
+                    OFFSET -2.88569 -0.067952 -0.308858
+                    CHANNELS 3 Zrotation Yrotation Xrotation
+                    JOINT RightHandRing2
+                    {
+                      OFFSET -5.854198 -0.48613 -1.373731
+                      CHANNELS 3 Zrotation Yrotation Xrotation
+                      JOINT RightHandRing3
+                      {
+                        OFFSET -4.33881 -4e-06 -0.0
+                        CHANNELS 3 Zrotation Yrotation Xrotation
+                        JOINT RightHandRing4
+                        {
+                          OFFSET -2.654903 -4e-06 4e-06
+                          CHANNELS 3 Zrotation Yrotation Xrotation
+                          JOINT RightHandRingEnd
+                          {
+                            OFFSET -1.933568 0.077527 -5.2e-05
+                            CHANNELS 3 Zrotation Yrotation Xrotation
+                          }
+                        }
+                      }
+                    }
+                  }
+                  JOINT RightHandPinky1
+                  {
+                    OFFSET -2.866425 -0.342796 -1.584145
+                    CHANNELS 3 Zrotation Yrotation Xrotation
+                    JOINT RightHandPinky2
+                    {
+                      OFFSET -5.091371 -1.332055 -1.772385
+                      CHANNELS 3 Zrotation Yrotation Xrotation
+                      JOINT RightHandPinky3
+                      {
+                        OFFSET -3.062664 -4e-06 1e-06
+                        CHANNELS 3 Zrotation Yrotation Xrotation
+                        JOINT RightHandPinky4
+                        {
+                          OFFSET -1.546529 4e-06 -2e-06
+                          CHANNELS 3 Zrotation Yrotation Xrotation
+                          JOINT RightHandPinkyEnd
+                          {
+                            OFFSET -1.945119 -0.157718 0.057211
+                            CHANNELS 3 Zrotation Yrotation Xrotation
+                          }
+                        }
+                      }
+                    }
+                  }
+                }
+              }
+            }
+          }
+        }
+      }
+    }
+    JOINT LeftLeg
+    {
+      OFFSET 10.043214 -8.434526 2.595655
+      CHANNELS 3 Zrotation Yrotation Xrotation
+      JOINT LeftShin
+      {
+        OFFSET -1e-06 -43.221752 -0.802913
+        CHANNELS 3 Zrotation Yrotation Xrotation
+        JOINT LeftFoot
+        {
+          OFFSET 1e-06 -42.155094 -3.481523
+          CHANNELS 3 Zrotation Yrotation Xrotation
+          JOINT LeftToeBase
+          {
+            OFFSET 0.0 -5.059472 13.231529
+            CHANNELS 3 Zrotation Yrotation Xrotation
+            JOINT LeftToeEnd
+            {
+              OFFSET -0.009607 -1.647619 6.513017
+              CHANNELS 3 Zrotation Yrotation Xrotation
+            }
+          }
+        }
+      }
+    }
+    JOINT RightLeg
+    {
+      OFFSET -10.047278 -8.29526 2.620317
+      CHANNELS 3 Zrotation Yrotation Xrotation
+      JOINT RightShin
+      {
+        OFFSET 1e-06 -43.362206 -0.805556
+        CHANNELS 3 Zrotation Yrotation Xrotation
+        JOINT RightFoot
+        {
+          OFFSET 2e-06 -42.117393 -3.478398
+          CHANNELS 3 Zrotation Yrotation Xrotation
+          JOINT RightToeBase
+          {
+            OFFSET -0.0 -5.079609 13.284196
+            CHANNELS 3 Zrotation Yrotation Xrotation
+            JOINT RightToeEnd
+            {
+              OFFSET 0.009532 -1.634378 6.460591
+              CHANNELS 3 Zrotation Yrotation Xrotation
+            }
+          }
+        }
+      }
+    }
+  }
+}
+MOTION
+Frames: 1
+Frame Time: 0.03333333333333333
+0.0 0.0 0.0 0.0 0.0 0.0 0.0 100.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0

kimodo/constraints.py

@@ -0,0 +1,625 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Constraint sets for conditioning motion generation (root 2D, full body, end-effectors)."""
+
+from typing import Optional, Union
+
+import torch
+from torch import Tensor
+
+from kimodo.motion_rep.feature_utils import compute_heading_angle
+from kimodo.skeleton import SkeletonBase, SOMASkeleton30, SOMASkeleton77
+from kimodo.tools import ensure_batched, load_json, save_json
+
+from .geometry import axis_angle_to_matrix, matrix_to_axis_angle
+
+
+def _convert_constraint_local_rots_to_skeleton(local_rot_mats: Tensor, skeleton: SkeletonBase) -> Tensor:
+    """Convert loaded local rotation matrices to match the skeleton's joint count.
+
+    Handles SOMA 30↔77: constraint files may have been saved with 30 or 77 joints while the session
+    skeleton (e.g. from the SOMA30 model) uses SOMASkeleton77.
+    """
+    n_joints = local_rot_mats.shape[-3]
+    skeleton_joints = skeleton.nbjoints
+    if n_joints == skeleton_joints:
+        return local_rot_mats
+    if n_joints == 77 and skeleton_joints == 30 and isinstance(skeleton, SOMASkeleton30):
+        return skeleton.from_SOMASkeleton77(local_rot_mats)
+    if n_joints == 30 and skeleton_joints == 77 and isinstance(skeleton, SOMASkeleton77):
+        skel30 = SOMASkeleton30()
+        return skel30.to_SOMASkeleton77(local_rot_mats)
+    raise ValueError(
+        f"Constraint joint count ({n_joints}) does not match skeleton joint count "
+        f"({skeleton_joints}). Only SOMA 30↔77 conversion is supported."
+    )
+
+
+def create_pairs(tensor_A: Tensor, tensor_B: Tensor) -> Tensor:
+    """Form all (a, b) pairs from two 1D tensors; output shape (len(A)*len(B), 2)."""
+    pairs = torch.stack(
+        (
+            tensor_A[:, None].expand(-1, len(tensor_B)),
+            tensor_B.expand(len(tensor_A), -1),
+        ),
+        dim=-1,
+    ).reshape(-1, 2)
+    return pairs
+
+
+def compute_global_heading(global_joints_positions: Tensor, skeleton: SkeletonBase) -> Tensor:
+    """Compute global root heading (cos, sin) from global joint positions using skeleton."""
+    root_heading_angle = compute_heading_angle(global_joints_positions, skeleton)
+    global_root_heading = torch.stack([torch.cos(root_heading_angle), torch.sin(root_heading_angle)], dim=-1)
+    return global_root_heading
+
+
+def _tensor_to(
+    t: Tensor,
+    device: Optional[Union[str, torch.device]] = None,
+    dtype: Optional[torch.dtype] = None,
+) -> Tensor:
+    """Move tensor to device and/or dtype.
+
+    Returns same tensor if no args.
+    """
+    if device is not None and dtype is not None:
+        return t.to(device=device, dtype=dtype)
+    if device is not None:
+        return t.to(device=device)
+    if dtype is not None:
+        return t.to(dtype=dtype)
+    return t
+
+
+class Root2DConstraintSet:
+    """Constraint set fixing root (x, z) trajectory and optionally global heading on given
+    frames."""
+
+    name = "root2d"
+
+    def __init__(
+        self,
+        skeleton: SkeletonBase,
+        frame_indices: Tensor,
+        smooth_root_2d: Tensor,
+        to_crop: bool = False,
+        global_root_heading: Optional[Tensor] = None,
+    ) -> None:
+        self.skeleton = skeleton
+
+        # if we pass the full smooth root 3D as input
+        if smooth_root_2d.shape[-1] == 3:
+            smooth_root_2d = smooth_root_2d[..., [0, 1]]
+
+        if to_crop:
+            smooth_root_2d = smooth_root_2d[frame_indices]
+            if global_root_heading is not None:
+                global_root_heading = global_root_heading[frame_indices]
+        else:
+            assert len(smooth_root_2d) == len(
+                frame_indices
+            ), "The number of smooth root 2d should be match the number of frames"
+            if global_root_heading is not None:
+                assert len(global_root_heading) == len(
+                    frame_indices
+                ), "The number of global root heading should be match the number of frames"
+
+        self.smooth_root_2d = smooth_root_2d
+        self.global_root_heading = global_root_heading
+        self.frame_indices = frame_indices
+
+    def update_constraints(self, data_dict: dict, index_dict: dict) -> None:
+        """Append this constraint's smooth_root_2d (and optional global_root_heading) to data/index
+        dicts."""
+        data_dict["smooth_root_2d"].append(self.smooth_root_2d)
+        index_dict["smooth_root_2d"].append(self.frame_indices)
+
+        if self.global_root_heading is not None:
+            # constraint the global heading
+            data_dict["global_root_heading"].append(self.global_root_heading)
+            index_dict["global_root_heading"].append(self.frame_indices)
+
+    def crop_move(self, start: int, end: int) -> "Root2DConstraintSet":
+        """Return a new constraint set for the cropped frame range [start, end)."""
+        mask = (self.frame_indices >= start) & (self.frame_indices < end)
+
+        if self.global_root_heading is not None:
+            masked_global_root_heading = self.global_root_heading[mask]
+        else:
+            masked_global_root_heading = None
+
+        return Root2DConstraintSet(
+            self.skeleton,
+            self.frame_indices[mask] - start,
+            self.smooth_root_2d[mask],
+            global_root_heading=masked_global_root_heading,
+        )
+
+    def get_save_info(self) -> dict:
+        """Return a dict suitable for JSON serialization (frame_indices, smooth_root_2d, optional
+        global_root_heading)."""
+        out = {
+            "type": self.name,
+            "frame_indices": self.frame_indices,
+            "smooth_root_2d": self.smooth_root_2d,
+        }
+        if self.global_root_heading is not None:
+            out["global_root_heading"] = self.global_root_heading
+        return out
+
+    def to(
+        self,
+        device: Optional[Union[str, torch.device]] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> "Root2DConstraintSet":
+        self.smooth_root_2d = _tensor_to(self.smooth_root_2d, device, dtype)
+        self.frame_indices = _tensor_to(self.frame_indices, device, dtype)
+        if self.global_root_heading is not None:
+            self.global_root_heading = _tensor_to(self.global_root_heading, device, dtype)
+        if device is not None and hasattr(self.skeleton, "to"):
+            self.skeleton = self.skeleton.to(device)
+        return self
+
+    @classmethod
+    def from_dict(cls, skeleton: SkeletonBase, dico: dict) -> "Root2DConstraintSet":
+        """Build a Root2DConstraintSet from a dict (e.g. loaded from JSON)."""
+        device = skeleton.device if hasattr(skeleton, "device") else "cpu"
+
+        if "global_root_heading" in dico:
+            global_root_heading = torch.tensor(dico["global_root_heading"], device=device)
+        else:
+            global_root_heading = None
+
+        return cls(
+            skeleton,
+            frame_indices=torch.tensor(dico["frame_indices"]),
+            smooth_root_2d=torch.tensor(dico["smooth_root_2d"], device=device),
+            global_root_heading=global_root_heading,
+        )
+
+
+class FullBodyConstraintSet:
+    """Constraint set fixing full-body global positions and rotations on given keyframes."""
+
+    name = "fullbody"
+
+    def __init__(
+        self,
+        skeleton: SkeletonBase,
+        frame_indices: Tensor,
+        global_joints_positions: Tensor,
+        global_joints_rots: Tensor,
+        smooth_root_2d: Optional[Tensor] = None,
+        to_crop: bool = False,
+    ):
+        self.skeleton = skeleton
+        self.frame_indices = frame_indices
+
+        # if we pass the full smooth root 3D as input
+        if smooth_root_2d is not None and smooth_root_2d.shape[-1] == 3:
+            smooth_root_2d = smooth_root_2d[..., [0, 1]]
+
+        if to_crop:
+            global_joints_positions = global_joints_positions[frame_indices]
+            global_joints_rots = global_joints_rots[frame_indices]
+            if smooth_root_2d is not None:
+                smooth_root_2d = smooth_root_2d[frame_indices]
+        else:
+            assert len(global_joints_positions) == len(
+                frame_indices
+            ), "The number of global positions should be match the number of frames"
+            assert len(global_joints_rots) == len(
+                frame_indices
+            ), "The number of global joint rotations should be match the number of frames"
+
+            if smooth_root_2d is not None:
+                assert len(smooth_root_2d) == len(
+                    frame_indices
+                ), "The number of smooth root 2d (if specified) should be match the number of frames"
+
+        if smooth_root_2d is None:
+            # substitute the smooth root 2d with the real root
+            smooth_root_2d = global_joints_positions[:, skeleton.root_idx, [0, 2]]
+
+        # root y: from smooth or pelvis is the same
+        self.root_y_pos = global_joints_positions[:, skeleton.root_idx, 1]
+
+        self.global_joints_positions = global_joints_positions
+        self.global_joints_rots = global_joints_rots
+        self.global_root_heading = compute_global_heading(global_joints_positions, skeleton)
+        self.smooth_root_2d = smooth_root_2d
+
+    def update_constraints(self, data_dict: dict, index_dict: dict) -> None:
+        """Append global positions, smooth root 2D, root y, and global heading to data/index
+        dicts."""
+        nbjoints = self.skeleton.nbjoints
+        indices_lst = create_pairs(
+            self.frame_indices,
+            torch.arange(nbjoints, device=self.frame_indices.device),
+        )
+        data_dict["global_joints_positions"].append(
+            self.global_joints_positions.reshape(-1, 3)
+        )  # flatten the global positions
+        index_dict["global_joints_positions"].append(indices_lst)
+
+        # global rotations are not used here
+
+        # as we use smooth root, also constraint the smooth root to get the same full body
+        # maybe keep storing the hips offset, if we smooth it ourselves
+        data_dict["smooth_root_2d"].append(self.smooth_root_2d)
+        index_dict["smooth_root_2d"].append(self.frame_indices)
+
+        # constraint the y pos of the root
+        data_dict["root_y_pos"].append(self.root_y_pos)
+        index_dict["root_y_pos"].append(self.frame_indices)
+
+        # constraint the global heading
+        data_dict["global_root_heading"].append(self.global_root_heading)
+        index_dict["global_root_heading"].append(self.frame_indices)
+
+    def crop_move(self, start: int, end: int) -> "FullBodyConstraintSet":
+        """Return a new FullBodyConstraintSet for the cropped frame range [start, end)."""
+        mask = (self.frame_indices >= start) & (self.frame_indices < end)
+        return FullBodyConstraintSet(
+            self.skeleton,
+            self.frame_indices[mask] - start,
+            self.global_joints_positions[mask],
+            self.global_joints_rots[mask],
+            self.smooth_root_2d[mask],
+        )
+
+    def get_save_info(self) -> dict:
+        """Return a dict for JSON save: type, frame_indices, local_joints_rot, root_positions, smooth_root_2d."""
+        local_joints_rot = self.skeleton.global_rots_to_local_rots(self.global_joints_rots)
+        if isinstance(self.skeleton, SOMASkeleton30):
+            local_joints_rot = self.skeleton.to_SOMASkeleton77(local_joints_rot)
+        local_joints_rot = matrix_to_axis_angle(local_joints_rot)
+
+        root_positions = self.global_joints_positions[:, self.skeleton.root_idx]
+        return {
+            "type": self.name,
+            "frame_indices": self.frame_indices,
+            "local_joints_rot": local_joints_rot,
+            "root_positions": root_positions,
+            "smooth_root_2d": self.smooth_root_2d,
+        }
+
+    def to(
+        self,
+        device: Optional[Union[str, torch.device]] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> "FullBodyConstraintSet":
+        self.frame_indices = _tensor_to(self.frame_indices, device, dtype)
+        self.global_joints_positions = _tensor_to(self.global_joints_positions, device, dtype)
+        self.global_joints_rots = _tensor_to(self.global_joints_rots, device, dtype)
+        self.root_y_pos = _tensor_to(self.root_y_pos, device, dtype)
+        self.global_root_heading = _tensor_to(self.global_root_heading, device, dtype)
+        self.smooth_root_2d = _tensor_to(self.smooth_root_2d, device, dtype)
+        if device is not None and hasattr(self.skeleton, "to"):
+            self.skeleton = self.skeleton.to(device)
+        return self
+
+    @classmethod
+    def from_dict(cls, skeleton: SkeletonBase, dico: dict) -> "FullBodyConstraintSet":
+        """Build a FullBodyConstraintSet from a dict (e.g. loaded from JSON)."""
+        frame_indices = torch.tensor(dico["frame_indices"])
+        device = skeleton.device if hasattr(skeleton, "device") else "cpu"
+        local_rot = torch.tensor(dico["local_joints_rot"], device=device)
+        local_rot_mats = axis_angle_to_matrix(local_rot)
+        local_rot_mats = _convert_constraint_local_rots_to_skeleton(local_rot_mats, skeleton)
+        global_joints_rots, global_joints_positions, _ = skeleton.fk(
+            local_rot_mats,
+            torch.tensor(dico["root_positions"], device=device),
+        )
+        smooth_root_2d = None
+        if "smooth_root_2d" in dico:
+            smooth_root_2d = torch.tensor(dico["smooth_root_2d"], device=device)
+
+        return cls(
+            skeleton,
+            frame_indices=frame_indices,
+            global_joints_positions=global_joints_positions,
+            global_joints_rots=global_joints_rots,
+            smooth_root_2d=smooth_root_2d,
+        )
+
+
+class EndEffectorConstraintSet:
+    """Constraint set fixing selected end-effector positions and rotations on given frames."""
+
+    name = "end-effector"
+
+    def __init__(
+        self,
+        skeleton: SkeletonBase,
+        frame_indices: Tensor,
+        global_joints_positions: Tensor,
+        global_joints_rots: Tensor,
+        smooth_root_2d: Optional[Tensor],
+        *,
+        joint_names: list[str],
+        to_crop: bool = False,
+    ) -> None:
+        self.skeleton = skeleton
+        self.frame_indices = frame_indices
+        self.joint_names = joint_names
+
+        # joint_names are constant for all the frames
+        rot_joint_names, pos_joint_names = self.skeleton.expand_joint_names(self.joint_names)
+        # indexing works for motion_rep with smooth root only (contains pelvis index)
+        self.pos_indices = torch.tensor([self.skeleton.bone_index[jname] for jname in pos_joint_names])
+        self.rot_indices = torch.tensor([self.skeleton.bone_index[jname] for jname in rot_joint_names])
+
+        # if we pass the full smooth root 3D as input
+        if smooth_root_2d is not None and smooth_root_2d.shape[-1] == 3:
+            smooth_root_2d = smooth_root_2d[..., [0, 1]]
+
+        if to_crop:
+            global_joints_positions = global_joints_positions[frame_indices]
+            global_joints_rots = global_joints_rots[frame_indices]
+            if smooth_root_2d is not None:
+                smooth_root_2d = smooth_root_2d[frame_indices]
+        else:
+            assert len(global_joints_positions) == len(
+                frame_indices
+            ), "The number of global positions should be match the number of frames"
+            assert len(global_joints_rots) == len(
+                frame_indices
+            ), "The number of global joint rotations should be match the number of frames"
+            if smooth_root_2d is not None:
+                assert len(smooth_root_2d) == len(
+                    frame_indices
+                ), "The number of smooth root 2d (if specified) should be match the number of frames"
+
+        if smooth_root_2d is None:
+            # substitute the smooth root 2d with the real root
+            smooth_root_2d = global_joints_positions[:, skeleton.root_idx, [0, 2]]
+
+        # root y: from smooth or pelvis is the same
+        self.root_y_pos = global_joints_positions[:, skeleton.root_idx, 1]
+
+        self.global_joints_positions = global_joints_positions
+        self.global_root_heading = compute_global_heading(global_joints_positions, skeleton)
+        self.global_joints_rots = global_joints_rots
+        self.smooth_root_2d = smooth_root_2d
+
+    def update_constraints(self, data_dict: dict, index_dict: dict) -> None:
+        """Append constrained joint positions/rots, smooth root 2D, root y, and heading to
+        data/index dicts."""
+        crop_frames_indexing = torch.arange(len(self.frame_indices), device=self.frame_indices.device)
+
+        # constraint positions
+        pos_indices_real = create_pairs(
+            self.frame_indices,
+            self.pos_indices,
+        )
+        pos_indices_crop = create_pairs(
+            crop_frames_indexing,
+            self.pos_indices,
+        )
+        data_dict["global_joints_positions"].append(self.global_joints_positions[tuple(pos_indices_crop.T)])
+        index_dict["global_joints_positions"].append(pos_indices_real)
+
+        # constraint rotations
+        rot_indices_real = create_pairs(
+            self.frame_indices,
+            self.rot_indices,
+        )
+        rot_indices_crop = create_pairs(
+            crop_frames_indexing,
+            self.rot_indices,
+        )
+        data_dict["global_joints_rots"].append(self.global_joints_rots[tuple(rot_indices_crop.T)])
+        index_dict["global_joints_rots"].append(rot_indices_real)
+
+        # as we use smooth root, also constraint the smooth root to get the same full body
+        # maybe keep storing the hips offset, if we smooth it ourselves
+        data_dict["smooth_root_2d"].append(self.smooth_root_2d)
+        index_dict["smooth_root_2d"].append(self.frame_indices)
+
+        # constraint the y pos of the root
+        data_dict["root_y_pos"].append(self.root_y_pos)
+        index_dict["root_y_pos"].append(self.frame_indices)
+
+        # constraint the global heading
+        data_dict["global_root_heading"].append(self.global_root_heading)
+        index_dict["global_root_heading"].append(self.frame_indices)
+
+    def crop_move(self, start: int, end: int) -> "EndEffectorConstraintSet":
+        """Return a new EndEffectorConstraintSet for the cropped frame range [start, end)."""
+        mask = (self.frame_indices >= start) & (self.frame_indices < end)
+
+        cls = type(self)
+        kwargs = {}
+        if not hasattr(cls, "joint_names"):
+            kwargs["joint_names"] = self.joint_names
+
+        return cls(
+            self.skeleton,
+            self.frame_indices[mask] - start,
+            self.global_joints_positions[mask],
+            self.global_joints_rots[mask],
+            self.smooth_root_2d[mask],
+            **kwargs,
+        )
+
+    def get_save_info(self) -> dict:
+        """Return a dict for JSON save: type, frame_indices, local_joints_rot, root_positions, smooth_root_2d, joint_names."""
+        local_joints_rot = self.skeleton.global_rots_to_local_rots(self.global_joints_rots)
+        if isinstance(self.skeleton, SOMASkeleton30):
+            local_joints_rot = self.skeleton.to_SOMASkeleton77(local_joints_rot)
+        local_joints_rot = matrix_to_axis_angle(local_joints_rot)
+
+        root_positions = self.global_joints_positions[:, self.skeleton.root_idx]
+        output = {
+            "type": self.name,
+            "frame_indices": self.frame_indices,
+            "local_joints_rot": local_joints_rot,
+            "root_positions": root_positions,
+            "smooth_root_2d": self.smooth_root_2d,
+        }
+        if not hasattr(self.__class__, "joint_names"):
+            # save the joint_names for this base class
+            # but not for children
+            output["joint_names"] = self.joint_names
+        return output
+
+    def to(
+        self,
+        device: Optional[Union[str, torch.device]] = None,
+        dtype: Optional[torch.dtype] = None,
+    ) -> "EndEffectorConstraintSet":
+        self.frame_indices = _tensor_to(self.frame_indices, device, dtype)
+        self.pos_indices = _tensor_to(self.pos_indices, device, dtype)
+        self.rot_indices = _tensor_to(self.rot_indices, device, dtype)
+        self.root_y_pos = _tensor_to(self.root_y_pos, device, dtype)
+        self.global_joints_positions = _tensor_to(self.global_joints_positions, device, dtype)
+        self.global_root_heading = _tensor_to(self.global_root_heading, device, dtype)
+        self.global_joints_rots = _tensor_to(self.global_joints_rots, device, dtype)
+        self.smooth_root_2d = _tensor_to(self.smooth_root_2d, device, dtype)
+        if device is not None and hasattr(self.skeleton, "to"):
+            self.skeleton = self.skeleton.to(device)
+        return self
+
+    @classmethod
+    def from_dict(cls, skeleton: SkeletonBase, dico: dict) -> "EndEffectorConstraintSet":
+        """Build an EndEffectorConstraintSet from a dict (e.g. loaded from JSON)."""
+        frame_indices = torch.tensor(dico["frame_indices"])
+        device = skeleton.device if hasattr(skeleton, "device") else "cpu"
+        local_rot = torch.tensor(dico["local_joints_rot"], device=device)
+        local_rot_mats = axis_angle_to_matrix(local_rot)
+        local_rot_mats = _convert_constraint_local_rots_to_skeleton(local_rot_mats, skeleton)
+        global_joints_rots, global_joints_positions, _ = skeleton.fk(
+            local_rot_mats,
+            torch.tensor(dico["root_positions"], device=device),
+        )
+        smooth_root_2d = None
+        if "smooth_root_2d" in dico:
+            smooth_root_2d = torch.tensor(dico["smooth_root_2d"], device=device)
+
+        kwargs = {}
+        if not hasattr(cls, "joint_names"):
+            kwargs["joint_names"] = dico["joint_names"]
+
+        return cls(
+            skeleton,
+            frame_indices=frame_indices,
+            global_joints_positions=global_joints_positions,
+            global_joints_rots=global_joints_rots,
+            smooth_root_2d=smooth_root_2d,
+            **kwargs,
+        )
+
+
+class LeftHandConstraintSet(EndEffectorConstraintSet):
+    """End-effector constraint for the left hand only."""
+
+    name = "left-hand"
+    joint_names: list[str] = ["LeftHand"]
+
+    def __init__(self, *args, **kwargs: dict):
+        super().__init__(*args, joint_names=self.joint_names, **kwargs)
+
+
+class RightHandConstraintSet(EndEffectorConstraintSet):
+    """End-effector constraint for the right hand only."""
+
+    name = "right-hand"
+    joint_names: list[str] = ["RightHand"]
+
+    def __init__(self, *args, **kwargs: dict):
+        super().__init__(*args, joint_names=self.joint_names, **kwargs)
+
+
+class LeftFootConstraintSet(EndEffectorConstraintSet):
+    """End-effector constraint for the left foot only."""
+
+    name = "left-foot"
+    joint_names: list[str] = ["LeftFoot"]
+
+    def __init__(self, *args, **kwargs: dict):
+        super().__init__(*args, joint_names=self.joint_names, **kwargs)
+
+
+class RightFootConstraintSet(EndEffectorConstraintSet):
+    """End-effector constraint for the right foot only."""
+
+    name = "right-foot"
+    joint_names: list[str] = ["RightFoot"]
+
+    def __init__(self, *args, **kwargs: dict):
+        super().__init__(*args, joint_names=self.joint_names, **kwargs)
+
+
+TYPE_TO_CLASS = {
+    "root2d": Root2DConstraintSet,
+    "fullbody": FullBodyConstraintSet,
+    "left-hand": LeftHandConstraintSet,
+    "right-hand": RightHandConstraintSet,
+    "left-foot": LeftFootConstraintSet,
+    "right-foot": RightFootConstraintSet,
+    "end-effector": EndEffectorConstraintSet,
+}
+
+
+def load_constraints_lst(
+    path_or_data: str | list,
+    skeleton: SkeletonBase,
+    device: Optional[Union[str, torch.device]] = None,
+    dtype: Optional[torch.dtype] = None,
+):
+    """Load a list of constraints from JSON path or list of dicts.
+
+    Args:
+        path_or_data: Path to constraints.json or list of constraint dicts.
+        skeleton: Skeleton instance (used for from_dict).
+        device: If set, move all constraint tensors and skeleton to this device.
+        dtype: If set, cast constraint tensors to this dtype.
+    """
+    if isinstance(path_or_data, str):
+        saved = load_json(path_or_data)
+    else:
+        saved = path_or_data
+
+    constraints_lst = []
+    for el in saved:
+        cls = TYPE_TO_CLASS[el["type"]]
+        c = cls.from_dict(skeleton, el)
+        if device is not None or dtype is not None:
+            c.to(device=device, dtype=dtype)
+        constraints_lst.append(c)
+    return constraints_lst
+
+
+def save_constraints_lst(path: str, constraints_lst: list) -> list | None:
+    """Save a list of constraint sets to a JSON file.
+
+    Returns None if list is empty.
+    """
+    if not constraints_lst:
+        print("The constraints lst is empty. Skip saving")
+        return
+
+    to_save = []
+
+    def tensor_to_list(obj):
+        """Recursively convert tensors to lists for JSON serialization."""
+        if isinstance(obj, Tensor):
+            return obj.cpu().tolist()
+        elif isinstance(obj, dict):
+            return {k: tensor_to_list(v) for k, v in obj.items()}
+        elif isinstance(obj, list):
+            return [tensor_to_list(v) for v in obj]
+        else:
+            return obj
+
+    for constraint in constraints_lst:
+        constraint_info = constraint.get_save_info()
+        # Convert all tensors to lists for JSON serialization
+        constraint_info = tensor_to_list(constraint_info)
+        to_save.append(constraint_info)
+
+    save_json(path, to_save)
+    print(f"Saved constraints to {path}")
+    return to_save

kimodo/demo/__init__.py

@@ -0,0 +1,29 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# ruff: noqa: I001
+import argparse
+
+from kimodo.model import DEFAULT_MODEL
+from kimodo.model.registry import resolve_model_name
+
+from .app import Demo
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Run the kimodo demo UI.")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default=DEFAULT_MODEL,
+        help="Default model to load (e.g. Kimodo-SOMA-RP-v1, kimodo-soma-rp, or SOMA).",
+    )
+    args = parser.parse_args()
+
+    resolved = resolve_model_name(args.model, "Kimodo")
+    demo = Demo(default_model_name=resolved)
+    demo.run()
+
+
+if __name__ == "__main__":
+    main()

kimodo/demo/__main__.py

@@ -0,0 +1,8 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Entry point for `python -m kimodo.demo`."""
+
+from kimodo.demo import main
+
+if __name__ == "__main__":
+    main()

kimodo/demo/app.py

@@ -0,0 +1,690 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import base64
+import os
+import shutil
+import threading
+import time
+from typing import Optional
+
+import numpy as np
+import torch
+
+import viser
+from kimodo.assets import DEMO_ASSETS_ROOT
+from kimodo.model.load_model import load_model
+from kimodo.model.registry import resolve_model_name
+from kimodo.skeleton import SkeletonBase, SOMASkeleton30
+from kimodo.tools import load_json
+from kimodo.viz import viser_utils
+from kimodo.viz.viser_utils import (
+    Character,
+    CharacterMotion,
+    EEJointsKeyframeSet,
+    FullbodyKeyframeSet,
+    RootKeyframe2DSet,
+)
+from viser.theme import TitlebarButton, TitlebarConfig, TitlebarImage
+
+from . import generation, ui
+from .config import (
+    DARK_THEME,
+    DEFAULT_CUR_DURATION,
+    DEFAULT_MODEL,
+    DEFAULT_PLAYBACK_SPEED,
+    DEFAULT_PROMPT,
+    DEMO_UI_QUICK_START_MODAL_MD,
+    EXAMPLES_ROOT_DIR,
+    HF_MODE,
+    LIGHT_THEME,
+    MAX_ACTIVE_USERS,
+    MAX_DURATION,
+    MAX_SESSION_MINUTES,
+    MIN_DURATION,
+    MODEL_EXAMPLES_DIRS,
+    MODEL_NAMES,
+    SERVER_NAME,
+    SERVER_PORT,
+)
+from .embedding_cache import CachedTextEncoder
+from .queue_manager import QueueManager, UserQueue
+from .state import ClientSession, ModelBundle
+
+
+class Demo:
+    def __init__(self, default_model_name: str = DEFAULT_MODEL):
+        self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        print(f"Using device: {self.device}")
+        self.models: dict[str, ModelBundle] = {}
+        self._text_encoder = None
+        resolved = resolve_model_name(default_model_name, "Kimodo")
+        if resolved not in MODEL_NAMES:
+            raise ValueError(f"Unknown model '{default_model_name}'. Expected one of: {MODEL_NAMES}")
+        self.default_model_name = resolved
+        self.ensure_examples_layout()
+        self.load_model(self.default_model_name)
+
+        # Serialize GPU-bound generation across all clients
+        self._generation_lock = threading.Lock()
+        self._cuda_healthy = True
+
+        # Per-client sessions
+        self.client_sessions: dict[int, ClientSession] = {}
+        self.start_direction_markers: dict[int, viser_utils.WaypointMesh] = {}
+        self.grid_handles: dict[int, viser.GridHandle] = {}
+
+        self.server = viser.ViserServer(
+            host=SERVER_NAME,
+            port=SERVER_PORT,
+            label="Kimodo",
+            enable_camera_keyboard_controls=False,  # don't move the camera with the arrow keys
+        )
+        self.server.scene.world_axes.visible = False  # used for debugging
+        self.server.scene.set_up_direction("+y")
+
+        # Register callbacks for session handling
+        self.server.on_client_connect(self.on_client_connect)
+        self.server.on_client_disconnect(self.on_client_disconnect)
+
+        # HF mode: queue and session limit
+        if HF_MODE:
+            self.user_queue = UserQueue(MAX_ACTIVE_USERS, MAX_SESSION_MINUTES)
+            self.queue_manager = QueueManager(
+                queue=self.user_queue,
+                server=self.server,
+                setup_demo_for_client=self._setup_demo_for_client,
+                cleanup_session=self._cleanup_session_for_client,
+            )
+        else:
+            self.user_queue = None
+            self.queue_manager = None
+
+        # create grid and floor
+        self.floor_len = 20.0  # meters
+
+    def ensure_examples_layout(self) -> None:
+        os.makedirs(EXAMPLES_ROOT_DIR, exist_ok=True)
+        for model_dir in MODEL_EXAMPLES_DIRS.values():
+            os.makedirs(model_dir, exist_ok=True)
+
+        for entry in os.listdir(EXAMPLES_ROOT_DIR):
+            if entry in MODEL_EXAMPLES_DIRS:
+                continue
+            src = os.path.join(EXAMPLES_ROOT_DIR, entry)
+            if not os.path.isdir(src):
+                continue
+            dst = os.path.join(
+                MODEL_EXAMPLES_DIRS.get(DEFAULT_MODEL, next(iter(MODEL_EXAMPLES_DIRS.values()))),
+                entry,
+            )
+            if not os.path.exists(dst):
+                shutil.move(src, dst)
+
+    def get_examples_base_dir(self, model_name: str, absolute: bool = True) -> str:
+        return MODEL_EXAMPLES_DIRS[model_name]
+
+    def load_model(self, model_name: str) -> ModelBundle:
+        if model_name in self.models:
+            return self.models[model_name]
+
+        print(f"Loading model {model_name}...")
+        try:
+            model = load_model(
+                modelname=model_name,
+                device=self.device,
+                text_encoder=self._text_encoder,
+            )
+        except Exception as e:
+            print(f"Error loading model: {e}\nMake sure text encoder server is running!")
+            raise e
+
+        if hasattr(model, "text_encoder"):
+            if self._text_encoder is None:
+                self._text_encoder = model.text_encoder
+            model.text_encoder = CachedTextEncoder(model.text_encoder, model_name=model_name)
+
+        skeleton = model.motion_rep.skeleton
+        if isinstance(skeleton, SOMASkeleton30):
+            skeleton = skeleton.somaskel77.to(model.device)
+        bundle = ModelBundle(
+            model=model,
+            motion_rep=model.motion_rep,
+            skeleton=skeleton,
+            model_fps=model.motion_rep.fps,
+        )
+        self.models[model_name] = bundle
+        print(f"Model {model_name} loaded successfully")
+        self.prewarm_embedding_cache(model_name, bundle.model)
+        return bundle
+
+    def prewarm_embedding_cache(self, model_name: str, model: object) -> None:
+        encoder = getattr(model, "text_encoder", None)
+        if not isinstance(encoder, CachedTextEncoder):
+            return
+
+        prompt_set = set()
+        prompt_set.add(DEFAULT_PROMPT)
+
+        examples_dir = MODEL_EXAMPLES_DIRS.get(model_name)
+        if examples_dir and os.path.isdir(examples_dir):
+            for entry in os.listdir(examples_dir):
+                example_dir = os.path.join(examples_dir, entry)
+                if not os.path.isdir(example_dir):
+                    continue
+                meta_path = os.path.join(example_dir, "meta.json")
+                if not os.path.exists(meta_path):
+                    continue
+                try:
+                    meta = load_json(meta_path)
+                except Exception:
+                    continue
+                for prompt in meta.get("prompts_text", []):
+                    if isinstance(prompt, str):
+                        prompt_set.add(prompt)
+
+        if prompt_set:
+            encoder.prewarm(list(prompt_set))
+
+    def build_constraint_tracks(
+        self, client: viser.ClientHandle, skeleton: SkeletonBase
+    ) -> dict[str, viser_utils.ConstraintSet]:
+        return {
+            "Full-Body": FullbodyKeyframeSet(
+                name="Full-Body",
+                server=client,
+                skeleton=skeleton,
+            ),
+            "End-Effectors": EEJointsKeyframeSet(
+                name="End-Effectors",
+                server=client,
+                skeleton=skeleton,
+            ),
+            "2D Root": RootKeyframe2DSet(
+                name="2D Root",
+                server=client,
+                skeleton=skeleton,
+            ),
+        }
+
+    def set_timeline_defaults(self, timeline, model_fps: float) -> None:
+        timeline.set_defaults(
+            default_text=DEFAULT_PROMPT,
+            default_duration=int(DEFAULT_CUR_DURATION * model_fps - 1),
+            min_duration=int(MIN_DURATION * model_fps - 1),  # 2 seconds minimum,
+            max_duration=int(
+                MAX_DURATION * model_fps - 1  # - NB_TRANSITION_FRAMES
+            ),  # 10 seconds maximum, minus the transition frames, if needed
+            default_num_frames_zoom=int(1.10 * 10 * model_fps),  # a bit more than the max
+            max_frames_zoom=1000,
+            fps=model_fps,
+        )
+
+    def _apply_constraint_overlay_visibility(self, session: ClientSession) -> None:
+        """Apply show-all vs show-only-current-frame to constraint overlays."""
+        only_frame = session.frame_idx if session.show_only_current_constraint else None
+        for constraint in session.constraints.values():
+            constraint.set_overlay_visibility(only_frame)
+
+    def set_constraint_tracks_visible(self, session: ClientSession, visible: bool) -> None:
+        timeline = session.client.timeline
+        timeline_data = session.timeline_data
+        if timeline_data.get("constraint_tracks_visible", True) == visible:
+            return
+
+        with timeline_data["keyframe_update_lock"]:
+            if visible:
+                for track_id, track_info in timeline_data["tracks"].items():
+                    timeline.add_track(
+                        track_info["name"],
+                        track_type=track_info.get("track_type", "keyframe"),
+                        color=track_info.get("color"),
+                        height_scale=track_info.get("height_scale", 1.0),
+                        uuid=track_id,
+                    )
+
+                for keyframe_id, keyframe_data in timeline_data["keyframes"].items():
+                    timeline.add_keyframe(
+                        track_id=keyframe_data["track_id"],
+                        frame=keyframe_data["frame"],
+                        value=keyframe_data.get("value"),
+                        opacity=keyframe_data.get("opacity", 1.0),
+                        locked=keyframe_data.get("locked", False),
+                        uuid=keyframe_id,
+                    )
+
+                for interval_id, interval_data in timeline_data["intervals"].items():
+                    timeline.add_interval(
+                        track_id=interval_data["track_id"],
+                        start_frame=interval_data["start_frame_idx"],
+                        end_frame=interval_data["end_frame_idx"],
+                        value=interval_data.get("value"),
+                        opacity=interval_data.get("opacity", 1.0),
+                        locked=interval_data.get("locked", False),
+                        uuid=interval_id,
+                    )
+            else:
+                for track_id in list(timeline_data["tracks"].keys()):
+                    timeline.remove_track(track_id)
+
+        timeline_data["constraint_tracks_visible"] = visible
+
+    def _cleanup_session_for_client(self, client_id: int) -> None:
+        """Remove session and scene state for a client (e.g. on session expiry)."""
+        if client_id in self.client_sessions:
+            del self.client_sessions[client_id]
+        self.start_direction_markers.pop(client_id, None)
+        self.grid_handles.pop(client_id, None)
+
+    def _setup_demo_for_client(self, client: viser.ClientHandle) -> None:
+        """Initialize scene, GUI, and session state for a client (no modals)."""
+        self.setup_scene(client)
+
+        model_bundle = self.load_model(self.default_model_name)
+
+        # Initialize each empty constraint track
+        constraint_tracks = self.build_constraint_tracks(client, model_bundle.skeleton)
+
+        # Create GUI elements for this client
+        (
+            gui_elements,
+            timeline_tracks,
+            example_dict,
+            gui_examples_dropdown,
+            gui_save_example_path_text,
+            gui_model_selector,
+        ) = ui.create_gui(
+            demo=self,
+            client=client,
+            model_name=self.default_model_name,
+            model_fps=model_bundle.model_fps,
+        )
+        timeline_data = {
+            "tracks": timeline_tracks,
+            "tracks_ids": {val["name"]: key for key, val in timeline_tracks.items()},
+            "keyframes": {},
+            "intervals": {},
+            "keyframe_update_lock": threading.Lock(),
+            "keyframe_move_timers": {},
+            "pending_keyframe_moves": {},  # keyframe_id -> new_frame
+            "constraint_tracks_visible": True,
+            "dense_path_after_release_timer": None,
+        }
+
+        # Initialize session state
+        cur_duration = DEFAULT_CUR_DURATION
+        max_frame_idx = int(cur_duration * model_bundle.model_fps - 1)
+
+        session = ClientSession(
+            client=client,
+            gui_elements=gui_elements,
+            motions={},
+            constraints=constraint_tracks,
+            timeline_data=timeline_data,
+            frame_idx=0,
+            playing=False,
+            playback_speed=DEFAULT_PLAYBACK_SPEED,
+            cur_duration=cur_duration,
+            max_frame_idx=max_frame_idx,
+            updating_motions=False,
+            edit_mode=False,
+            model_name=self.default_model_name,
+            model_fps=model_bundle.model_fps,
+            skeleton=model_bundle.skeleton,
+            motion_rep=model_bundle.motion_rep,
+            examples_base_dir=self.get_examples_base_dir(self.default_model_name, absolute=True),
+            example_dict=example_dict,
+            gui_examples_dropdown=gui_examples_dropdown,
+            gui_save_example_path_text=gui_save_example_path_text,
+            gui_model_selector=gui_model_selector,
+        )
+
+        self.client_sessions[client.client_id] = session
+
+        # Initialize default character for this client
+        self.add_character_motion(client, session.skeleton)
+
+    def on_client_connect(self, client: viser.ClientHandle) -> None:
+        """Initialize GUI and state for each new client."""
+        print(f"Client {client.client_id} connected")
+
+        if HF_MODE and self.queue_manager is not None:
+            self.queue_manager.on_client_connect(client)
+        else:
+            # Show quick start popup when a browser client connects (non-HF mode).
+            with client.gui.add_modal(
+                "Welcome — Quick Start",
+                size="xl",
+                show_close_button=True,
+                save_choice="kimodo.demo.quick_start_ack",
+            ) as modal:
+                client.gui.add_markdown(DEMO_UI_QUICK_START_MODAL_MD)
+                client.gui.add_button("Got it (don't remind me again)").on_click(lambda _event: modal.close())
+            self._setup_demo_for_client(client)
+
+    def setup_scene(self, client: viser.ClientHandle) -> None:
+        self.configure_theme(client)
+        client.camera.position = np.array(
+            [2.7417358737841426, 1.8790455698853281, 7.675741569777456],
+            dtype=np.float64,
+        )
+        client.camera.look_at = np.array([0.0, 0.0, 0.0], dtype=np.float64)
+        client.camera.up_direction = np.array(
+            [-1.1102230246251568e-16, 1.0, 1.3596310734468913e-32],
+            dtype=np.float64,
+        )
+        client.camera.fov = np.deg2rad(45.0)
+        grid_handle = client.scene.add_grid(
+            "/grid",
+            width=self.floor_len,
+            height=self.floor_len,
+            wxyz=viser.transforms.SO3.from_x_radians(-np.pi / 2.0).wxyz,
+            position=(0.0, 0.0001, 0.0),
+            fade_distance=3 * self.floor_len,
+            section_color=LIGHT_THEME["grid"],
+            infinite_grid=True,
+        )
+        self.grid_handles[client.client_id] = grid_handle
+        # marker for origin
+        origin_waypoint = viser_utils.WaypointMesh(
+            "/origin_waypoint",
+            client,
+            position=np.array([0.0, 0.0, 0.0]),
+            heading=np.array([0.0, 1.0]),
+            color=(0, 0, 255),
+        )
+        self.start_direction_markers[client.client_id] = origin_waypoint
+
+    def on_client_disconnect(self, client: viser.ClientHandle) -> None:
+        """Clean up when client disconnects."""
+        print(f"Client {client.client_id} disconnected")
+        client_id = client.client_id
+
+        if HF_MODE and self.queue_manager is not None:
+            self.queue_manager.on_client_disconnect(client_id)
+
+        self._cleanup_session_for_client(client_id)
+
+    def set_start_direction_visible(self, client_id: int, visible: bool) -> None:
+        marker = self.start_direction_markers.get(client_id)
+        if marker is None:
+            return
+        marker.set_visible(visible)
+
+    def client_active(self, client_id: int) -> bool:
+        return client_id in self.client_sessions
+
+    def add_character_motion(
+        self,
+        client: viser.ClientHandle,
+        skeleton: SkeletonBase,
+        joints_pos: Optional[torch.Tensor] = None,
+        joints_rot: Optional[torch.Tensor] = None,
+        foot_contacts: Optional[torch.Tensor] = None,
+    ) -> None:
+        client_id = client.client_id
+        if not self.client_active(client_id):
+            return
+        session = self.client_sessions[client_id]
+
+        ci = len(session.motions)
+        character_name = f"character{ci}"
+        # build character skeleton and skinning mesh
+        if "g1" in session.model_name:
+            mesh_mode = "g1_stl"
+        elif "smplx" in session.model_name:
+            mesh_mode = "smplx_skin"
+        elif "soma" in session.model_name:
+            if session.gui_elements.gui_use_soma_layer_checkbox.value:
+                mesh_mode = "soma_layer_skin"
+            else:
+                mesh_mode = "soma_skin"
+        else:
+            raise ValueError("The model name is not recognized for skinning.")
+
+        new_character = Character(
+            character_name,
+            client,
+            skeleton,
+            create_skeleton_mesh=True,
+            create_skinned_mesh=True,
+            visible_skeleton=False,  # don't show immediately
+            visible_skinned_mesh=False,  # don't show immediately
+            skinned_mesh_opacity=session.gui_elements.gui_viz_skinned_mesh_opacity_slider.value,
+            show_foot_contacts=session.gui_elements.gui_viz_foot_contacts_checkbox.value,
+            dark_mode=session.gui_elements.gui_dark_mode_checkbox.value,
+            mesh_mode=mesh_mode,
+            gui_use_soma_layer_checkbox=session.gui_elements.gui_use_soma_layer_checkbox,
+        )
+
+        # if no motion given, initialize to character default (rest) pose for one frame
+        init_joints_pos, init_joints_rot = new_character.get_pose()
+        if joints_pos is None:
+            joints_pos = init_joints_pos[None].repeat(session.max_frame_idx + 1, 1, 1)
+        if joints_rot is None:
+            joints_rot = init_joints_rot[None].repeat(session.max_frame_idx + 1, 1, 1, 1)
+
+        new_motion = CharacterMotion(new_character, joints_pos, joints_rot, foot_contacts)
+        # save the motion in our dict
+        session.motions[character_name] = new_motion
+
+        # put the character at the right frame
+        new_motion.set_frame(session.frame_idx)
+
+        # put them visible with a small delay
+        # so that the set_frame function has time to finish
+        def _set_visibility():
+            new_motion.character.set_skinned_mesh_visibility(session.gui_elements.gui_viz_skinned_mesh_checkbox.value)
+            new_motion.character.set_skeleton_visibility(session.gui_elements.gui_viz_skeleton_checkbox.value)
+
+        timer = threading.Timer(
+            0.2,  # 0.2s delay
+            _set_visibility,
+        )
+        timer.start()
+
+    def clear_motions(self, client_id: int) -> None:
+        if not self.client_active(client_id):
+            return
+        session = self.client_sessions[client_id]
+        for motion in list(session.motions.values()):
+            motion.clear()
+        session.motions.clear()
+
+    def compute_model_constraints_lst(
+        self,
+        session: ClientSession,
+        model_bundle: ModelBundle,
+        num_frames: int,
+    ):
+        return generation.compute_model_constraints_lst(session, model_bundle, num_frames, self.device)
+
+    def check_cuda_health(self) -> bool:
+        """Check if CUDA is still functional.
+
+        Trigger auto-restart if corrupted.
+        """
+        if self.device == "cpu":
+            return True
+        try:
+            torch.tensor([1.0], device=self.device) + torch.tensor([1.0], device=self.device)
+            return True
+        except RuntimeError as e:
+            if "device-side assert" in str(e) or "CUDA error" in str(e):
+                if self._cuda_healthy:
+                    self._cuda_healthy = False
+                    print("FATAL: CUDA context is corrupted (device-side assert). " "The process must be restarted.")
+                    self._trigger_restart()
+                return False
+            raise
+
+    def _trigger_restart(self) -> None:
+        """Exit the process so the HF Space (or systemd/Docker) can restart it."""
+        import sys
+
+        print("Initiating automatic restart due to unrecoverable CUDA error...")
+        sys.stdout.flush()
+        sys.stderr.flush()
+        os._exit(1)
+
+    def generate(
+        self,
+        client: viser.ClientHandle,
+        prompts: list[str],
+        num_frames: list[int],
+        num_samples: int,
+        seed: int,
+        diffusion_steps: int,
+        cfg_weight: Optional[list[float]] = None,
+        cfg_type: Optional[str] = None,
+        postprocess_parameters: Optional[dict] = None,
+        transitions_parameters: Optional[dict] = None,
+        real_robot_rotations: bool = False,
+    ) -> None:
+        if not self._cuda_healthy:
+            raise RuntimeError("CUDA is in a corrupted state. The space is restarting...")
+
+        locked = self._generation_lock.acquire(blocking=False)
+        if not locked:
+            waiting_notif = client.add_notification(
+                title="Waiting for GPU...",
+                body="Another generation is in progress. Yours will start automatically.",
+                loading=True,
+                with_close_button=False,
+            )
+            self._generation_lock.acquire()
+            waiting_notif.remove()
+
+        try:
+            session = self.client_sessions[client.client_id]
+            model_bundle = self.load_model(session.model_name)
+            generation.generate(
+                client=client,
+                session=session,
+                model_bundle=model_bundle,
+                prompts=prompts,
+                num_frames=num_frames,
+                num_samples=num_samples,
+                seed=seed,
+                diffusion_steps=diffusion_steps,
+                cfg_weight=cfg_weight,
+                cfg_type=cfg_type,
+                postprocess_parameters=postprocess_parameters,
+                transitions_parameters=transitions_parameters,
+                real_robot_rotations=real_robot_rotations,
+                device=self.device,
+                clear_motions=self.clear_motions,
+                add_character_motion=self.add_character_motion,
+            )
+        finally:
+            self._generation_lock.release()
+
+    def set_frame(self, client_id: int, frame_idx: int, update_timeline: bool = True):
+        if not self.client_active(client_id):
+            return
+
+        session = self.client_sessions[client_id]
+
+        session.frame_idx = frame_idx
+        if update_timeline:
+            session.client.timeline.set_current_frame(frame_idx)
+        for motion in list(session.motions.values()):
+            motion.set_frame(frame_idx)
+        self._apply_constraint_overlay_visibility(session)
+
+    def run(self) -> None:
+        update_counter = 0
+        cuda_check_interval = 300
+        while True:
+            last_update_time = time.time()
+            if self.models:
+                # the max playback speed is 2x the model fps (from gui_playback_speed_buttons)
+                playback_fps = max(bundle.model_fps for bundle in self.models.values()) * 2.0
+            else:
+                playback_fps = 60.0
+
+            # update each client session independently
+            #   copy to a list first to avoid changing size if client disconnects
+            for client_id, session in list(self.client_sessions.items()):
+                update_interval = int(playback_fps / (session.playback_speed * session.model_fps))
+                new_frame_idx = session.frame_idx
+                if session.playing and update_counter % update_interval == 0:
+                    if session.frame_idx >= session.max_frame_idx:
+                        new_frame_idx = 0
+                    else:
+                        new_frame_idx = session.frame_idx + 1
+
+                    # make sure the client is still active before updating the frame
+                    if self.client_active(client_id):
+                        self.set_frame(client_id, new_frame_idx)
+
+            if update_counter % cuda_check_interval == 0:
+                self.check_cuda_health()
+
+            time_remaining = max(0, 1.0 / playback_fps - (time.time() - last_update_time))
+            time.sleep(time_remaining)
+            update_counter += 1
+            update_counter %= playback_fps  # wrap around to 0 every second
+
+    def configure_theme(
+        self,
+        client: viser.ClientHandle,
+        dark_mode: bool = False,
+        titlebar_dark_mode_checkbox_uuid: str | None = None,
+    ):
+        # Sync grid color with theme (light vs dark)
+        theme = DARK_THEME if dark_mode else LIGHT_THEME
+        grid_handle = self.grid_handles.get(client.client_id)
+        if grid_handle is not None:
+            grid_handle.section_color = theme["grid"]
+
+        #
+        # setup theme
+        #
+        buttons = (
+            TitlebarButton(
+                text="Documentation",
+                icon="Description",
+                href="https://research.nvidia.com/labs/sil/projects/kimodo/docs/interactive_demo/index.html",
+            ),
+            TitlebarButton(
+                text="Project Page",
+                icon=None,
+                href="https://research.nvidia.com/labs/sil/projects/kimodo/",
+            ),
+            TitlebarButton(
+                text="Github",
+                icon="GitHub",
+                href="https://github.com/nv-tlabs/kimodo",
+            ),
+        )
+        assets_dir = DEMO_ASSETS_ROOT
+        logo_light_path = assets_dir / "nvidia_logo.png"
+        logo_dark_path = assets_dir / "nvidia_logo_dark.png"
+        if logo_light_path.exists():
+            light_b64 = base64.standard_b64encode(logo_light_path.read_bytes()).decode("ascii")
+            dark_b64 = (
+                base64.standard_b64encode(logo_dark_path.read_bytes()).decode("ascii")
+                if logo_dark_path.exists()
+                else None
+            )
+            image = TitlebarImage(
+                image_url_light=f"data:image/png;base64,{light_b64}",
+                image_url_dark=(f"data:image/png;base64,{dark_b64}" if dark_b64 else None),
+                image_alt="NVIDIA",
+                href="https://www.nvidia.com/",
+            )
+        else:
+            image = None
+        titlebar_theme = TitlebarConfig(buttons=buttons, image=image, title_text="Kimodo")
+        client.gui.set_panel_label("Kimodo")
+        client.gui.configure_theme(
+            titlebar_content=titlebar_theme,
+            control_layout="floating",  # "floating",  # ['floating', 'collapsible', 'fixed']
+            control_width="large",  # ['small', 'medium', 'large']
+            dark_mode=dark_mode,
+            show_logo=False,  # hide viser logo on bottom left corner
+            show_share_button=False,
+            titlebar_dark_mode_checkbox_uuid=titlebar_dark_mode_checkbox_uuid,
+            brand_color=(152, 189, 255),  # (60, 131, 0),  # (R, G, B) tuple
+        )

kimodo/demo/config.py

@@ -0,0 +1,163 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+
+from kimodo.assets import DEMO_EXAMPLES_ROOT
+from kimodo.model.registry import (
+    AVAILABLE_MODELS,
+    DEFAULT_MODEL,
+    FRIENDLY_NAMES,
+    get_datasets,
+    get_model_info,
+    get_models_for_dataset_skeleton,
+    get_short_key_from_display_name,
+    get_skeleton_display_name,
+    get_skeleton_display_names_for_dataset,
+    get_skeleton_key_from_display_name,
+    get_skeletons_for_dataset,
+    get_versions_for_dataset_skeleton,
+    resolve_to_short_key,
+)
+
+SERVER_NAME = os.environ.get("SERVER_NAME", "0.0.0.0")
+SERVER_PORT = int(os.environ.get("SERVER_PORT", "7860"))
+HF_MODE = os.environ.get("HF_MODE", False)
+
+# HF mode: user queue and session limit (override via env in Spaces)
+MAX_ACTIVE_USERS = int(os.environ.get("MAX_ACTIVE_USERS", "5"))
+MAX_SESSION_MINUTES = float(os.environ.get("MAX_SESSION_MINUTES", "5.0"))
+
+DEFAULT_PLAYBACK_SPEED = 1.0
+# default start duration is 6.0 sec, but model can handle up to 10 sec
+DEFAULT_CUR_DURATION = 6.0
+DEFAULT_PROMPT = "A person walks forward."
+MIN_DURATION = 2.0
+MAX_DURATION = 10.0
+
+SHOW_TRANSITION_PARAMS = True
+INIT_POSTPROCESSING = True
+NB_TRANSITION_FRAMES = 5
+
+LIGHT_THEME = dict(
+    floor=(220, 220, 220),
+    grid=(180, 180, 180),
+)
+
+# Dark theme: slightly lighter grid and floor for better visibility and less flat black
+DARK_THEME = dict(
+    floor=(48, 48, 52),
+    grid=(105, 105, 110),
+)
+
+EXAMPLES_ROOT_DIR = str(DEMO_EXAMPLES_ROOT)
+
+# Model list and paths from kimodo registry (all models: Kimodo + TMR)
+MODEL_NAMES = tuple(AVAILABLE_MODELS)
+MODEL_EXAMPLES_DIRS = {name: os.path.join(EXAMPLES_ROOT_DIR, name) for name in MODEL_NAMES}
+# Display labels for backward compatibility (short_key -> display name)
+MODEL_LABELS = {name: FRIENDLY_NAMES.get(name, f"Model ({name})") for name in MODEL_NAMES}
+MODEL_LABEL_TO_NAME = {label: name for name, label in MODEL_LABELS.items()}
+
+# -----------------------------------------------------------------------------
+# Demo UI copy
+# -----------------------------------------------------------------------------
+
+DEMO_UI_QUICK_START_CORE_MD = """
+### Camera
+- **Left-drag**: rotate
+- **Right-drag**: pan
+- **Scroll**: zoom
+
+### Playback
+- **Space** to play/pause
+- **←/→** to step frames, or click the frame number.
+- **Scroll up/down** in the timeline: move left/right
+- **Shift + scroll** in the timeline: zoom in/out
+
+### Prompts
+- **Double-click** a text prompt to edit it.
+- **Click and drag** the right edge of a prompt box to extend/shorten it.
+- **Click empty space** to add a prompt.
+- **Right-click** a prompt to delete it.
+
+### Generate
+- Go to the **Generate** tab to modify options
+- It is also possible to **load** examples
+- Click **Generate** to generate a motion
+
+### Constraints
+- This is **optional**: should be use after a first generation
+- **Click** in the timeline tracks (Full-Body / 2D root etc) to add a constraint.
+- **Right-click** on a constraint to delete it.
+- To **edit** a constraint:
+    - Move playback to the target frame
+    - Click **Enter Editing Mode** in the Constraints tab.
+"""
+
+DEMO_UI_QUICK_START_MODAL_MD = (
+    DEMO_UI_QUICK_START_CORE_MD
+    + """
+
+See the **Instructions** tab for the full user manual.
+"""
+)
+
+DEMO_UI_INSTRUCTIONS_TAB_MD = (
+    """
+## How to Use This Demo
+
+"""
+    + DEMO_UI_QUICK_START_CORE_MD
+    + """
+
+---
+
+### Generating Motion (step-by-step)
+
+1. **Edit the text prompts** in the timeline (e.g., "A person walks forward.")
+2. **Modify the duration** by moving the right edge of each prompts (2–10 seconds)
+3. **Add constraints** (optional) to control the motion:
+   - Click **Enter Editing Mode** to adjust the character pose
+   - Use the timeline to place keyframes or intervals in constraint tracks (see below)
+4. **Click Generate** to create the motion
+5. If generating multiple samples, **click on a mesh** to select which one to keep
+
+### Timeline Editing
+
+**Adding Constraints:**
+1. Click anywhere on the timeline to add a keyframe at that frame. The keyframe is created based on the current character motion.
+2. Ctrl/Cmd+click+drag to add an interval constraint, or expand a keyframe into an interval
+3. Enter editing mode with the **Enter Editing Mode** button to adjust character pose before/after adding constraints.
+
+**Constraint Types:**
+- **Full-Body**: constrains the entire character pose
+- **2D Root**: constrains the character's path on the ground plane
+  - Enable **Densify** to create a continuous path
+- **End-Effectors**: constrains hands and feet positions
+  - Use separate tracks for Left/Right Hand/Foot
+
+
+**Moving & Deleting:**
+- **Drag keyframes/intervals** to move them to different frames
+- **Right-click** a keyframe or interval to delete it
+- Use **Clear All Constraints** to remove everything
+
+**Tips:**
+- The posing skeleton becomes visible in editing mode for precise positioning
+- Use **Snap to constraint** to align the current frame to a constraint
+
+### Saving & Loading
+
+You can save the current constraints or current motion to load in later from the Load/Save menu.
+Saving an **Example** will save the full constraints, motion, and generation metadata.
+
+### Visualization Options
+
+Switch to the **Visualize** tab to:
+- Toggle mesh and skeleton visibility
+- Adjust mesh opacity
+- Show/hide foot contact indicators
+- Switch between light and dark modes
+"""
+)

kimodo/demo/embedding_cache.py

@@ -0,0 +1,253 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import contextlib
+import contextvars
+import hashlib
+import json
+import os
+import threading
+import time
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import Iterable, Optional
+
+import numpy as np
+import torch
+
+from kimodo.sanitize import sanitize_texts
+
+_ACTIVE_SESSION = contextvars.ContextVar("kimodo_demo_active_session", default=None)
+
+
+@dataclass
+class CacheStats:
+    hits: int = 0
+    misses: int = 0
+    disk_hits: int = 0
+
+
+class EmbeddingCache:
+    """Disk-backed text embedding cache with a small in-memory LRU."""
+
+    def __init__(
+        self,
+        *,
+        model_name: str,
+        encoder_id: str,
+        base_dir: Optional[str] = None,
+        max_mem_entries: int = 128,
+    ) -> None:
+        cache_root = base_dir or os.environ.get(
+            "kimodo_EMBED_CACHE_DIR",
+            os.path.join("~", ".cache", "kimodo_demo", "embeddings"),
+        )
+        self.base_dir = os.path.expanduser(cache_root)
+        self.model_name = model_name
+        self.encoder_id = encoder_id
+        self.max_mem_entries = max_mem_entries
+        self.stats = CacheStats()
+
+        self._lock = threading.Lock()
+        self._mem_cache: OrderedDict[str, np.ndarray] = OrderedDict()
+        self._index = {}
+        self._index_loaded = False
+
+    def _model_dir(self) -> str:
+        return os.path.join(self.base_dir, self.model_name)
+
+    def _index_path(self) -> str:
+        return os.path.join(self._model_dir(), "index.json")
+
+    def _prewarm_marker_path(self, key: str) -> str:
+        return os.path.join(self._model_dir(), f"prewarm_{key}.json")
+
+    def has_prewarm_marker(self, key: str) -> bool:
+        return os.path.exists(self._prewarm_marker_path(key))
+
+    def write_prewarm_marker(self, key: str, *, prompt_count: int) -> None:
+        os.makedirs(self._model_dir(), exist_ok=True)
+        payload = {"prompt_count": prompt_count, "updated_at": time.time()}
+        tmp_path = f"{self._prewarm_marker_path(key)}.tmp"
+        with open(tmp_path, "w", encoding="utf-8") as f:
+            json.dump(payload, f)
+        os.replace(tmp_path, self._prewarm_marker_path(key))
+
+    def _load_index(self) -> None:
+        if self._index_loaded:
+            return
+        index_path = self._index_path()
+        if os.path.exists(index_path):
+            try:
+                with open(index_path, "r", encoding="utf-8") as f:
+                    self._index = json.load(f)
+            except json.JSONDecodeError:
+                self._index = {}
+        self._index_loaded = True
+
+    def _save_index(self) -> None:
+        os.makedirs(self._model_dir(), exist_ok=True)
+        tmp_path = f"{self._index_path()}.tmp"
+        with open(tmp_path, "w", encoding="utf-8") as f:
+            json.dump(self._index, f)
+        os.replace(tmp_path, self._index_path())
+
+    def _make_key(self, text: str) -> str:
+        key_src = f"{self.model_name}|{self.encoder_id}|{text}"
+        return hashlib.sha256(key_src.encode("utf-8")).hexdigest()
+
+    def _entry_path(self, key: str) -> str:
+        return os.path.join(self._model_dir(), f"{key}.npy")
+
+    def _mem_get(self, key: str) -> Optional[np.ndarray]:
+        if key in self._mem_cache:
+            self._mem_cache.move_to_end(key)
+            return self._mem_cache[key]
+        return None
+
+    def _mem_put(self, key: str, value: np.ndarray) -> None:
+        self._mem_cache[key] = value
+        self._mem_cache.move_to_end(key)
+        while len(self._mem_cache) > self.max_mem_entries:
+            self._mem_cache.popitem(last=False)
+
+    def _disk_load(self, key: str) -> Optional[np.ndarray]:
+        path = self._entry_path(key)
+        if not os.path.exists(path):
+            return None
+        try:
+            return np.load(path)
+        except Exception:
+            return None
+
+    def _disk_save(self, key: str, value: np.ndarray) -> None:
+        os.makedirs(self._model_dir(), exist_ok=True)
+        np.save(self._entry_path(key), value)
+        self._index[key] = {
+            "length": int(value.shape[0]),
+            "dtype": str(value.dtype),
+            "updated_at": time.time(),
+        }
+
+    def _maybe_use_session_cache(self, texts: list[str]):
+        session = _ACTIVE_SESSION.get()
+        if session is None:
+            return None
+        if session.last_prompt_texts == texts and session.last_prompt_embeddings is not None:
+            return session.last_prompt_embeddings, session.last_prompt_lengths
+        return None
+
+    def _update_session_cache(self, texts: list[str], tensor: torch.Tensor, lengths: list[int]) -> None:
+        session = _ACTIVE_SESSION.get()
+        if session is None:
+            return
+        session.last_prompt_texts = texts
+        session.last_prompt_embeddings = tensor
+        session.last_prompt_lengths = lengths
+
+    def get_or_encode(self, texts: Iterable[str], encoder):
+        if isinstance(texts, str):
+            texts = [texts]
+        texts = sanitize_texts(list(texts))
+        if len(texts) == 0:
+            empty = torch.empty()
+            return empty, []
+
+        session_cache = self._maybe_use_session_cache(texts)
+        if session_cache is not None:
+            return session_cache
+
+        arrays: list[Optional[np.ndarray]] = [None] * len(texts)
+        lengths: list[int] = [0] * len(texts)
+        misses: list[tuple[int, str, str]] = []
+
+        with self._lock:
+            self._load_index()
+            for idx, text in enumerate(texts):
+                key = self._make_key(text)
+                cached = self._mem_get(key)
+                if cached is not None:
+                    arrays[idx] = cached
+                    lengths[idx] = cached.shape[0]
+                    self.stats.hits += 1
+                    continue
+
+                cached = self._disk_load(key)
+                if cached is not None:
+                    arrays[idx] = cached
+                    lengths[idx] = cached.shape[0]
+                    self._mem_put(key, cached)
+                    self.stats.disk_hits += 1
+                    continue
+
+                misses.append((idx, text, key))
+                self.stats.misses += 1
+
+        if misses:
+            miss_texts = [text for _, text, _ in misses]
+            miss_tensor, miss_lengths = encoder(miss_texts)
+            miss_tensor = miss_tensor.detach().cpu()
+            miss_tensor_np = miss_tensor.numpy()
+
+            with self._lock:
+                self._load_index()
+                for miss_idx, length in enumerate(miss_lengths):
+                    idx, _text, key = misses[miss_idx]
+                    arr = miss_tensor_np[miss_idx, :length].copy()
+                    arrays[idx] = arr
+                    lengths[idx] = int(length)
+                    self._mem_put(key, arr)
+                    self._disk_save(key, arr)
+                self._save_index()
+
+        max_len = max(lengths) if lengths else 0
+        feat_dim = arrays[0].shape[-1] if arrays[0] is not None else 0
+        dtype = arrays[0].dtype if arrays[0] is not None else np.float32
+        padded = np.zeros((len(texts), max_len, feat_dim), dtype=dtype)
+        for idx, arr in enumerate(arrays):
+            if arr is None:
+                continue
+            padded[idx, : arr.shape[0]] = arr
+
+        result = torch.from_numpy(padded)
+        self._update_session_cache(texts, result, lengths)
+        return result, lengths
+
+
+class CachedTextEncoder:
+    """Wrapper around a text encoder to add disk-backed caching."""
+
+    def __init__(self, encoder, *, model_name: str, base_dir: Optional[str] = None):
+        self.encoder = encoder
+        self.model_name = model_name
+        encoder_id = f"{type(encoder).__name__}"
+        self.cache = EmbeddingCache(model_name=model_name, encoder_id=encoder_id, base_dir=base_dir)
+
+    def __call__(self, texts):
+        return self.cache.get_or_encode(texts, self.encoder)
+
+    def prewarm(self, texts) -> None:
+        if isinstance(texts, str):
+            texts = [texts]
+        texts = sanitize_texts(list(texts))
+        prewarm_key = hashlib.sha256("|".join(texts).encode("utf-8")).hexdigest()
+        if self.cache.has_prewarm_marker(prewarm_key):
+            return
+        self.cache.get_or_encode(texts, self.encoder)
+        self.cache.write_prewarm_marker(prewarm_key, prompt_count=len(texts))
+
+    def to(self, device=None, dtype=None):
+        if hasattr(self.encoder, "to"):
+            self.encoder.to(device=device, dtype=dtype)
+        return self
+
+    @contextlib.contextmanager
+    def session_context(self, session):
+        token = _ACTIVE_SESSION.set(session)
+        try:
+            yield
+        finally:
+            _ACTIVE_SESSION.reset(token)
+
+    def __getattr__(self, name):
+        return getattr(self.encoder, name)

kimodo/demo/generation.py

@@ -0,0 +1,218 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from collections import defaultdict
+from typing import Optional
+
+import numpy as np
+import torch
+
+import viser
+from kimodo.constraints import (
+    TYPE_TO_CLASS,
+    FullBodyConstraintSet,
+    Root2DConstraintSet,
+)
+from kimodo.exports.mujoco import apply_g1_real_robot_projection
+from kimodo.skeleton import G1Skeleton34, SOMASkeleton30
+from kimodo.tools import seed_everything
+
+from .embedding_cache import CachedTextEncoder
+from .state import ClientSession, ModelBundle
+
+
+def compute_model_constraints_lst(
+    session: ClientSession,
+    model_bundle: ModelBundle,
+    num_frames: int,
+    device: str,
+):
+    """Compute the lst of constraints for the model based on the constraints in viser."""
+    assert len(session.motions) == 1, "Only one motion allowed for constrained generation"
+    if not session.constraints:
+        return []
+
+    model_skeleton = model_bundle.model.skeleton
+    # For SOMA, UI uses somaskel77; extract 30-joint subset for the model
+    use_skel_slice = isinstance(model_skeleton, SOMASkeleton30) and session.skeleton.nbjoints != model_skeleton.nbjoints
+    skel_slice = model_skeleton.get_skel_slice(session.skeleton) if use_skel_slice else None
+
+    dense_smooth_root_pos_2d = None
+    if session.constraints["2D Root"].dense_path:
+        # get the full 2d root
+        dense_smooth_root_pos_2d = session.constraints["2D Root"].get_constraint_info(device=device)["root_pos"][
+            :, [0, 2]
+        ]
+
+    model_constraints = []
+    for track_name, constraint in session.constraints.items():
+        constraint_info = constraint.get_constraint_info(device=device)
+        frame_idx = constraint_info["frame_idx"]
+        # drop any constraints outside the generation range
+        valid_info = [(i, fi) for i, fi in enumerate(frame_idx) if fi < num_frames]
+        valid_idx = [i for i, _ in valid_info]
+        valid_frame_idx = [fi for _, fi in valid_info]
+
+        if len(valid_frame_idx) == 0:
+            continue
+
+        frame_indices = torch.tensor(valid_frame_idx)
+        if track_name == "2D Root":
+            smooth_root_pos_2d = constraint_info["root_pos"][valid_idx][:, [0, 2]].to(device)
+            # same as "smooth_root_2d"
+            model_constraints.append(
+                Root2DConstraintSet(
+                    model_skeleton,
+                    frame_indices,
+                    smooth_root_pos_2d,
+                )
+            )
+        elif track_name == "Full-Body":
+            constraint_joints_pos = constraint_info["joints_pos"][valid_idx].to(device)
+            constraint_joints_rot = constraint_info["joints_rot"][valid_idx].to(device)
+            if skel_slice is not None:
+                constraint_joints_pos = constraint_joints_pos[:, skel_slice]
+                constraint_joints_rot = constraint_joints_rot[:, skel_slice]
+
+            smooth_root_pos_2d = None
+            if dense_smooth_root_pos_2d is not None:
+                smooth_root_pos_2d = dense_smooth_root_pos_2d[frame_indices]
+
+            model_constraints.append(
+                FullBodyConstraintSet(
+                    model_skeleton,
+                    frame_indices,
+                    constraint_joints_pos,
+                    constraint_joints_rot,
+                    smooth_root_2d=smooth_root_pos_2d,
+                )
+            )
+        elif track_name == "End-Effectors":
+            constraint_joints_pos = constraint_info["joints_pos"][valid_idx].to(device)
+            constraint_joints_rot = constraint_info["joints_rot"][valid_idx].to(device)
+            if skel_slice is not None:
+                constraint_joints_pos = constraint_joints_pos[:, skel_slice]
+                constraint_joints_rot = constraint_joints_rot[:, skel_slice]
+
+            end_effector_type_set_lst = [
+                end_effector_type_set
+                for i, end_effector_type_set in enumerate(constraint_info["end_effector_type"])
+                if i in valid_idx
+            ]
+
+            # regroup the end effector data by type
+            cls_idx = defaultdict(list)
+            for idx, end_effector_type_set in enumerate(end_effector_type_set_lst):
+                for end_effector_type in end_effector_type_set:
+                    cls_idx[TYPE_TO_CLASS[end_effector_type]].append(idx)
+
+            for cls, lst_idx in cls_idx.items():
+                frame_indices_cls = frame_indices[lst_idx]
+                smooth_root_pos_2d = None
+                if dense_smooth_root_pos_2d is not None:
+                    smooth_root_pos_2d = dense_smooth_root_pos_2d[frame_indices_cls]
+
+                constraint_joints_pos_el = constraint_joints_pos[lst_idx]
+                constraint_joints_rot_el = constraint_joints_rot[lst_idx]
+
+                model_constraints.append(
+                    cls(
+                        model_skeleton,
+                        frame_indices_cls,
+                        constraint_joints_pos_el,
+                        constraint_joints_rot_el,
+                        smooth_root_2d=smooth_root_pos_2d,
+                    )
+                )
+        else:
+            raise ValueError(f"Unsupported constraint type: {constraint.display_name}")
+    return model_constraints
+
+
+def generate(
+    *,
+    client: viser.ClientHandle,
+    session: ClientSession,
+    model_bundle: ModelBundle,
+    prompts: list[str],
+    num_frames: list[int],
+    num_samples: int,
+    seed: int,
+    diffusion_steps: int,
+    cfg_weight: Optional[list[float]] = None,
+    cfg_type: Optional[str] = None,
+    postprocess_parameters: Optional[dict] = None,
+    transitions_parameters: Optional[dict] = None,
+    real_robot_rotations: bool = False,
+    device: str,
+    clear_motions,
+    add_character_motion,
+) -> None:
+    client_id = client.client_id
+    print(
+        f"Generating {num_samples} samples for a total of {sum(num_frames)} frames with those prompt: {prompts} (client {client_id})"
+    )
+
+    seed_everything(seed)
+
+    model_constraints = compute_model_constraints_lst(session, model_bundle, sum(num_frames), device)
+    cfg_weight = cfg_weight or [2.0, 2.0]
+    postprocess_parameters = postprocess_parameters or {}
+    transitions_parameters = transitions_parameters or {}
+
+    encoder = getattr(model_bundle.model, "text_encoder", None)
+    if isinstance(encoder, CachedTextEncoder):
+        with encoder.session_context(session):
+            pred_joints_output = model_bundle.model(
+                prompts,
+                num_frames,
+                diffusion_steps,
+                multi_prompt=True,
+                constraint_lst=model_constraints,
+                cfg_weight=cfg_weight,
+                num_samples=num_samples,
+                cfg_type=cfg_type,
+                **(postprocess_parameters | transitions_parameters),
+            )  # [B, T, motion_rep_dim]
+    else:
+        pred_joints_output = model_bundle.model(
+            prompts,
+            num_frames,
+            diffusion_steps,
+            multi_prompt=True,
+            constraint_lst=model_constraints,
+            cfg_weight=cfg_weight,
+            num_samples=num_samples,
+            cfg_type=cfg_type,
+            **(postprocess_parameters | transitions_parameters),
+        )  # [B, T, motion_rep_dim]
+
+    joints_pos = pred_joints_output["posed_joints"]  # [B, T, J, 3]
+    joints_rot = pred_joints_output["global_rot_mats"]
+    foot_contacts = pred_joints_output.get("foot_contacts")
+
+    # Optionally project G1 to real robot DoF (1-DoF per joint, clamped) for display.
+    if real_robot_rotations and isinstance(session.skeleton, G1Skeleton34):
+        joints_pos, joints_rot = apply_g1_real_robot_projection(
+            session.skeleton,
+            pred_joints_output["posed_joints"],
+            pred_joints_output["global_rot_mats"],
+            clamp_to_limits=True,
+        )
+
+    # Display on characters (callbacks keep this module UI-agnostic).
+    clear_motions(client_id)
+    # Keep one sample centered at the origin so constraints align.
+    spread_factor = 1.0  # meters
+    center_idx = num_samples // 2
+    x_trans = (np.arange(num_samples) - center_idx) * spread_factor
+    for i in range(num_samples):
+        cur_joints_pos = joints_pos[i]
+        cur_joints_pos[..., 0] += x_trans[i]
+        add_character_motion(
+            client,
+            session.skeleton,
+            cur_joints_pos,
+            joints_rot[i],
+            foot_contacts[i],
+        )

kimodo/demo/queue_manager.py

@@ -0,0 +1,336 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""HF mode user queue and session time limit."""
+
+import math
+import threading
+import time
+from collections.abc import Callable
+from typing import Any
+
+import viser
+
+from .config import DEMO_UI_QUICK_START_MODAL_MD, MAX_SESSION_MINUTES
+
+# Link for "Duplicate this Space" on Hugging Face (used in queue and expiry modals).
+DUPLICATE_SPACE_URL = "https://huggingface.co/spaces/nvidia/Kimodo?duplicate=true"
+GITHUB_REPO_URL = "https://github.com/nv-tlabs/kimodo"
+
+# How often to refresh queue modal content (position, total, estimated wait).
+QUEUE_MODAL_REFRESH_INTERVAL_SEC = 15
+
+
+class UserQueue:
+    """Thread-safe queue: active users (with activation timestamp) and waiting queue."""
+
+    def __init__(self, max_active: int, max_minutes: float) -> None:
+        self._max_active = max_active
+        self._max_minutes = max_minutes
+        self._max_seconds = max_minutes * 60.0
+        self._active: dict[int, float] = {}  # client_id -> activation timestamp
+        self._queued: list[int] = []
+        self._lock = threading.Lock()
+
+    def try_activate(self, client_id: int) -> bool:
+        """If a slot is free, add client as active and return True.
+
+        Else return False.
+        """
+        with self._lock:
+            if len(self._active) < self._max_active:
+                self._active[client_id] = time.time()
+                return True
+            return False
+
+    def enqueue(self, client_id: int) -> None:
+        with self._lock:
+            if client_id not in self._queued:
+                self._queued.append(client_id)
+
+    def remove(self, client_id: int) -> bool:
+        """Remove from active or queue.
+
+        Returns True if was active.
+        """
+        with self._lock:
+            was_active = client_id in self._active
+            self._active.pop(client_id, None)
+            if client_id in self._queued:
+                self._queued.remove(client_id)
+            return was_active
+
+    def promote_next(self) -> int | None:
+        """If queue non-empty, pop first, activate them, return their client_id.
+
+        Else None.
+        """
+        with self._lock:
+            if not self._queued:
+                return None
+            client_id = self._queued.pop(0)
+            self._active[client_id] = time.time()
+            return client_id
+
+    def get_queue_position(self, client_id: int) -> tuple[int, int] | None:
+        """(1-based position, total_in_queue) or None if not queued."""
+        with self._lock:
+            if client_id not in self._queued:
+                return None
+            pos = self._queued.index(client_id)
+            return (pos + 1, len(self._queued))
+
+    def get_estimated_wait_seconds(self, client_id: int) -> float:
+        """Estimated seconds until this queued client gets a slot."""
+        with self._lock:
+            if client_id not in self._queued:
+                return 0.0
+            pos = self._queued.index(client_id) + 1  # 1-based
+            # Expiry times of active users (when they free a slot)
+            now = time.time()
+            expiries = sorted(now + self._max_seconds - (now - t) for t in self._active.values())
+            if not expiries:
+                return 0.0
+            # Nth slot to free (1-indexed) wraps over expiries
+            idx = (pos - 1) % len(expiries)
+            cycles = (pos - 1) // len(expiries)
+            slot_free_time = expiries[idx] + cycles * self._max_seconds
+            return max(0.0, slot_free_time - now)
+
+    def is_active(self, client_id: int) -> bool:
+        with self._lock:
+            return client_id in self._active
+
+    def was_active(self, client_id: int) -> bool:
+        """True if client is currently active (for use when already holding lock)."""
+        return client_id in self._active
+
+
+def _format_wait(seconds: float) -> str:
+    if seconds < 60:
+        return "less than a minute"
+    mins = int(math.ceil(seconds / 60))
+    return f"~{mins} minute{'s' if mins != 1 else ''}"
+
+
+def _queue_modal_markdown(position: int, total: int, estimated_wait_sec: float) -> str:
+    wait_str = _format_wait(estimated_wait_sec)
+    mins = int(MAX_SESSION_MINUTES) if MAX_SESSION_MINUTES == int(MAX_SESSION_MINUTES) else MAX_SESSION_MINUTES
+    return f"""## Kimodo Demo — Please Wait
+
+This demo runs with limited capacity.
+Each user gets **{mins} minute{"s" if mins != 1 else ""}** of interactive time.
+
+**Your position in queue:** {position} / {total}
+
+**Estimated wait:** {wait_str}
+
+Please keep this tab open — the demo will start automatically when it's your turn.
+
+---
+*Want unlimited access? [Duplicate this Space]({DUPLICATE_SPACE_URL}) or clone the [GitHub repo]({GITHUB_REPO_URL}) to run locally!*
+"""
+
+
+def _welcome_modal_markdown() -> str:
+    mins = int(MAX_SESSION_MINUTES) if MAX_SESSION_MINUTES == int(MAX_SESSION_MINUTES) else MAX_SESSION_MINUTES
+    return f"""## Welcome to Kimodo Demo
+
+You have been granted a **{mins}-minute** demo session.
+Your session timer has started.
+
+Click the button below to begin!
+"""
+
+
+def _expiry_modal_markdown() -> str:
+    mins = int(MAX_SESSION_MINUTES) if MAX_SESSION_MINUTES == int(MAX_SESSION_MINUTES) else MAX_SESSION_MINUTES
+    return f"""## Session Expired
+
+Your {mins}-minute demo session has ended.
+Thank you for trying Kimodo!
+
+Refresh this page to rejoin the queue, or [duplicate this Space]({DUPLICATE_SPACE_URL}) for unlimited access.
+"""
+
+
+class QueueManager:
+    """Orchestrates HF mode: queue modals, welcome modal, session timer, promotion."""
+
+    def __init__(
+        self,
+        queue: UserQueue,
+        server: viser.ViserServer,
+        setup_demo_for_client: Callable[[viser.ClientHandle], None],
+        cleanup_session: Callable[[int], None],
+    ) -> None:
+        self._queue = queue
+        self._server = server
+        self._setup_demo_for_client = setup_demo_for_client
+        self._cleanup_session = cleanup_session
+        self._max_seconds = queue._max_seconds
+
+        self._queue_modal_handles: dict[int, tuple[Any, Any]] = {}
+        self._welcome_modal_handles: dict[int, Any] = {}
+        self._expiry_timers: dict[int, threading.Timer] = {}
+        self._lock = threading.Lock()
+        self._refresh_stop = threading.Event()
+        self._refresh_thread = threading.Thread(
+            target=self._queue_modal_refresh_loop,
+            name="queue-modal-refresh",
+            daemon=True,
+        )
+        self._refresh_thread.start()
+
+    def _queue_modal_refresh_loop(self) -> None:
+        """Periodically refresh queue modals so position, total, and estimated wait stay current."""
+        while not self._refresh_stop.wait(timeout=QUEUE_MODAL_REFRESH_INTERVAL_SEC):
+            self._update_all_queue_modals()
+
+    def on_client_connect(self, client: viser.ClientHandle) -> None:
+        """Handle new connection: activate if slot free, else enqueue and show queue modal."""
+        client_id = client.client_id
+        if self._queue.try_activate(client_id):
+            try:
+                self._setup_demo_for_client(client)
+            except RuntimeError as e:
+                if "CUDA error" in str(e):
+                    print(f"CUDA error while setting up client {client_id}: {e}")
+                    return
+                raise
+            self._start_session_timer(client_id)
+            self._show_welcome_modal(client)
+        else:
+            self._queue.enqueue(client_id)
+            self._show_queue_modal(client)
+            self._update_all_queue_modals()
+
+    def on_client_disconnect(self, client_id: int) -> None:
+        """Remove from queue/active, cancel timer, promote next if was active.
+
+        Session/scene cleanup is done by the demo's on_client_disconnect.
+        """
+        with self._lock:
+            self._expiry_timers.pop(client_id, None)
+            self._queue_modal_handles.pop(client_id, None)
+            self._welcome_modal_handles.pop(client_id, None)
+        was_active = self._queue.remove(client_id)
+        if was_active:
+            self._promote_next_user()
+        else:
+            self._update_all_queue_modals()
+
+    def _show_queue_modal(self, client: viser.ClientHandle) -> None:
+        client_id = client.client_id
+        pos, total = self._queue.get_queue_position(client_id) or (0, 0)
+        wait_sec = self._queue.get_estimated_wait_seconds(client_id)
+        md_content = _queue_modal_markdown(pos, total, wait_sec)
+
+        modal = client.gui.add_modal(
+            "Kimodo Demo — Please Wait",
+            size="xl",
+            show_close_button=False,
+        )
+        with modal:
+            md_handle = client.gui.add_markdown(md_content)
+        with self._lock:
+            self._queue_modal_handles[client_id] = (modal, md_handle)
+
+    def _show_quick_start_modal(self, client: viser.ClientHandle) -> None:
+        """Show the quick start instructions modal (same as non-HF mode)."""
+        with client.gui.add_modal(
+            "Welcome — Quick Start",
+            size="xl",
+            show_close_button=True,
+            save_choice="kimodo.demo.quick_start_ack",
+        ) as quick_start_modal:
+            client.gui.add_markdown(DEMO_UI_QUICK_START_MODAL_MD)
+            client.gui.add_button("Got it (don't remind me again)").on_click(lambda _: quick_start_modal.close())
+
+    def _show_welcome_modal(self, client: viser.ClientHandle) -> None:
+        client_id = client.client_id
+
+        def _on_start_demo(_: Any) -> None:
+            modal.close()
+            self._show_quick_start_modal(client)
+
+        modal = client.gui.add_modal(
+            "Welcome to Kimodo Demo",
+            size="xl",
+            show_close_button=True,
+        )
+        with modal:
+            client.gui.add_markdown(_welcome_modal_markdown())
+            client.gui.add_button("Start Demo").on_click(_on_start_demo)
+        with self._lock:
+            self._welcome_modal_handles[client_id] = modal
+
+    def _update_all_queue_modals(self) -> None:
+        with self._lock:
+            handles = list(self._queue_modal_handles.items())
+        for client_id, (modal, md_handle) in handles:
+            pos_total = self._queue.get_queue_position(client_id)
+            if pos_total is None:
+                continue
+            pos, total = pos_total
+            wait_sec = self._queue.get_estimated_wait_seconds(client_id)
+            try:
+                md_handle.content = _queue_modal_markdown(pos, total, wait_sec)
+            except Exception:
+                pass
+
+    def _promote_next_user(self) -> None:
+        promoted_id = self._queue.promote_next()
+        if promoted_id is None:
+            return
+        clients = self._server.get_clients()
+        client = clients.get(promoted_id)
+        if client is None:
+            return
+        with self._lock:
+            old = self._queue_modal_handles.pop(promoted_id, None)
+        if old is not None:
+            try:
+                old[0].close()
+            except Exception:
+                pass
+        try:
+            self._setup_demo_for_client(client)
+        except RuntimeError as e:
+            if "CUDA error" in str(e):
+                print(f"CUDA error while setting up client {promoted_id}: {e}")
+                return
+            raise
+        self._start_session_timer(promoted_id)
+        self._show_welcome_modal(client)
+        self._update_all_queue_modals()
+
+    def _start_session_timer(self, client_id: int) -> None:
+        def on_expiry() -> None:
+            self._on_session_expired(client_id)
+
+        t = threading.Timer(self._max_seconds, on_expiry)
+        t.daemon = True
+        with self._lock:
+            self._expiry_timers[client_id] = t
+        t.start()
+
+    def _on_session_expired(self, client_id: int) -> None:
+        with self._lock:
+            self._expiry_timers.pop(client_id, None)
+        if not self._queue.is_active(client_id):
+            return
+        self._queue.remove(client_id)
+        clients = self._server.get_clients()
+        client = clients.get(client_id)
+        if client is not None:
+            try:
+                with client.gui.add_modal(
+                    "Session Expired",
+                    size="lg",
+                    show_close_button=False,
+                ) as modal_ctx:
+                    client.gui.add_markdown(_expiry_modal_markdown())
+            except Exception:
+                pass
+        self._cleanup_session(client_id)
+        self._promote_next_user()

kimodo/demo/state.py

@@ -0,0 +1,59 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from dataclasses import dataclass, field
+from typing import Optional
+
+import torch
+
+import kimodo.viz.viser_utils as viser_utils
+import viser
+from kimodo.skeleton import SkeletonBase
+from kimodo.viz.viser_utils import GuiElements
+
+from .config import (
+    DEFAULT_CUR_DURATION,
+    DEFAULT_MODEL,
+    DEFAULT_PLAYBACK_SPEED,
+)
+
+
+@dataclass(frozen=True)
+class ModelBundle:
+    model: object
+    motion_rep: object
+    skeleton: SkeletonBase
+    model_fps: float
+
+
+@dataclass
+class ClientSession:
+    """Per-client session data."""
+
+    client: viser.ClientHandle
+    gui_elements: GuiElements
+    motions: dict  # character_name -> CharacterMotion
+    constraints: dict[str, viser_utils.ConstraintSet] = field(default_factory=dict)
+    timeline_data: object = None
+    frame_idx: int = 0
+    playing: bool = False
+    playback_speed: float = DEFAULT_PLAYBACK_SPEED
+    cur_duration: float = DEFAULT_CUR_DURATION
+    max_frame_idx: int = 100  # will be updated based on model_fps
+    updating_motions: bool = False
+    edit_mode: bool = False
+    model_name: str = DEFAULT_MODEL
+    model_fps: float = 0.0
+    skeleton: SkeletonBase | None = None
+    motion_rep: object | None = None
+    examples_base_dir: str = ""
+    example_dict: dict[str, str] = field(default_factory=dict)
+    gui_examples_dropdown: Optional[viser.GuiInputHandle] = None
+    gui_save_example_path_text: Optional[viser.GuiInputHandle] = None
+    gui_model_selector: Optional[viser.GuiInputHandle] = None
+    last_prompt_texts: Optional[list[str]] = None
+    last_prompt_embeddings: Optional[torch.Tensor] = None
+    last_prompt_lengths: Optional[list[int]] = None
+    edit_mode_snapshot: Optional[dict[int, dict[str, object]]] = None
+    undo_drag_snapshot: Optional[dict[str, object]] = None
+    show_only_current_constraint: bool = False  # False = Show All, True = Show only Current

kimodo/exports/__init__.py

@@ -0,0 +1,65 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Export utilities: MuJoCo, BVH, SMPLX/AMASS, and motion I/O helpers."""
+
+from .bvh import bvh_to_kimodo_motion, motion_to_bvh_bytes, read_bvh_frame_time_seconds, save_motion_bvh
+from .motion_convert_lib import convert_motion_files
+from .motion_formats import (
+    infer_npz_kind,
+    infer_source_format_from_path,
+    infer_target_format_from_path,
+    resolve_source_fps,
+)
+from .motion_io import (
+    KIMODO_CONVERT_TARGET_FPS,
+    amass_npz_to_bytes,
+    complete_motion_dict,
+    g1_csv_to_bytes,
+    kimodo_npz_to_bytes,
+    load_amass_npz,
+    load_g1_csv,
+    load_kimodo_npz,
+    load_kimodo_npz_as_torch,
+    load_motion_file,
+    motion_dict_to_numpy,
+    save_kimodo_npz,
+    save_kimodo_npz_at_target_fps,
+)
+from .mujoco import MujocoQposConverter, apply_g1_real_robot_projection
+from .smplx import (
+    AMASSConverter,
+    amass_npz_to_kimodo_motion,
+    get_amass_parameters,
+    kimodo_y_up_to_amass_coord_rotation_matrix,
+)
+
+__all__ = [
+    "AMASSConverter",
+    "KIMODO_CONVERT_TARGET_FPS",
+    "MujocoQposConverter",
+    "amass_npz_to_bytes",
+    "amass_npz_to_kimodo_motion",
+    "apply_g1_real_robot_projection",
+    "bvh_to_kimodo_motion",
+    "complete_motion_dict",
+    "convert_motion_files",
+    "g1_csv_to_bytes",
+    "get_amass_parameters",
+    "infer_npz_kind",
+    "infer_source_format_from_path",
+    "infer_target_format_from_path",
+    "kimodo_npz_to_bytes",
+    "kimodo_y_up_to_amass_coord_rotation_matrix",
+    "load_amass_npz",
+    "load_g1_csv",
+    "load_kimodo_npz",
+    "load_kimodo_npz_as_torch",
+    "load_motion_file",
+    "motion_dict_to_numpy",
+    "motion_to_bvh_bytes",
+    "read_bvh_frame_time_seconds",
+    "resolve_source_fps",
+    "save_kimodo_npz",
+    "save_kimodo_npz_at_target_fps",
+    "save_motion_bvh",
+]

kimodo/exports/bvh.py

@@ -0,0 +1,298 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Export utilities for converting internal motion representations into common file formats.
+
+This module is intended to hold lightweight serialization / export helpers that can be reused
+outside of interactive demos.
+"""
+
+import os
+import tempfile
+from pathlib import Path
+from typing import Tuple, Union
+
+import numpy as np
+import torch
+
+from kimodo.geometry import matrix_to_quaternion as _matrix_to_quaternion
+
+
+def _strip_end_site_blocks(bvh_text: str) -> str:
+    """Remove all 'End Site { ... }' blocks from BVH text so output matches original format.
+
+    bvhio adds an End Site for every leaf joint when writing; we do not set EndSite on joints, so we
+    post-process the string to remove these blocks for Blender/original compatibility.
+    """
+    lines = bvh_text.splitlines(keepends=True)
+    result = []
+    i = 0
+    while i < len(lines):
+        line = lines[i]
+        if "End Site" in line:
+            # Skip this line and the following block { ... }; brace-count to find closing }
+            i += 1
+            if i < len(lines) and "{" in lines[i]:
+                i += 1
+                depth = 1
+                while i < len(lines) and depth > 0:
+                    if "{" in lines[i]:
+                        depth += 1
+                    if "}" in lines[i]:
+                        depth -= 1
+                    i += 1
+            continue
+        result.append(line)
+        i += 1
+    return "".join(result)
+
+
+def _coerce_batch(name: str, x: torch.Tensor, *, expected_ndim: int) -> torch.Tensor:
+    """Coerce (T, ...) or (1, T, ...) into (T, ...)."""
+    if x.ndim == expected_ndim:
+        return x
+    if x.ndim == expected_ndim + 1:
+        if int(x.shape[0]) != 1:
+            raise ValueError(
+                f"{name} has batch dimension B={int(x.shape[0])}, but BVH export " "only supports a single clip (B==1)."
+            )
+        return x[0]
+    raise ValueError(f"{name} must have shape (T, ...) or (1, T, ...); got {tuple(x.shape)}")
+
+
+def motion_to_bvh(
+    local_rot_mats: torch.Tensor,
+    root_positions: torch.Tensor,
+    *,
+    skeleton,
+    fps: float,
+    standard_tpose: bool = False,
+) -> str:
+    """Convert local rotations and root positions to BVH format; return UTF-8 string.
+
+    Args:
+        local_rot_mats: (T, J, 3, 3) or (1, T, J, 3, 3) local rotation matrices.
+        root_positions: (T, 3) or (1, T, 3) root joint positions (e.g. from posed joints).
+        skeleton: Skeleton with bone_order_names, bvh_neutral_joints, etc.
+        fps: Frames per second for the motion.
+        standard_tpose: If True, export with the rest pose being the standard T-pose rather than the rest pose consistent with the BONES-SEED dataset.
+    Notes:
+        BVH is plain-text. Root is named "Root" with ZYX rotation order; leaf joints
+        have no End Site block.
+    """
+    try:
+        import bvhio  # type: ignore[import-not-found]
+        import glm  # type: ignore[import-not-found]
+        from SpatialTransform import Pose  # type: ignore[import-not-found]
+    except Exception as e:  # pragma: no cover
+        raise ImportError(
+            "BVH export requires `bvhio` (and its deps `PyGLM` + `SpatialTransform`). "
+            "Install with: `pip install bvhio`."
+        ) from e
+
+    local_rot_mats = local_rot_mats.detach()
+    root_positions = root_positions.detach()
+    # SOMA: accept either somaskel30 (convert to 77) or somaskel77 (use as-is)
+    if skeleton.name == "somaskel30":
+        local_rot_mats = skeleton.to_SOMASkeleton77(local_rot_mats)
+        skeleton = skeleton.somaskel77
+
+    if standard_tpose:
+        neutral = skeleton.neutral_joints.detach().cpu().numpy()
+    else:
+        # transform local rots to the original rest pose consistent with the BONES-SEED dataset
+        local_rot_mats, _ = skeleton.from_standard_tpose(local_rot_mats)
+        neutral = skeleton.bvh_neutral_joints.detach().cpu().numpy()
+
+    joint_names = list(skeleton.bone_order_names)
+    parents = skeleton.joint_parents.detach().cpu().numpy().astype(int)
+    root_idx = int(skeleton.root_idx)
+
+    local_rot_mats = _coerce_batch("local_rot_mats", local_rot_mats, expected_ndim=4)
+    T, J = local_rot_mats.shape[:2]
+    q_wxyz = _matrix_to_quaternion(local_rot_mats).detach().cpu().numpy()  # [T, J, 4]
+
+    root_xyz = _coerce_batch("root_positions", root_positions, expected_ndim=2)
+    root_xyz = root_xyz.cpu().numpy()  # [T, 3]
+
+    # Build BVH hierarchy: Root (wrapper at origin) -> Hips (pelvis with offset in meters) -> ...
+    # Offsets are in meters to match the original format.
+    children: dict[int, list[int]] = {i: [] for i in range(J)}
+    for i, p in enumerate(parents):
+        if p >= 0:
+            children[int(p)].append(int(i))
+
+    _ROOT_CHANNELS = [
+        "Xposition",
+        "Yposition",
+        "Zposition",
+        "Zrotation",
+        "Yrotation",
+        "Xrotation",
+    ]
+    _JOINT_CHANNELS = ["Zrotation", "Yrotation", "Xrotation"]
+
+    # Scale from meters to centimeters (match original SEED data BVH scale).
+    neutral = neutral * 100
+    root_xyz = root_xyz * 100
+
+    # Hips offset from Root: use skeleton neutral; if root is at origin (zeros), use a
+    # nominal pelvis height so the hierarchy is non-degenerate in Blender.
+    hips_offset = neutral[root_idx]
+    if (hips_offset == 0).all():
+        hips_offset = np.array([0.0, 100.0, 0.0], dtype=neutral.dtype)  # 1 m in cm
+
+    def _make_joint(i: int) -> "bvhio.BvhJoint":
+        name = joint_names[i]
+        j = bvhio.BvhJoint(name, offset=glm.vec3(0, 0, 0))
+        if i == root_idx:
+            # Hips: offset from Root (origin) in cm
+            off = hips_offset
+            j.Offset = glm.vec3(float(off[0]), float(off[1]), float(off[2]))
+            j.Channels = _ROOT_CHANNELS.copy()
+        else:
+            p = int(parents[i])
+            off = neutral[i] - neutral[p]
+            j.Offset = glm.vec3(float(off[0]), float(off[1]), float(off[2]))
+            j.Channels = _JOINT_CHANNELS.copy()
+
+        for c in children[i]:
+            j.Children.append(_make_joint(c))
+        return j
+
+    # Wrapper Root at origin; single child is Hips (skeleton root).
+    root_wrapper = bvhio.BvhJoint("Root", offset=glm.vec3(0.0, 0.0, 0.0))
+    root_wrapper.Channels = _ROOT_CHANNELS.copy()
+    root_wrapper.Children.append(_make_joint(root_idx))
+    root_joint = root_wrapper
+
+    # Populate keyframes: Root = identity/zero, Hips = root motion, others = local rotation.
+    bvh_layout = root_joint.layout()
+    name_to_id = {n: idx for idx, n in enumerate(joint_names)}
+    ordered_joint_ids = []
+    for bj, _, _ in bvh_layout:
+        if bj.Name == "Root":
+            ordered_joint_ids.append(None)
+        else:
+            ordered_joint_ids.append(name_to_id[bj.Name])
+
+    bvh_joints = [bj for bj, _, _ in bvh_layout]
+    for bj in bvh_joints:
+        bj.Keyframes = [None] * T  # type: ignore[list-item]
+
+    identity_quat = glm.quat(1.0, 0.0, 0.0, 0.0)
+    zero_vec = glm.vec3(0.0, 0.0, 0.0)
+    for t in range(T):
+        for bj, jid in zip(bvh_joints, ordered_joint_ids):
+            if jid is None:
+                position = zero_vec
+                rotation = identity_quat
+            elif jid == root_idx:
+                pos = root_xyz[t]
+                position = glm.vec3(float(pos[0]), float(pos[1]), float(pos[2]))
+                qw, qx, qy, qz = q_wxyz[t, jid]
+                rotation = glm.quat(float(qw), float(qx), float(qy), float(qz))
+            else:
+                position = zero_vec
+                qw, qx, qy, qz = q_wxyz[t, jid]
+                rotation = glm.quat(float(qw), float(qx), float(qy), float(qz))
+            bj.Keyframes[t] = Pose(position, rotation)  # type: ignore[index]
+
+    container = bvhio.BvhContainer(root_joint, frameCount=T, frameTime=1.0 / float(fps))
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".bvh", delete=False, encoding="utf-8") as f:
+        tmp_path = f.name
+    try:
+        bvhio.writeBvh(tmp_path, container, percision=6)
+        bvh_text = Path(tmp_path).read_text(encoding="utf-8")
+        return _strip_end_site_blocks(bvh_text)
+    finally:
+        try:
+            os.remove(tmp_path)
+        except Exception:
+            pass
+
+
+def motion_to_bvh_bytes(
+    local_rot_mats: torch.Tensor,
+    root_positions: torch.Tensor,
+    *,
+    skeleton,
+    fps: float,
+    standard_tpose: bool = False,
+) -> bytes:
+    """Convert local rotations and root positions to BVH bytes (UTF-8).
+
+    Convenience wrapper around :func:`motion_to_bvh`.
+    """
+    return motion_to_bvh(
+        local_rot_mats,
+        root_positions,
+        skeleton=skeleton,
+        fps=fps,
+        standard_tpose=standard_tpose,
+    ).encode("utf-8")
+
+
+def save_motion_bvh(
+    path: Union[str, Path],
+    local_rot_mats: torch.Tensor,
+    root_positions: torch.Tensor,
+    *,
+    skeleton,
+    fps: float,
+    standard_tpose: bool = False,
+) -> None:
+    """Write local rotations and root positions to a BVH file at the given path."""
+    Path(path).write_text(
+        motion_to_bvh(local_rot_mats, root_positions, skeleton=skeleton, fps=fps, standard_tpose=standard_tpose),
+        encoding="utf-8",
+    )
+
+
+def read_bvh_frame_time_seconds(path: Union[str, Path]) -> float:
+    """Read ``Frame Time`` from a BVH file (seconds per frame)."""
+    with open(path, encoding="utf-8") as f:
+        for line in f:
+            if "Frame Time:" in line:
+                parts = line.split()
+                return float(parts[-1])
+    raise ValueError(f"Could not find 'Frame Time:' in {path}")
+
+
+def bvh_to_kimodo_motion(
+    path: Union[str, Path],
+    skeleton=None,
+    *,
+    standard_tpose: bool = False,
+) -> Tuple:
+    """Load a Kimodo-style SOMA BVH into a Kimodo motion dict.
+
+    Expects the same hierarchy as :func:`save_motion_bvh` (``Root`` wrapper + SOMA77 joints).
+    The frame rate is always read from the BVH ``Frame Time`` header.  Callers
+    that need a different playback rate should resample the returned motion dict
+    (see :func:`~kimodo.exports.motion_io.resample_motion_dict_to_kimodo_fps`).
+
+    Returns:
+        ``(motion_dict, source_fps)`` where ``source_fps`` is the native BVH
+        frame rate read from the file header.
+    """
+    from kimodo.exports.motion_io import complete_motion_dict
+    from kimodo.skeleton.bvh import parse_bvh_motion
+    from kimodo.skeleton.registry import build_skeleton
+
+    if skeleton is None:
+        skeleton = build_skeleton(77)
+    device = skeleton.neutral_joints.device
+
+    local_rot_mats, root_trans, bvh_fps = parse_bvh_motion(str(path))
+    local_rot_mats = local_rot_mats.to(device=device)
+    root_trans = root_trans.to(device=device)
+
+    if int(local_rot_mats.shape[1]) != int(skeleton.nbjoints):
+        raise ValueError(
+            f"BVH has {local_rot_mats.shape[1]} joints but skeleton has {skeleton.nbjoints}; "
+            "use a Kimodo-exported SOMA BVH or matching skeleton."
+        )
+    if not standard_tpose:
+        local_rot_mats, _ = skeleton.to_standard_tpose(local_rot_mats)
+
+    return complete_motion_dict(local_rot_mats, root_trans, skeleton, float(bvh_fps)), bvh_fps

kimodo/exports/motion_convert_lib.py

@@ -0,0 +1,159 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Library API for converting between Kimodo NPZ, AMASS NPZ, SOMA BVH, and G1 MuJoCo CSV."""
+
+from __future__ import annotations
+
+import warnings
+
+import numpy as np
+
+from kimodo.exports.bvh import bvh_to_kimodo_motion, save_motion_bvh
+from kimodo.exports.motion_formats import (
+    infer_source_format_from_path,
+    infer_target_format_from_path,
+    resolve_source_fps,
+)
+from kimodo.exports.motion_io import (
+    load_amass_npz,
+    load_g1_csv,
+    load_kimodo_npz_as_torch,
+    save_kimodo_npz_at_target_fps,
+)
+from kimodo.exports.mujoco import MujocoQposConverter
+from kimodo.exports.smplx import AMASSConverter
+from kimodo.skeleton.registry import build_skeleton
+
+
+def convert_motion_files(
+    input_path: str,
+    output_path: str,
+    *,
+    from_fmt: str | None = None,
+    to_fmt: str | None = None,
+    source_fps: float | None = None,
+    z_up: bool = True,
+    mujoco_rest_zero: bool = False,
+    bvh_standard_tpose: bool = False,
+) -> None:
+    """Convert a motion file between Kimodo-supported formats.
+
+    Supported pairs (hub-and-spoke through Kimodo NPZ):
+
+    - amass <-> kimodo
+    - soma-bvh <-> kimodo
+    - g1-csv <-> kimodo
+
+    Args:
+        input_path: Source file (``.npz``, ``.bvh``, or ``.csv``).
+        output_path: Destination file.
+        from_fmt: Source format; inferred from extension/contents when ``None``.
+        to_fmt: Target format; inferred from extension when ``None``.
+        source_fps: Source motion frame rate (Hz).  If provided, trusted as-is.
+            If ``None``, auto-detected from BVH ``Frame Time``, AMASS
+            ``mocap_frame_rate``, or default 30.
+        z_up: For AMASS conversions, apply the Z-up <-> Kimodo Y-up transform.
+        mujoco_rest_zero: For G1 CSV, joint angles relative to MuJoCo rest pose.
+        bvh_standard_tpose: If input or output is BVH: the BVH file uses the standard T-pose 
+            as its rest pose instead of the BONES-SEED rest pose.
+    """
+    from_fmt = from_fmt or infer_source_format_from_path(input_path)
+    to_fmt = to_fmt or infer_target_format_from_path(output_path, from_fmt)
+
+    _validate_output_extension(to_fmt, output_path)
+
+    pair = (from_fmt, to_fmt)
+
+    if pair == ("amass", "kimodo"):
+        sk = build_skeleton(22)
+        effective_source = source_fps
+        if effective_source is None:
+            with np.load(input_path, allow_pickle=True) as z:
+                effective_source = float(z["mocap_frame_rate"]) if "mocap_frame_rate" in z.files else 30.0
+        motion = load_amass_npz(input_path, source_fps=effective_source, z_up=z_up)
+        save_kimodo_npz_at_target_fps(motion, sk, effective_source, output_path)
+        return
+
+    if pair == ("kimodo", "amass"):
+        data, J = load_kimodo_npz_as_torch(input_path, ensure_complete=False)
+        if J != 22:
+            raise ValueError(f"Kimodo→AMASS requires 22 joints (SMPL-X); this file has J={J}.")
+        sk = build_skeleton(22)
+        effective_source = resolve_source_fps(source_fps, "kimodo", input_path, None)
+        converter = AMASSConverter(fps=effective_source, skeleton=sk)
+        converter.convert_save_npz(data, output_path, z_up=z_up)
+        return
+
+    if pair == ("soma-bvh", "kimodo"):
+        sk = build_skeleton(77)
+        motion, bvh_fps = bvh_to_kimodo_motion(input_path, skeleton=sk, standard_tpose=bvh_standard_tpose)
+        effective_source = source_fps if source_fps is not None else bvh_fps
+        save_kimodo_npz_at_target_fps(motion, sk, effective_source, output_path)
+        return
+
+    if pair == ("kimodo", "soma-bvh"):
+        data, J = load_kimodo_npz_as_torch(input_path, ensure_complete=False)
+        if J == 30:
+            warnings.warn(
+                f"Input has 30 joints (somaskel30); expanding to somaskel77 for BVH export.",
+                UserWarning,
+                stacklevel=2,
+            )
+            sk = build_skeleton(30)
+        elif J == 77:
+            sk = build_skeleton(77)
+        else:
+            raise ValueError(f"Kimodo→BVH requires a SOMA skeleton (30 or 77 joints); this file has J={J}.")
+        effective_source = resolve_source_fps(source_fps, "kimodo", input_path, None)
+        save_motion_bvh(
+            output_path,
+            data["local_rot_mats"],
+            data["root_positions"],
+            skeleton=sk,
+            fps=effective_source,
+            standard_tpose=bvh_standard_tpose,
+        )
+        return
+
+    if pair == ("g1-csv", "kimodo"):
+        sk = build_skeleton(34)
+        effective_source = resolve_source_fps(source_fps, "g1-csv", input_path, None)
+        motion = load_g1_csv(input_path, source_fps=effective_source, mujoco_rest_zero=mujoco_rest_zero)
+        save_kimodo_npz_at_target_fps(motion, sk, effective_source, output_path)
+        return
+
+    if pair == ("kimodo", "g1-csv"):
+        data, J = load_kimodo_npz_as_torch(input_path, ensure_complete=False)
+        if J != 34:
+            raise ValueError(f"Kimodo→CSV requires G1 with 34 joints; this file has J={J}.")
+        sk = build_skeleton(34)
+        effective_source = resolve_source_fps(source_fps, "kimodo", input_path, None)
+        converter = MujocoQposConverter(sk)
+        qpos = converter.dict_to_qpos(
+            {k: v for k, v in data.items() if k in ("local_rot_mats", "root_positions")},
+            device=str(sk.neutral_joints.device),
+            numpy=True,
+            mujoco_rest_zero=mujoco_rest_zero,
+        )
+        converter.save_csv(qpos, output_path)
+        return
+
+    raise ValueError(
+        f"Unsupported conversion {from_fmt!r} → {to_fmt!r}. "
+        "Supported: amass↔kimodo (SMPL-X NPZ), soma-bvh↔kimodo, g1-csv↔kimodo."
+    )
+
+
+def _validate_output_extension(to_fmt: str, output_path: str) -> None:
+    lower = output_path.lower()
+    if to_fmt == "kimodo" and lower.endswith(".npz"):
+        return
+    if to_fmt == "amass":
+        if not lower.endswith(".npz"):
+            raise ValueError("AMASS output must use a .npz path.")
+    elif to_fmt == "soma-bvh":
+        if not lower.endswith(".bvh"):
+            raise ValueError("SOMA BVH output must use a .bvh path.")
+    elif to_fmt == "g1-csv":
+        if not lower.endswith(".csv"):
+            raise ValueError("G1 CSV output must use a .csv path.")

kimodo/exports/motion_formats.py

@@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Infer motion file formats from paths and NPZ contents."""
+
+from __future__ import annotations
+
+import os
+from typing import Literal
+
+import numpy as np
+
+MotionSourceFormat = Literal["amass", "kimodo", "soma-bvh", "g1-csv"]
+MotionTargetFormat = Literal["amass", "kimodo", "soma-bvh", "g1-csv"]
+NpzMotionKind = Literal["amass", "kimodo"]
+
+
+def infer_npz_kind(path: str) -> NpzMotionKind:
+    """Classify a ``.npz`` as AMASS SMPL-X or Kimodo from required array keys."""
+    with np.load(path, allow_pickle=False) as z:
+        keys = set(z.files)
+    if "trans" in keys and "pose_body" in keys and "root_orient" in keys:
+        return "amass"
+    if "local_rot_mats" in keys or "posed_joints" in keys:
+        return "kimodo"
+    raise ValueError(
+        f"Unrecognized NPZ {path!r}: expected AMASS keys (trans, pose_body, ...) "
+        "or Kimodo keys (local_rot_mats, posed_joints, ...)."
+    )
+
+
+def infer_source_format_from_path(path: str) -> MotionSourceFormat:
+    """Infer converter input format from file extension and NPZ contents when needed."""
+    ext = os.path.splitext(path)[1].lower()
+    if ext == ".bvh":
+        return "soma-bvh"
+    if ext == ".csv":
+        return "g1-csv"
+    if ext == ".npz":
+        return infer_npz_kind(path)  # type: ignore[return-value]
+    raise ValueError(f"Cannot infer format from extension of {path!r}")
+
+
+def infer_target_format_from_path(path: str, from_fmt: MotionSourceFormat) -> MotionTargetFormat:
+    """Infer converter output format from destination path and source format."""
+    ext = os.path.splitext(path)[1].lower()
+    if ext == ".bvh":
+        return "soma-bvh"
+    if ext == ".csv":
+        return "g1-csv"
+    if ext == ".npz":
+        if from_fmt == "amass":
+            return "kimodo"
+        if from_fmt == "kimodo":
+            return "amass"
+        if from_fmt in ("g1-csv", "soma-bvh"):
+            return "kimodo"
+        raise ValueError(
+            "Ambiguous .npz output: set --to to 'kimodo' or 'amass' when the input format is not amass/kimodo."
+        )
+    raise ValueError(f"Cannot infer output format from extension of {path!r}")
+
+
+def resolve_source_fps(
+    fps: float | None,
+    from_kind: str,
+    input_path: str,
+    data: dict | None,
+) -> float:
+    """Resolve source frame rate (Hz) for conversion when ``fps`` is not overridden."""
+    if fps is not None:
+        return float(fps)
+    if data is not None and "mocap_frame_rate" in data:
+        return float(np.asarray(data["mocap_frame_rate"]).item())
+    if from_kind == "soma-bvh":
+        from kimodo.exports.bvh import read_bvh_frame_time_seconds
+
+        return 1.0 / read_bvh_frame_time_seconds(input_path)
+    return 30.0

kimodo/exports/motion_io.py

@@ -0,0 +1,443 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Assemble Kimodo NPZ-compatible motion dicts from local rotations + root trajectory."""
+
+from __future__ import annotations
+
+import os
+import warnings
+from typing import Any, Dict, Tuple
+
+import numpy as np
+import torch
+
+from kimodo.geometry import matrix_to_quaternion, quaternion_to_matrix
+from kimodo.motion_rep.feature_utils import compute_heading_angle, compute_vel_xyz
+from kimodo.motion_rep.feet import foot_detect_from_pos_and_vel
+from kimodo.motion_rep.smooth_root import get_smooth_root_pos
+from kimodo.skeleton import SkeletonBase
+from kimodo.skeleton.registry import build_skeleton
+from kimodo.tools import to_numpy
+
+# Default motion rate for Kimodo NPZ produced by format conversion (matches common model FPS).
+KIMODO_CONVERT_TARGET_FPS = 30.0
+
+
+def _quaternion_slerp(q0: torch.Tensor, q1: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+    """Spherical linear interpolation; ``q0``, ``q1`` (..., 4) wxyz; ``t`` broadcastable to (...,
+    1)."""
+    if t.dim() < q0.dim():
+        t = t.unsqueeze(-1)
+    dot = (q0 * q1).sum(dim=-1, keepdim=True)
+    q1 = torch.where(dot < 0, -q1, q1)
+    dot = torch.abs(dot).clamp(-1.0, 1.0)
+    theta_0 = torch.acos(dot)
+    sin_theta = torch.sin(theta_0)
+    s0 = torch.sin((1.0 - t) * theta_0) / sin_theta.clamp(min=1e-8)
+    s1 = torch.sin(t * theta_0) / sin_theta.clamp(min=1e-8)
+    q = s0 * q0 + s1 * q1
+    return q / torch.linalg.norm(q, dim=-1, keepdim=True).clamp(min=1e-8)
+
+
+def resample_motion_dict_to_kimodo_fps(
+    motion_dict: Dict[str, torch.Tensor],
+    skeleton: SkeletonBase,
+    source_fps: float,
+    target_fps: float = KIMODO_CONVERT_TARGET_FPS,
+) -> Tuple[Dict[str, torch.Tensor], bool]:
+    """Resample a Kimodo motion dict to ``target_fps``.
+
+    When the fps ratio is close to an integer (e.g. 120 / 30 = 4), the faster
+    stepping method is used (take every *step*-th frame).  Otherwise falls back
+    to linear interp (root) + quaternion slerp (joints).
+
+    Re-runs :func:`complete_motion_dict` at the target rate so derived channels stay consistent.
+
+    Returns:
+        The motion dict and ``True`` if time resampling was applied, else ``False`` (already at
+        ``target_fps`` with matching frame count; only re-derived via FK).
+    """
+    local_rot_mats = motion_dict["local_rot_mats"]
+    root_positions = motion_dict["root_positions"]
+    local_rot_mats, root_positions = _coerce_time_local_root(local_rot_mats, root_positions)
+    t_in = int(local_rot_mats.shape[0])
+    if t_in < 1:
+        raise ValueError("Motion must have at least one frame.")
+    if source_fps <= 0:
+        raise ValueError(f"source_fps must be positive; got {source_fps}")
+
+    t_out = max(1, int(round(t_in * target_fps / source_fps)))
+    if t_out == t_in and abs(float(source_fps) - float(target_fps)) < 1e-3:
+        return complete_motion_dict(local_rot_mats, root_positions, skeleton, float(target_fps)), False
+
+    ratio = source_fps / target_fps
+    step = round(ratio)
+    if step >= 2 and abs(ratio - step) < 0.05:
+        local_out = local_rot_mats[::step]
+        root_out = root_positions[::step]
+    else:
+        device = local_rot_mats.device
+        dtype = local_rot_mats.dtype
+        u = torch.linspace(0, t_in - 1, t_out, device=device, dtype=dtype)
+        i0 = u.floor().long().clamp(0, t_in - 1)
+        i1 = torch.minimum(i0 + 1, torch.tensor(t_in - 1, device=device))
+        tau_1d = (u - i0.float()).unsqueeze(-1)
+        rp0 = root_positions[i0]
+        rp1 = root_positions[i1]
+        root_out = (1.0 - tau_1d) * rp0 + tau_1d * rp1
+
+        quats = matrix_to_quaternion(local_rot_mats)
+        q0 = quats[i0]
+        q1 = quats[i1]
+        tau_q = (u - i0.float()).view(t_out, 1, 1)
+        quat_out = _quaternion_slerp(q0, q1, tau_q)
+        local_out = quaternion_to_matrix(quat_out)
+
+    return complete_motion_dict(local_out, root_out, skeleton, float(target_fps)), True
+
+
+def warn_kimodo_npz_framerate(source_fps: float, t_before: int, t_after: int) -> None:
+    """Emit a warning after time resampling for Kimodo NPZ (linear root, quaternion slerp per
+    joint)."""
+    warnings.warn(
+        f"Resampled motion to {KIMODO_CONVERT_TARGET_FPS:.0f} Hz for Kimodo NPZ "
+        f"(source ~{source_fps:.4g} Hz, {t_before} input frames → {t_after} output frames). "
+        "Pass --source-fps if the detected source rate is wrong.",
+        UserWarning,
+        stacklevel=3,
+    )
+
+
+def _coerce_time_local_root(
+    local_rot_mats: torch.Tensor,
+    root_positions: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Normalize to shapes (T, J, 3, 3) and (T, 3)."""
+    if local_rot_mats.dim() == 5:
+        if int(local_rot_mats.shape[0]) != 1:
+            raise ValueError(f"local_rot_mats batch size must be 1 for single clip; got {local_rot_mats.shape[0]}")
+        local_rot_mats = local_rot_mats[0]
+    if root_positions.dim() == 3:
+        if int(root_positions.shape[0]) != 1:
+            raise ValueError(f"root_positions batch size must be 1; got {root_positions.shape[0]}")
+        root_positions = root_positions[0]
+    if local_rot_mats.dim() != 4:
+        raise ValueError(f"local_rot_mats must be (T,J,3,3); got {tuple(local_rot_mats.shape)}")
+    if root_positions.dim() != 2 or int(root_positions.shape[-1]) != 3:
+        raise ValueError(f"root_positions must be (T,3); got {tuple(root_positions.shape)}")
+    if int(local_rot_mats.shape[0]) != int(root_positions.shape[0]):
+        raise ValueError("local_rot_mats and root_positions must have the same number of frames")
+    return local_rot_mats, root_positions
+
+
+def complete_motion_dict(
+    local_rot_mats: torch.Tensor,
+    root_positions: torch.Tensor,
+    skeleton: SkeletonBase,
+    fps: float,
+) -> Dict[str, torch.Tensor]:
+    """Build the Kimodo motion output dict from local rotations and root positions.
+
+    Matches keys written by CLI generation (see docs/source/user_guide/output_formats.md).
+
+    Args:
+        local_rot_mats: (T, J, 3, 3) or (1, T, J, 3, 3) local rotation matrices.
+        root_positions: (T, 3) or (1, T, 3) root / pelvis world positions (meters).
+        skeleton: Skeleton instance (SOMA77, G1, SMPL-X, etc.).
+        fps: Sampling rate (Hz).
+
+    Returns:
+        Dict with tensors ``posed_joints``, ``global_rot_mats``, ``local_rot_mats``,
+        ``foot_contacts``, ``smooth_root_pos``, ``root_positions``, ``global_root_heading``.
+    """
+    device = local_rot_mats.device
+    dtype = local_rot_mats.dtype
+    local_rot_mats, root_positions = _coerce_time_local_root(
+        local_rot_mats.to(device=device, dtype=dtype),
+        root_positions.to(device=device, dtype=dtype),
+    )
+
+    global_rot_mats, posed_joints, _ = skeleton.fk(local_rot_mats, root_positions)
+
+    smooth_root_pos = get_smooth_root_pos(root_positions.unsqueeze(0)).squeeze(0)
+
+    lengths = torch.tensor([posed_joints.shape[0]], device=device)
+    velocities = compute_vel_xyz(posed_joints.unsqueeze(0), fps, lengths=lengths).squeeze(0)
+
+    heading_angle = compute_heading_angle(posed_joints.unsqueeze(0), skeleton).squeeze(0)
+    global_root_heading = torch.stack([torch.cos(heading_angle), torch.sin(heading_angle)], dim=-1)
+
+    foot_contacts = foot_detect_from_pos_and_vel(
+        posed_joints.unsqueeze(0),
+        velocities.unsqueeze(0),
+        skeleton,
+        0.15,
+        0.10,
+    ).squeeze(0)
+
+    return {
+        "posed_joints": posed_joints,
+        "global_rot_mats": global_rot_mats,
+        "local_rot_mats": local_rot_mats,
+        "foot_contacts": foot_contacts,
+        "smooth_root_pos": smooth_root_pos,
+        "root_positions": root_positions,
+        "global_root_heading": global_root_heading,
+    }
+
+
+def motion_dict_to_numpy(d: Dict[str, Any]) -> Dict[str, np.ndarray]:
+    """Convert motion dict values to numpy arrays for ``np.savez``."""
+    out: Dict[str, np.ndarray] = {}
+    for k, v in d.items():
+        if hasattr(v, "detach"):
+            out[k] = to_numpy(v)
+        elif isinstance(v, np.ndarray):
+            out[k] = v
+        else:
+            out[k] = np.asarray(v)
+    return out
+
+
+def save_kimodo_npz(path: str, motion_dict: Dict[str, Any]) -> None:
+    """Save a Kimodo-compatible motion dict to ``.npz`` (numpy arrays)."""
+    np.savez(path, **motion_dict_to_numpy(motion_dict))
+
+
+def load_kimodo_npz(path: str) -> Dict[str, np.ndarray]:
+    """Load arrays from a Kimodo ``.npz`` file."""
+    with np.load(path, allow_pickle=False) as data:
+        return {k: np.asarray(data[k]) for k in data.files}
+
+
+def load_g1_csv(
+    path: str,
+    source_fps: float = KIMODO_CONVERT_TARGET_FPS,
+    *,
+    mujoco_rest_zero: bool = False,
+) -> Dict[str, torch.Tensor]:
+    """Load a G1 MuJoCo ``qpos`` CSV (``(T, 36)``) into a Kimodo motion dict.
+
+    Args:
+        path: CSV path (comma-separated, no header).
+        source_fps: Source frame rate (Hz) of the CSV data.
+        mujoco_rest_zero: Must match how the CSV was written (see :class:`MujocoQposConverter`).
+    """
+    from kimodo.exports.mujoco import MujocoQposConverter
+
+    qpos = np.loadtxt(path, delimiter=",")
+    if qpos.ndim != 2 or qpos.shape[-1] != 36:
+        raise ValueError(f"Expected G1 CSV with shape (T, 36); got {qpos.shape}")
+    sk = build_skeleton(34)
+    converter = MujocoQposConverter(sk)
+    return converter.qpos_to_motion_dict(qpos, float(source_fps), mujoco_rest_zero=mujoco_rest_zero)
+
+
+def load_amass_npz(
+    path: str,
+    source_fps: float | None = None,
+    *,
+    z_up: bool = True,
+) -> Dict[str, torch.Tensor]:
+    """Load an AMASS-style SMPL-X ``.npz`` into a Kimodo motion dict (22 joints).
+
+    Args:
+        path: NPZ with ``trans``, ``root_orient``, ``pose_body``, etc.
+        source_fps: Source frame rate (Hz); if ``None``, uses ``mocap_frame_rate``
+            from the file when present, else 30 Hz.
+        z_up: If ``True``, apply AMASS Z-up to Kimodo Y-up transform (same as CLI).
+    """
+    from kimodo.exports.smplx import amass_npz_to_kimodo_motion
+
+    sk = build_skeleton(22)
+    return amass_npz_to_kimodo_motion(path, sk, source_fps=source_fps, z_up=z_up)
+
+
+def load_kimodo_npz_as_torch(
+    path: str,
+    source_fps: float = KIMODO_CONVERT_TARGET_FPS,
+    *,
+    ensure_complete: bool = True,
+) -> tuple[Dict[str, torch.Tensor], int]:
+    """Load a Kimodo NPZ and return all arrays as torch tensors on the skeleton device.
+
+    Args:
+        path: Kimodo NPZ file path.
+        source_fps: Source frame rate (Hz) used for derived channels when
+            ``ensure_complete=True``.
+        ensure_complete: If ``True`` and the NPZ lacks derived channels
+            (``posed_joints``, ``global_rot_mats``, …), run :func:`complete_motion_dict`
+            to fill them from ``local_rot_mats`` + ``root_positions``.
+            If ``False``, load all arrays verbatim (requires ``local_rot_mats``).
+
+    Returns:
+        ``(tensor_dict, num_joints)``
+    """
+    raw = load_kimodo_npz(path)
+    if "local_rot_mats" in raw:
+        j = int(raw["local_rot_mats"].shape[1])
+    elif "posed_joints" in raw:
+        j = int(raw["posed_joints"].shape[1])
+    else:
+        raise ValueError("Kimodo NPZ must contain 'local_rot_mats' or 'posed_joints'.")
+    sk = build_skeleton(j)
+    device = sk.neutral_joints.device
+    dtype = torch.float32
+
+    if not ensure_complete:
+        if "local_rot_mats" not in raw:
+            raise ValueError("Kimodo NPZ must contain 'local_rot_mats' (and typically 'root_positions').")
+        out: Dict[str, torch.Tensor] = {}
+        for k, v in raw.items():
+            out[k] = torch.from_numpy(np.asarray(v)).to(device=device, dtype=dtype)
+        return out, j
+
+    if "posed_joints" in raw and "global_rot_mats" in raw:
+        out = {}
+        for k, v in raw.items():
+            out[k] = torch.from_numpy(np.asarray(v)).to(device=device, dtype=dtype)
+        return out, j
+
+    if "local_rot_mats" not in raw or "root_positions" not in raw:
+        raise ValueError("Kimodo NPZ must contain posed_joints+global_rot_mats, or local_rot_mats+root_positions.")
+    local = torch.from_numpy(np.asarray(raw["local_rot_mats"])).to(device=device, dtype=dtype)
+    root = torch.from_numpy(np.asarray(raw["root_positions"])).to(device=device, dtype=dtype)
+    return complete_motion_dict(local, root, sk, float(source_fps)), j
+
+
+def save_kimodo_npz_at_target_fps(
+    motion: Dict[str, torch.Tensor],
+    skeleton: SkeletonBase,
+    source_fps: float,
+    output_path: str,
+    target_fps: float = KIMODO_CONVERT_TARGET_FPS,
+) -> None:
+    """Resample a motion dict to ``target_fps`` when needed, then save Kimodo NPZ."""
+    t_before = int(motion["local_rot_mats"].shape[0])
+    motion, did_resample = resample_motion_dict_to_kimodo_fps(motion, skeleton, source_fps, target_fps)
+    t_after = int(motion["local_rot_mats"].shape[0])
+    if did_resample:
+        warn_kimodo_npz_framerate(source_fps, t_before, t_after)
+    save_kimodo_npz(output_path, motion)
+
+
+def kimodo_npz_to_bytes(motion_dict: Dict[str, Any]) -> bytes:
+    """Serialize a Kimodo motion dict to in-memory NPZ bytes."""
+    import io
+
+    buf = io.BytesIO()
+    np.savez(buf, **motion_dict_to_numpy(motion_dict))
+    return buf.getvalue()
+
+
+def g1_csv_to_bytes(motion_dict: Dict[str, Any], skeleton: SkeletonBase, device: Any) -> bytes:
+    """Convert a motion dict to G1 MuJoCo CSV bytes via :class:`MujocoQposConverter`."""
+    import io
+
+    from kimodo.exports.mujoco import MujocoQposConverter
+
+    converter = MujocoQposConverter(skeleton)
+    qpos = converter.dict_to_qpos(
+        {k: v for k, v in motion_dict.items() if k in ("local_rot_mats", "root_positions")},
+        device,
+        numpy=True,
+    )
+    buf = io.StringIO()
+    np.savetxt(buf, qpos, delimiter=",")
+    return buf.getvalue().encode("utf-8")
+
+
+def amass_npz_to_bytes(motion_dict: Dict[str, Any], skeleton: SkeletonBase, fps: float) -> bytes:
+    """Convert a motion dict to AMASS NPZ bytes via :class:`AMASSConverter`."""
+    import io
+
+    from kimodo.exports.smplx import AMASSConverter
+
+    converter = AMASSConverter(skeleton=skeleton, fps=fps)
+    buf = io.BytesIO()
+    converter.convert_save_npz(
+        {k: v for k, v in motion_dict.items() if k in ("local_rot_mats", "root_positions")},
+        buf,
+    )
+    return buf.getvalue()
+
+
+def _read_amass_source_fps(path: str) -> float:
+    """Read the source frame rate from an AMASS NPZ, defaulting to 30 Hz."""
+    with np.load(path, allow_pickle=True) as z:
+        if "mocap_frame_rate" in z.files:
+            return float(z["mocap_frame_rate"])
+    return 30.0
+
+
+def load_motion_file(
+    path: str,
+    source_fps: float | None = None,
+    target_fps: float | None = None,
+    *,
+    z_up: bool = True,
+    mujoco_rest_zero: bool = False,
+) -> tuple[Dict[str, torch.Tensor], int]:
+    """Load a motion file and return a Kimodo motion dict plus joint count.
+
+    Supports SOMA BVH (``.bvh``), G1 MuJoCo CSV (``.csv``), Kimodo NPZ, and AMASS SMPL-X NPZ
+    (``.npz``).
+
+    The motion is loaded at its native (or overridden) source rate, then
+    resampled to ``target_fps`` when they differ.
+
+    Args:
+        path: Path to ``.bvh``, ``.csv``, or ``.npz``.
+        source_fps: Source frame rate (Hz).  If provided, trusted as-is.
+            If ``None``, auto-detected per format: BVH ``Frame Time`` header,
+            AMASS ``mocap_frame_rate``, or :data:`KIMODO_CONVERT_TARGET_FPS`
+            (30 Hz) for CSV / Kimodo NPZ.
+        target_fps: Desired output frame rate (Hz).  Defaults to
+            :data:`KIMODO_CONVERT_TARGET_FPS` (30 Hz).  The motion is
+            resampled when ``source_fps`` and ``target_fps`` differ.
+        z_up: AMASS NPZ only; passed to :func:`load_amass_npz`.
+        mujoco_rest_zero: G1 CSV only; passed to :func:`load_g1_csv`.
+
+    Returns:
+        ``(motion_dict, num_joints)`` with the same keys as :func:`complete_motion_dict`.
+    """
+    from kimodo.exports.motion_formats import infer_npz_kind
+
+    if target_fps is None:
+        target_fps = KIMODO_CONVERT_TARGET_FPS
+
+    ext = os.path.splitext(path)[1].lower()
+    if ext == ".bvh":
+        from kimodo.exports.bvh import bvh_to_kimodo_motion
+
+        motion_dict, bvh_fps = bvh_to_kimodo_motion(path)
+        effective_source = source_fps if source_fps is not None else bvh_fps
+        num_joints = int(motion_dict["local_rot_mats"].shape[1])
+    elif ext == ".csv":
+        effective_source = source_fps if source_fps is not None else KIMODO_CONVERT_TARGET_FPS
+        motion_dict = load_g1_csv(path, source_fps=effective_source, mujoco_rest_zero=mujoco_rest_zero)
+        num_joints = 34
+    elif ext == ".npz":
+        kind = infer_npz_kind(path)
+        if kind == "amass":
+            effective_source = source_fps if source_fps is not None else _read_amass_source_fps(path)
+            motion_dict = load_amass_npz(path, source_fps=effective_source, z_up=z_up)
+            num_joints = 22
+        else:
+            effective_source = source_fps if source_fps is not None else KIMODO_CONVERT_TARGET_FPS
+            motion_dict, num_joints = load_kimodo_npz_as_torch(path, source_fps=effective_source)
+    else:
+        raise ValueError(f"Unsupported motion file {path!r}; expected .bvh, .csv, or .npz")
+
+    if abs(effective_source - target_fps) > 0.5:
+        sk = build_skeleton(num_joints)
+        motion_dict, did_resample = resample_motion_dict_to_kimodo_fps(motion_dict, sk, effective_source, target_fps)
+        if did_resample:
+            t_out = int(motion_dict["local_rot_mats"].shape[0])
+            warnings.warn(
+                f"Resampled motion from {effective_source:.4g} Hz to " f"{target_fps:.0f} Hz ({t_out} frames).",
+                UserWarning,
+                stacklevel=2,
+            )
+
+    return motion_dict, num_joints

kimodo/exports/mujoco.py

@@ -0,0 +1,588 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Convert kimodo motion (y-up, z-forward) to MuJoCo qpos (z-up, x-forward) for G1 skeleton."""
+
+import os
+import xml.etree.ElementTree as ET
+from typing import Optional
+
+import numpy as np
+import torch
+from scipy.spatial.transform import Rotation
+
+from kimodo.assets import skeleton_asset_path
+from kimodo.geometry import (
+    axis_angle_to_matrix,
+    matrix_to_axis_angle,
+    matrix_to_quaternion,
+    quaternion_to_matrix,
+)
+from kimodo.skeleton import G1Skeleton34, SkeletonBase, global_rots_to_local_rots
+from kimodo.tools import ensure_batched, to_numpy, to_torch
+
+# Cache so that the same (skeleton, xml_path) returns the same converter instance.
+_converter_cache: dict[tuple[int, str], "MujocoQposConverter"] = {}
+
+
+class MujocoQposConverter:
+    """Fast batch converter from our dictionary format to mujoco qpos with precomputed transforms.
+
+    In mujoco, the coordination is z up and x forward, right handed.
+
+    Features (30 joints):
+    - root (pelvis, 7 = translation + rotation) + 29 dof joints (29)
+
+    In kimodo, the coordinate system is y up and z forward, right handed.
+    Features (34 joints):
+    - root (pelvis) + (34 - 1) joints; among these joints, 4 are end-effector joints added by kimodo.
+
+    Cached by (input_skeleton id, xml_path); repeated calls with the same args return the same instance.
+    """
+
+    def __new__(
+        cls,
+        input_skeleton: SkeletonBase,
+        xml_path: str = str(skeleton_asset_path("g1skel34", "xml", "g1.xml")),
+    ):
+        key = (id(input_skeleton), xml_path)
+        if key not in _converter_cache:
+            inst = object.__new__(cls)
+            _converter_cache[key] = inst
+        return _converter_cache[key]
+
+    def __init__(
+        self,
+        input_skeleton: SkeletonBase,
+        xml_path: str = str(skeleton_asset_path("g1skel34", "xml", "g1.xml")),
+    ):
+        """Initialize converter with precomputed transforms.
+
+        Args:
+            xml_path: Path to the mujoco XML file containing joint definitions
+        """
+        if getattr(self, "_initialized", False):
+            return
+        self.xml_path = xml_path
+        self.skeleton = input_skeleton
+        self._prepare_transforms()
+        self._subtree_joints = {}
+        self._initialized = True
+
+    def _prepare_transforms(self):
+        """Precompute all necessary transforms for efficient batch processing."""
+        # Define coordinate transformations between mujoco and kimodo space
+        # 1) R_zup_to_yup: rotation around x-axis by -90 degrees
+        # 2) x_forward_to_y_forward: rotation around z-axis by -90 degrees
+        # Combined transformation matrix: mujoco_to_kimodo = R_zup_to_yup * x_forward_to_y_forward
+        self.mujoco_to_kimodo_matrix = torch.tensor(
+            [[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]], dtype=torch.float32
+        )
+        self.kimodo_to_mujoco_matrix = self.mujoco_to_kimodo_matrix.T  # Inverse transformation: kimodo_to_mujoco
+
+        # Parse XML once and extract joint information
+        tree = ET.parse(self.xml_path)
+        root = tree.getroot()
+
+        xml_classes = [x for x in tree.findall(".//default") if "class" in x.attrib]
+        joint_axes = dict()
+        class_ranges: dict[str, tuple[float, float]] = {}
+        for xml_class in xml_classes:
+            j = xml_class.findall("joint")
+            if j:
+                joint_axes[xml_class.get("class")] = j[0].get("axis")
+                range_str = j[0].get("range")
+                if range_str:
+                    range_vals = [float(x) for x in range_str.split()]
+                    if len(range_vals) == 2:
+                        class_ranges[xml_class.get("class")] = (
+                            range_vals[0],
+                            range_vals[1],
+                        )
+
+        mujoco_hinge_joints = root.find("worldbody").findall(".//joint")  # skip the base joint
+        self._mujoco_joint_axis_values_kimodo_space = torch.zeros(
+            (len(mujoco_hinge_joints), 3), dtype=torch.float32
+        )  # mujoco order but kimodo space
+        self._mujoco_joint_axis_values_mujoco_space = torch.zeros(
+            (len(mujoco_hinge_joints), 3), dtype=torch.float32
+        )  # mujoco order but mujoco space
+
+        # for the below indices, mujoco_indices_to_kimodo_indices does not include mujoco root (30 - 1 = 29 elements),
+        # while kimodo_indices_to_mujoco_indices inclues the kimodo root (32 elements).
+        self._mujoco_indices_to_kimodo_indices = torch.zeros((len(mujoco_hinge_joints),), dtype=torch.int32)
+        self._kimodo_indices_to_mujoco_indices = (
+            torch.ones((self.skeleton.nbjoints,), dtype=torch.int32) * -1
+        )  # -1 means not in the csv skeleton
+
+        self._nb_joints_mujoco = len(mujoco_hinge_joints) + 1
+        self._nb_joints_kimodo = self.skeleton.nbjoints
+        self._mujoco_joint_including_root_parent_list = torch.full(
+            (len(mujoco_hinge_joints) + 1,), -1, dtype=torch.int32
+        )
+        self._mujoco_joint_including_root_list = ["pelvis_skel"]
+
+        for joint_id_in_csv, joint in enumerate(mujoco_hinge_joints):
+            joint_name_in_skeleton = joint.get("name").replace("_joint", "_skel")
+            joint_parent_name_in_skeleton = self.skeleton.bone_parents[joint_name_in_skeleton]
+
+            self._mujoco_joint_including_root_list.append(joint_name_in_skeleton)
+            self._mujoco_joint_including_root_parent_list[joint_id_in_csv + 1] = (
+                self._mujoco_joint_including_root_list.index(joint_parent_name_in_skeleton)
+            )
+
+            joint_idx_in_kimodo_skeleton = self.skeleton.bone_order_names.index(joint_name_in_skeleton)
+            axis_values = [float(x) for x in (joint.get("axis") or joint_axes[joint.get("class")]).split(" ")]
+
+            # the mapped axis in kimodo skeleton space is calculated as bones_axis = mujoco_to_kimodo.apply(axis_values)
+            # [1, 0, 0] -> [0, 0, 1]; [0, 1, 0] -> [1, 0, 0]; [0, 0, 1] -> [0, 1, 0]
+            mujoco_joint_axis_mapping_kimodo_space = [
+                torch.tensor([0, 0, 1]),
+                torch.tensor([1, 0, 0]),
+                torch.tensor([0, 1, 0]),
+            ][np.argmax(axis_values)]
+
+            self._mujoco_joint_axis_values_kimodo_space[joint_id_in_csv] = mujoco_joint_axis_mapping_kimodo_space
+            self._mujoco_joint_axis_values_mujoco_space[joint_id_in_csv] = torch.tensor(axis_values)
+
+            self._mujoco_indices_to_kimodo_indices[joint_id_in_csv] = joint_idx_in_kimodo_skeleton
+            self._kimodo_indices_to_mujoco_indices[joint_idx_in_kimodo_skeleton] = (
+                joint_id_in_csv + 1
+            )  # +1 for the root
+        self._kimodo_indices_to_mujoco_indices[0] = 0  # the root joint mapping
+
+        # Joint limits (min, max) in radians for each mujoco hinge, for clamping
+        self._joint_limits_min = torch.full((len(mujoco_hinge_joints),), float("-inf"), dtype=torch.float32)
+        self._joint_limits_max = torch.full((len(mujoco_hinge_joints),), float("inf"), dtype=torch.float32)
+        for joint_id_in_csv, joint in enumerate(mujoco_hinge_joints):
+            range_vals = None
+            if joint.get("range"):
+                range_vals = [float(x) for x in joint.get("range").split()]
+            elif joint.get("class") and joint.get("class") in class_ranges:
+                lo, hi = class_ranges[joint.get("class")]
+                range_vals = [lo, hi]
+            if range_vals is not None and len(range_vals) == 2:
+                self._joint_limits_min[joint_id_in_csv] = range_vals[0]
+                self._joint_limits_max[joint_id_in_csv] = range_vals[1]
+
+        # load the offset matrices from the xml
+        R_zup_to_yup = Rotation.from_euler("x", -90, degrees=True)
+        x_forward_to_y_forward = Rotation.from_euler("z", -90, degrees=True)
+        mujoco_to_kimodo = R_zup_to_yup * x_forward_to_y_forward
+
+        self._rot_offsets_q2t = torch.zeros(len(self._kimodo_indices_to_mujoco_indices), 3, 3, dtype=torch.float32)
+        self._rot_offsets_q2t[...] = torch.eye(3)[None]
+
+        self._rot_offsets_f2q = torch.zeros(len(self._kimodo_indices_to_mujoco_indices), 3, 3, dtype=torch.float32)
+        self._rot_offsets_f2q[...] = torch.eye(3)[None]
+        parent_map = {child: parent for parent in root.iter() for child in parent}
+        for i, joint in enumerate(mujoco_hinge_joints):
+            body = parent_map[joint]
+            if "quat" in body.attrib:
+                rot = Rotation.from_quat(
+                    [float(x) for x in body.get("quat").strip().split(" ")],
+                    scalar_first=True,
+                )
+                idx = self._mujoco_indices_to_kimodo_indices[i]
+                self._rot_offsets_q2t[idx] = torch.from_numpy(rot.as_matrix())
+                rot = mujoco_to_kimodo * rot * mujoco_to_kimodo.inv()
+                self._rot_offsets_f2q[idx] = torch.from_numpy(rot.as_matrix().T)
+
+        # Hinge axis in f2q space so extraction uses the same frame as joint_rot_f2q.
+        # Then extract(offset) gives the angle s.t. axis_angle(angle * axis_f2q) = offset, and
+        # reconstruction R_local = offset.T @ axis_angle(angle * axis_f2q) = I when input is identity.
+        axis_kimodo = self._mujoco_joint_axis_values_kimodo_space
+        self._mujoco_joint_axis_values_f2q_space = torch.zeros_like(axis_kimodo)
+        for i in range(len(mujoco_hinge_joints)):
+            j = self._mujoco_indices_to_kimodo_indices[i].item()
+            axis_f2q = torch.mv(self._rot_offsets_f2q[j], axis_kimodo[i])
+            n = axis_f2q.norm()
+            if n > 1e-8:
+                axis_f2q = axis_f2q / n
+            self._mujoco_joint_axis_values_f2q_space[i] = axis_f2q
+
+        # Rest-pose DOFs: angle we extract when R_local = I (t-pose). MuJoCo limits are
+        # relative to joint zero (rest pose), so we must clamp in MuJoCo space: convert
+        # joint_dofs to mujoco_angle = joint_dofs - rest_dofs, clamp, then back.
+        rest_rot_f2q = self._rot_offsets_f2q[self._mujoco_indices_to_kimodo_indices]
+        rest_rot_f2q = rest_rot_f2q.unsqueeze(0).unsqueeze(0)
+        self._rest_dofs = self._local_rots_f2q_to_joint_dofs(rest_rot_f2q).squeeze(0).squeeze(0)
+        # Axis-angle rest DOFs: angle s.t. axis_angle(angle * axis_f2q) = offset. Used in
+        # project_to_real_robot_rotations so extract+reconstruct round-trip and t-pose is preserved.
+        rest_rot_f2q_flat = self._rot_offsets_f2q[self._mujoco_indices_to_kimodo_indices]
+        full_aa = matrix_to_axis_angle(rest_rot_f2q_flat)
+        self._rest_dofs_axis_angle = (full_aa * self._mujoco_joint_axis_values_f2q_space).sum(dim=-1)
+
+    def dict_to_qpos(
+        self,
+        output: dict,
+        device: Optional[str] = None,
+        root_quat_w_first: bool = True,
+        numpy: bool = True,
+        mujoco_rest_zero: bool = False,
+    ):
+        """Convert kimodo output dict to mujoco qpos format.
+
+        Args:
+            output: dict with keys "local_rot_mats" and "root_positions".
+            device: device to use for the output.
+            root_quat_w_first: If True, quaternion in qpos is (w,x,y,z).
+            numpy: If True, convert the output to numpy array.
+            mujoco_rest_zero: If True, joint angles are written so that kimodo rest (t-pose)
+                maps to q=0 in MuJoCo. If False, write raw joint_dofs.
+
+        Returns:
+            qpos: (B, T, 7+J) mujoco qpos format.
+        """
+        local_rot_mats = to_torch(output["local_rot_mats"], device)
+        root_positions = to_torch(output["root_positions"], device)
+
+        qpos = self.to_qpos(
+            local_rot_mats,
+            root_positions,
+            root_quat_w_first=root_quat_w_first,
+            mujoco_rest_zero=mujoco_rest_zero,
+        )
+        if numpy:
+            qpos = to_numpy(qpos)
+        return qpos
+
+    def qpos_to_motion_dict(
+        self,
+        qpos: torch.Tensor | np.ndarray,
+        source_fps: float,
+        *,
+        root_quat_w_first: bool = True,
+        mujoco_rest_zero: bool = False,
+    ):
+        """Inverse of :meth:`to_qpos` / :meth:`dict_to_qpos` for MuJoCo CSV ``(T, 36)`` rows.
+
+        Args:
+            qpos: Shape ``(T, 36)`` or ``(1, T, 36)`` (root xyz, root quat wxyz, 29 joint angles).
+            source_fps: Source frame rate (Hz) of the qpos data.
+            root_quat_w_first: Must match how the CSV was written (default ``True``).
+            mujoco_rest_zero: Must match :meth:`dict_to_qpos` / :meth:`to_qpos`.
+
+        Returns:
+            Kimodo motion dict (see :func:`kimodo.exports.motion_io.complete_motion_dict`).
+        """
+        from kimodo.exports.motion_io import complete_motion_dict
+
+        qpos = to_torch(qpos, None)
+        if qpos.dim() == 2:
+            qpos = qpos.unsqueeze(0)
+        device = qpos.device
+        dtype = qpos.dtype
+        batch_size, num_frames, ncols = qpos.shape
+        if ncols != 36:
+            raise ValueError(f"Expected qpos last dim 36; got {ncols}")
+
+        kimodo_to_mujoco_matrix = self.kimodo_to_mujoco_matrix.to(device=device, dtype=dtype)
+        mujoco_to_kimodo_matrix = kimodo_to_mujoco_matrix.T
+
+        root_mujoco = qpos[..., :3]
+        root_positions = torch.matmul(mujoco_to_kimodo_matrix[None, None, ...], root_mujoco[..., None]).squeeze(-1)
+
+        quat = qpos[..., 3:7]
+        if root_quat_w_first:
+            root_rot_mujoco = quaternion_to_matrix(quat)
+        else:
+            quat_wxyz = quat[..., [3, 0, 1, 2]]
+            root_rot_mujoco = quaternion_to_matrix(quat_wxyz)
+
+        O0 = self._rot_offsets_f2q[0].to(device=device, dtype=dtype)
+        # root_rot_mujoco is (..., 3, 3) after optional batch unsqueeze (e.g. (1, T, 3, 3)).
+        # Use ``...il`` so ``k`` sums with ``kl``; ``...ik`` incorrectly keeps ``k`` in the output.
+        R_f2q_root = torch.einsum(
+            "ij,...jk,kl->...il",
+            mujoco_to_kimodo_matrix,
+            root_rot_mujoco,
+            kimodo_to_mujoco_matrix,
+        )
+        R_kimodo_root = torch.einsum("ij,...jk->...ik", O0.T, R_f2q_root)
+
+        joint_dofs = qpos[..., 7:]
+        if mujoco_rest_zero:
+            rest_dofs = self._rest_dofs.to(device=device, dtype=dtype)
+            angles = joint_dofs + rest_dofs[None, None, :]
+            use_relative = True
+        else:
+            angles = joint_dofs
+            use_relative = False
+
+        nb_joints = self.skeleton.nbjoints
+        template = torch.eye(3, device=device, dtype=dtype).expand(batch_size, num_frames, nb_joints, 3, 3).contiguous()
+        template[:, :, 0] = R_kimodo_root
+
+        local_rot_mats = self._joint_dofs_to_local_rot_mats(
+            angles,
+            template,
+            device,
+            dtype,
+            use_relative=use_relative,
+        )
+
+        if batch_size != 1:
+            raise ValueError(f"Only a single clip is supported; got batch_size={batch_size}")
+
+        return complete_motion_dict(local_rot_mats[0], root_positions[0], self.skeleton, source_fps)
+
+    def save_csv(self, qpos: torch.Tensor | np.ndarray, csv_path):
+        # comment this
+        qpos = to_numpy(qpos)
+        shape = qpos.shape
+        if len(shape) == 2:
+            # only one motion: save it
+            np.savetxt(csv_path, qpos, delimiter=",")
+        if len(shape) == 3:
+            # batch of motions
+            if shape[0] == 1:
+                # if only one motion, just save it
+                np.savetxt(csv_path, qpos[0], delimiter=",")
+            else:
+                csv_path_base, ext = os.path.splitext(csv_path)
+                for i in range(shape[0]):
+                    self.save_csv(qpos[i], csv_path_base + "_" + str(i).zfill(2) + ext)
+
+    def _local_rots_to_joint_dofs(
+        self,
+        local_rot_mats: torch.Tensor,
+        axis_vals: torch.Tensor,
+    ) -> torch.Tensor:
+        """Extract per-joint single-DoF angles (radians) via Euler projection (for to_qpos/f2q)."""
+        x_joint_dof = torch.atan2(local_rot_mats[..., 2, 1], local_rot_mats[..., 2, 2])
+        y_joint_dof = torch.atan2(local_rot_mats[..., 0, 2], local_rot_mats[..., 0, 0])
+        z_joint_dof = torch.atan2(local_rot_mats[..., 1, 0], local_rot_mats[..., 1, 1])
+        xyz_joint_dofs = torch.stack([x_joint_dof, y_joint_dof, z_joint_dof], dim=-1)
+        axis_vals = axis_vals.to(device=local_rot_mats.device, dtype=local_rot_mats.dtype)
+        joint_dofs = (xyz_joint_dofs * axis_vals[None, None, :, :]).sum(dim=-1)
+        return joint_dofs
+
+    def _local_rots_to_joint_dofs_axis_angle(
+        self,
+        local_rot_mats: torch.Tensor,
+        axis_vals: torch.Tensor,
+    ) -> torch.Tensor:
+        """Extract per-joint single-DoF angles (radians) via axis-angle; round-trips with
+        axis_angle_to_matrix.
+
+        Args:
+            local_rot_mats: (..., num_hinges, 3, 3) in same frame as axis_vals.
+            axis_vals: (num_hinges, 3) unit axis per hinge.
+        Returns:
+            joint_dofs: (..., num_hinges) signed angle = dot(axis_angle(R), axis).
+        """
+        axis_vals = axis_vals.to(device=local_rot_mats.device, dtype=local_rot_mats.dtype)
+        full_aa = matrix_to_axis_angle(local_rot_mats)
+        joint_dofs = (full_aa * axis_vals).sum(dim=-1)
+        return joint_dofs
+
+    def _local_rots_f2q_to_joint_dofs(self, local_rot_mats_f2q: torch.Tensor) -> torch.Tensor:
+        """Extract per-joint single-DoF angles from local rotations in f2q space (for to_qpos)."""
+        axis_vals = self._mujoco_joint_axis_values_f2q_space
+        return self._local_rots_to_joint_dofs(local_rot_mats_f2q, axis_vals)
+
+    def _clamp_to_limits(self, joint_dofs: torch.Tensor) -> torch.Tensor:
+        """Clamp joint angles to XML limits (radians).
+
+        Angles are in kimodo convention (0 = rest).
+        """
+        device = joint_dofs.device
+        lo = self._joint_limits_min.to(device=device, dtype=joint_dofs.dtype)
+        hi = self._joint_limits_max.to(device=device, dtype=joint_dofs.dtype)
+        return torch.clamp(joint_dofs, lo[None, None, :], hi[None, None, :])
+
+    def _clamp_joint_dofs(self, joint_dofs: torch.Tensor, rest_dofs: torch.Tensor) -> torch.Tensor:
+        """Clamp joint angles to MuJoCo limits (radians), with rest_dofs conversion."""
+        device = joint_dofs.device
+        rest_dofs = rest_dofs.to(device=device, dtype=joint_dofs.dtype)
+        mujoco_dofs = joint_dofs - rest_dofs[None, None, :]
+        lo = self._joint_limits_min.to(device=device, dtype=joint_dofs.dtype)
+        hi = self._joint_limits_max.to(device=device, dtype=joint_dofs.dtype)
+        mujoco_dofs = torch.clamp(mujoco_dofs, lo[None, None, :], hi[None, None, :])
+        return mujoco_dofs + rest_dofs[None, None, :]
+
+    def _joint_dofs_to_local_rot_mats(
+        self,
+        joint_dofs: torch.Tensor,
+        original_local_rot_mats: torch.Tensor,
+        device: torch.device,
+        dtype: torch.dtype,
+        use_relative: bool = False,
+    ) -> torch.Tensor:
+        """Reconstruct full local rotation matrices from 1-DoF angles."""
+        out = original_local_rot_mats.clone()
+        axis_kimodo = self._mujoco_joint_axis_values_kimodo_space.to(device=device, dtype=dtype)
+        for i in range(joint_dofs.shape[-1]):
+            j = self._mujoco_indices_to_kimodo_indices[i].item()
+            angle = joint_dofs[..., i]
+            axis = axis_kimodo[i]
+            if use_relative:
+                axis_angle = angle[..., None] * axis[None, None, :]
+                R_local = axis_angle_to_matrix(axis_angle)
+            else:
+                rot_offsets_f2q = self._rot_offsets_f2q.to(device=device, dtype=dtype)
+                axis_in_f2q = torch.mv(rot_offsets_f2q[j], axis)
+                axis_angle = angle[..., None] * axis_in_f2q[None, None, :]
+                R_f2q = axis_angle_to_matrix(axis_angle)
+                R_local = torch.einsum("ij,btjk->btik", rot_offsets_f2q[j].T, R_f2q)
+            out[:, :, j, :, :] = R_local
+        return out
+
+    @ensure_batched(local_rot_mats=5, root_positions=3, lengths=1)
+    def project_to_real_robot_rotations(
+        self,
+        local_rot_mats: torch.Tensor,
+        root_positions: torch.Tensor,
+        clamp_to_limits: bool = True,
+        mujoco_rest_zero: bool = False,
+    ) -> dict:
+        """Project full 3D local rotations to G1 real robot DoF and back to 3D for viz.
+
+        Joint angles are extracted along each hinge axis, optionally clamped to XML limits, then
+        reconstructed to 3D rotations. When mujoco_rest_zero=False (default), raw angles are used
+        (baked-with-quat). When True, angles are relative to rest (0 = T-pose in MuJoCo).
+        """
+        device = local_rot_mats.device
+        dtype = local_rot_mats.dtype
+
+        # Transform to f2q frame and extract 1-DoF angles (axis-angle projection).
+        local_rot_f2q = torch.matmul(self._rot_offsets_f2q.to(device=device, dtype=dtype), local_rot_mats)
+        hinge_rots = local_rot_f2q[:, :, self._mujoco_indices_to_kimodo_indices, :, :]
+        axis_f2q = self._mujoco_joint_axis_values_f2q_space.to(device=device, dtype=dtype)
+        joint_dofs = self._local_rots_to_joint_dofs_axis_angle(hinge_rots, axis_f2q)
+
+        # Optionally express angles relative to rest (MuJoCo q=0 at T-pose).
+        if mujoco_rest_zero:
+            rest_dofs = self._rest_dofs_axis_angle.to(device=device, dtype=dtype)
+            angles = joint_dofs - rest_dofs[None, None, :]
+            use_relative = True
+        else:
+            angles = joint_dofs
+            use_relative = False
+
+        if clamp_to_limits:
+            if mujoco_rest_zero:
+                angles = self._clamp_to_limits(angles)
+            else:
+                rest_dofs_aa = self._rest_dofs_axis_angle.to(device=device, dtype=dtype)
+                angles = self._clamp_joint_dofs(angles, rest_dofs_aa)
+
+        # Reconstruct 3D local rotations from 1-DoF angles and run FK.
+        local_rot_mats_proj = self._joint_dofs_to_local_rot_mats(
+            angles, local_rot_mats, device, dtype, use_relative=use_relative
+        )
+        global_rot_mats, posed_joints, _ = self.skeleton.fk(local_rot_mats_proj, root_positions)
+        return {
+            "local_rot_mats": local_rot_mats_proj,
+            "global_rot_mats": global_rot_mats,
+            "posed_joints": posed_joints,
+            "root_positions": root_positions,
+        }
+
+    @ensure_batched(local_rot_mats=5, root_positions=3, lengths=1)
+    def to_qpos(
+        self,
+        local_rot_mats: torch.Tensor,
+        root_positions: torch.Tensor,
+        root_quat_w_first: bool = True,
+        mujoco_rest_zero: bool = False,
+    ) -> torch.Tensor:
+        """Fast batch conversion from kimodo features to mujoco qpos format.
+
+        Args:
+            local_rot_mats: (B, T, J, 3, 3) local rotation matrices (kimodo convention).
+            root_positions: (B, T, 3) root positions.
+            root_quat_w_first: If True, quaternion in qpos is (w,x,y,z).
+            mujoco_rest_zero: If True, joint angles are written so that kimodo rest (t-pose)
+                maps to q=0 in MuJoCo. If False, write raw joint_dofs.
+
+        Returns:
+            torch.Tensor of shape [batch, numFrames, 36] containing mujoco qpos data:
+            - root_trans (3) + root_quat (4) + joint_dofs (29) = 36 columns
+        """
+
+        batch_size, num_frames, nb_joints = local_rot_mats.shape[:3]
+        device, dtype = local_rot_mats.device, local_rot_mats.dtype
+
+        local_rot_mats = torch.matmul(self._rot_offsets_f2q.to(device), local_rot_mats)
+
+        batch_size, num_frames = root_positions.shape[0], root_positions.shape[1]
+
+        # Move precomputed matrices to the same device/dtype
+        kimodo_to_mujoco_matrix = self.kimodo_to_mujoco_matrix.to(device=device, dtype=dtype)
+
+        # Initialize output tensor: [batch, numFrames, 36]
+        qpos = torch.zeros((batch_size, num_frames, 36), dtype=dtype, device=device)
+
+        # Convert root translation: apply coordinate transformation
+        root_positions_mujoco = torch.matmul(kimodo_to_mujoco_matrix[None, None, ...], root_positions[..., None])
+        qpos[:, :, :3] = root_positions_mujoco.view(batch_size, num_frames, 3)
+
+        # Convert root rotation: apply coordinate transformation to rotation matrix
+        root_rot = local_rot_mats[:, :, 0, :]  # [batch, numFrames, 3, 3]
+
+        # Apply coordinate transformation: R_mujoco = kimodo_to_mujoco * R_kimodo * kimodo_to_mujoco^T
+        mujoco_to_kimodo_matrix = kimodo_to_mujoco_matrix.T
+        root_rot_mujoco = torch.matmul(
+            torch.matmul(kimodo_to_mujoco_matrix[None, None, ...], root_rot),
+            mujoco_to_kimodo_matrix[None, None, ...],
+        )
+        root_rot_quat = matrix_to_quaternion(root_rot_mujoco)  # [w, x, y, z]
+        if root_quat_w_first:
+            qpos[:, :, 3:7] = root_rot_quat[:, :, [0, 1, 2, 3]]  # [w, x, y, z]
+        else:
+            qpos[:, :, 3:7] = root_rot_quat[:, :, [1, 2, 3, 0]]  # [w, x, y, z] -> [x, y, z, w]
+
+        # Joint DOFs: raw angles or relative to rest (rest = q=0 in MuJoCo).
+        joint_rot_f2q = local_rot_mats[:, :, self._mujoco_indices_to_kimodo_indices, :, :]
+        joint_dofs = self._local_rots_f2q_to_joint_dofs(joint_rot_f2q)
+        if mujoco_rest_zero:
+            rest_dofs = self._rest_dofs.to(device=device, dtype=dtype)
+            qpos[:, :, 7:] = joint_dofs - rest_dofs[None, None, :]
+        else:
+            qpos[:, :, 7:] = joint_dofs
+        return qpos
+
+
+def apply_g1_real_robot_projection(
+    skeleton: G1Skeleton34,
+    joints_pos: torch.Tensor,
+    joints_rot: torch.Tensor,
+    clamp_to_limits: bool = True,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Project G1 motion to real robot DoF (1-DoF per joint) with optional axis limits.
+
+    Extracts a single angle per hinge along its axis (1-DoF), optionally clamps to
+    joint limits from the MuJoCo XML (when clamp_to_limits=True), then reconstructs
+    3D rotations and runs FK. T-pose (identity local rotations) is preserved.
+
+    Args:
+        skeleton: G1 skeleton instance.
+        joints_pos: (T, J, 3) or (B, T, J, 3) joint positions in global space.
+        joints_rot: (T, J, 3, 3) or (B, T, J, 3, 3) global rotation matrices.
+        clamp_to_limits: If True, clamp joint angles to XML axis limits (default True).
+
+    Returns:
+        (posed_joints, global_rot_mats) as tensors, same shape as inputs (batch preserved).
+    """
+
+    local_rot_mats = global_rots_to_local_rots(joints_rot, skeleton)
+    root_positions = joints_pos[..., skeleton.root_idx, :]
+
+    # Converter expects batch dim (B, T, ...); add and remove if single sequence.
+    single_sequence = local_rot_mats.dim() == 4
+    if single_sequence:
+        local_rot_mats = local_rot_mats.unsqueeze(0)
+        root_positions = root_positions.unsqueeze(0)
+
+    converter = MujocoQposConverter(skeleton)
+    projected = converter.project_to_real_robot_rotations(
+        local_rot_mats, root_positions, clamp_to_limits=clamp_to_limits
+    )
+
+    out_pos = projected["posed_joints"]
+    out_rot = projected["global_rot_mats"]
+    if single_sequence:
+        out_pos = out_pos.squeeze(0)
+        out_rot = out_rot.squeeze(0)
+    return out_pos, out_rot

kimodo/exports/smplx.py

@@ -0,0 +1,251 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Convert kimodo motion to AMASS/SMPL-X compatible parameters (axis-angle, Y-up or Z-up)."""
+
+import os
+from typing import Optional
+
+import einops
+import numpy as np
+import torch
+
+from kimodo.assets import skeleton_asset_path
+from kimodo.geometry import axis_angle_to_matrix, matrix_to_axis_angle
+from kimodo.tools import ensure_batched, to_numpy, to_torch
+
+
+def kimodo_y_up_to_amass_coord_rotation_matrix() -> np.ndarray:
+    """3x3 rotation mapping Kimodo Y-up (+Z forward) to AMASS Z-up (+Y forward).
+
+    Used by :func:`get_amass_parameters` and :func:`amass_arrays_to_kimodo_motion` (inverse).
+    """
+    y_up_to_z_up = np.array(
+        [
+            [1.0, 0.0, 0.0],
+            [0.0, 0.0, -1.0],
+            [0.0, 1.0, 0.0],
+        ],
+        dtype=np.float32,
+    )
+    rot_z_180 = np.array(
+        [
+            [-1.0, 0.0, 0.0],
+            [0.0, -1.0, 0.0],
+            [0.0, 0.0, 1.0],
+        ],
+        dtype=np.float32,
+    )
+    return np.matmul(rot_z_180, y_up_to_z_up).astype(np.float32)
+
+
+@ensure_batched(local_rot_mats=5, root_positions=3, lengths=1)
+def get_amass_parameters(
+    local_rot_mats,
+    root_positions,
+    skeleton,
+    z_up=True,
+):
+    """Convert local rot mats and root positions to AMASS-style trans and pose_body; optional z_up
+    coordinate transform.
+
+    Our method generates motions with Y-up and +Z forward; if z_up=True, transform to Z-up and +Y
+    forward as in AMASS.
+    """
+    # Our method generate motions with Y-up and +Z forward
+    # if z_up = True, we transform this to: Z-up with +Y forward, as in AMASS
+    # Remove the root offset; SMPL-X FK adds pelvis offset back.
+    pelvis_offset = skeleton.neutral_joints[skeleton.root_idx].cpu().numpy()
+    trans = root_positions - pelvis_offset
+
+    root_rot_mats = to_numpy(local_rot_mats[:, :, 0])
+    local_rot_axis_angle = to_numpy(matrix_to_axis_angle(to_torch(local_rot_mats)))
+    pose_body = einops.rearrange(local_rot_axis_angle[:, :, 1:], "b t j d -> b t (j d)")
+
+    # Optionally convert from Y-up to Z-up coordinates.
+    if z_up:
+        y_up_to_z_up = kimodo_y_up_to_amass_coord_rotation_matrix()
+        root_rot_mats = np.matmul(y_up_to_z_up, root_rot_mats)
+        trans = np.matmul(trans + pelvis_offset, y_up_to_z_up.T) - pelvis_offset
+
+    root_orient = to_numpy(matrix_to_axis_angle(to_torch(root_rot_mats)))
+    return trans, root_orient, pose_body
+
+
+def amass_arrays_to_kimodo_motion(
+    trans: np.ndarray,
+    root_orient: np.ndarray,
+    pose_body: np.ndarray,
+    skeleton,
+    source_fps: float,
+    *,
+    z_up: bool = True,
+):
+    """Inverse of :func:`get_amass_parameters` for a single sequence (AMASS → Kimodo motion dict).
+
+    Args:
+        trans: ``(T, 3)`` AMASS root translation (same as ``trans`` in AMASS NPZ).
+        root_orient: ``(T, 3)`` axis-angle root orientation in AMASS coordinates (z-up when ``z_up``).
+        pose_body: ``(T, 63)`` body pose axis-angle (21 joints × 3).
+        skeleton: :class:`~kimodo.skeleton.definitions.SMPLXSkeleton22` instance.
+        source_fps: Source frame rate (Hz) of the AMASS recording.
+        z_up: If ``True``, invert the same Y-up↔Z-up transform as ``get_amass_parameters(..., z_up=True)``.
+
+    Returns:
+        Motion dict compatible with :func:`kimodo.exports.motion_io.save_kimodo_npz`.
+    """
+    from kimodo.exports.motion_io import complete_motion_dict
+
+    trans = np.asarray(trans, dtype=np.float32)
+    root_orient = np.asarray(root_orient, dtype=np.float32)
+    pose_body = np.asarray(pose_body, dtype=np.float32)
+    if trans.ndim != 2 or trans.shape[-1] != 3:
+        raise ValueError(f"trans must be (T, 3); got {trans.shape}")
+    if root_orient.shape != trans.shape:
+        raise ValueError(f"root_orient shape {root_orient.shape} must match trans {trans.shape}")
+    t = trans.shape[0]
+    if pose_body.shape != (t, 63):
+        raise ValueError(f"pose_body must be (T, 63); got {pose_body.shape}")
+
+    pelvis_offset = skeleton.neutral_joints[skeleton.root_idx].detach().cpu().numpy().astype(np.float32)
+    device = skeleton.neutral_joints.device
+    dtype = torch.float32
+
+    Y_np = kimodo_y_up_to_amass_coord_rotation_matrix()
+    if z_up:
+        y_up_to_z_up = torch.from_numpy(Y_np).to(device=device, dtype=dtype)
+        # trans_amass = root_kimodo @ Y.T - pelvis_offset  =>  root_kimodo = (trans_amass + pelvis_offset) @ Y
+        root_positions_np = (trans + pelvis_offset) @ Y_np
+    else:
+        root_positions_np = trans + pelvis_offset
+
+    root_positions = torch.from_numpy(root_positions_np).to(device=device, dtype=dtype)
+
+    R_amass_root = axis_angle_to_matrix(torch.from_numpy(root_orient).to(device=device, dtype=dtype))
+    if z_up:
+        R_kimodo_root = torch.einsum("ij,tjk->tik", y_up_to_z_up.T, R_amass_root)
+    else:
+        R_kimodo_root = R_amass_root
+
+    nb = skeleton.nbjoints
+    if nb != 22:
+        raise ValueError(f"Expected SMPL-X body skeleton with 22 joints; got {nb}")
+
+    local_rot_mats = torch.zeros((t, nb, 3, 3), device=device, dtype=dtype)
+    local_rot_mats[:, 0] = R_kimodo_root
+
+    pose_aa = torch.from_numpy(pose_body.reshape(t, 21, 3)).to(device=device, dtype=dtype)
+    local_rot_mats[:, 1:] = axis_angle_to_matrix(pose_aa.reshape(-1, 3)).reshape(t, 21, 3, 3)
+
+    return complete_motion_dict(local_rot_mats, root_positions, skeleton, source_fps)
+
+
+def amass_npz_to_kimodo_motion(npz_path: str, skeleton, source_fps: Optional[float] = None, *, z_up: bool = True):
+    """Load an AMASS-style ``.npz`` and return a Kimodo motion dict.
+
+    Args:
+        npz_path: Path to AMASS NPZ (``trans``, ``root_orient``, ``pose_body``, ...).
+        skeleton: SMPL-X skeleton instance.
+        source_fps: Source frame rate (Hz); if ``None``, uses ``mocap_frame_rate``
+            from the file when present, else ``30.0``.
+        z_up: Same meaning as :func:`amass_arrays_to_kimodo_motion`.
+    """
+    with np.load(npz_path, allow_pickle=True) as data:
+        trans = np.asarray(data["trans"], dtype=np.float32)
+        root_orient = np.asarray(data["root_orient"], dtype=np.float32)
+        pose_body = np.asarray(data["pose_body"], dtype=np.float32)
+        if source_fps is None:
+            source_fps = float(data["mocap_frame_rate"]) if "mocap_frame_rate" in data.files else 30.0
+
+    return amass_arrays_to_kimodo_motion(trans, root_orient, pose_body, skeleton, source_fps, z_up=z_up)
+
+
+class AMASSConverter:
+    def __init__(
+        self,
+        fps,
+        skeleton,
+        beta_path=str(skeleton_asset_path("smplx22", "beta.npy")),
+        mean_hands_path=str(skeleton_asset_path("smplx22", "mean_hands.npy")),
+    ):
+        self.fps = fps
+        self.skeleton = skeleton
+        # Load betas
+        if os.path.exists(beta_path):
+            # only use first 16 betas to match AMASS
+            betas = np.load(beta_path)[:16]
+        else:
+            betas = np.zeros(16)
+
+        # Load mean hands
+        if os.path.exists(mean_hands_path):
+            mean_hands = np.load(mean_hands_path)
+        else:
+            mean_hands = np.zeros(90)
+
+        self.default_frame_params = {
+            "pose_jaw": np.zeros(3),
+            "pose_eye": np.zeros(6),
+            "pose_hand": mean_hands,
+        }
+        self.output_dict_base = {
+            "gender": "neutral",
+            "surface_model_type": "smplx",
+            "betas": betas,
+            "num_betas": len(betas),
+            "mocap_frame_rate": float(fps),
+        }
+
+    def convert_save_npz(self, output: dict, npz_path, z_up=True):
+        trans, root_orient, pose_body = get_amass_parameters(
+            output["local_rot_mats"],
+            output["root_positions"],
+            self.skeleton,
+            z_up=z_up,
+        )
+        nb_frames = trans.shape[-2]
+
+        amass_output_base = self.output_dict_base.copy()
+        for key, val in self.default_frame_params.items():
+            amass_output_base[key] = einops.repeat(val, "d -> t d", t=nb_frames)
+
+        amass_output_base["mocap_time_length"] = nb_frames / self.fps
+        self.save_npz(trans, root_orient, pose_body, amass_output_base, npz_path)
+
+    def save_npz(self, trans, root_orient, pose_body, base_output, npz_path):
+        shape = trans.shape
+        if len(shape) == 3 and shape[0] == 1:
+            # if only one motion, squeeze the data
+            trans = trans[0]
+            root_orient = root_orient[0]
+            pose_body = pose_body[0]
+            shape = trans.shape
+        if len(shape) == 2:
+            amass_output = {
+                "trans": trans,
+                "root_orient": root_orient,
+                "pose_body": pose_body,
+            } | base_output
+            np.savez(npz_path, **amass_output)
+
+        elif len(shape) == 3:
+            # real batch of motions
+            npz_path_base, ext = os.path.splitext(npz_path)
+            for i in range(shape[0]):
+                npz_path_i = npz_path_base + "_" + str(i).zfill(2) + ext
+                self.save_npz(trans[i], root_orient[i], pose_body[i], base_output, npz_path_i)
+
+
+# amass_output = {
+#     "gender": "neutral",
+#     "surface_model_type": "smplx",
+#     "mocap_frame_rate": float(fps),
+#     "mocap_time_length": len(motion) / float(fps)
+#     "trans": trans,
+#     "betas": betas,
+#     "num_betas": len(betas),
+#     "root_orient": np.array([T, 3]), # axis angle
+#     "pose_body": np.array([T, 63]), # 63=21*3, axis angle 21 = 22 - root
+#     "pose_hand": np.array([T, 90]), # 90=30*3=15*2*3 axis angle (load from mean_hands)
+#     "pose_jaw": np.array([T, 3]), # all zeros is fine
+#     "pose_eye": np.array([T, 6]), # all zeros is fine`
+# }

kimodo/geometry.py

@@ -0,0 +1,216 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Rotation and representation conversions: axis-angle, quaternion, matrix, 6D continuous."""
+
+import torch
+import torch.nn.functional as F
+
+
+def angle_to_Y_rotation_matrix(angle: torch.Tensor) -> torch.Tensor:
+    """Build a rotation matrix around the Y axis from a scalar angle (radians).
+
+    Shape: angle.shape + (3, 3).
+    """
+    cos, sin = torch.cos(angle), torch.sin(angle)
+    one, zero = torch.ones_like(angle), torch.zeros_like(angle)
+    mat = torch.stack((cos, zero, sin, zero, one, zero, -sin, zero, cos), -1)
+    mat = mat.reshape(angle.shape + (3, 3))
+    return mat
+
+
+def matrix_to_cont6d(matrix: torch.Tensor) -> torch.Tensor:
+    """Convert rotation matrix to 6D continuous representation (first two columns).
+
+    Shape: (..., 3, 3) -> (..., 6).
+    """
+    cont_6d = torch.concat([matrix[..., 0], matrix[..., 1]], dim=-1)
+    return cont_6d
+
+
+def cont6d_to_matrix(cont6d: torch.Tensor) -> torch.Tensor:
+    """Convert 6D continuous representation to rotation matrix (Gram–Schmidt on two columns).
+
+    Last dim must be 6.
+    """
+    assert cont6d.shape[-1] == 6, "The last dimension must be 6"
+    x_raw = cont6d[..., 0:3]
+    y_raw = cont6d[..., 3:6]
+
+    x = x_raw / torch.norm(x_raw, dim=-1, keepdim=True)
+    z = torch.cross(x, y_raw, dim=-1)
+    z = z / torch.norm(z, dim=-1, keepdim=True)
+
+    y = torch.cross(z, x, dim=-1)
+
+    x = x[..., None]
+    y = y[..., None]
+    z = z[..., None]
+
+    mat = torch.cat([x, y, z], dim=-1)
+    return mat
+
+
+def axis_angle_to_matrix(axis_angle: torch.Tensor) -> torch.Tensor:
+    """Convert axis-angle to rotation matrix.
+
+    Args:
+        axis_angle: (..., 3) axis-angle vectors (angle = norm, axis = normalized)
+    Returns:
+        rotmat: (..., 3, 3) rotation matrices
+    """
+    eps = 1e-6
+    angle = torch.norm(axis_angle, dim=-1, keepdim=True)  # (..., 1)
+    axis = axis_angle / (angle + eps)
+
+    x, y, z = axis.unbind(-1)
+
+    zero = torch.zeros_like(x)
+    K = torch.stack([zero, -z, y, z, zero, -x, -y, x, zero], dim=-1).reshape(*axis.shape[:-1], 3, 3)
+
+    eye = torch.eye(3, device=axis.device, dtype=axis.dtype)
+    eye = eye.expand(*axis.shape[:-1], 3, 3)
+
+    sin = torch.sin(angle)[..., None]
+    cos = torch.cos(angle)[..., None]
+
+    R = eye + sin * K + (1 - cos) * (K @ K)
+    return R
+
+
+def matrix_to_axis_angle(R: torch.Tensor) -> torch.Tensor:
+    """Convert rotation matrix to axis-angle via quaternions (more numerically stable).
+
+    Args:
+        R: (..., 3, 3) rotation matrices
+    Returns:
+        axis_angle: (..., 3)
+    """
+    # Go through quaternions for numerical stability
+    quat = matrix_to_quaternion(R)  # (..., 4) with (w, x, y, z)
+    return quaternion_to_axis_angle(quat)
+
+
+def quaternion_to_axis_angle(quat: torch.Tensor) -> torch.Tensor:
+    """Convert quaternion to axis-angle representation.
+
+    Args:
+        quat: (..., 4) quaternions with real part first (w, x, y, z)
+    Returns:
+        axis_angle: (..., 3)
+    """
+    eps = 1e-6
+
+    # Ensure canonical form to avoid sign ambiguity.
+    # Primary: prefer w > 0. When w ≈ 0 (angle ≈ π), prefer first nonzero xyz > 0.
+    w = quat[..., 0:1]
+    xyz = quat[..., 1:]
+
+    # Find first significant component of xyz for tie-breaking when w ≈ 0
+    first_significant = xyz[..., 0:1]  # use x component as tie-breaker
+
+    # Flip if: w < 0, OR (w ≈ 0 AND first xyz component < 0)
+    should_flip = (w < -eps) | ((w.abs() <= eps) & (first_significant < 0))
+    quat = torch.where(should_flip, -quat, quat)
+
+    w = quat[..., 0]
+    xyz = quat[..., 1:]
+
+    # sin(angle/2) = ||xyz||
+    sin_half_angle = xyz.norm(dim=-1)
+
+    # angle = 2 * atan2(sin(angle/2), cos(angle/2))
+    # This is more stable than 2 * acos(w) near angle=0
+    angle = 2.0 * torch.atan2(sin_half_angle, w)
+
+    # axis = xyz / sin(angle/2), but handle small angles
+    # For small angles: axis-angle ≈ 2 * xyz (since sin(x) ≈ x for small x)
+    small_angle = sin_half_angle.abs() < eps
+
+    # Safe division
+    scale = torch.where(
+        small_angle,
+        2.0 * torch.ones_like(angle),  # small angle: axis_angle ≈ 2 * xyz
+        angle / sin_half_angle.clamp(min=eps),
+    )
+
+    return xyz * scale.unsqueeze(-1)
+
+
+def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor:
+    """Returns torch.sqrt(torch.max(0, x)) subgradient is zero where x is 0."""
+    return torch.sqrt(x * (x > 0).to(x.dtype))
+
+
+def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor:
+    """Convert rotations given as rotation matrices to quaternions.
+
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+    Returns:
+        quaternions with real part first, as tensor of shape (..., 4).
+    """
+    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
+        raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.")
+
+    batch_dim = matrix.shape[:-2]
+    m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind(matrix.reshape(batch_dim + (9,)), dim=-1)
+
+    q_abs = _sqrt_positive_part(
+        torch.stack(
+            [
+                1.0 + m00 + m11 + m22,
+                1.0 + m00 - m11 - m22,
+                1.0 - m00 + m11 - m22,
+                1.0 - m00 - m11 + m22,
+            ],
+            dim=-1,
+        )
+    )
+
+    quat_by_rijk = torch.stack(
+        [
+            torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1),
+            torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1),
+            torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1),
+            torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1),
+        ],
+        dim=-2,
+    )
+
+    flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device)
+    quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr))
+
+    return (
+        (F.one_hot(q_abs.argmax(dim=-1), num_classes=4)[..., None] * quat_candidates)
+        .sum(dim=-2)
+        .reshape(batch_dim + (4,))
+    )
+
+
+def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor:
+    """Convert rotations given as quaternions to rotation matrices.
+
+    Args:
+        quaternions: quaternions with real part first,
+            as tensor of shape (..., 4).
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    r, i, j, k = torch.unbind(quaternions, -1)
+    two_s = 2.0 / (quaternions * quaternions).sum(-1)
+
+    o = torch.stack(
+        (
+            1 - two_s * (j * j + k * k),
+            two_s * (i * j - k * r),
+            two_s * (i * k + j * r),
+            two_s * (i * j + k * r),
+            1 - two_s * (i * i + k * k),
+            two_s * (j * k - i * r),
+            two_s * (i * k - j * r),
+            two_s * (j * k + i * r),
+            1 - two_s * (i * i + j * j),
+        ),
+        -1,
+    )
+    return o.reshape(quaternions.shape[:-1] + (3, 3))

kimodo/meta.py

@@ -0,0 +1,80 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Parse and normalize prompt text/duration data from meta dicts."""
+
+import os
+from typing import Any, Optional
+
+from kimodo.tools import load_json
+
+from .sanitize import sanitize_text, sanitize_texts
+
+
+def load_prompts_from_meta(meta_path: str, **kwargs):
+    """Load prompts from a meta dict or file. If fps is provided, the durations are converted to
+    frames.
+
+    Args:
+        meta_path: Path to the meta file.
+        **kwargs: Additional arguments to pass to parse_prompts_from_meta.
+
+    Returns:
+        texts: List of texts.
+        durations: List of durations in seconds or frames.
+    """
+    if not os.path.exists(meta_path):
+        raise FileNotFoundError(f"meta.json not found in input folder: {meta_path}")
+
+    meta = load_json(meta_path)
+    return parse_prompts_from_meta(meta, **kwargs)
+
+
+def parse_prompts_from_meta(
+    meta: dict[str, Any],
+    fps: Optional[float] = None,
+    sanitize: bool = False,
+) -> tuple[list[str], list[float]]:
+    """Parse prompt texts and durations from a meta dict into normalized lists. If fps is provided,
+    the durations are converted to frames.
+
+    Accepts either:
+    - Single prompt: "text" (str) and "duration" (float) in seconds.
+    - Multiple prompts: "texts" (list of str) and "durations" (list of float) in seconds.
+
+    Returns:
+        (texts, durations): texts as list of str, durations as list of float (seconds or frames).
+        Lengths of both lists are equal.
+
+    Raises:
+        ValueError: If meta does not contain a recognized format.
+    """
+    # Single prompt
+    if "text" in meta and "duration" in meta:
+        text = meta["text"]
+        duration = float(meta["duration"])
+        if fps is not None:
+            duration = int(duration * fps)
+        if isinstance(text, list):
+            raise ValueError("meta has 'text' but it is a list; use 'texts' for multiple prompts")
+
+        if sanitize:
+            text = sanitize_text(text)
+        return ([text], [duration])
+
+    # Multiple prompts
+    if "texts" in meta and "durations" in meta:
+        texts = meta["texts"]
+        durations = meta["durations"]
+        if not isinstance(texts, list) or not isinstance(durations, list):
+            raise ValueError("meta 'texts' and 'durations' must be lists")
+        if len(texts) != len(durations):
+            raise ValueError(f"meta 'texts' and 'durations' length mismatch: {len(texts)} vs {len(durations)}")
+        durations = [float(d) for d in durations]
+        if fps is not None:
+            durations = [int(d * fps) for d in durations]
+
+        if sanitize:
+            texts = sanitize_texts(texts)
+        return texts, durations
+
+    raise ValueError("meta must contain either 'text' and 'duration', or 'texts' and 'durations'.")

kimodo/metrics/__init__.py

@@ -0,0 +1,39 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Evaluation metrics for motion quality (foot skate, contact consistency, constraint following)."""
+
+from .base import (
+    Metric,
+    aggregate_metrics,
+    clear_metrics,
+    compute_metrics,
+)
+from .constraints import ContraintFollow
+from .foot_skate import (
+    FootContactConsistency,
+    FootSkateFromContacts,
+    FootSkateFromHeight,
+    FootSkateRatio,
+)
+from .tmr import (
+    TMR_EmbeddingMetric,
+    TMR_Metric,
+    compute_tmr_per_sample_retrieval,
+    compute_tmr_retrieval_metrics,
+)
+
+__all__ = [
+    "Metric",
+    "ContraintFollow",
+    "FootContactConsistency",
+    "FootSkateFromContacts",
+    "FootSkateFromHeight",
+    "FootSkateRatio",
+    "TMR_EmbeddingMetric",
+    "TMR_Metric",
+    "aggregate_metrics",
+    "clear_metrics",
+    "compute_metrics",
+    "compute_tmr_per_sample_retrieval",
+    "compute_tmr_retrieval_metrics",
+]

kimodo/metrics/base.py

@@ -0,0 +1,66 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Base metric class and batch/aggregate helpers."""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import Dict, List
+
+import torch
+
+
+class Metric:
+    """Base class for metrics that accumulate results over multiple __call__ and expose
+    aggregate()."""
+
+    def __init__(self, **kwargs):
+        self.clear()
+
+    def __call__(self, *args, **kwargs):
+        """Compute metric for current batch, append to saved_metrics, and return the batch
+        result."""
+        metrics = self._compute(*args, **kwargs)
+        for key, val in metrics.items():
+            self.saved_metrics[key].append(val.detach().cpu().float())
+        return metrics
+
+    def _compute(self, **kwargs):
+        """Subclasses implement this to compute metric dict from batch inputs."""
+        raise NotImplementedError()
+
+    def clear(self):
+        """Reset all accumulated metric values."""
+        self.saved_metrics = defaultdict(list)
+
+    def aggregate(self):
+        """Return a dict of concatenated/stacked tensors over all accumulated batches."""
+        output = {}
+        for key, lst in self.saved_metrics.items():
+            try:
+                output[key] = torch.cat(lst)
+            except RuntimeError:
+                output[key] = torch.stack(lst)
+        return output
+
+
+def compute_metrics(metrics_list: List[Metric], metrics_in: Dict) -> Dict:
+    """Run each metric on metrics_in and return the combined dict of batch results."""
+    metrics_out = {}
+    for metric in metrics_list:
+        metrics_out.update(metric(**metrics_in))
+    return metrics_out
+
+
+def aggregate_metrics(metrics_list: List[Metric]) -> Dict:
+    """Return combined aggregated results (concatenated over batches) for all metrics."""
+    metrics_out = {}
+    for metric in metrics_list:
+        metrics_out.update(metric.aggregate())
+    return metrics_out
+
+
+def clear_metrics(metrics_list: List[Metric]) -> None:
+    """Clear accumulated values for all metrics in the list."""
+    for metric in metrics_list:
+        metric.clear()

kimodo/metrics/constraints.py

@@ -0,0 +1,87 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Constraint-following metrics."""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import Dict, List, Optional
+
+import torch
+from torch import Tensor
+
+from kimodo.constraints import (
+    EndEffectorConstraintSet,
+    FullBodyConstraintSet,
+    Root2DConstraintSet,
+)
+from kimodo.tools import ensure_batched
+
+from .base import Metric
+
+
+class ContraintFollow(Metric):
+    """Constraint-following metric dispatcher for kimodo constraint sets."""
+
+    def __init__(
+        self,
+        skeleton,
+        root_threshold: float = 0.10,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.skeleton = skeleton
+        self.root_threshold = root_threshold
+
+    @ensure_batched(posed_joints=4, constraints_lst=2, lengths=1)
+    def _compute(
+        self,
+        posed_joints: Tensor,
+        constraints_lst: Optional[List],
+        lengths: Optional[Tensor] = None,
+        **kwargs,
+    ) -> Dict:
+        if not constraints_lst:
+            return {}
+
+        root_idx = self.skeleton.root_idx
+        output = defaultdict(list)
+
+        for posed_joints_s, constraint_lst_s, lengths_s in zip(posed_joints, constraints_lst, lengths):
+            output_seq = defaultdict(list)
+            for constraint in constraint_lst_s:
+                frame_idx = constraint.frame_indices.to(device=posed_joints_s.device, dtype=torch.long)
+                assert frame_idx.max() < lengths_s, "The constraint is defined outsite the lenght of the motion."
+                if frame_idx.numel() == 0:
+                    continue
+
+                if isinstance(constraint, Root2DConstraintSet):
+                    pred_root2d = posed_joints_s[frame_idx, root_idx][:, [0, 2]]
+                    target = constraint.smooth_root_2d.to(posed_joints_s.device)
+
+                    dist = torch.norm(pred_root2d - target, dim=-1)
+                    output_seq["constraint_root2d_err"].append(dist)
+                    hit = (dist <= self.root_threshold).float()
+                    output_seq["constraint_root2d_acc"].append(hit)
+
+                elif isinstance(constraint, FullBodyConstraintSet):
+                    pred = posed_joints_s[frame_idx]
+                    target = constraint.global_joints_positions.to(posed_joints_s.device)
+                    err = torch.norm(pred - target, dim=-1)
+                    output_seq["constraint_fullbody_keyframe"].append(err)
+
+                elif isinstance(constraint, EndEffectorConstraintSet):
+                    pos_idx = constraint.pos_indices.to(device=posed_joints_s.device, dtype=torch.long)
+                    pred = posed_joints_s[frame_idx].index_select(1, pos_idx)
+                    target = constraint.global_joints_positions.to(posed_joints_s.device).index_select(1, pos_idx)
+                    err = torch.norm(pred - target, dim=-1)
+                    output_seq["constraint_end_effector"].append(err)
+
+            # in case we have several same constraints in the list
+            for key, val in output_seq.items():
+                output[key].append(torch.cat(val).mean())
+
+        reduced = {}
+        for key, vals in output.items():
+            reduced[key] = torch.stack(vals, dim=0)
+        return reduced

kimodo/metrics/foot_skate.py

@@ -0,0 +1,253 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Foot skate and contact consistency metrics."""
+
+from __future__ import annotations
+
+from typing import Dict, Optional
+
+import torch
+from torch import Tensor
+
+from kimodo.motion_rep.feature_utils import compute_vel_xyz
+from kimodo.motion_rep.feet import foot_detect_from_pos_and_vel
+from kimodo.skeleton import SkeletonBase
+from kimodo.tools import ensure_batched
+
+from .base import Metric
+
+
+def get_four_contacts(fidx: list):
+    if len(fidx) == 4:
+        return fidx
+    if len(fidx) == 6:
+        # For soma77
+        # remove "LeftToeEnd" and "RightToeEnd"
+        fidx = fidx[:2] + fidx[3:5]
+        return fidx
+    raise ValueError("Expects 4 or 6 foot joints (heel/toe per foot)")
+
+
+class FootSkateFromHeight(Metric):
+    """When toe joint is near the floor, measures mean velocity of the toes."""
+
+    def __init__(
+        self,
+        skeleton: SkeletonBase,
+        fps: float,
+        height_thresh: float = 0.05,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.height_thresh = height_thresh
+        self.skeleton = skeleton
+        self.fps = fps
+
+    @ensure_batched(posed_joints=4, lengths=1)
+    def _compute(
+        self,
+        posed_joints: Tensor,
+        lengths: Optional[Tensor] = None,
+        **kwargs,
+    ) -> Dict:
+        fidx = self.skeleton.foot_joint_idx
+        fidx = get_four_contacts(fidx)
+
+        feet_pos = posed_joints[:, :, fidx]
+        toe_pos = feet_pos[:, :, [1, 3]]
+
+        toe_on_floor = (toe_pos[..., 1] < self.height_thresh)[:, :-1]  # y-up [B, T, 2] where [left right]
+
+        dt = 1.0 / self.fps
+        toe_vel = torch.norm(toe_pos[:, 1:] - toe_pos[:, :-1], dim=-1) / dt  # [B, nframes-1, 2]
+
+        # compute err
+        contact_toe_vel = toe_vel * toe_on_floor  # vel when corresponding toe is on ground
+
+        # account for generated length
+        # since they are velocities use length-1 to avoid inaccurate vel going one frame past len
+        device = toe_on_floor.device
+        len_mask = torch.arange(toe_on_floor.shape[1], device=device)[None, :, None].expand(toe_on_floor.shape) < (
+            lengths[:, None, None] - 1
+        )
+        toe_on_floor = toe_on_floor * len_mask
+        contact_toe_vel = contact_toe_vel * len_mask
+
+        mean_vel = torch.sum(contact_toe_vel, (1, 2)) / (torch.sum(toe_on_floor, (1, 2)) + 1e-6)
+        return {"foot_skate_from_height": mean_vel}
+
+
+class FootSkateFromContacts(Metric):
+    """Measures velocity of the toes and ankles when predicted to be in contact."""
+
+    def __init__(
+        self,
+        skeleton: SkeletonBase,
+        fps: float,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.skeleton = skeleton
+        self.fps = fps
+
+    @ensure_batched(posed_joints=4, foot_contacts=3, lengths=1)
+    def _compute(
+        self,
+        posed_joints: Tensor,
+        foot_contacts: Tensor,
+        lengths: Optional[Tensor] = None,
+        **kwargs,
+    ) -> Dict:
+        fidx = self.skeleton.foot_joint_idx
+        fidx = get_four_contacts(fidx)
+
+        feet_pos = posed_joints[:, :, fidx]
+        dt = 1.0 / self.fps
+        foot_vel = torch.norm(feet_pos[:, 1:] - feet_pos[:, :-1], dim=-1) / dt
+
+        if foot_contacts.shape[-1] == 6:
+            # For soma77
+            # remove "LeftToeEnd" and "RightToeEnd"
+            foot_contacts = foot_contacts[..., [0, 1, 3, 4]]
+
+        foot_contacts = foot_contacts[:, :-1]
+        vel_err = foot_vel * foot_contacts
+
+        # account for generated length
+        # since they are velocities use length-1 to avoid inaccurate vel going one frame past len
+        device = foot_contacts.device
+        len_mask = torch.arange(foot_contacts.shape[1], device=device)[None, :, None].expand(foot_contacts.shape) < (
+            lengths[:, None, None] - 1
+        )
+        foot_contacts = foot_contacts * len_mask
+        vel_err = vel_err * len_mask
+
+        mean_vel = torch.sum(vel_err, (1, 2)) / (torch.sum(foot_contacts, (1, 2)) + 1e-6)  # mean over contacting frames
+
+        # Compute max velocity error across all feet and frames (per batch)
+        max_vel = vel_err.amax(dim=(1, 2))  # [B]
+
+        return {
+            "foot_skate_from_pred_contacts": mean_vel,
+            "foot_skate_max_vel": max_vel,
+        }
+
+
+class FootSkateRatio(Metric):
+    """Compute fraction of frames where the foot skates when it is on the ground.
+
+    Inspired by GMD: https://github.com/korrawe/guided-motion-diffusion/blob/main/data_loaders/humanml/utils/metrics.py#L204
+    """
+
+    def __init__(
+        self,
+        skeleton: SkeletonBase,
+        fps: float,
+        height_thresh=0.05,
+        vel_thresh=0.2,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.height_thresh = height_thresh
+        self.vel_thresh = vel_thresh
+
+        self.skeleton = skeleton
+        self.fps = fps
+
+    @ensure_batched(posed_joints=4, foot_contacts=3, lengths=1)
+    def _compute(
+        self,
+        posed_joints: Tensor,
+        foot_contacts: Tensor,
+        lengths: Optional[Tensor] = None,
+        **kwargs,
+    ) -> Dict:
+        fidx = self.skeleton.foot_joint_idx
+        fidx = get_four_contacts(fidx)
+
+        feet_pos = posed_joints[:, :, fidx]
+        toe_pos = feet_pos[:, :, [1, 3]]
+
+        toe_on_floor = toe_pos[..., 1] < self.height_thresh  # y-up [B, T, 2] where [left right]
+        # current and next frame on floor to consider it in contact
+        toe_on_floor = torch.logical_and(toe_on_floor[:, :-1], toe_on_floor[:, 1:])  # [B, T-1, 2]
+
+        dt = 1.0 / self.fps
+        toe_vel = torch.norm(toe_pos[:, 1:] - toe_pos[:, :-1], dim=-1) / dt  # [B, nframes-1, 2]
+
+        # compute err
+        contact_toe_vel = toe_vel * toe_on_floor  # vel when corresponding toe is on ground
+
+        # account for generated length
+        # since they are velocities use length-1 to avoid inaccurate vel going one frame past len
+        device = toe_on_floor.device
+        len_mask = torch.arange(toe_on_floor.shape[1], device=device)[None, :, None].expand(toe_on_floor.shape) < (
+            lengths[:, None, None] - 1
+        )
+        toe_on_floor = toe_on_floor * len_mask
+        contact_toe_vel = contact_toe_vel * len_mask
+
+        # skating if velocity during contact > thresh
+        toe_skate = contact_toe_vel > self.vel_thresh
+        skate_ratio = torch.sum(toe_skate, (1, 2)) / (torch.sum(toe_on_floor, (1, 2)) + 1e-6)
+        return {"foot_skate_ratio": skate_ratio}
+
+
+class FootContactConsistency(Metric):
+    """Measures consistency between heuristic detected foot contacts (from height and velocity) and
+    predicted foot contacts.
+
+    i.e. accuracy of how well predicted matches heuristic.
+    """
+
+    def __init__(
+        self,
+        skeleton: SkeletonBase,
+        fps: float,
+        vel_thresh: float = 0.15,
+        height_thresh: float = 0.10,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.vel_thresh = vel_thresh
+        self.height_thresh = height_thresh
+
+        self.skeleton = skeleton
+        self.fps = fps
+
+    @ensure_batched(posed_joints=4, foot_contacts=3, lengths=1)
+    def _compute(
+        self,
+        posed_joints: Tensor,
+        foot_contacts: Tensor,
+        lengths: Optional[Tensor] = None,
+        **kwargs,
+    ) -> Dict:
+        velocity = compute_vel_xyz(posed_joints, float(self.fps), lengths=lengths)
+        heuristic_contacts = foot_detect_from_pos_and_vel(
+            posed_joints,
+            velocity,
+            self.skeleton,
+            self.vel_thresh,
+            self.height_thresh,
+        )
+
+        if foot_contacts.shape[-1] == 6:
+            # For soma77
+            # remove "LeftToeEnd" and "RightToeEnd"
+            foot_contacts = foot_contacts[..., [0, 1, 3, 4]]
+
+        num_contacts = foot_contacts.shape[-1]
+        incorrect = torch.logical_xor(heuristic_contacts, foot_contacts)
+        # account for generated length
+        # since they are velocities, use length-1 to avoid inaccurate vel going one frame past len
+        device = foot_contacts.device
+        len_mask = torch.arange(foot_contacts.shape[1], device=device)[None, :, None].expand(foot_contacts.shape) < (
+            lengths[:, None, None] - 1
+        )
+        incorrect = incorrect * len_mask
+
+        incorrect_ratio = torch.sum(incorrect, (1, 2)) / (num_contacts * (lengths - 1))
+        accuracy = 1 - incorrect_ratio
+
+        return {"foot_contact_consistency": accuracy}

kimodo/metrics/tmr.py

@@ -0,0 +1,545 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""TMR evaluation metrics: text-motion retrieval, R-Precision, and related scores."""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+import torch
+from scipy import linalg
+from torch import Tensor
+
+from kimodo.model.tmr import TMR
+
+from .base import Metric
+
+
+# Scores are between 0 and 1
+def get_score_matrix_unit(x, y):
+    sim_matrix = np.einsum("b i, c i -> b c", x, y)
+    scores = sim_matrix / 2 + 0.5
+    return scores
+
+
+def get_scores_unit(x, y):
+    similarity = np.einsum("... i, ... i", x, y)
+    scores = similarity / 2 + 0.5
+    return scores
+
+
+def compute_tmr_per_sample_retrieval(
+    motion_emb: np.ndarray,
+    text_emb: np.ndarray,
+    sample_ids: List[str],
+    texts: List[str],
+    top_k: int = 5,
+) -> List[Dict[str, Any]]:
+    """For each sample (text query i), compute t2m rank of motion i and top-k retrieved motions with
+    ids and texts.
+
+    Returns list of dicts: [{"rank": int, "top_k": [{"id": str, "text": str}, ...]}, ...].
+    """
+    motion_emb = np.asarray(motion_emb).squeeze()
+    text_emb = np.asarray(text_emb).squeeze()
+    if motion_emb.ndim == 1:
+        motion_emb = motion_emb[np.newaxis, :]
+    if text_emb.ndim == 1:
+        text_emb = text_emb[np.newaxis, :]
+    n = motion_emb.shape[0]
+    assert text_emb.shape[0] == n and len(sample_ids) == n and len(texts) == n
+    scores = get_score_matrix_unit(text_emb, motion_emb)
+    out: List[Dict[str, Any]] = []
+    for i in range(n):
+        row = np.asarray(scores[i])
+        order = np.argsort(row)[::-1]
+        rank = int(np.where(order == i)[0][0]) + 1
+        top_indices = order[:top_k]
+        top_k_list = [{"id": sample_ids[j], "text": texts[j]} for j in top_indices]
+        out.append({"rank": rank, "top_k": top_k_list})
+    return out
+
+
+class TMR_Metric(Metric):
+    def __init__(
+        self,
+        tmr_model: TMR,
+        ranks: List = [1, 2, 3, 5, 10],
+        ranks_rounding=2,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.tmr_model = tmr_model
+        self.ranks = ranks
+        self.ranks_rounding = ranks_rounding
+
+    def clear(self):
+        self.saved_metrics = defaultdict(list)
+        self.saved_text_latents = []
+        self.saved_motion_gen_latents = []
+        self.saved_motion_gt_latents = []
+
+    def _compute(
+        self,
+        motion_rep,
+        pred_joints_output: Dict,
+        gt_joints_output: Dict,
+        text_x_dict: Dict,
+        lengths: Tensor,
+        **kwargs,
+    ) -> Dict:
+        pred_posed_joints = pred_joints_output["posed_joints"]
+        original_skeleton = motion_rep.skeleton if motion_rep is not None else None
+        latents_motion = self.tmr_model.encode_motion(
+            pred_posed_joints,
+            lengths=lengths,
+            original_skeleton=original_skeleton,
+            unit_vector=True,
+        )
+        latents_motion = latents_motion.cpu().numpy()
+
+        if isinstance(text_x_dict, dict) and "texts" in text_x_dict:
+            latents_text = self.tmr_model.encode_raw_text(text_x_dict["texts"], unit_vector=True)
+        else:
+            latents_text = self.tmr_model.encode_text(text_x_dict, unit_vector=True)
+        if latents_text.dim() == 1:
+            latents_text = latents_text.unsqueeze(0)
+        latents_text = latents_text.cpu().numpy()
+
+        self.saved_text_latents.append(latents_text)
+        self.saved_motion_gen_latents.append(latents_motion)
+
+        scores_text = get_scores_unit(latents_motion, latents_text)
+        output = {"TMR/t2m_sim": scores_text}
+
+        if gt_joints_output is not None and "posed_joints" in gt_joints_output:
+            gt_posed_joints = gt_joints_output["posed_joints"]
+            gt_latents_motion = self.tmr_model.encode_motion(
+                gt_posed_joints,
+                lengths=lengths,
+                original_skeleton=original_skeleton,
+                unit_vector=True,
+            )
+            gt_latents_motion = gt_latents_motion.cpu().numpy()
+            self.saved_motion_gt_latents.append(gt_latents_motion)
+
+            gt_scores_text = get_scores_unit(gt_latents_motion, latents_text)
+            scores_motion = get_scores_unit(latents_motion, gt_latents_motion)
+
+            output["TMR/t2m_gt_sim"] = gt_scores_text
+            output["TMR/m2m_sim"] = scores_motion
+
+        # pytorch tensors
+        for key, val in output.items():
+            output[key] = torch.tensor(val)
+        return output
+
+    def aggregate(self):
+        output = {}
+        for key, lst in self.saved_metrics.items():
+            output[key] = np.concatenate(lst)
+
+        assert self.saved_text_latents, "Should call the metric at least once."
+
+        text_latents = np.concatenate(self.saved_text_latents)
+        motion_gen_latents = np.concatenate(self.saved_motion_gen_latents)
+
+        batch_size = len(text_latents)
+        assert text_latents.shape == motion_gen_latents.shape
+
+        scores_t2m = get_score_matrix_unit(text_latents, motion_gen_latents)
+        scores_t2t = get_score_matrix_unit(text_latents, text_latents)
+
+        t2m_metrics = contrastive_metrics(
+            scores=scores_t2m,
+            scores_t2t=scores_t2t,
+            threshold=0.99,
+            rounding=2,
+        )
+
+        for key, val in t2m_metrics.items():
+            output["TMR/t2m_R/" + key] = val
+
+        mu_gen, cov_gen = calculate_activation_statistics(motion_gen_latents)
+        mu_text, cov_text = calculate_activation_statistics(text_latents)
+
+        fid_gen_text = calculate_frechet_distance(mu_gen, cov_gen, mu_text, cov_text)
+        output["TMR/FID/gen_text"] = fid_gen_text
+
+        if self.saved_motion_gt_latents:
+            motion_gt_latents = np.concatenate(self.saved_motion_gt_latents)
+            assert motion_gt_latents.shape == motion_gen_latents.shape
+
+            scores_m2gm = get_score_matrix_unit(motion_gen_latents, motion_gt_latents)
+            scores_t2gm = get_score_matrix_unit(text_latents, motion_gt_latents)
+
+            m2gm_metrics = contrastive_metrics(
+                scores=scores_m2gm,
+                scores_t2t=scores_t2t,
+                threshold=0.99,
+                rounding=2,
+            )
+            for key, val in m2gm_metrics.items():
+                output["TMR/m2m_R/" + key] = val
+
+            t2gm_metrics = contrastive_metrics(
+                scores=scores_t2gm,
+                scores_t2t=scores_t2t,
+                threshold=0.99,
+                rounding=2,
+            )
+            for key, val in t2gm_metrics.items():
+                output["TMR/t2m_gt_R/" + key] = val
+
+            mu_gt_motion, cov_gt_motion = calculate_activation_statistics(motion_gt_latents)
+            fid_gen_motion = calculate_frechet_distance(
+                mu_gen,
+                cov_gen,
+                mu_gt_motion,
+                cov_gt_motion,
+            )
+            output["TMR/FID/gen_gt"] = fid_gen_motion
+
+            fid_gt_text = calculate_frechet_distance(
+                mu_gt_motion,
+                cov_gt_motion,
+                mu_text,
+                cov_text,
+            )
+            output["TMR/FID/gt_text"] = fid_gt_text
+
+        for key, val in output.items():
+            if isinstance(val, (int, float, np.integer, np.floating)):
+                val = torch.tensor([val for _ in range(batch_size)])
+
+            if isinstance(val, np.ndarray):
+                val = torch.from_numpy(val)
+
+            output[key] = val.cpu().float()
+        return output
+
+
+class TMR_EmbeddingMetric(Metric):
+    """TMR metrics from precomputed motion and text embeddings (no model load).
+
+    Use in the loop: pass motion_emb and text_emb per sample; aggregate() computes retrieval metrics.
+    """
+
+    def __init__(self, ranks_rounding: int = 2, **kwargs):
+        super().__init__(**kwargs)
+        self.ranks_rounding = ranks_rounding
+
+    def clear(self):
+        self.saved_metrics = defaultdict(list)
+        self.saved_text_latents = []
+        self.saved_motion_gen_latents = []
+        self.saved_motion_gt_latents = []
+
+    def _compute(
+        self,
+        motion_emb=None,
+        text_emb=None,
+        gt_motion_emb=None,
+        **kwargs,
+    ) -> Dict:
+        if motion_emb is None or text_emb is None:
+            return {}
+        motion_emb = np.asarray(motion_emb)
+        text_emb = np.asarray(text_emb)
+        if motion_emb.ndim == 1:
+            motion_emb = motion_emb[np.newaxis, :]
+        if text_emb.ndim == 1:
+            text_emb = text_emb[np.newaxis, :]
+        self.saved_text_latents.append(text_emb)
+        self.saved_motion_gen_latents.append(motion_emb)
+        if gt_motion_emb is not None:
+            gt_motion_emb = np.asarray(gt_motion_emb)
+            if gt_motion_emb.ndim == 1:
+                gt_motion_emb = gt_motion_emb[np.newaxis, :]
+            self.saved_motion_gt_latents.append(gt_motion_emb)
+        scores = get_scores_unit(motion_emb, text_emb)
+        return {"TMR/t2m_sim": torch.tensor(scores, dtype=torch.float32)}
+
+    def aggregate(self):
+        output = {}
+        for key, lst in self.saved_metrics.items():
+            output[key] = np.concatenate(lst)
+        if not self.saved_text_latents:
+            return output
+        text_latents = np.concatenate(self.saved_text_latents)
+        motion_gen_latents = np.concatenate(self.saved_motion_gen_latents)
+        batch_size = len(text_latents)
+        assert text_latents.shape == motion_gen_latents.shape
+        scores_t2m = get_score_matrix_unit(text_latents, motion_gen_latents)
+        scores_t2t = get_score_matrix_unit(text_latents, text_latents)
+        t2m_metrics = contrastive_metrics(
+            scores=scores_t2m,
+            scores_t2t=scores_t2t,
+            threshold=0.99,
+            rounding=self.ranks_rounding,
+        )
+        for key, val in t2m_metrics.items():
+            output["TMR/t2m_R/" + key] = val
+        if batch_size >= 2:
+            mu_gen, cov_gen = calculate_activation_statistics(motion_gen_latents)
+            mu_text, cov_text = calculate_activation_statistics(text_latents)
+            output["TMR/FID/gen_text"] = calculate_frechet_distance(mu_gen, cov_gen, mu_text, cov_text)
+        else:
+            output["TMR/FID/gen_text"] = float("nan")
+        if self.saved_motion_gt_latents:
+            motion_gt_latents = np.concatenate(self.saved_motion_gt_latents)
+            assert motion_gt_latents.shape == motion_gen_latents.shape
+            scores_m2gm = get_score_matrix_unit(motion_gen_latents, motion_gt_latents)
+            scores_t2gm = get_score_matrix_unit(text_latents, motion_gt_latents)
+            m2gm_metrics = contrastive_metrics(
+                scores=scores_m2gm,
+                scores_t2t=scores_t2t,
+                threshold=0.99,
+                rounding=self.ranks_rounding,
+            )
+            for key, val in m2gm_metrics.items():
+                output["TMR/m2m_R/" + key] = val
+            t2gm_metrics = contrastive_metrics(
+                scores=scores_t2gm,
+                scores_t2t=scores_t2t,
+                threshold=0.99,
+                rounding=self.ranks_rounding,
+            )
+            for key, val in t2gm_metrics.items():
+                output["TMR/t2m_gt_R/" + key] = val
+            if batch_size >= 2:
+                mu_gt_motion, cov_gt_motion = calculate_activation_statistics(motion_gt_latents)
+                output["TMR/FID/gen_gt"] = calculate_frechet_distance(mu_gen, cov_gen, mu_gt_motion, cov_gt_motion)
+                output["TMR/FID/gt_text"] = calculate_frechet_distance(mu_gt_motion, cov_gt_motion, mu_text, cov_text)
+            else:
+                output["TMR/FID/gen_gt"] = float("nan")
+                output["TMR/FID/gt_text"] = float("nan")
+        for key, val in output.items():
+            if isinstance(val, (int, float, np.integer, np.floating)):
+                val = torch.tensor([val for _ in range(batch_size)])
+            if isinstance(val, np.ndarray):
+                val = torch.from_numpy(val)
+            output[key] = val.cpu().float()
+        return output
+
+
+def compute_tmr_retrieval_metrics(
+    motion_emb: np.ndarray,
+    text_emb: np.ndarray,
+    gt_motion_emb: Optional[np.ndarray] = None,
+    rounding: int = 2,
+) -> Dict[str, float]:
+    """Compute TMR retrieval metrics from precomputed embeddings."""
+    if motion_emb.shape != text_emb.shape:
+        raise ValueError(f"Expected same shape for motion/text embeddings, got {motion_emb.shape} vs {text_emb.shape}")
+
+    scores_t2m = get_score_matrix_unit(text_emb, motion_emb)
+    scores_t2t = get_score_matrix_unit(text_emb, text_emb)
+
+    output: Dict[str, float] = {}
+    t2m_metrics = contrastive_metrics(
+        scores=scores_t2m,
+        scores_t2t=scores_t2t,
+        threshold=0.99,
+        rounding=rounding,
+    )
+    for key, val in t2m_metrics.items():
+        output[f"TMR/t2m_R/{key}"] = float(val)
+
+    n_samples = len(motion_emb)
+    if n_samples >= 2:
+        mu_gen, cov_gen = calculate_activation_statistics(motion_emb)
+        mu_text, cov_text = calculate_activation_statistics(text_emb)
+        output["TMR/FID/gen_text"] = float(calculate_frechet_distance(mu_gen, cov_gen, mu_text, cov_text))
+    else:
+        output["TMR/FID/gen_text"] = float("nan")
+
+    if gt_motion_emb is not None:
+        if gt_motion_emb.shape != motion_emb.shape:
+            raise ValueError(f"Expected gt motion embeddings shape {motion_emb.shape}, got {gt_motion_emb.shape}")
+
+        scores_m2gm = get_score_matrix_unit(motion_emb, gt_motion_emb)
+        scores_t2gm = get_score_matrix_unit(text_emb, gt_motion_emb)
+
+        m2gm_metrics = contrastive_metrics(
+            scores=scores_m2gm,
+            scores_t2t=scores_t2t,
+            threshold=0.99,
+            rounding=rounding,
+        )
+        for key, val in m2gm_metrics.items():
+            output[f"TMR/m2m_R/{key}"] = float(val)
+
+        t2gm_metrics = contrastive_metrics(
+            scores=scores_t2gm,
+            scores_t2t=scores_t2t,
+            threshold=0.99,
+            rounding=rounding,
+        )
+        for key, val in t2gm_metrics.items():
+            output[f"TMR/t2m_gt_R/{key}"] = float(val)
+
+        if n_samples >= 2:
+            mu_gt_motion, cov_gt_motion = calculate_activation_statistics(gt_motion_emb)
+            output["TMR/FID/gen_gt"] = float(calculate_frechet_distance(mu_gen, cov_gen, mu_gt_motion, cov_gt_motion))
+            output["TMR/FID/gt_text"] = float(calculate_frechet_distance(mu_gt_motion, cov_gt_motion, mu_text, cov_text))
+        else:
+            output["TMR/FID/gen_gt"] = float("nan")
+            output["TMR/FID/gt_text"] = float("nan")
+
+    return output
+
+
+def all_contrastive_metrics(sims, emb=None, threshold=None, rounding=2, return_cols=False):
+    text_selfsim = None
+    if emb is not None:
+        text_selfsim = emb @ emb.T
+
+    t2m_m, t2m_cols = contrastive_metrics(sims, text_selfsim, threshold, return_cols=True, rounding=rounding)
+    m2t_m, m2t_cols = contrastive_metrics(sims.T, text_selfsim, threshold, return_cols=True, rounding=rounding)
+
+    all_m = {}
+    for key in t2m_m:
+        all_m[f"t2m/{key}"] = t2m_m[key]
+        all_m[f"m2t/{key}"] = m2t_m[key]
+
+    all_m["t2m/len"] = float(len(sims))
+    all_m["m2t/len"] = float(len(sims[0]))
+    if return_cols:
+        return all_m, t2m_cols, m2t_cols
+    return all_m
+
+
+def contrastive_metrics(
+    scores,
+    scores_t2t=None,
+    threshold=None,
+    rounding=2,
+):
+    n, m = scores.shape
+    assert n == m
+    num_queries = n
+
+    dists = -scores
+    sorted_dists = np.sort(dists, axis=1)
+    # GT is in the diagonal
+    gt_dists = np.diag(dists)[:, None]
+
+    if scores_t2t is not None and threshold is not None:
+        real_threshold = 2 * threshold - 1
+        idx = np.argwhere(scores_t2t > real_threshold)
+        partition = np.unique(idx[:, 0], return_index=True)[1]
+        # take as GT the minimum score of similar values
+        gt_dists = np.minimum.reduceat(dists[tuple(idx.T)], partition)
+        gt_dists = gt_dists[:, None]
+
+    rows, cols = np.where((sorted_dists - gt_dists) == 0)  # find column position of GT
+
+    # if there are ties
+    if rows.size > num_queries:
+        assert np.unique(rows).size == num_queries, "issue in metric evaluation"
+        avg_cols = break_ties_average(sorted_dists, gt_dists)
+        cols = avg_cols
+
+    msg = "expected ranks to match queries ({} vs {}) "
+    assert cols.size == num_queries, msg
+
+    metrics = {}
+    vals = [str(x).zfill(2) for x in [1, 2, 3, 5, 10]]
+    for val in vals:
+        metrics[f"R{val}"] = 100 * float(np.sum(cols < int(val))) / num_queries
+
+    metrics["MedR"] = float(np.median(cols) + 1)
+    metrics["len"] = num_queries
+
+    if rounding is not None:
+        for key in metrics:
+            metrics[key] = round(metrics[key], rounding)
+    return metrics
+
+
+def break_ties_average(sorted_dists, gt_dists):
+    # fast implementation, based on this code:
+    # https://stackoverflow.com/a/49239335
+    locs = np.argwhere((sorted_dists - gt_dists) == 0)
+
+    # Find the split indices
+    steps = np.diff(locs[:, 0])
+    splits = np.nonzero(steps)[0] + 1
+    splits = np.insert(splits, 0, 0)
+
+    # Compute the result columns
+    summed_cols = np.add.reduceat(locs[:, 1], splits)
+    counts = np.diff(np.append(splits, locs.shape[0]))
+    avg_cols = summed_cols / counts
+    return avg_cols
+
+
+def calculate_activation_statistics(activations):
+    """
+    Params:
+    -- activation: num_samples x dim_feat
+    Returns:
+    -- mu: dim_feat
+    -- sigma: dim_feat x dim_feat
+    """
+    mu = np.mean(activations, axis=0)
+    cov = np.cov(activations, rowvar=False)
+    return mu, cov
+
+
+def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
+    """Numpy implementation of the Frechet Distance. The Frechet distance between two multivariate
+    Gaussians X_1 ~ N(mu_1, C_1)
+
+    and X_2 ~ N(mu_2, C_2) is
+            d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
+    Stable version by Dougal J. Sutherland.
+    Params:
+    -- mu1   : Numpy array containing the activations of a layer of the
+               inception net (like returned by the function 'get_predictions')
+               for generated samples.
+    -- mu2   : The sample mean over activations, precalculated on an
+               representative dataset set.
+    -- sigma1: The covariance matrix over activations for generated samples.
+    -- sigma2: The covariance matrix over activations, precalculated on an
+               representative dataset set.
+    Returns:
+    --   : The Frechet Distance.
+    """
+
+    mu1 = np.atleast_1d(mu1)
+    mu2 = np.atleast_1d(mu2)
+
+    sigma1 = np.atleast_2d(sigma1)
+    sigma2 = np.atleast_2d(sigma2)
+
+    assert mu1.shape == mu2.shape, "Training and test mean vectors have different lengths"
+    assert sigma1.shape == sigma2.shape, "Training and test covariances have different dimensions"
+
+    diff = mu1 - mu2
+
+    # Product might be almost singular
+    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
+    if not np.isfinite(covmean).all():
+        msg = ("fid calculation produces singular product; " "adding %s to diagonal of cov estimates") % eps
+        print(msg)
+        offset = np.eye(sigma1.shape[0]) * eps
+        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
+
+    # Numerical error might give slight imaginary component
+    if np.iscomplexobj(covmean):
+        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
+            # try again with diagonal %s
+            offset = np.eye(sigma1.shape[0]) * eps
+            covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
+            if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
+                m = np.max(np.abs(covmean.imag))
+                raise ValueError("Imaginary component {}".format(m))
+        covmean = covmean.real
+
+    tr_covmean = np.trace(covmean)
+
+    return diff.dot(diff) + np.trace(sigma1) + np.trace(sigma2) - 2 * tr_covmean

kimodo/model/__init__.py

@@ -0,0 +1,31 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Kimodo model package: main model class, text encoders, and loading utilities."""
+
+from .common import resolve_target
+from .kimodo_model import Kimodo
+from .llm2vec import LLM2VecEncoder
+from .load_model import load_model
+from .loading import (
+    AVAILABLE_MODELS,
+    DEFAULT_MODEL,
+    DEFAULT_TEXT_ENCODER_URL,
+    MODEL_NAMES,
+    load_checkpoint_state_dict,
+)
+from .tmr import TMR
+from .twostage_denoiser import TwostageDenoiser
+
+__all__ = [
+    "Kimodo",
+    "LLM2VecEncoder",
+    "TMR",
+    "TwostageDenoiser",
+    "load_model",
+    "load_checkpoint_state_dict",
+    "resolve_target",
+    "AVAILABLE_MODELS",
+    "DEFAULT_MODEL",
+    "DEFAULT_TEXT_ENCODER_URL",
+    "MODEL_NAMES",
+]

kimodo/model/backbone.py

@@ -0,0 +1,312 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Transformer backbone: padding, masking, and encoder stack for the denoiser."""
+
+import logging
+from typing import Optional, Union
+
+import torch
+from omegaconf import ListConfig
+from pydantic.dataclasses import dataclass
+from torch import Tensor, nn
+from torch.nn import TransformerEncoder, TransformerEncoderLayer
+
+from kimodo.tools import validate
+
+log = logging.getLogger(__name__)
+
+
+def pad_x_and_mask_to_fixed_size(x: Tensor, mask: Tensor, size: int):
+    """Pad a feature vector x and the mask to always have the same size.
+
+    Args:
+        x (torch.Tensor): [B, T, D]
+        mask (torch.Tensor): [B, T]
+        size (int)
+    Returns:
+        torch.Tensor: [B, size, D]
+        torch.Tensor: [B, size]
+    """
+
+    batch_size, cur_max_size, dim = x.shape[0], x.shape[1], x.shape[2]
+
+    if cur_max_size == size:
+        # already padded to this size, probably in the collate function
+        return x, mask
+
+    if cur_max_size > size:
+        # This issue should have been handled in the collate function
+        # usefull as a check for test time
+        log.warn("The size of the tensor is larger than the maximum size. Cropping the input..")
+        cur_max_size = size
+
+    new_x = torch.zeros(
+        (batch_size, size, dim),
+        dtype=x.dtype,
+        device=x.device,
+    )
+    new_x[:, :cur_max_size] = x
+
+    # same for the mask
+    new_mask = torch.zeros(
+        (batch_size, size),
+        dtype=mask.dtype,
+        device=mask.device,
+    )
+    new_mask[:, :cur_max_size] = mask
+    return new_x, new_mask
+
+
+@dataclass(frozen=True, config=dict(extra="forbid", arbitrary_types_allowed=True))
+class TransformerEncoderBlockConfig:
+    """Configuration for the transformer encoder backbone."""
+
+    # input features dimension
+    input_dim: int
+    # output features dimension
+    output_dim: int
+
+    # skeleton object
+    skeleton: object
+
+    # dimension of the text embeddings
+    llm_shape: Union[list[int], ListConfig]
+
+    # mask the text or not
+    use_text_mask: bool
+
+    # latent dimension of the model
+    latent_dim: int
+    # dimension of the feedforward network in transformer
+    ff_size: int
+    # num layers in transformer
+    num_layers: int
+    # num heads in transformer
+    num_heads: int
+    # activation in transformer
+    activation: str
+    # dropout rate for the transformer
+    dropout: float
+    # dropout rate for the positional embeddings
+    pe_dropout: float
+    # use norm first or not
+    norm_first: bool = False
+    # artificially extend the number of text tokens
+    num_text_tokens_override: Optional[int] = None
+
+    # Input first heading angle
+    input_first_heading_angle: bool = False
+
+
+class TransformerEncoderBlock(nn.Module):
+    @validate(TransformerEncoderBlockConfig, save_args=True, super_init=True)
+    def __init__(self, conf):
+        self.nbjoints = self.skeleton.nbjoints
+        llm_dim = self.llm_shape[-1]
+        self.embed_text = nn.Linear(llm_dim, self.latent_dim)
+
+        self.sequence_pos_encoder = PositionalEncoding(self.latent_dim, self.pe_dropout)
+
+        # maximum number of tokens
+        self.num_text_tokens = self.llm_shape[0]
+        if self.num_text_tokens_override is not None:
+            self.num_text_tokens = self.num_text_tokens_override
+
+        self.embed_timestep = TimestepEmbedder(self.latent_dim, self.sequence_pos_encoder)
+
+        self.input_linear = nn.Linear(self.input_dim, self.latent_dim)
+        self.output_linear = nn.Linear(self.latent_dim, self.output_dim)
+        self.linear_first_heading_angle = nn.Linear(2, self.latent_dim)
+
+        trans_enc_layer = TransformerEncoderLayer(
+            d_model=self.latent_dim,
+            nhead=self.num_heads,
+            dim_feedforward=self.ff_size,
+            dropout=self.dropout,
+            activation=self.activation,
+            batch_first=True,
+            norm_first=self.norm_first,
+        )
+        self.seqTransEncoder = TransformerEncoder(
+            trans_enc_layer,
+            num_layers=self.num_layers,
+            enable_nested_tensor=False,
+        )
+
+    def forward(
+        self,
+        x: Tensor,
+        x_pad_mask: torch.Tensor,
+        text_feat: torch.Tensor,
+        text_feat_pad_mask: torch.Tensor,
+        timesteps: Tensor,
+        first_heading_angle: Optional[Tensor] = None,
+    ) -> Tensor:
+        """
+        Args:
+            x (torch.Tensor): [B, T, dim_motion] current noisy motion
+            x_pad_mask (torch.Tensor): [B, T] attention mask, positions with True are allowed to attend, False are not
+            text_feat (torch.Tensor): [B, max_text_len, llm_dim] embedded text prompts
+            text_feat_pad_mask (torch.Tensor): [B, max_text_len] attention mask, positions with True are allowed to attend, False are not
+            timesteps (torch.Tensor): [B,] current denoising step
+
+        Returns:
+            torch.Tensor: [B, T, output_dim]
+        """
+        batch_size = len(x)
+        x = self.input_linear(x)  # [B, T, D]
+
+        # Pad the text tokens + mask to always have the same size == self.num_text_tokens
+        # done here if it was not done in the collate function
+        if self.num_text_tokens is not None:
+            text_feat, text_feat_pad_mask = pad_x_and_mask_to_fixed_size(
+                text_feat,
+                text_feat_pad_mask,
+                self.num_text_tokens,
+            )
+
+        # Encode the text features and the time information
+        emb_text = self.embed_text(text_feat)  # [B, max_text_len, D]
+        emb_time = self.embed_timestep(timesteps)  # [B, 1, D]
+
+        # Create mask for the time information
+        time_mask = torch.ones((batch_size, 1), dtype=bool, device=x.device)
+
+        # Create the prefix features (text, time, etc): [B, max_text_len + 1 + etc]
+        prefix_feats = torch.cat((emb_text, emb_time), axis=1)
+
+        # Behavior from old code: not use text mask -> True for all the tokens
+        if not self.use_text_mask:
+            text_feat_pad_mask = torch.ones(
+                (batch_size, emb_text.shape[1]),
+                dtype=torch.bool,
+                device=x.device,
+            )
+
+        prefix_mask = torch.cat((text_feat_pad_mask, time_mask), axis=1)
+
+        # add the input first heading angle
+        if self.input_first_heading_angle:
+            assert first_heading_angle is not None, "The first heading angle is mandatory for this model"
+            # cos(angle) / sin(angle)
+            first_heading_angle_feats = torch.stack(
+                [
+                    torch.cos(first_heading_angle),
+                    torch.sin(first_heading_angle),
+                ],
+                axis=-1,
+            )
+
+            first_heading_angle_feats = self.linear_first_heading_angle(first_heading_angle_feats)
+            first_heading_angle_feats = first_heading_angle_feats[:, None]  # for cat
+            first_heading_angle_mask = torch.ones(
+                (batch_size, 1),
+                dtype=bool,
+                device=x.device,
+            )
+            prefix_feats = torch.cat((prefix_feats, first_heading_angle_feats), axis=1)
+            prefix_mask = torch.cat((prefix_mask, first_heading_angle_mask), axis=1)
+
+        # compute the number of prefix features
+        pose_start_ind = prefix_feats.shape[1]
+
+        # Concatenate prefix and x: [B, len(prefix) + T, D]
+        xseq = torch.cat((prefix_feats, x), axis=1)
+
+        # Concatenate the masks and negate them: [B, len(prefix) + T]
+        src_key_padding_mask = ~torch.cat((prefix_mask, x_pad_mask), axis=1)
+
+        # Add positional encoding
+        xseq = self.sequence_pos_encoder(xseq)
+
+        # Input to the transformer and keep the motion indexes
+        if isinstance(self.seqTransEncoder, nn.TransformerEncoder):
+            assert not self.seqTransEncoder.use_nested_tensor, "Flash attention should be disabled due to bug!"
+
+        output = self.seqTransEncoder(
+            xseq,
+            src_key_padding_mask=src_key_padding_mask,
+        )
+        output = output[:, pose_start_ind:]  # [B, T, D]
+        output = self.output_linear(output)  # [B, T, OD]
+        return output
+
+
+class PositionalEncoding(nn.Module):
+    """Non-learned positional encoding."""
+
+    def __init__(
+        self,
+        d_model: int,
+        dropout: Optional[float] = 0.1,
+        max_len: Optional[int] = 5000,
+    ):
+        """
+        Args:
+            d_model (int): input dim
+            dropout (Optional[float] = 0.1): dropout probability on output
+            max_len (Optional[int] = 5000): maximum sequence length
+        """
+        super(PositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+
+        # Note: have to replace torch.exp() and math.log() with torch.pow()
+        # due to MKL exp() and ln() throws floating point exceptions on certain CPUs
+        # see corresponding commit and MR
+        div_term = torch.pow(10000.0, -torch.arange(0, d_model, 2).float() / d_model)
+        # div_term = torch.exp(
+        #     torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model)
+        # )
+
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)  # [1, T, D]
+
+        self.register_buffer("pe", pe, persistent=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply positional encoding to input sequence.
+
+        Args:
+            x (torch.Tensor): [B, T, D] input motion sequence
+
+        Returns:
+            torch.Tensor: [B, T, D] input motion with PE added to it (and optionally dropout)
+        """
+        x = x + self.pe[:, : x.shape[1], :]
+        return self.dropout(x)
+
+
+class TimestepEmbedder(nn.Module):
+    """Encoder for diffusion step."""
+
+    def __init__(self, latent_dim: int, sequence_pos_encoder: PositionalEncoding):
+        """
+        Args:
+            latent_dim (int): dim to encode to
+            sequence_pos_encoder (PositionalEncoding): the PE to use on timesteps
+        """
+        super().__init__()
+        self.latent_dim = latent_dim
+        self.sequence_pos_encoder = sequence_pos_encoder
+
+        time_embed_dim = self.latent_dim
+        self.time_embed = nn.Sequential(
+            nn.Linear(self.latent_dim, time_embed_dim),
+            nn.SiLU(),
+            nn.Linear(time_embed_dim, time_embed_dim),
+        )
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        """Embed timesteps by adding PE then going through linear layers.
+
+        Args:
+            timesteps (torch.Tensor): [B]
+
+        Returns:
+            torch.Tensor: [B, 1, D]
+        """
+        return self.time_embed(self.sequence_pos_encoder.pe.transpose(0, 1)[timesteps])

kimodo/model/cfg.py

@@ -0,0 +1,133 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Classifier-free guidance wrapper for the denoiser at sampling time."""
+
+from typing import Dict, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+CFG_TYPES = ["nocfg", "regular", "separated"]
+
+
+class ClassifierFreeGuidedModel(nn.Module):
+    """Wrapper around denoiser to use classifier-free guidance at sampling time."""
+
+    def __init__(self, model: nn.Module, cfg_type: Optional[str] = "separated"):
+        """Wrap the denoiser for classifier-free guidance; cfg_type in CFG_TYPES (e.g. 'regular',
+        'nocfg')."""
+        super().__init__()
+        self.model = model
+        assert cfg_type in CFG_TYPES, f"Invalid cfg_type: {cfg_type}"
+        self.cfg_type_default = cfg_type
+
+    def forward(
+        self,
+        cfg_weight: Union[float, Tuple[float, float]],
+        x: torch.Tensor,
+        x_pad_mask: torch.Tensor,
+        text_feat: torch.Tensor,
+        text_feat_pad_mask: torch.Tensor,
+        timesteps: torch.Tensor,
+        first_heading_angle: Optional[torch.Tensor] = None,
+        motion_mask: Optional[torch.Tensor] = None,
+        observed_motion: Optional[torch.Tensor] = None,
+        cfg_type: Optional[str] = None,
+    ) -> torch.Tensor:
+        """
+        Args:
+            cfg_weight (float): guidance weight float or tuple of floats with (text, constraint) weights if using separated cfg
+            x (torch.Tensor): [B, T, dim_motion] current noisy motion
+            x_pad_mask (torch.Tensor): [B, T] attention mask, positions with True are allowed to attend, False are not
+            text_feat (torch.Tensor): [B, max_text_len, llm_dim] embedded text prompts
+            text_feat_pad_mask (torch.Tensor): [B, max_text_len] attention mask, positions with True are allowed to attend, False are not
+            timesteps (torch.Tensor): [B,] current denoising step
+            motion_mask
+            observed_motion
+            neutral_joints (torch.Tensor): [B, nbjoints] The neutral joints of the motions
+
+        Returns:
+            torch.Tensor: same size as input x
+        """
+
+        if cfg_type is None:
+            cfg_type = self.cfg_type_default
+
+        assert cfg_type in CFG_TYPES, f"Invalid cfg_type: {cfg_type}"
+
+        # batched conditional and uncond pass together
+        if cfg_type == "nocfg":
+            return self.model(
+                x,
+                x_pad_mask,
+                text_feat,
+                text_feat_pad_mask,
+                timesteps,
+                first_heading_angle=first_heading_angle,
+                motion_mask=motion_mask,
+                observed_motion=observed_motion,
+            )
+        elif cfg_type == "regular":
+            assert isinstance(cfg_weight, (float, int)), "cfg_weight must be a single float for regular CFG"
+            # out_uncond + w * (out_text_and_constraint - out_uncond)
+            text_feat = torch.concatenate([text_feat, 0 * text_feat], dim=0)
+            if motion_mask is not None:
+                motion_mask = torch.concatenate([motion_mask, 0 * motion_mask], dim=0)
+            if observed_motion is not None:
+                observed_motion = torch.concatenate([observed_motion, observed_motion], dim=0)
+            if first_heading_angle is not None:
+                first_heading_angle = torch.concatenate([first_heading_angle, first_heading_angle], dim=0)
+
+            out_cond_uncond = self.model(
+                torch.concatenate([x, x], dim=0),
+                torch.concatenate([x_pad_mask, x_pad_mask], dim=0),
+                text_feat,
+                torch.concatenate([text_feat_pad_mask, False * text_feat_pad_mask], dim=0),
+                torch.concatenate([timesteps, timesteps], dim=0),
+                first_heading_angle=first_heading_angle,
+                motion_mask=motion_mask,
+                observed_motion=observed_motion,
+            )
+
+            out, out_uncond = torch.chunk(out_cond_uncond, 2)
+            out_new = out_uncond + (cfg_weight * (out - out_uncond))
+        elif cfg_type == "separated":
+            assert len(cfg_weight) == 2, "cfg_weight must be a tuple of two floats for separated CFG"
+            # out_uncond + w_text * (out_text - out_uncond) + w_constraint * (out_constraint - out_uncond)
+            text_feat = torch.concatenate([text_feat, 0 * text_feat, 0 * text_feat], dim=0)
+            if motion_mask is not None:
+                motion_mask = torch.concatenate([0 * motion_mask, motion_mask, 0 * motion_mask], dim=0)
+            if observed_motion is not None:
+                observed_motion = torch.concatenate([observed_motion, observed_motion, observed_motion], dim=0)
+            if first_heading_angle is not None:
+                first_heading_angle = torch.concatenate(
+                    [first_heading_angle, first_heading_angle, first_heading_angle],
+                    dim=0,
+                )
+
+            out_cond_uncond = self.model(
+                torch.concatenate([x, x, x], dim=0),
+                torch.concatenate([x_pad_mask, x_pad_mask, x_pad_mask], dim=0),
+                text_feat,
+                torch.concatenate(
+                    [
+                        text_feat_pad_mask,
+                        False * text_feat_pad_mask,
+                        False * text_feat_pad_mask,
+                    ],
+                    dim=0,
+                ),
+                torch.concatenate([timesteps, timesteps, timesteps], dim=0),
+                first_heading_angle=first_heading_angle,
+                motion_mask=motion_mask,
+                observed_motion=observed_motion,
+            )
+
+            out_text, out_constraint, out_uncond = torch.chunk(out_cond_uncond, 3)
+            out_new = (
+                out_uncond + (cfg_weight[0] * (out_text - out_uncond)) + (cfg_weight[1] * (out_constraint - out_uncond))
+            )
+        else:
+            raise ValueError(f"Invalid cfg_type: {cfg_type}")
+
+        return out_new

kimodo/model/common.py

@@ -0,0 +1,48 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Config hydration: env vars, _target_ resolution, and recursive instantiation."""
+
+import importlib
+import os
+
+
+def get_env_var(name: str, default=None):
+    """Read env var by name and by lowercased name; return default if neither set."""
+    return os.getenv(name, os.getenv(name.lower(), default))
+
+
+def resolve_target(target: str):
+    """Import module and return the attribute named by a dotted path (e.g. 'pkg.mod.Class')."""
+    module_name, attr_name = target.rsplit(".", 1)
+    module = importlib.import_module(module_name)
+    return getattr(module, attr_name)
+
+
+def materialize_value(value):
+    """Recursively turn dicts with '_target_' into instances; lists/dicts traversed; leaves
+    unchanged."""
+    if isinstance(value, dict):
+        if "_target_" in value:
+            return instantiate_from_dict(value)
+        return {k: materialize_value(v) for k, v in value.items()}
+    if isinstance(value, list):
+        return [materialize_value(v) for v in value]
+    return value
+
+
+def instantiate_from_dict(node, overrides=None):
+    """Build an instance from a config dict: '_target_' gives the class, other keys are kwargs; overrides merged in."""
+    if not isinstance(node, dict) or "_target_" not in node:
+        raise ValueError("Config node must be a dict with a '_target_' key.")
+
+    target = resolve_target(node["_target_"])
+    kwargs = {}
+    for key, value in node.items():
+        if key == "_target_":
+            continue
+        kwargs[key] = materialize_value(value)
+
+    if overrides:
+        kwargs.update({k: v for k, v in overrides.items() if v is not None})
+
+    return target(**kwargs)

kimodo/model/diffusion.py

@@ -0,0 +1,133 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Diffusion process and DDIM sampling for motion generation."""
+
+import math
+from typing import Optional, Tuple
+
+import torch
+from torch import nn
+
+
+def get_beta_schedule(
+    num_diffusion_timesteps: int,
+    max_beta: Optional[float] = 0.999,
+) -> torch.Tensor:
+    """Get cosine beta schedule."""
+
+    def alpha_bar(t):
+        return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
+
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
+    return torch.tensor(betas, dtype=torch.float)
+
+
+class Diffusion(torch.nn.Module):
+    """Cosine-schedule diffusion process: betas, alphas, and DDIM step mapping."""
+
+    def __init__(self, num_base_steps: int):
+        """Set up cosine beta schedule and precompute diffusion variables for num_base_steps."""
+        super().__init__()
+        self.num_base_steps = num_base_steps
+        betas_base = get_beta_schedule(self.num_base_steps)
+        self.register_buffer("betas_base", betas_base, persistent=False)
+        alphas_cumprod_base = torch.cumprod(1.0 - self.betas_base, dim=0)
+        self.register_buffer("alphas_cumprod_base", alphas_cumprod_base, persistent=False)
+        use_timesteps, _ = self.space_timesteps(self.num_base_steps)
+        self.calc_diffusion_vars(use_timesteps)
+
+    def extra_repr(self) -> str:
+        return f"num_base_steps={self.num_base_steps}"
+
+    @property
+    def device(self):
+        return self.betas_base.device
+
+    def space_timesteps(self, num_denoising_steps: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Return (use_timesteps, map_tensor) for a subsampled denoising schedule of
+        num_denoising_steps."""
+        nsteps_train = self.num_base_steps
+        frac_stride = (nsteps_train - 1) / max(1, num_denoising_steps - 1)
+        use_timesteps = torch.round(torch.arange(nsteps_train, device=self.device) * frac_stride).to(torch.long)
+        use_timesteps = torch.clamp(use_timesteps, max=nsteps_train - 1)
+        map_tensor = torch.arange(nsteps_train, device=self.device, dtype=torch.long)[use_timesteps]
+        return use_timesteps, map_tensor
+
+    def calc_diffusion_vars(self, use_timesteps: torch.Tensor) -> None:
+        """Update buffers (betas, alphas, alphas_cumprod, etc.) for the given subsampled
+        timesteps."""
+        alphas_cumprod = self.alphas_cumprod_base[use_timesteps]
+        last_alpha_cumprod = torch.cat([torch.tensor([1.0]).to(alphas_cumprod), alphas_cumprod[:-1]])
+        betas = 1.0 - alphas_cumprod / last_alpha_cumprod
+        self.register_buffer("betas", betas, persistent=False)
+
+        alphas = 1.0 - self.betas
+        self.register_buffer("alphas", alphas, persistent=False)
+        alphas_cumprod = torch.cumprod(self.alphas, dim=0)
+        alphas_cumprod = torch.clamp(alphas_cumprod, min=1e-9)
+        self.register_buffer("alphas_cumprod", alphas_cumprod, persistent=False)
+
+        alphas_cumprod_prev = torch.cat([torch.tensor([1.0]).to(self.alphas_cumprod), self.alphas_cumprod[:-1]])
+        self.register_buffer("alphas_cumprod_prev", alphas_cumprod_prev, persistent=False)
+
+        sqrt_recip_alphas_cumprod = torch.rsqrt(self.alphas_cumprod)
+        self.register_buffer("sqrt_recip_alphas_cumprod", sqrt_recip_alphas_cumprod, persistent=False)
+
+        sqrt_recipm1_alphas_cumprod = torch.rsqrt(self.alphas_cumprod / (1.0 - self.alphas_cumprod))
+        self.register_buffer("sqrt_recipm1_alphas_cumprod", sqrt_recipm1_alphas_cumprod, persistent=False)
+
+        posterior_variance = self.betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        self.register_buffer("posterior_variance", posterior_variance, persistent=False)
+
+        sqrt_alphas_cumprod = torch.rsqrt(1.0 / self.alphas_cumprod)
+        self.register_buffer("sqrt_alphas_cumprod", sqrt_alphas_cumprod, persistent=False)
+
+        sqrt_one_minus_alphas_cumprod = torch.rsqrt(1.0 / (1.0 - self.alphas_cumprod))
+        self.register_buffer(
+            "sqrt_one_minus_alphas_cumprod",
+            sqrt_one_minus_alphas_cumprod,
+            persistent=False,
+        )
+
+    def q_sample(
+        self,
+        x_start: torch.Tensor,
+        t: torch.Tensor,
+        noise: torch.Tensor = None,
+    ):
+        if noise is None:
+            noise = torch.randn_like(x_start)
+        assert noise.shape == x_start.shape
+
+        xt = (
+            self.sqrt_alphas_cumprod[t, None, None] * x_start
+            + self.sqrt_one_minus_alphas_cumprod[t, None, None] * noise
+        )
+        return xt
+
+
+class DDIMSampler(nn.Module):
+    """Deterministic DDIM sampler (eta = 0)."""
+
+    def __init__(self, diffusion: Diffusion):
+        super().__init__()
+        self.diffusion = diffusion
+
+    def __call__(
+        self,
+        use_timesteps: torch.Tensor,
+        x_t: torch.Tensor,
+        pred_xstart: torch.Tensor,
+        t: torch.Tensor,
+    ) -> torch.Tensor:
+        self.diffusion.calc_diffusion_vars(use_timesteps)
+        eps = (
+            self.diffusion.sqrt_recip_alphas_cumprod[t, None, None] * x_t - pred_xstart
+        ) / self.diffusion.sqrt_recipm1_alphas_cumprod[t, None, None]
+        alpha_bar_prev = self.diffusion.alphas_cumprod_prev[t, None, None]
+        x = pred_xstart * torch.sqrt(alpha_bar_prev) + torch.sqrt(1 - alpha_bar_prev) * eps
+        return x

kimodo/model/kimodo_model.py

@@ -0,0 +1,634 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Kimodo model: denoiser, text encoder, diffusion sampling, and post-processing."""
+
+import logging
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+from torch import nn
+from tqdm.auto import tqdm
+
+from kimodo.constraints import EndEffectorConstraintSet, FullBodyConstraintSet
+from kimodo.motion_rep.feature_utils import compute_heading_angle, length_to_mask
+from kimodo.postprocess import post_process_motion
+from kimodo.sanitize import sanitize_texts
+from kimodo.skeleton import SOMASkeleton30
+from kimodo.tools import to_numpy
+
+from .cfg import ClassifierFreeGuidedModel
+from .diffusion import DDIMSampler, Diffusion
+
+log = logging.getLogger(__name__)
+
+
+class Kimodo(nn.Module):
+    """Helper class for test time."""
+
+    def __init__(
+        self,
+        denoiser: nn.Module,
+        text_encoder: nn.Module,
+        num_base_steps: int,
+        device: Optional[Union[str, torch.device]] = None,
+        cfg_type: Optional[str] = "separated",
+    ):
+        super().__init__()
+
+        self.denoiser = denoiser.eval()
+
+        if cfg_type is None:
+            cfg_type = "nocfg"
+
+        # Add Classifier-free guidance to the model if needed
+        self.denoiser = ClassifierFreeGuidedModel(self.denoiser, cfg_type=cfg_type)
+
+        self.motion_rep = denoiser.motion_rep
+        self.skeleton = self.motion_rep.skeleton
+
+        self.fps = denoiser.motion_rep.fps
+
+        self.diffusion = Diffusion(num_base_steps=num_base_steps)
+        self.sampler = DDIMSampler(self.diffusion)
+        self.text_encoder = text_encoder
+
+        self.device = device
+        # for classifier-free guidance
+
+        self.to(device)
+
+    @property
+    def output_skeleton(self):
+        """Skeleton used for model output (somaskel77 for SOMA, else unchanged)."""
+        if isinstance(self.skeleton, SOMASkeleton30):
+            return self.skeleton.somaskel77
+        return self.skeleton
+
+    def train(self, mode: bool):
+        self.denoiser.train(mode)
+        return self
+
+    def eval(self):
+        self.denoiser.eval()
+        return self
+
+    def denoising_step(
+        self,
+        motion: torch.Tensor,
+        pad_mask: torch.Tensor,
+        text_feat: torch.Tensor,
+        text_pad_mask: torch.Tensor,
+        t: torch.Tensor,
+        first_heading_angle: Optional[torch.Tensor],
+        motion_mask: torch.Tensor,
+        observed_motion: torch.Tensor,
+        num_denoising_steps: torch.Tensor,
+        cfg_weight: Union[float, Tuple[float, float]],
+        guide_masks: Optional[Dict] = None,
+        cfg_type: Optional[str] = None,
+    ) -> torch.Tensor:
+        """Single denoising step.
+
+        Returns:
+            torch.Tensor: [B, T, D] noisy motion input to t-1
+        """
+        # subsample timesteps
+        #   NOTE: do this at every step due to ONNX export, i.e. num_samp_stepsmay change dynamically when
+        #       running onnx version so need to account for that.
+        num_denoising_steps = num_denoising_steps[0]
+        use_timesteps, map_tensor = self.diffusion.space_timesteps(num_denoising_steps)
+        self.diffusion.calc_diffusion_vars(use_timesteps)
+
+        # first compute initial clean prediction from denoiser
+        t_map = map_tensor[t]
+
+        with torch.inference_mode():
+            pred_clean = self.denoiser(
+                cfg_weight,
+                motion,
+                pad_mask,
+                text_feat,
+                text_pad_mask,
+                t_map,
+                first_heading_angle,
+                motion_mask,
+                observed_motion,
+                cfg_type=cfg_type,
+            )
+
+        # sampler computes next step noisy motion
+        x_tm1 = self.sampler(use_timesteps, motion, pred_clean, t)
+        return x_tm1
+
+    def _multiprompt(
+        self,
+        prompts: list[str],
+        num_frames: int | list[int],
+        num_denoising_steps: int,
+        constraint_lst: Optional[list] = [],
+        cfg_weight: Optional[float] = [2.0, 2.0],
+        num_samples: Optional[int] = None,
+        cfg_type: Optional[str] = None,
+        return_numpy: bool = False,
+        first_heading_angle: Optional[torch.Tensor] = None,
+        # for transitioning
+        num_transition_frames: int = 5,
+        # for postprocess
+        post_processing: bool = False,
+        root_margin: float = 0.04,
+        # progress bar
+        progress_bar=tqdm,
+    ) -> torch.Tensor:
+        device = self.device
+
+        bs = num_samples
+        texts = sanitize_texts(prompts)
+
+        if isinstance(num_frames, int):
+            # same duration for all the segments
+            num_frames = [num_frames for _ in range(num_samples)]
+
+        tosqueeze = False
+        if num_samples is None:
+            num_samples = 1
+            tosqueeze = True
+
+        if constraint_lst is None:
+            constraint_lst = []
+
+        # Generate one chunck at a time
+        current_frame = 0
+        generated_motions = []
+
+        for idx, (text, num_frame) in enumerate(zip(texts, num_frames)):
+            texts_bs = [text for _ in range(num_samples)]
+
+            lengths = torch.tensor(
+                [num_frame for _ in range(num_samples)],
+                device=device,
+            )
+
+            is_first_motion = not generated_motions
+
+            observed_motion, motion_mask = None, None
+
+            # filter the constraint_lst to only keep the relevent ones
+            constraint_lst_base = [
+                constraint.crop_move(current_frame, current_frame + num_frame) for constraint in constraint_lst
+            ]  # this move temporally but not spatially
+
+            observed_motion, motion_mask = self.motion_rep.create_conditions_from_constraints_batched(
+                constraint_lst_base,
+                lengths,
+                to_normalize=False,  # don't normalize yet, it needs to be moved around
+                device=device,
+            )
+
+            if not is_first_motion:
+                nb_transition_frames = num_transition_frames
+
+                if nb_transition_frames < 1:
+                    raise ValueError(f"num_transition_frames must be at least 1, got {nb_transition_frames}")
+
+                latest_motions = generated_motions.pop()
+                # remove the transition part of A (will be put back afterward)
+                generated_motions.append(latest_motions[:, :-nb_transition_frames])
+                latest_frames = latest_motions[:, -nb_transition_frames:]
+
+                last_output = self.motion_rep.inverse(
+                    latest_frames,
+                    is_normalized=False,
+                    return_numpy=False,
+                )
+                smooth_root_2d = last_output["smooth_root_pos"][..., [0, 2]]
+
+                # add constraints at the begining to allow natural transitions
+                constraint_lst_transition = []
+                for batch_id in range(bs):
+                    new_constraint = FullBodyConstraintSet(
+                        self.skeleton,
+                        torch.arange(num_transition_frames),
+                        last_output["posed_joints"][batch_id, :num_transition_frames],
+                        last_output["global_rot_mats"][batch_id, :num_transition_frames],
+                        smooth_root_2d[batch_id, :num_transition_frames],
+                    )
+                    # separate end-effector constraint to capture hand/feet rotations
+                    new_ee_constraint = EndEffectorConstraintSet(
+                        self.skeleton,
+                        torch.arange(num_transition_frames),
+                        last_output["posed_joints"][batch_id, :num_transition_frames],
+                        last_output["global_rot_mats"][batch_id, :num_transition_frames],
+                        smooth_root_2d[batch_id, :num_transition_frames],
+                        joint_names=["LeftHand", "RightHand", "LeftFoot", "RightFoot"],
+                    )
+
+                    constraint_lst_transition.append([new_constraint, new_ee_constraint])
+
+                transition_lengths = torch.tensor(
+                    [nb_transition_frames for _ in range(num_samples)],
+                    device=device,
+                )
+
+                observed_motion_transition, motion_mask_transition = (
+                    self.motion_rep.create_conditions_from_constraints_batched(
+                        constraint_lst_transition,
+                        transition_lengths,
+                        to_normalize=False,  # don't normalize yet
+                        device=device,
+                    )
+                )
+
+                # concatenate the obversed motion / motion mask
+                observed_motion = torch.cat([observed_motion_transition, observed_motion], axis=1)
+                motion_mask = torch.cat([motion_mask_transition, motion_mask], axis=1)
+
+                # we need to move each observed motion in the batch to the new starting points
+                last_smooth_root_2d = smooth_root_2d[:, 0]
+                observed_motion = self.motion_rep.translate_2d(
+                    observed_motion, -last_smooth_root_2d
+                )  # equivalent to:  self.motion_rep.translate_2d_to_zero(observed_motion)
+
+                # remove dummy values after moving
+                observed_motion = observed_motion * motion_mask
+
+                lengths = lengths + transition_lengths
+                first_heading_angle = compute_heading_angle(last_output["posed_joints"], self.skeleton)[:, 0]
+            else:
+                if first_heading_angle is None:
+                    # Start at 0 angle, but this will change afterward
+                    first_heading_angle = torch.tensor([0.0] * bs, device=device)
+                else:
+                    first_heading_angle = torch.as_tensor(first_heading_angle, device=device)
+                    if first_heading_angle.numel() == 1:
+                        first_heading_angle = first_heading_angle.repeat(bs)
+
+            observed_motion = self.motion_rep.normalize(observed_motion)
+
+            max_frames = max(lengths)
+            motion_pad_mask = length_to_mask(lengths)
+
+            motion = self._generate(
+                texts_bs,
+                max_frames,
+                num_denoising_steps=num_denoising_steps,
+                pad_mask=motion_pad_mask,
+                first_heading_angle=first_heading_angle,
+                motion_mask=motion_mask,
+                observed_motion=observed_motion,
+                cfg_weight=cfg_weight,
+                cfg_type=cfg_type,
+            )
+
+            motion = self.motion_rep.unnormalize(motion)
+
+            if not is_first_motion:
+                motion_with_transition = self.motion_rep.translate_2d(
+                    motion,
+                    last_smooth_root_2d,
+                )
+
+                if post_processing:
+                    # Per-segment postprocessing: inverse, postprocess, re-encode.
+                    # The full transition+segment is postprocessed together so the
+                    # transition constraints keep the junction smooth.
+                    seg_output = self.motion_rep.inverse(
+                        motion_with_transition, is_normalized=False, return_numpy=False,
+                    )
+                    seg_constraints = [list(cl) for cl in constraint_lst_transition]
+                    for bi in range(bs):
+                        seg_constraints[bi].extend(
+                            [c.crop_move(current_frame - nb_transition_frames,
+                                         current_frame - nb_transition_frames + num_frame + nb_transition_frames)
+                             for c in constraint_lst]
+                        )
+                    corrected = post_process_motion(
+                        seg_output["local_rot_mats"],
+                        seg_output["root_positions"],
+                        seg_output["foot_contacts"],
+                        self.skeleton,
+                        seg_constraints,
+                        root_margin=root_margin,
+                    )
+                    seg_output.update(corrected)
+                    motion = self.motion_rep(
+                        seg_output["local_rot_mats"],
+                        seg_output["root_positions"],
+                        to_normalize=False,
+                        lengths=lengths,
+                    )
+                else:
+                    motion = motion_with_transition[:, num_transition_frames:]
+                    transition_frames = motion_with_transition[:, :num_transition_frames]
+
+                    # linearly combine the previously generated transitions with the newly generated ones
+                    alpha = torch.linspace(1, 0, num_transition_frames, device=device)[:, None]
+                    new_transition_frames = (
+                        latest_frames[:, :num_transition_frames] * alpha + (1 - alpha) * transition_frames
+                    )
+
+                    # add new transitions frames for A (merging with B prediction of the history)
+                    generated_motions.append(new_transition_frames)
+
+            elif post_processing:
+                # First segment: postprocess immediately
+                seg_output = self.motion_rep.inverse(
+                    motion, is_normalized=False, return_numpy=False,
+                )
+                seg_constraints = constraint_lst_base if constraint_lst_base else []
+                corrected = post_process_motion(
+                    seg_output["local_rot_mats"],
+                    seg_output["root_positions"],
+                    seg_output["foot_contacts"],
+                    self.skeleton,
+                    seg_constraints,
+                    root_margin=root_margin,
+                )
+                seg_output.update(corrected)
+                motion = self.motion_rep(
+                    seg_output["local_rot_mats"],
+                    seg_output["root_positions"],
+                    to_normalize=False,
+                    lengths=lengths,
+                )
+
+            generated_motions.append(motion)
+            current_frame += num_frame
+
+        generated_motions = torch.cat(generated_motions, axis=1)  # temporal axis (b, t, d)
+
+        if tosqueeze:
+            generated_motions = generated_motions[0]
+
+        output = self.motion_rep.inverse(
+            generated_motions,
+            is_normalized=False,
+            return_numpy=False,
+        )
+
+        # Post-processing: already applied per-segment inside the loop above,
+        # so no additional post-processing pass is needed here.
+
+        # Convert SOMA output to somaskel77 for external API
+        if isinstance(self.skeleton, SOMASkeleton30):
+            output = self.skeleton.output_to_SOMASkeleton77(output)
+
+        # Convert to numpy if requested
+        if return_numpy:
+            output = to_numpy(output)
+        return output
+
+    def __call__(
+        self,
+        prompts: str | list[str],
+        num_frames: int | list[int],
+        num_denoising_steps: int,
+        multi_prompt: bool = False,
+        constraint_lst: Optional[list] = [],
+        cfg_weight: Optional[float] = [2.0, 2.0],
+        num_samples: Optional[int] = None,
+        cfg_type: Optional[str] = None,
+        return_numpy: bool = False,
+        first_heading_angle: Optional[torch.Tensor] = None,
+        # for transitioning
+        num_transition_frames: int = 5,
+        # for postprocess
+        post_processing: bool = False,
+        root_margin: float = 0.04,
+        # progress bar
+        progress_bar=tqdm,
+    ) -> dict:
+        """Generate motion from text prompts and optional kinematic constraints.
+
+        When a single prompt/num_frames pair is given, one motion is generated.
+        Passing lists of prompts and/or num_frames produces a batch of
+        independent motions. With ``multi_prompt=True``, the prompts are
+        treated as sequential segments that are generated and stitched together
+        with smooth transitions.
+
+        Args:
+            prompts: One or more text descriptions of the desired motion.
+                A single string generates one sample; a list generates a batch
+                (or sequential segments when ``multi_prompt=True``).
+            num_frames: Duration of the generated motion in frames.  Can be a
+                single int applied to every prompt or a per-prompt list.
+            num_denoising_steps: Number of DDIM denoising steps.  More steps
+                generally improve quality at the cost of speed.
+            multi_prompt: If ``True``, treat ``prompts`` as an ordered sequence
+                of segments and concatenate them with transitions.
+            constraint_lst: Per-sample list of kinematic constraints (e.g.
+                keyframe poses, end-effector targets, 2-D paths).  Pass an
+                empty list for unconstrained generation.
+            cfg_weight: Classifier-free guidance scale(s).  A two-element list
+                ``[text_cfg, constraint_cfg]`` controls text and constraint
+                guidance independently.
+            num_samples: Number of samples to generate.
+            cfg_type: Override the default CFG strategy set at init
+                (e.g. ``"separated"``).
+            return_numpy: If ``True``, convert all output tensors to numpy
+                arrays.
+            first_heading_angle: Initial body heading in radians.  Shape
+                ``(B,)`` or scalar.  Defaults to ``0`` (facing +Z).
+            num_transition_frames: Number of overlapping frames used to blend
+                consecutive segments in multi-prompt mode.
+            post_processing: If ``True``, apply post-processing
+                (foot-skate cleanup and constraint enforcement).
+            root_margin: Horizontal margin (in meters) used by the post-processor
+                to determine when to correct root motion. When root deviates more than
+                margin from the constraint, the post-processor will correct it.
+            progress_bar: Callable wrapping an iterable to display progress
+                (default: ``tqdm``).  Pass a no-op to silence output.
+
+        Returns:
+            dict: A dictionary of motion tensors (or numpy arrays if
+            ``return_numpy=True``) with the following keys:
+
+            - ``local_rot_mats`` – Local joint rotations as rotation matrices.
+            - ``global_rot_mats`` – Global joint rotations as rotation matrices.
+            - ``posed_joints`` – Joint positions in world space.
+            - ``root_positions`` – Root joint positions.
+            - ``smooth_root_pos`` – Smoothed root trajectory.
+            - ``foot_contacts`` – Boolean foot-contact labels [left heel, left toe, right heel, right toe].
+            - ``global_root_heading`` – Root heading angle over time.
+        """
+        device = self.device
+
+        if multi_prompt:
+            # multi prompt generation
+            return self._multiprompt(
+                prompts,
+                num_frames,
+                num_denoising_steps,
+                constraint_lst,
+                cfg_weight,
+                num_samples,
+                cfg_type,
+                return_numpy,
+                first_heading_angle,
+                num_transition_frames,
+                post_processing,
+                root_margin,
+                progress_bar,
+            )
+
+        # Input checking
+        tosqueeze = False
+        if isinstance(prompts, list) and isinstance(num_frames, list):
+            assert len(prompts) == len(num_frames), "The number of prompts should match the number of num_frames."
+            num_samples = len(prompts)
+        elif isinstance(prompts, list):
+            num_samples = len(prompts)
+            num_frames = [num_frames for _ in range(num_samples)]
+        elif isinstance(num_frames, list):
+            num_samples = len(num_frames)
+            prompts = [prompts for _ in range(num_samples)]
+        else:
+            if num_samples is None:
+                tosqueeze = True
+                num_samples = 1
+            prompts = [prompts for _ in range(num_samples)]
+            num_frames = [num_frames for _ in range(num_samples)]
+
+        bs = num_samples
+        texts = sanitize_texts(prompts)
+
+        lengths = torch.tensor(
+            num_frames,
+            device=device,
+        )
+        max_frames = max(lengths)
+        motion_pad_mask = length_to_mask(lengths)
+
+        if first_heading_angle is None:
+            # Start at 0 angle
+            first_heading_angle = torch.tensor([0.0] * bs, device=device)
+        else:
+            first_heading_angle = torch.as_tensor(first_heading_angle, device=device)
+            if first_heading_angle.numel() == 1:
+                first_heading_angle = first_heading_angle.repeat(bs)
+
+        observed_motion, motion_mask = None, None
+        if constraint_lst:
+            observed_motion, motion_mask = self.motion_rep.create_conditions_from_constraints_batched(
+                constraint_lst,
+                lengths,
+                to_normalize=True,
+                device=device,
+            )
+
+        motion = self._generate(
+            texts,
+            max_frames,
+            num_denoising_steps=num_denoising_steps,
+            pad_mask=motion_pad_mask,
+            first_heading_angle=first_heading_angle,
+            motion_mask=motion_mask,
+            observed_motion=observed_motion,
+            cfg_weight=cfg_weight,
+            cfg_type=cfg_type,
+            progress_bar=progress_bar,
+        )
+
+        if tosqueeze:
+            motion = motion[0]
+
+        output = self.motion_rep.inverse(
+            motion,
+            is_normalized=True,
+            return_numpy=False,  # Keep as tensor for potential post-processing
+        )
+
+        # Apply post-processing if requested
+        if post_processing:
+            corrected = post_process_motion(
+                output["local_rot_mats"],
+                output["root_positions"],
+                output["foot_contacts"],
+                self.skeleton,
+                constraint_lst,
+                root_margin=root_margin,
+            )
+            # key frame outputs / foot contacts are not changed
+            output.update(corrected)
+
+        # Convert SOMA output to somaskel77 for external API
+        if isinstance(self.skeleton, SOMASkeleton30):
+            output = self.skeleton.output_to_SOMASkeleton77(output)
+
+        # Convert to numpy if requested
+        if return_numpy:
+            output = to_numpy(output)
+        return output
+
+    def _generate(
+        self,
+        texts: List[str],
+        max_frames: int,
+        num_denoising_steps: int,
+        pad_mask: torch.Tensor,
+        first_heading_angle: Optional[torch.Tensor],
+        motion_mask: torch.Tensor,
+        observed_motion: torch.Tensor,
+        cfg_weight: Optional[float] = 2.0,
+        text_feat: Optional[torch.Tensor] = None,
+        text_pad_mask: Optional[torch.Tensor] = None,
+        guide_masks: Optional[Dict] = None,
+        cfg_type: Optional[str] = None,
+        progress_bar=tqdm,
+    ) -> torch.Tensor:
+        """Sample full denoising loop.
+
+        Args:
+            texts (List[str]): batch of text prompts to use for sampling (if text_feat is not passed in)
+        """
+
+        device = self.device
+        if text_feat is None:
+            assert text_pad_mask is None
+            log.info("Encoding text...")
+            text_feat, text_length = self.text_encoder(texts)
+            text_feat = text_feat.to(device)
+
+            # handle empty string (set to zero)
+            empty_text_mask = [len(text.strip()) == 0 for text in texts]
+            text_feat[empty_text_mask] = 0
+
+            # Create the pad mask for the text
+            batch_size, maxlen = text_feat.shape[:2]
+            tensor_text_length = torch.tensor(text_length, device=device)
+            tensor_text_length[empty_text_mask] = 0
+            text_pad_mask = torch.arange(maxlen, device=device).expand(batch_size, maxlen) < tensor_text_length[:, None]
+
+        if motion_mask is not None:
+            if motion_mask.dtype == torch.bool:
+                motion_mask = 1 * motion_mask
+
+        batch_size = text_feat.shape[0]
+
+        # sample loop
+        indices = list(range(num_denoising_steps))[::-1]
+        shape = (batch_size, max_frames, self.motion_rep.motion_rep_dim)
+        cur_mot = torch.randn(shape, device=self.device)
+        num_denoising_steps = torch.tensor(
+            [num_denoising_steps], device=self.device
+        )  # this and t need to be tensor for onnx export
+        # init diffusion with correct num steps before looping
+        use_timesteps = self.diffusion.space_timesteps(num_denoising_steps[0])[0]
+        self.diffusion.calc_diffusion_vars(use_timesteps)
+        for i in progress_bar(indices):
+            t = torch.tensor([i] * cur_mot.size(0), device=self.device)
+            with torch.inference_mode():
+                cur_mot = self.denoising_step(
+                    cur_mot,
+                    pad_mask,
+                    text_feat,
+                    text_pad_mask,
+                    t,
+                    first_heading_angle,
+                    motion_mask,
+                    observed_motion,
+                    num_denoising_steps,
+                    cfg_weight,
+                    guide_masks=guide_masks,
+                    cfg_type=cfg_type,
+                )
+        return cur_mot

kimodo/model/llm2vec/README.md

@@ -0,0 +1 @@
+This is a patched version of the original [LLM2Vec](https://github.com/McGill-NLP/llm2vec) codebase so that `McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntp-supervised` works with `transformers==5.0.0rc3`.

kimodo/model/llm2vec/__init__.py

@@ -0,0 +1,11 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""LLM2Vec text encoder and wrapper for Kimodo."""
+
+from .llm2vec import LLM2Vec
+from .llm2vec_wrapper import LLM2VecEncoder
+
+__all__ = [
+    "LLM2Vec",
+    "LLM2VecEncoder",
+]

kimodo/model/llm2vec/llm2vec.py

@@ -0,0 +1,477 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 McGill NLP
+# SPDX-License-Identifier: MIT
+#
+# Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"),
+# to deal in the Software without restriction, including without limitation
+# the rights to use, copy, modify, merge, publish, distribute, sublicense,
+# and/or sell copies of the Software, and to permit persons to whom the
+# Software is furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+# DEALINGS IN THE SOFTWARE.
+
+
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import logging
+import os
+from functools import partial
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+import torch
+import torch.multiprocessing as mp
+from peft import PeftModel
+from torch import Tensor, device, nn
+from tqdm.autonotebook import tqdm, trange
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoTokenizer,
+    GemmaConfig,
+    LlamaConfig,
+    MistralConfig,
+    PretrainedConfig,
+    Qwen2Config,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def batch_to_device(batch, target_device: device):
+    """Send a pytorch batch to a device (CPU/GPU)"""
+    for key in batch:
+        if isinstance(batch[key], Tensor):
+            batch[key] = batch[key].to(target_device)
+    return batch
+
+
+class LLM2Vec(nn.Module):
+    def __init__(
+        self,
+        model: AutoModel,
+        tokenizer: AutoTokenizer,
+        pooling_mode: str = "mean",
+        max_length: int = 512,
+        doc_max_length: int = 400,
+        skip_instruction: bool = True,
+    ):
+        super().__init__()
+        self.model = model
+        self.tokenizer = tokenizer
+        self.pooling_mode = pooling_mode
+        self.skip_instruction = skip_instruction
+        self.max_length = max_length
+        self.doc_max_length = doc_max_length
+        self.config = model.config
+
+    @classmethod
+    def _get_model_class(cls, config_class_name, enable_bidirectional):
+        if not enable_bidirectional:
+            return AutoModel
+        if config_class_name == "MistralConfig":
+            from .models.bidirectional_mistral import MistralBiModel
+
+            return MistralBiModel
+        elif config_class_name == "LlamaConfig":
+            from .models.bidirectional_llama import LlamaBiModel
+
+            return LlamaBiModel
+        elif config_class_name == "GemmaConfig":
+            from .models.bidirectional_gemma import GemmaBiModel
+
+            return GemmaBiModel
+        elif config_class_name == "Qwen2Config":
+            from .models.bidirectional_qwen2 import Qwen2BiModel
+
+            return Qwen2BiModel
+        else:
+            raise ValueError(f"{config_class_name} is not supported yet with bidirectional models.")
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        base_model_name_or_path,
+        peft_model_name_or_path=None,
+        merge_peft=False,
+        enable_bidirectional=True,
+        **kwargs,
+    ):
+        # pop out encoder args
+        keys = ["pooling_mode", "max_length", "doc_max_length", "skip_instruction"]
+        encoder_args = {key: kwargs.pop(key, None) for key in keys if kwargs.get(key) is not None}
+
+        tokenizer = AutoTokenizer.from_pretrained(base_model_name_or_path)
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.padding_side = "left"
+
+        config = AutoConfig.from_pretrained(base_model_name_or_path)
+        config_class_name = config.__class__.__name__
+
+        model_class = cls._get_model_class(config_class_name, enable_bidirectional=enable_bidirectional)
+
+        model = model_class.from_pretrained(base_model_name_or_path, **kwargs)
+
+        if os.path.isdir(base_model_name_or_path) and os.path.exists(f"{base_model_name_or_path}/config.json"):
+            with open(f"{base_model_name_or_path}/config.json", "r") as fIn:
+                config_dict = json.load(fIn)
+            config = PretrainedConfig.from_dict(config_dict)
+            model.config._name_or_path = config._name_or_path
+
+        # For special case where config.json and adapter weights are in the same directory
+        if hasattr(model, "peft_config"):
+            model = PeftModel.from_pretrained(
+                model,
+                base_model_name_or_path,
+            )
+            model = model.merge_and_unload()
+
+        if peft_model_name_or_path is not None:
+            model = PeftModel.from_pretrained(
+                model,
+                peft_model_name_or_path,
+            )
+            if merge_peft:
+                model = model.merge_and_unload()
+
+        config = {}
+        config_addr = peft_model_name_or_path if peft_model_name_or_path is not None else base_model_name_or_path
+        if os.path.exists(f"{config_addr}/llm2vec_config.json"):
+            with open(f"{config_addr}/llm2vec_config.json", "r") as fIn:
+                llm2vec_config = json.load(fIn)
+            config.update(llm2vec_config)
+
+        for key, value in encoder_args.items():
+            config[key] = value
+
+        return cls(model=model, tokenizer=tokenizer, **config)
+
+    def prepare_for_tokenization(self, text):
+        if self.model.config._name_or_path == "meta-llama/Meta-Llama-3-8B-Instruct":
+            text = "<|start_header_id|>user<|end_header_id|>\n\n" + text.strip() + "<|eot_id|>"
+            return text
+        if self.model.config._name_or_path in [
+            "mistralai/Mistral-7B-Instruct-v0.2",
+            "meta-llama/Llama-2-7b-chat-hf",
+        ]:
+            text = "[INST] " + text.strip() + " [/INST]"
+        if self.model.config._name_or_path in [
+            "google/gemma-2-9b-it",
+        ]:
+            text = "<bos><start_of_turn>user\n" + text.strip() + "<end_of_turn>"
+        if self.model.config._name_or_path in [
+            "Qwen/Qwen2-1.5B-Instruct",
+            "Qwen/Qwen2-7B-Instruct",
+        ]:
+            text = "<|im_start|>user\n" + text.strip() + "<|im_end|>"
+        if self.pooling_mode == "eos_token":
+            if self.model.config._name_or_path == "meta-llama/Meta-Llama-3-8B":
+                text = text.strip() + "<|end_of_text|>"
+            elif isinstance(self.model.config, LlamaConfig) or isinstance(self.model.config, MistralConfig):
+                text = text.strip() + " </s>"
+            elif isinstance(self.model.config, GemmaConfig):
+                text = text.strip() + "<eos>"
+            elif isinstance(self.model.config, Qwen2Config):
+                text = text.strip() + "<|endoftext|>"
+        return text
+
+    def tokenize(self, texts):
+        texts_2 = []
+        original_texts = []
+        for text in texts:
+            t = text.split("!@#$%^&*()")
+            texts_2.append(t[1] if len(t) > 1 else "")
+            original_texts.append("".join(t))
+
+        original = self.tokenizer(
+            original_texts,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=self.max_length,
+        )
+        embed_mask = None
+        for t_i, t in enumerate(texts_2):
+            ids = self.tokenizer(
+                [t],
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=self.max_length,
+                add_special_tokens=False,
+            )
+            if embed_mask is None:
+                e_m = torch.zeros_like(original["attention_mask"][t_i])
+                if len(ids["input_ids"][0]) > 0:
+                    e_m[-len(ids["input_ids"][0]) :] = torch.ones(len(ids["input_ids"][0]))
+                embed_mask = e_m.unsqueeze(0)
+            else:
+                e_m = torch.zeros_like(original["attention_mask"][t_i])
+                if len(ids["input_ids"][0]) > 0:
+                    e_m[-len(ids["input_ids"][0]) :] = torch.ones(len(ids["input_ids"][0]))
+                embed_mask = torch.cat((embed_mask, e_m.unsqueeze(0)), dim=0)
+
+        original["embed_mask"] = embed_mask
+        return original
+
+    def _skip_instruction(self, sentence_feature):
+        assert sentence_feature["attention_mask"].shape == sentence_feature["embed_mask"].shape
+        sentence_feature["attention_mask"] = sentence_feature["embed_mask"]
+
+    def forward(self, sentence_feature: Dict[str, Tensor]):
+        embed_mask = None
+        if "embed_mask" in sentence_feature:
+            embed_mask = sentence_feature.pop("embed_mask")
+        reps = self.model(**sentence_feature)
+        sentence_feature["embed_mask"] = embed_mask
+
+        return self.get_pooling(sentence_feature, reps.last_hidden_state)
+
+    def get_pooling(self, features, last_hidden_states):  # All models padded from left
+        assert self.tokenizer.padding_side == "left", "Pooling modes are implemented for padding from left."
+        if self.skip_instruction:
+            self._skip_instruction(features)
+        seq_lengths = features["attention_mask"].sum(dim=-1)
+        if self.pooling_mode == "mean":
+            return torch.stack(
+                [last_hidden_states[i, -length:, :].mean(dim=0) for i, length in enumerate(seq_lengths)],
+                dim=0,
+            )
+        elif self.pooling_mode == "weighted_mean":
+            bs, l, _ = last_hidden_states.shape
+            complete_weights = torch.zeros(bs, l, device=last_hidden_states.device)
+            for i, seq_l in enumerate(seq_lengths):
+                if seq_l > 0:
+                    complete_weights[i, -seq_l:] = torch.arange(seq_l) + 1
+                    complete_weights[i] /= torch.clamp(complete_weights[i].sum(), min=1e-9)
+            return torch.sum(last_hidden_states * complete_weights.unsqueeze(-1), dim=1)
+        elif self.pooling_mode == "eos_token" or self.pooling_mode == "last_token":
+            return last_hidden_states[:, -1]
+        elif self.pooling_mode == "bos_token":
+            return last_hidden_states[features["input_ids"] == self.tokenizer.bos_token_id]
+        else:
+            raise ValueError(f"{self.pooling_mode} is not implemented yet.")
+
+    def _convert_to_str(self, instruction, text):
+        tokenized_q = self.tokenizer(
+            text,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=self.max_length,
+            add_special_tokens=False,
+        )
+        tokenized_q_length = len(tokenized_q["input_ids"][0])
+
+        while tokenized_q_length > self.doc_max_length:
+            reduction_ratio = self.doc_max_length / tokenized_q_length
+            reduced_length = int(len(text.split()) * reduction_ratio)
+            text = " ".join(text.split()[:reduced_length])
+            tokenized_q = self.tokenizer(
+                text,
+                return_tensors="pt",
+                padding=True,
+                truncation=True,
+                max_length=self.max_length,
+                add_special_tokens=False,
+            )
+            tokenized_q_length = len(tokenized_q["input_ids"][0])
+
+        return f"{instruction.strip()} !@#$%^&*(){text}" if instruction else f"!@#$%^&*(){text}"
+
+    def encode(
+        self,
+        sentences: Union[str, List[str]],
+        batch_size: int = 32,
+        show_progress_bar: bool = True,
+        convert_to_numpy: bool = False,
+        convert_to_tensor: bool = False,
+        device: Optional[str] = None,
+    ):
+        """
+        Encode a list of sentences to their respective embeddings. The sentences can be a list of strings or a string.
+        Args:
+            sentences: sentence or sentences to encode.
+            batch_size: batch size for turning sentence tokens into embeddings.
+            show_progress_bar: whether to show progress bars during encoding steps.
+            convert_to_numpy: If true, return numpy arrays instead of torch tensors.
+            convert_to_tensor: If true, return torch tensors (default).
+            device: torch backend device identifier (e.g., 'cuda', 'cpu','mps' etc.). If not specified,
+            the default is to use cuda when available, otherwise cpu. Note that only the choice of 'cuda' supports
+            multiprocessing as currently implemented.
+
+        Returns: embeddings of the sentences. Embeddings are detached and always on the CPU (see _encode implementation).
+
+        """
+        if isinstance(sentences[0], str) and isinstance(sentences[-1], int):
+            sentences = [sentences]
+        # required for MEDI version of MTEB
+        if isinstance(sentences[0], str):
+            sentences = [[""] + [sentence] for sentence in sentences]
+
+        if device is None:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        concatenated_input_texts = []
+        for sentence in sentences:
+            assert isinstance(sentence[0], str)
+            assert isinstance(sentence[1], str)
+            concatenated_input_texts.append(self._convert_to_str(sentence[0], sentence[1]))
+        sentences = concatenated_input_texts
+
+        self.eval()
+
+        if convert_to_tensor:
+            convert_to_numpy = False
+
+        length_sorted_idx = np.argsort([-self._text_length(sen) for sen in sentences])
+        sentences_sorted = [sentences[idx] for idx in length_sorted_idx]
+        all_embeddings = []
+
+        if torch.cuda.device_count() <= 1:
+            # This branch also support mps devices
+            self.to(device)
+            for start_index in trange(
+                0,
+                len(sentences),
+                batch_size,
+                desc="Batches",
+                disable=not show_progress_bar,
+            ):
+                sentences_batch = sentences_sorted[start_index : start_index + batch_size]
+                embeddings = self._encode(sentences_batch, device=device, convert_to_numpy=convert_to_numpy)
+                all_embeddings.append(embeddings)
+        else:
+            num_proc = torch.cuda.device_count()
+            cuda_compatible_multiprocess = mp.get_context("spawn")
+            with cuda_compatible_multiprocess.Pool(num_proc) as p:
+                sentences_batches = [
+                    sentences_sorted[start_index : start_index + batch_size]
+                    for start_index in range(0, len(sentences), batch_size)
+                ]
+
+                progress_bar = tqdm(
+                    total=len(sentences_batches),
+                    desc="Batches",
+                    disable=not show_progress_bar,
+                )
+                results = []
+
+                def update(*args):
+                    progress_bar.update()
+
+                for batch in sentences_batches:
+                    results.append(
+                        p.apply_async(
+                            self._encode,
+                            args=(batch, None, convert_to_numpy, True),
+                            callback=update,
+                        )
+                    )
+
+                all_embeddings = [result.get() for result in results]
+                progress_bar.close()
+
+        all_embeddings = torch.cat(all_embeddings, dim=0)
+        all_embeddings = all_embeddings[np.argsort(length_sorted_idx)]
+        all_embeddings = all_embeddings.to(torch.float32)
+        if convert_to_numpy:
+            all_embeddings = np.asarray([emb.numpy() for emb in all_embeddings])
+        return all_embeddings
+
+    def save(self, output_path, merge_before_save=False, save_config=True):
+        if merge_before_save and isinstance(self.model, PeftModel):
+            self.model = self.model.merge_and_unload()
+            # Fixes the issue of saving - https://huggingface.co/McGill-NLP/LLM2Vec-Mistral-7B-Instruct-v2-mntp-unsup-simcse/discussions/1
+            if hasattr(self.model, "_hf_peft_config_loaded"):
+                self.model._hf_peft_config_loaded = False
+
+        self.model.save_pretrained(output_path)
+        self.tokenizer.save_pretrained(output_path)
+
+        llm2vec_config = {
+            "pooling_mode": self.pooling_mode,
+            "max_length": self.max_length,
+            "doc_max_length": self.doc_max_length,
+            "skip_instruction": self.skip_instruction,
+        }
+
+        if save_config:
+            os.makedirs(output_path, exist_ok=True)
+            with open(f"{output_path}/llm2vec_config.json", "w") as fOut:
+                json.dump(llm2vec_config, fOut, indent=4)
+
+    def _encode(
+        self,
+        sentences_batch,
+        device: Optional[str] = None,
+        convert_to_numpy: bool = False,
+        multiprocessing=False,
+    ):
+        if multiprocessing:
+            # multiprocessing only supports CUDA devices at this time, so we ignore the value of device
+            # and use cuda:rank for the device
+            rank = mp.current_process()._identity[0]
+            if device is None and torch.cuda.is_available():
+                device = f"cuda:{rank % torch.cuda.device_count()}"
+
+        self.to(device)
+        features = self.tokenize([self.prepare_for_tokenization(sentence) for sentence in sentences_batch])
+        features = batch_to_device(features, device)
+
+        with torch.no_grad():
+            embeddings = self.forward(features)
+            embeddings = embeddings.detach()
+            embeddings = embeddings.cpu()
+
+        return embeddings
+
+    def _text_length(self, text: Union[List[int], List[List[int]]]):
+        """Help function to get the length for the input text.
+
+        Text can be either a string (which means a single text) a list of ints (which means a single
+        tokenized text), or a tuple of list of ints (representing several text inputs to the model).
+        """
+        if (
+            isinstance(text, str) or (isinstance(text, list) and isinstance(text[0], int)) or len(text) == 0
+        ):  # Single text, list of ints, or empty
+            return len(text)
+        if isinstance(text, dict):  # {key: value} case
+            return len(next(iter(text.values())))
+        elif not hasattr(text, "__len__"):  # Object has no len() method
+            return 1
+        else:
+            return sum([len(t) for t in text])
+
+    def resize_token_embeddings(
+        self,
+        new_num_tokens: Optional[int] = None,
+        pad_to_multiple_of: Optional[int] = None,
+    ) -> nn.Embedding:
+        return self.model.resize_token_embeddings(new_num_tokens=new_num_tokens, pad_to_multiple_of=pad_to_multiple_of)
+
+    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
+        self.model.gradient_checkpointing_enable(gradient_checkpointing_kwargs=gradient_checkpointing_kwargs)

kimodo/model/llm2vec/llm2vec_wrapper.py

@@ -0,0 +1,95 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""LLM2Vec encoder wrapper for Kimodo text conditioning."""
+
+import os
+
+import numpy as np
+import torch
+
+from .llm2vec import LLM2Vec
+
+
+class LLM2VecEncoder:
+    """LLM2Vec text embeddings."""
+
+    def __init__(
+        self,
+        base_model_name_or_path: str,
+        peft_model_name_or_path: str,
+        dtype: str,
+        llm_dim: int,
+        device: str = "auto",
+    ) -> None:
+        torch_dtype = getattr(torch, dtype)
+        self.llm_dim = llm_dim
+
+        cache_dir = os.environ.get("HUGGINGFACE_CACHE_DIR")
+
+        if "TEXT_ENCODERS_DIR" in os.environ:
+            base_model_name_or_path = os.path.join(os.environ["TEXT_ENCODERS_DIR"], base_model_name_or_path)
+            peft_model_name_or_path = os.path.join(os.environ["TEXT_ENCODERS_DIR"], peft_model_name_or_path)
+
+        self.model = LLM2Vec.from_pretrained(
+            base_model_name_or_path=base_model_name_or_path,
+            peft_model_name_or_path=peft_model_name_or_path,
+            torch_dtype=torch_dtype,
+            cache_dir=cache_dir,
+        )
+
+        env_device = os.environ.get("TEXT_ENCODER_DEVICE")
+        if env_device:
+            device = env_device
+        if device == "auto":
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+        self._device = device
+        if device is not None:
+            self.model = self.model.to(device)
+
+        self.model.eval()
+        for p in self.model.parameters():
+            p.requires_grad = False
+
+    def to(self, device: torch.device):
+        self.model = self.model.to(device)
+        self._device = str(device) if not isinstance(device, str) else device
+        return self
+
+    def eval(self):
+        self.model.eval()
+        return self
+
+    def get_device(self):
+        return self.model.model.device
+
+    def __call__(self, text: list[str] | str):
+        is_string = False
+        if isinstance(text, str):
+            text = [text]
+            is_string = True
+
+        with torch.no_grad():
+            encoded_text = self.model.encode(
+                text,
+                # IMPORTANT: different batch sizes unexpectedly change the output embeddings, so we always set it to 1
+                #            here for repeatability no matter how many texts are being encoded. This
+                #            is a fundamental issue with transformers, and is especially bad at lower
+                #            precisions (https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535)
+                #            note: this is an internal batch size used by llm2vec - the text list can still be of arbitrary length.
+                batch_size=1,
+                show_progress_bar=False,
+                device=self._device,
+            )
+
+        assert len(encoded_text.shape)
+        assert self.llm_dim == encoded_text.shape[-1]
+
+        encoded_text = encoded_text[:, None]
+        lengths = np.ones(len(encoded_text), dtype=int).tolist()
+
+        if is_string:
+            encoded_text = encoded_text[0]
+            lengths = lengths[0]
+
+        encoded_text = torch.tensor(encoded_text).to(self._device)
+        return encoded_text, lengths