add blocks

LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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.

NOTICE

@@ -0,0 +1,11 @@
+MotionBricks SOMA Modular Diffusers
+Copyright 2026 The HuggingFace Team.
+
+The guide encoding, root and output feature decoding, and endpoint correction in motion_utils.py are adapted from
+NVIDIA MotionBricks (https://github.com/NVlabs/motionbricks):
+
+    SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+    SPDX-License-Identifier: Apache-2.0
+
+The MotionBricks model weights are not part of this repository. They are licensed under the NVIDIA Open Model
+Agreement: https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-agreement/

README.md

@@ -0,0 +1,118 @@
+---
+language:
+- en
+license: apache-2.0
+base_model:
+- nvidia/MotionBricks
+library_name: diffusers
+tags:
+- motion-generation
+- motion-inbetweening
+- motionbricks
+- modular-diffusers
+- soma
+---
+
+# MotionBricks SOMA Modular Diffusers
+
+<video controls loop muted playsinline src="https://huggingface.co/datasets/OzzyGT/diffusers-examples/resolve/main/motionbricks_soma_diffusers/motionbricks_soma_run_stop.mp4"></video>
+
+*The start of a Kimodo run, then the 30 frames generated with the code below, then the standing end of another Kimodo take, placed 1.4 m ahead. The generated frames are drawn on SOMA77 with a relaxed finger pose. At normal playback speed; the cameras follow the root.*
+
+Modular Diffusers blocks for [NVIDIA MotionBricks](https://huggingface.co/nvidia/MotionBricks) on the SOMA skeleton: generate the transition from the last poses of one motion to the first poses of another in a single forward pass, with no sampling loop. The model code and weights are in [OzzyGT/motionbricks_soma_diffusers](https://huggingface.co/OzzyGT/motionbricks_soma_diffusers).
+
+MotionBricks has no text input. It fills the gap between poses you provide, for example two motions from [Kimodo SOMA Modular](https://huggingface.co/OzzyGT/Kimodo_SOMA_Modular), and chooses the length of the transition itself.
+
+Requires PyTorch, `diffusers>=0.41.0` and `numpy`.
+
+## Licensing
+
+The blocks are [Apache-2.0](LICENSE), adapted from the [MotionBricks runtime](https://github.com/NVlabs/motionbricks). See [NOTICE](NOTICE). This repository contains no model weights; they use the [NVIDIA Open Model Agreement](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-agreement/), see the [model repository](https://huggingface.co/OzzyGT/motionbricks_soma_diffusers).
+
+## Usage
+
+The example uses `trust_remote_code=True`, which loads these blocks, and `trust_components_code=True`, which allows the model repository's custom classes.
+
+> **Note:** both flags run Python code downloaded from the Hub: `trust_remote_code=True` runs the code of the blocks repository, and `trust_components_code=True` runs the code shipped in the `components_repo`. Only use them with repositories you trust, and review their code before passing different ones.
+
+Bring a run to a stop, as in the video. Both clips come from [Kimodo SOMA Modular](https://huggingface.co/OzzyGT/Kimodo_SOMA_Modular): the start of a run, and a standing pose that the pipeline places 1.4 m ahead of it:
+
+```python
+import numpy as np
+import torch
+from diffusers import ModularPipelineBlocks
+from huggingface_hub import hf_hub_download
+
+blocks = ModularPipelineBlocks.from_pretrained(
+    "OzzyGT/motionbricks_soma_custom_blocks",
+    trust_remote_code=True,
+    components_repo="OzzyGT/motionbricks_soma_diffusers",
+    trust_components_code=True,
+)
+pipe = blocks.init_pipeline()
+pipe.load_components(dtype=torch.float32)
+pipe.to("cuda")
+
+first = hf_hub_download("OzzyGT/diffusers-examples", "motionbricks_soma_diffusers/first.npz", repo_type="dataset")
+second = hf_hub_download("OzzyGT/diffusers-examples", "motionbricks_soma_diffusers/second.npz", repo_type="dataset")
+context = dict(np.load(first))
+target = dict(np.load(second))
+
+motion = pipe(context=context, target=target, target_offset=(1.4, 0.0), num_target_poses=4, output="motion")[0]
+np.savez_compressed("transition.npz", **{key: value.cpu().numpy() for key, value in motion.items()})
+```
+
+Here the model picks a 30-frame transition.
+
+### Inputs
+
+- `context`: the motion to continue from.
+- `target`: the motion to arrive at.
+- `target_offset`: where `target` starts, `(forward, left)` in meters from the end of `context`.
+- `target_turn`: which way `target` faces, in degrees from the end of `context`, positive to the left.
+- `num_target_poses`: how many frames of `target` to match, 1 to 4. Default 1.
+- `num_tokens`: the lengths the model may choose. Default 26 to 50 frames.
+- `correct_endpoints`: start and end exactly on the given poses. Default `True`.
+
+### Notes
+
+- **Motions.** `context` and `target` are SOMA30 or SOMA77 motions at 30 FPS, such as Kimodo output loaded from an NPZ. The result is a SOMA30 motion.
+- **Joining.** The transition starts on the last frame of `context` and ends on the first frame of `target`, so cut them where they should join. Ask for `output=["motion", "placed_target"]` to also get `target` where the pipeline placed it, to join the three motions.
+- **Placement.** Without `target_offset` or `target_turn`, `target` stays where it is. With only one of them, the other is 0.
+- **Target poses.** The model always reads the last 4 frames of `context` and the first `num_target_poses` frames of `target`, consecutive frames, to see how the body moves at each end. Use 4 whenever `target` has at least 4 frames, so the transition also matches how `target` moves; use 1 when it is a single pose.
+- **Length.** `num_tokens` counts in steps of 4 frames, from 6 to 16; `n` tokens give `4 * n - 6` frames. Pass one value to fix the length.
+- **Longer transitions.** Feed the generated motion back as the next `context`.
+
+A call takes about 7.5 ms in float32 on an RTX 5090 and peaks at 0.8 GB of VRAM.
+
+## Precision
+
+The weights are float32, and float32 reproduces NVIDIA's ONNX graph with PyTorch's default settings. Enabling TF32 matrix multiplication (`torch.backends.cuda.matmul.allow_tf32 = True`) can change a few pose codes. `dtype=torch.float16` halves the memory and is a little faster, with small drift: in a test of 8 transitions, every length matched float32 and joints moved by about 4 cm in the median. Avoid `torch.bfloat16`, which drifts much more. NVIDIA notes the same drift in its mixed FP16 model, and it grows when the output is fed back as context.
+
+## Real-time use
+
+These blocks bring MotionBricks to apps and environments that already run on Diffusers. For a real-time product, such as a game or a robot, prefer NVIDIA's [MotionBricks runtime](https://github.com/NVlabs/motionbricks), which runs the ONNX models through TensorRT. With the blocks, about half of each call is Python and data preparation rather than GPU work.
+
+To generate transitions while the motion plays in a Diffusers app, keep the pipeline loaded and capture the three models in CUDA graphs. On an RTX 5090, a call then takes about 4.5 ms instead of 7.5 ms.
+
+Compile while loading, not in the loop: the first call compiles the models in about 15 seconds, and the second records the graphs.
+
+```python
+for name in ("root_model", "pose_model", "pose_decoder"):
+    getattr(pipe, name).compile(mode="reduce-overhead", fullgraph=True)
+
+for _ in range(2):
+    torch.compiler.cudagraph_mark_step_begin()
+    pipe(context=context, target=target, target_offset=(1.4, 0.0), num_target_poses=4, output="motion")
+```
+
+Then, whenever a transition is needed:
+
+```python
+torch.compiler.cudagraph_mark_step_begin()
+motion = pipe(context=context, target=target, target_offset=(1.4, 0.0), num_target_poses=4, output="motion")[0]
+```
+
+- Call `torch.compiler.cudagraph_mark_step_begin()` before every call, so the graphs can reuse their memory; without it, the second call fails.
+- New clips, target placements, `num_target_poses`, `num_tokens` and `correct_endpoints` reuse the graphs. The first new batch size compiles again, once, in about 20 seconds; later sizes don't.
+- Several characters, or several candidates for one character such as different stop distances, fit in one call: stack their motions on a leading batch dimension. Each gets its own transition and length, matching separate calls to within 0.01 mm. With CUDA graphs, a batch of 4 takes about 5.5 ms and a batch of 8 about 6 ms, against 4.5 ms for one.

decoders.py

@@ -0,0 +1,214 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0. See the LICENSE file in this repository.
+
+import torch
+
+from diffusers import AutoModel
+from diffusers.modular_pipelines import ModularPipelineBlocks, PipelineState
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+
+from .motion_utils import CONTEXT_FRAMES, decode_poses, integrate_local_root, select_root
+
+
+class MotionBricksDecodeStep(ModularPipelineBlocks):
+    model_name = "motionbricks"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Decode step that decodes the pose codes to frame features and integrates their local root into a global "
+            "one. When the decoded root height misses the target poses by more than 5 cm and the root step's is "
+            "closer, the root step's trajectory is kept."
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [ComponentSpec("root_model", AutoModel), ComponentSpec("pose_decoder", AutoModel)]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                "pose_codes",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size, 16, 5)` pose codes. Can be generated in the pose tokens step.",
+            ),
+            InputParam(
+                "pose_timeline",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size, 64, 268)` normalized guide poses on the frame timeline. Can be generated "
+                "in the pose tokens step.",
+            ),
+            InputParam(
+                "has_pose_timeline",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size, 64)` frames with a guide pose. Can be generated in the pose tokens step.",
+            ),
+            InputParam(
+                "local_root",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size, 64, 4)` normalized local root. Can be generated in the pose tokens step.",
+            ),
+            InputParam(
+                "motion_tokens",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size,)` motion lengths in tokens. Can be generated in the root step.",
+            ),
+            InputParam(
+                "root",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size, 64, 5)` normalized global root. Can be generated in the root step.",
+            ),
+            InputParam(
+                "guide_root",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size, 8, 5)` aligned global root of the guide poses. Can be generated in the "
+                "guides step.",
+            ),
+            InputParam(
+                "has_global_root",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size, 8)` guide poses with a global root. Can be generated in the guides step.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam(
+                "motion_features",
+                type_hint=torch.Tensor,
+                description="`(batch_size, 64, 366)` float32 features in the aligned space: global root, child joint "
+                "positions, 6D global rotations, joint velocities and foot contacts.",
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        root_model, pose_decoder = components.root_model, components.pose_decoder
+        num_tokens, guide_root = block_state.motion_tokens, block_state.guide_root
+        dtype = pose_decoder.dtype
+
+        features = pose_decoder(
+            block_state.pose_codes,
+            num_tokens,
+            block_state.pose_timeline.to(dtype),
+            block_state.has_pose_timeline,
+            block_state.local_root[..., 1:3].to(dtype),
+        ).features
+        features = pose_decoder.denormalize(features.float())
+        decoded_root = integrate_local_root(features[..., :4])
+        predicted_root = root_model.denormalize(block_state.root.float(), "global_root")
+        root = select_root(
+            decoded_root,
+            predicted_root,
+            guide_root[:, CONTEXT_FRAMES:, 1],
+            block_state.has_global_root[:, CONTEXT_FRAMES:],
+            num_tokens,
+        )
+        # Both roots start at the first context pose's ground position and heading.
+        angle = torch.atan2(root[..., 4], root[..., 3]) + torch.atan2(guide_root[:, :1, 4], guide_root[:, :1, 3])
+        origin = guide_root[:, :1]
+        root = torch.stack(
+            [
+                root[..., 0] + origin[..., 0],
+                root[..., 1],
+                root[..., 2] + origin[..., 2],
+                torch.cos(angle),
+                torch.sin(angle),
+            ],
+            dim=-1,
+        )
+        block_state.motion_features = torch.cat([root, features[..., 4:]], dim=-1)
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class MotionBricksPostprocessStep(ModularPipelineBlocks):
+    model_name = "motionbricks"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Postprocess step that drops the three boundary frames at each end, builds world-space SOMA30 poses by "
+            "forward kinematics on the skeleton's neutral bone offsets and, by default, blends corrections so the "
+            "motion starts at the last context pose and ends at the first target pose."
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [ComponentSpec("pose_decoder", AutoModel)]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                "correct_endpoints",
+                default=True,
+                type_hint=bool,
+                description="Blend the motion onto the last context pose and the first target pose. `False` keeps the "
+                "predicted endpoints.",
+            ),
+            InputParam(
+                "motion_features",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size, 64, 366)` frame features in the aligned space. Can be generated in the "
+                "decode step.",
+            ),
+            InputParam(
+                "motion_tokens",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size,)` motion lengths in tokens. Can be generated in the root step.",
+            ),
+            InputParam(
+                "alignment",
+                required=True,
+                type_hint=tuple,
+                description="`(origin, yaw)` from the aligned space back to the input space. Can be generated in the "
+                "guides step.",
+            ),
+            InputParam(
+                "endpoints",
+                required=True,
+                type_hint=dict,
+                description="Local rotations and root positions of the last context pose and the first target pose. "
+                "Can be generated in the guides step.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam(
+                "motion",
+                type_hint=list[dict[str, torch.Tensor]],
+                description="One SOMA30 motion per batch item, 30 FPS, meters, +Y up, in the input space: "
+                "`posed_joints` `(frames, 30, 3)`, `global_rot_mats` and `local_rot_mats` `(frames, 30, 3, 3)`, "
+                "`root_positions` `(frames, 3)` and `foot_contacts` `(frames, 4)` with the left heel, left toe, right "
+                "heel and right toe contact confidences.",
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        features = block_state.motion_features
+        neutral_joints = components.pose_decoder.neutral_joints.to(features.device, torch.float32)
+        endpoints = block_state.endpoints if block_state.correct_endpoints else None
+        block_state.motion = decode_poses(
+            features, block_state.motion_tokens, block_state.alignment, neutral_joints, endpoints
+        )
+        self.set_block_state(state, block_state)
+        return components, state

encoders.py

@@ -0,0 +1,176 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0. See the LICENSE file in this repository.
+
+import torch
+
+from diffusers import AutoModel
+from diffusers.modular_pipelines import ModularPipelineBlocks, PipelineState
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+
+from .motion_utils import (
+    CONTEXT_FRAMES,
+    TARGET_FRAMES,
+    encode_guides,
+    guide_masks,
+    local_rotations,
+    place_target,
+    soma30,
+)
+
+
+MOTION_DESCRIPTION = (
+    "`posed_joints` `(frames, joints, 3)` and `global_rot_mats` `(frames, joints, 3, 3)`, optionally with a leading "
+    "batch dimension: SOMA30 or SOMA77 joints at 30 FPS, in meters with +Y up, with the rotations of Kimodo SOMA "
+    "motions (the identity is the neutral pose)."
+)
+
+
+def _motion(motion: dict, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
+    positions = torch.as_tensor(motion["posed_joints"], device=device, dtype=torch.float32)
+    rotations = torch.as_tensor(motion["global_rot_mats"], device=device, dtype=torch.float32)
+    if positions.ndim == 3:
+        positions, rotations = positions[None], rotations[None]
+    return positions, rotations
+
+
+class MotionBricksGuidesStep(ModularPipelineBlocks):
+    model_name = "motionbricks"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Guide step that optionally places the target relative to the end of the context, then turns the last "
+            "four poses of the context and the first target poses into the model's normalized root and pose "
+            "features. The poses are turned so the last context pose faces +X."
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [ComponentSpec("root_model", AutoModel)]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                "context",
+                required=True,
+                type_hint=dict,
+                description="The motion to continue from; its last four frames are used. " + MOTION_DESCRIPTION,
+            ),
+            InputParam(
+                "target",
+                required=True,
+                type_hint=dict,
+                description="The motion to reach; its first `num_target_poses` frames are used. " + MOTION_DESCRIPTION,
+            ),
+            InputParam(
+                "num_target_poses",
+                default=1,
+                type_hint=int,
+                description="How many target poses to reach, 1 to 4. One pose sets where the motion ends; more also "
+                "set how it moves there.",
+            ),
+            InputParam(
+                "target_offset",
+                type_hint=tuple,
+                description="Where the first target pose stands, `(forward, left)` in meters from the last context "
+                "pose, along the direction the context faces; one pair, or one per batch item. With `target_offset` "
+                "or `target_turn`, `target` is moved and turned there; without either, it is used where it is.",
+            ),
+            InputParam(
+                "target_turn",
+                type_hint=float,
+                description="Which way the first target pose faces, in degrees relative to the last context pose, "
+                "positive to the left; one value, or one per batch item. Defaults to 0 when only `target_offset` is "
+                "given.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam("guide_root", type_hint=torch.Tensor, description="`(batch_size, 8, 5)` aligned global root."),
+            OutputParam("guide_global_root", type_hint=torch.Tensor, description="Normalized, centered global root."),
+            OutputParam("guide_local_root", type_hint=torch.Tensor, description="Normalized local root."),
+            OutputParam("guide_pose", type_hint=torch.Tensor, description="`(batch_size, 8, 268)` normalized pose."),
+            OutputParam("has_global_root", type_hint=torch.Tensor, description="`(batch_size, 8)` given roots."),
+            OutputParam("has_local_root", type_hint=torch.Tensor, description="`(batch_size, 8)` given local roots."),
+            OutputParam("has_pose", type_hint=torch.Tensor, description="`(batch_size, 8)` given poses."),
+            OutputParam("alignment", type_hint=tuple, description="`(origin, yaw)` back to the input space."),
+            OutputParam(
+                "placed_target",
+                type_hint=list,
+                description="`target` as used, one dict per batch item with all its frames and joints: "
+                "`posed_joints` and `global_rot_mats`, moved by `target_offset` and `target_turn`.",
+            ),
+            OutputParam(
+                "endpoints",
+                type_hint=dict,
+                description="Local rotations and root positions of the last context pose and the first target pose.",
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        root_model = components.root_model
+        device = components._execution_device
+        num_targets = block_state.num_target_poses
+        if not 1 <= num_targets <= TARGET_FRAMES:
+            raise ValueError(f"`num_target_poses` must be 1 to {TARGET_FRAMES}, got {num_targets}")
+
+        context_positions, context_rotations = _motion(block_state.context, device)
+        target_positions, target_rotations = _motion(block_state.target, device)
+        if context_positions.shape[1] < CONTEXT_FRAMES:
+            raise ValueError(f"`context` needs at least {CONTEXT_FRAMES} frames, got {context_positions.shape[1]}")
+        if context_positions.shape[0] != target_positions.shape[0]:
+            raise ValueError("`context` and `target` must have the same batch size")
+        if target_positions.shape[1] < num_targets:
+            raise ValueError(f"`target` has {target_positions.shape[1]} frames, fewer than `num_target_poses`")
+        context_positions = soma30(context_positions[:, -CONTEXT_FRAMES:], 2)
+        context_rotations = soma30(context_rotations[:, -CONTEXT_FRAMES:], 2)
+        if block_state.target_offset is not None or block_state.target_turn is not None:
+            batch_size = target_positions.shape[0]
+            offset = (0.0, 0.0) if block_state.target_offset is None else block_state.target_offset
+            turn = 0.0 if block_state.target_turn is None else block_state.target_turn
+            offset = torch.as_tensor(offset, device=device, dtype=torch.float32).reshape(-1, 2).expand(batch_size, 2)
+            turn = torch.deg2rad(torch.as_tensor(turn, device=device, dtype=torch.float32).reshape(-1))
+            target_positions, target_rotations = place_target(
+                target_positions, target_rotations, context_positions[:, -1], offset, turn.expand(batch_size)
+            )
+        block_state.placed_target = [
+            {"posed_joints": positions, "global_rot_mats": rotations}
+            for positions, rotations in zip(target_positions, target_rotations)
+        ]
+        target_positions = soma30(target_positions[:, :num_targets], 2)
+        target_rotations = soma30(target_rotations[:, :num_targets], 2)
+        # Unused target slots repeat the last target pose.
+        target_index = torch.arange(TARGET_FRAMES, device=device).clamp(max=num_targets - 1)
+        positions = torch.cat([context_positions, target_positions[:, target_index]], 1)
+        rotations = torch.cat([context_rotations, target_rotations[:, target_index]], 1)
+
+        global_root, local_root, pose, alignment = encode_guides(positions, rotations)
+        centered = global_root.clone()
+        centered[..., 0:3:2] -= global_root[:, :1, 0:3:2]
+        batch_size = positions.shape[0]
+        has_global_root, has_local_root, has_pose = guide_masks(torch.full((batch_size,), num_targets, device=device))
+
+        dtype = root_model.dtype
+        block_state.guide_root = global_root
+        block_state.guide_global_root = root_model.normalize(centered.to(dtype), "global_root")
+        block_state.guide_local_root = root_model.normalize(local_root.to(dtype), "local_root")
+        block_state.guide_pose = root_model.normalize(torch.cat([centered[..., 1:2], pose], -1).to(dtype), "pose")
+        block_state.has_global_root = has_global_root
+        block_state.has_local_root = has_local_root
+        block_state.has_pose = has_pose
+        block_state.alignment = alignment
+        block_state.endpoints = {
+            "context_local": local_rotations(rotations[:, CONTEXT_FRAMES - 1]),
+            "context_root": positions[:, CONTEXT_FRAMES - 1, 0],
+            "target_local": local_rotations(rotations[:, CONTEXT_FRAMES]),
+            "target_root": positions[:, CONTEXT_FRAMES, 0],
+        }
+
+        self.set_block_state(state, block_state)
+        return components, state

generation.py

@@ -0,0 +1,215 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0. See the LICENSE file in this repository.
+
+import functools
+
+import torch
+
+from diffusers import AutoModel
+from diffusers.modular_pipelines import ModularPipelineBlocks, PipelineState
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+
+from .motion_utils import FRAMES_PER_TOKEN, root_to_local
+
+
+@functools.cache
+def allowed_lengths(min_tokens: int, max_tokens: int, num_tokens: tuple, device: torch.device) -> torch.Tensor:
+    """`(1, lengths)` mask of the allowed lengths among `min_tokens` to `max_tokens`."""
+    lengths = torch.arange(min_tokens, max_tokens + 1)
+    return torch.isin(lengths, torch.tensor(num_tokens))[None].to(device)
+
+
+class MotionBricksRootStep(ModularPipelineBlocks):
+    model_name = "motionbricks"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Root step that chooses the length of the transition among the allowed ones and predicts the root "
+            "trajectory of every frame."
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [ComponentSpec("root_model", AutoModel)]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                "num_tokens",
+                default=list(range(8, 15)),
+                type_hint=int | list[int],
+                description="The transition lengths the model may choose, in tokens of 4 frames, from 6 to 16. A "
+                "length of `n` tokens gives `4 * n - 6` frames.",
+            ),
+            InputParam(
+                "guide_global_root",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size, 8, 5)` normalized global root of the guide poses. Can be generated in the "
+                "guides step.",
+            ),
+            InputParam(
+                "guide_local_root",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size, 8, 4)` normalized local root of the guide poses. Can be generated in the "
+                "guides step.",
+            ),
+            InputParam(
+                "guide_pose",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size, 8, 268)` normalized pose features of the guide poses. Can be generated in "
+                "the guides step.",
+            ),
+            InputParam(
+                "has_global_root",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size, 8)` guide poses with a global root. Can be generated in the guides step.",
+            ),
+            InputParam(
+                "has_local_root",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size, 8)` guide poses with a local root. Can be generated in the guides step.",
+            ),
+            InputParam(
+                "has_pose",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size, 8)` guide poses with pose features. Can be generated in the guides step.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam("motion_tokens", type_hint=torch.Tensor, description="`(batch_size,)` chosen lengths."),
+            OutputParam("root", type_hint=torch.Tensor, description="`(batch_size, 64, 5)` normalized global root."),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        root_model = components.root_model
+        config = root_model.config
+        num_tokens = block_state.num_tokens
+        num_tokens = [num_tokens] if isinstance(num_tokens, int) else list(num_tokens)
+        if not num_tokens or not all(config.min_tokens <= n <= config.max_tokens for n in num_tokens):
+            raise ValueError(f"`num_tokens` must be in {config.min_tokens} to {config.max_tokens}, got {num_tokens}")
+        allowed = allowed_lengths(
+            config.min_tokens, config.max_tokens, tuple(num_tokens), block_state.guide_pose.device
+        ).expand(block_state.guide_pose.shape[0], -1)
+
+        output = root_model(
+            block_state.guide_global_root,
+            block_state.guide_local_root,
+            block_state.guide_pose,
+            block_state.has_global_root,
+            block_state.has_local_root,
+            block_state.has_pose,
+            allowed,
+        )
+        block_state.motion_tokens = output.num_tokens
+        block_state.root = output.root
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class MotionBricksPoseTokensStep(ModularPipelineBlocks):
+    model_name = "motionbricks"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Pose step that places the guide poses on the frame timeline and predicts the pose codes of every token "
+            "in one pass of the masked token transformer."
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [ComponentSpec("root_model", AutoModel), ComponentSpec("pose_model", AutoModel)]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                "guide_local_root",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size, 8, 4)` normalized local root of the guide poses. Can be generated in the "
+                "guides step.",
+            ),
+            InputParam(
+                "guide_pose",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size, 8, 268)` normalized pose features of the guide poses. Can be generated in "
+                "the guides step.",
+            ),
+            InputParam(
+                "has_local_root",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size, 8)` guide poses with a local root. Can be generated in the guides step.",
+            ),
+            InputParam(
+                "has_pose",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size, 8)` guide poses with pose features. Can be generated in the guides step.",
+            ),
+            InputParam(
+                "motion_tokens",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size,)` motion lengths in tokens. Can be generated in the root step.",
+            ),
+            InputParam(
+                "root",
+                required=True,
+                type_hint=torch.Tensor,
+                description="`(batch_size, 64, 5)` normalized global root. Can be generated in the root step.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam("pose_codes", type_hint=torch.Tensor, description="`(batch_size, 16, 5)` pose codes."),
+            OutputParam("pose_timeline", type_hint=torch.Tensor, description="`(batch_size, 64, 268)` guide poses."),
+            OutputParam("has_pose_timeline", type_hint=torch.Tensor, description="`(batch_size, 64)` given poses."),
+            OutputParam(
+                "local_root", type_hint=torch.Tensor, description="`(batch_size, 64, 4)` normalized local root."
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        root_model, pose_model = components.root_model, components.pose_model
+        num_tokens, root = block_state.motion_tokens, block_state.root
+
+        local_root = root_to_local(root_model.denormalize(root.float(), "global_root"), num_tokens)
+        local_root = root_model.normalize(local_root.to(root.dtype), "local_root")
+        # A given local root at the last target pose replaces the derived one.
+        batch = torch.arange(len(num_tokens), device=num_tokens.device)
+        last = num_tokens * FRAMES_PER_TOKEN - 1
+        given = block_state.has_local_root[:, -1:]
+        local_root[batch, last] = torch.where(given, block_state.guide_local_root[:, -1], local_root[batch, last])
+
+        pose_timeline = root_model.place_guides(block_state.guide_pose, num_tokens, fill=0)
+        has_pose_timeline = root_model.place_guides(block_state.has_pose, num_tokens, fill=False)
+        ground_root = torch.cat([root[..., :1], root[..., 2:]], dim=-1)
+        logits = pose_model(pose_timeline, has_pose_timeline, ground_root, num_tokens).logits
+
+        block_state.pose_codes = logits.argmax(-1)
+        block_state.pose_timeline = pose_timeline
+        block_state.has_pose_timeline = has_pose_timeline
+        block_state.local_root = local_root
+        self.set_block_state(state, block_state)
+        return components, state

modular_blocks_motionbricks.py

@@ -0,0 +1,94 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0. See the LICENSE file in this repository.
+
+from dataclasses import replace
+
+from diffusers.modular_pipelines import ModularPipeline, SequentialPipelineBlocks
+from diffusers.modular_pipelines.modular_pipeline_utils import OutputParam
+from diffusers.utils.dynamic_modules_utils import get_class_from_dynamic_module
+
+from .decoders import MotionBricksDecodeStep, MotionBricksPostprocessStep
+from .encoders import MotionBricksGuidesStep
+from .generation import MotionBricksPoseTokensStep, MotionBricksRootStep
+
+# The remote-code loader only copies modules imported here directly, so the shared helpers are imported too.
+from .motion_utils import encode_guides  # noqa: F401
+
+
+class MotionBricksBlocks(SequentialPipelineBlocks):
+    """Motion in-betweening with NVIDIA MotionBricks on the SOMA skeleton.
+
+    Every component loads from `components_repo`, one subfolder per component, and so does its code, which requires
+    `trust_components_code=True`.
+    """
+
+    model_name = "motionbricks"
+    block_classes = [
+        MotionBricksGuidesStep(),
+        MotionBricksRootStep(),
+        MotionBricksPoseTokensStep(),
+        MotionBricksDecodeStep(),
+        MotionBricksPostprocessStep(),
+    ]
+    block_names = ["guides", "root", "pose_tokens", "decode", "postprocess"]
+
+    _requirements = {"diffusers": ">=0.41.0"}
+
+    def __init__(self, components_repo=None, components_revision=None, trust_components_code=False):
+        super().__init__()
+        self.components_repo = components_repo
+        self.components_revision = components_revision
+        self.component_types = {}
+        if components_repo is None:
+            return
+        if not trust_components_code:
+            raise ValueError("Set `trust_components_code=True` to load the model code from the components repository")
+        index = ModularPipeline.load_config(str(components_repo), revision=components_revision)
+        for spec in super().expected_components:
+            module_name, class_name, _ = index[spec.name]
+            self.component_types[spec.name] = get_class_from_dynamic_module(
+                str(components_repo),
+                subfolder=spec.name,
+                module_file=f"{module_name}.py",
+                class_name=class_name,
+                revision=components_revision,
+                trust_remote_code=True,
+            )
+
+    @property
+    def expected_components(self):
+        if self.components_repo is None:
+            return super().expected_components
+        return [
+            replace(
+                spec,
+                type_hint=self.component_types.get(spec.name, spec.type_hint),
+                pretrained_model_name_or_path=str(self.components_repo),
+                subfolder=spec.name,
+                revision=self.components_revision,
+            )
+            for spec in super().expected_components
+        ]
+
+    @property
+    def description(self) -> str:
+        return (
+            "Generates the transition from the last poses of a context motion to target poses, choosing its length "
+            "among the allowed ones."
+        )
+
+    @property
+    def outputs(self):
+        return [
+            OutputParam(
+                "motion",
+                type_hint=list,
+                description="One SOMA30 motion dict per batch item, 30 FPS, meters, +Y up.",
+            ),
+            OutputParam(
+                "placed_target",
+                type_hint=list,
+                description="`target` as used, one dict per batch item, moved by `target_offset` and `target_turn`.",
+            ),
+        ]

modular_config.json

@@ -0,0 +1,10 @@
+{
+  "_class_name": "MotionBricksBlocks",
+  "_diffusers_version": "0.41.0",
+  "auto_map": {
+    "ModularPipelineBlocks": "modular_blocks_motionbricks.MotionBricksBlocks"
+  },
+  "requirements": {
+    "diffusers": ">=0.41.0"
+  }
+}

modular_model_index.json

@@ -0,0 +1,4 @@
+{
+  "_class_name": "ModularPipeline",
+  "_blocks_class_name": "MotionBricksBlocks"
+}

motion_utils.py

@@ -0,0 +1,421 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0. See the LICENSE file in this repository.
+#
+# Feature encoding, decoding and endpoint correction adapted from NVIDIA MotionBricks
+# (https://github.com/NVlabs/motionbricks), Apache-2.0. See NOTICE.
+
+import functools
+import math
+
+import torch
+import torch.nn.functional as F
+
+
+FPS = 30.0
+CONTEXT_FRAMES = 4
+TARGET_FRAMES = 4
+BOUNDARY_FRAMES = 3
+FRAMES_PER_TOKEN = 4
+# Heading angle of +X, which the last context pose is turned to face.
+ALIGNED_HEADING = math.pi / 2
+
+# SOMA30 joint names in model order; the root is Hips.
+SOMA30_JOINTS = (
+    "Hips", "Spine1", "Spine2", "Chest", "Neck1", "Neck2", "Head", "Jaw", "LeftEye", "RightEye",
+    "LeftShoulder", "LeftArm", "LeftForeArm", "LeftHand", "LeftHandThumbEnd", "LeftHandMiddleEnd",
+    "RightShoulder", "RightArm", "RightForeArm", "RightHand", "RightHandThumbEnd", "RightHandMiddleEnd",
+    "LeftLeg", "LeftShin", "LeftFoot", "LeftToeBase", "RightLeg", "RightShin", "RightFoot", "RightToeBase",
+)  # fmt: skip
+SOMA30_PARENTS = (
+    -1, 0, 1, 2, 3, 4, 5, 6, 6, 6, 3, 10, 11, 12, 13, 13, 3, 16, 17, 18, 19, 19, 0, 22, 23, 24, 0, 26, 27, 28,
+)  # fmt: skip
+# SOMA77 index of each SOMA30 joint.
+SOMA77_TO_SOMA30 = (
+    0, 1, 2, 3, 4, 5, 6, 8, 9, 10, 11, 12, 13, 14, 18, 28, 39, 40, 41, 42, 46, 56, 67, 68, 69, 70, 72, 73, 74, 75,
+)  # fmt: skip
+RIGHT_LEG, LEFT_LEG = SOMA30_JOINTS.index("RightLeg"), SOMA30_JOINTS.index("LeftLeg")
+FOOT_CONTACT_JOINTS = ("LeftFoot", "LeftToeBase", "RightFoot", "RightToeBase")
+
+
+@functools.cache
+def soma77_index(device: torch.device) -> torch.Tensor:
+    return torch.tensor(SOMA77_TO_SOMA30, device=device)
+
+
+def on_ground(positions: torch.Tensor) -> torch.Tensor:
+    """`(..., 3)` positions with their height set to 0."""
+    return torch.stack([positions[..., 0], torch.zeros_like(positions[..., 1]), positions[..., 2]], dim=-1)
+
+
+@functools.cache
+def following_frames(num_frames: int, device: torch.device) -> torch.Tensor:
+    """Index of the frame after each guide frame; the last context frame and the last frame are their own."""
+    following = torch.arange(1, num_frames + 1).clamp(max=num_frames - 1)
+    following[CONTEXT_FRAMES - 1] = CONTEXT_FRAMES - 1
+    return following.to(device)
+
+
+def soma30(values: torch.Tensor, joint_dim: int) -> torch.Tensor:
+    """Select the SOMA30 joints of SOMA30 or SOMA77 values along `joint_dim`."""
+    count = values.shape[joint_dim]
+    if count == len(SOMA30_JOINTS):
+        return values
+    if count == 77:
+        return values.index_select(joint_dim, soma77_index(values.device))
+    raise ValueError(f"Expected SOMA30 or SOMA77 joints, got {count}")
+
+
+def yaw_matrix(angle: torch.Tensor) -> torch.Tensor:
+    """Rotation about +Y by `angle` radians, `(..., 3, 3)`."""
+    c, s = torch.cos(angle), torch.sin(angle)
+    zero, one = torch.zeros_like(angle), torch.ones_like(angle)
+    return torch.stack([c, zero, s, zero, one, zero, -s, zero, c], dim=-1).unflatten(-1, (3, 3))
+
+
+def heading_from_legs(right: torch.Tensor, left: torch.Tensor) -> torch.Tensor:
+    """Heading angle about +Y of the forward direction, from the right and left leg joint positions."""
+    across = right - left
+    length = torch.hypot(across[..., 0], across[..., 2])
+    angle = torch.atan2(across[..., 2], -across[..., 0])
+    return torch.where(length > 1e-6, angle, torch.zeros_like(angle))
+
+
+def place_target(
+    positions: torch.Tensor,
+    rotations: torch.Tensor,
+    anchor: torch.Tensor,
+    offset: torch.Tensor,
+    turn: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Move and turn a motion so its first pose stands relative to an anchor pose.
+
+    Args:
+        positions: `(batch_size, frames, joints, 3)` SOMA30 or SOMA77 joint positions of the motion.
+        rotations: `(batch_size, frames, joints, 3, 3)` global joint rotations of the motion.
+        anchor: `(batch_size, 30, 3)` SOMA30 joint positions of the anchor pose.
+        offset: `(batch_size, 2)` distance `(forward, left)` in meters from the anchor's ground point, along the
+            direction the anchor faces.
+        turn: `(batch_size,)` heading of the motion's first pose relative to the anchor's, in radians, positive to
+            the left.
+
+    Returns:
+        The moved `positions` and `rotations`.
+    """
+    first = soma30(positions[:, 0], 1)
+    anchor_angle = heading_from_legs(anchor[:, RIGHT_LEG], anchor[:, LEFT_LEG])
+    first_angle = heading_from_legs(first[:, RIGHT_LEG], first[:, LEFT_LEG])
+    zero = torch.zeros_like(anchor_angle)
+    forward = torch.stack([torch.sin(anchor_angle), zero, torch.cos(anchor_angle)], dim=-1)
+    left = torch.stack([torch.cos(anchor_angle), zero, -torch.sin(anchor_angle)], dim=-1)
+    destination = on_ground(anchor[:, 0]) + offset[:, :1] * forward + offset[:, 1:] * left
+    rotation = yaw_matrix(anchor_angle + turn - first_angle)[:, None, None]
+    origin = on_ground(first[:, 0])[:, None, None]
+    moved = (rotation @ (positions - origin)[..., None])[..., 0] + destination[:, None, None]
+    return moved, rotation @ rotations
+
+
+def rotation_6d(rotations: torch.Tensor) -> torch.Tensor:
+    """The first two columns of `(..., 3, 3)` rotation matrices, `(..., 6)`."""
+    return torch.cat([rotations[..., :, 0], rotations[..., :, 1]], dim=-1)
+
+
+def rotation_from_6d(features: torch.Tensor) -> torch.Tensor:
+    """Rotation matrices `(..., 3, 3)` from two `(..., 6)` column vectors, orthonormalized."""
+    x = F.normalize(features[..., :3], dim=-1)
+    z = F.normalize(torch.linalg.cross(x, features[..., 3:]), dim=-1)
+    y = torch.linalg.cross(z, x)
+    return torch.stack([x, y, z], dim=-1)
+
+
+def encode_guides(positions: torch.Tensor, rotations: torch.Tensor):
+    """Encode the context and target poses as the model's guide features.
+
+    Args:
+        positions: `(batch_size, 8, 30, 3)` SOMA30 joint positions, four context poses then four target poses, in
+            meters with +Y up.
+        rotations: `(batch_size, 8, 30, 3, 3)` SOMA30 global joint rotations, in the convention where the identity is
+            the neutral pose of the skeleton.
+
+    Returns:
+        `(global_root, local_root, pose, alignment)`: `(batch_size, 8, 5)` root position and heading cosine and sine,
+        `(batch_size, 8, 4)` heading rate, ground velocity and height, `(batch_size, 8, 267)` root-relative joint
+        positions and 6D rotations, and the `(origin, yaw)` that maps aligned poses back to the input space.
+    """
+    last = CONTEXT_FRAMES - 1
+    origin = on_ground(positions[:, last, 0])
+    current = heading_from_legs(positions[:, last, RIGHT_LEG], positions[:, last, LEFT_LEG])
+    yaw = current - ALIGNED_HEADING
+    inverse = yaw_matrix(yaw).transpose(-1, -2)[:, None, None]
+    aligned = (inverse @ (positions - origin[:, None, None])[..., None])[..., 0]
+    aligned_rotations = inverse @ rotations
+
+    num_frames = positions.shape[1]
+    following = following_frames(num_frames, positions.device)
+    hips, next_hips = aligned[:, :, 0], aligned[:, following, 0]
+    angle = heading_from_legs(aligned[:, :, RIGHT_LEG], aligned[:, :, LEFT_LEG])
+    next_angle = heading_from_legs(aligned[:, following, RIGHT_LEG], aligned[:, following, LEFT_LEG])
+    delta = torch.atan2(torch.sin(next_angle - angle), torch.cos(next_angle - angle))
+
+    global_root = torch.cat([hips, torch.cos(angle)[..., None], torch.sin(angle)[..., None]], dim=-1)
+    velocity = (next_hips - hips) * FPS
+    local_root = torch.stack([delta * FPS, velocity[..., 0], velocity[..., 2], hips[..., 1]], dim=-1)
+    ground = on_ground(hips)
+    joint_positions = (aligned[:, :, 1:] - ground[:, :, None]).flatten(2)
+    pose = torch.cat([joint_positions, rotation_6d(aligned_rotations).flatten(2)], dim=-1)
+    return global_root, local_root, pose, (origin, yaw)
+
+
+def guide_masks(num_targets: torch.Tensor):
+    """`(has_global_root, has_local_root, has_pose)` masks of the guide frames, `(batch_size, 8)` each.
+
+    The context poses and the first `num_targets` target poses are given. The local root, a derivative, is unknown
+    at the last context pose and the last target pose.
+    """
+    index = torch.arange(CONTEXT_FRAMES + TARGET_FRAMES, device=num_targets.device)[None]
+    given = index < CONTEXT_FRAMES + num_targets[:, None]
+    last = (index == CONTEXT_FRAMES - 1) | (index == CONTEXT_FRAMES + num_targets[:, None] - 1)
+    return given, given & ~last, given
+
+
+def root_to_local(root: torch.Tensor, num_tokens: torch.Tensor) -> torch.Tensor:
+    """Local root `(heading rate, x velocity, z velocity, height)` of a `(batch_size, num_frames, 5)` global root.
+
+    The last frame of the motion repeats the derivatives of the frame before it.
+    """
+    angle = torch.atan2(root[..., 4], root[..., 3])
+    cos, sin = torch.cos(angle), torch.sin(angle)
+    rate = torch.atan2(
+        sin[:, 1:] * cos[:, :-1] - cos[:, 1:] * sin[:, :-1], cos[:, 1:] * cos[:, :-1] + sin[:, 1:] * sin[:, :-1]
+    )
+    rate = F.pad(rate * FPS, (0, 1))
+    velocity = F.pad((root[:, 1:, :3] - root[:, :-1, :3]) * FPS, (0, 0, 0, 1))
+    local = torch.cat([rate[..., None], velocity[..., 0:3:2], root[..., 1:2]], dim=-1)
+    batch = torch.arange(root.shape[0], device=root.device)
+    last = num_tokens * FRAMES_PER_TOKEN - 1
+    local[batch, last] = local[batch, last - 1]
+    local[batch, last, 3] = root[batch, last, 1]
+    return local
+
+
+def integrate_local_root(local_root: torch.Tensor) -> torch.Tensor:
+    """Global root `(x, y, z, heading cosine, heading sine)` from a local root, starting at the origin."""
+    rate, vx, vz, height = local_root.unbind(-1)
+    angle = F.pad(torch.cumsum(rate[:, :-1] / FPS, dim=1), (1, 0))
+    x = F.pad(torch.cumsum(vx[:, :-1] / FPS, dim=1), (1, 0))
+    z = F.pad(torch.cumsum(vz[:, :-1] / FPS, dim=1), (1, 0))
+    return torch.stack([x, height, z, torch.cos(angle), torch.sin(angle)], dim=-1)
+
+
+def select_root(
+    decoded_root: torch.Tensor,
+    predicted_root: torch.Tensor,
+    target_height: torch.Tensor,
+    has_target: torch.Tensor,
+    num_tokens: torch.Tensor,
+) -> torch.Tensor:
+    """Keep the decoded root, unless its height misses the targets by over 5 cm and the root model's is closer."""
+    frames = (
+        num_tokens[:, None] * FRAMES_PER_TOKEN - TARGET_FRAMES + torch.arange(TARGET_FRAMES, device=num_tokens.device)
+    )
+    batch = torch.arange(decoded_root.shape[0], device=decoded_root.device)[:, None]
+    weight = has_target.to(decoded_root.dtype)
+    count = weight.sum(-1)
+    decoded_error = ((decoded_root[batch, frames, 1] - target_height).abs() * weight).sum(-1) / count
+    predicted_error = ((predicted_root[batch, frames, 1] - target_height).abs() * weight).sum(-1) / count
+    use_predicted = (predicted_error < decoded_error) & (decoded_error > 0.05)
+    return torch.where(use_predicted[:, None, None], predicted_root, decoded_root)
+
+
+def decode_poses(
+    motion: torch.Tensor,
+    num_tokens: torch.Tensor,
+    alignment: tuple[torch.Tensor, torch.Tensor],
+    neutral_joints: torch.Tensor,
+    endpoints: dict | None = None,
+):
+    """World-space SOMA30 poses of the valid frames of each motion.
+
+    Args:
+        motion: `(batch_size, num_frames, 366)` features: global root, child joint positions, 6D global rotations,
+            joint velocities and foot contacts, in the aligned space.
+        num_tokens: `(batch_size,)` motion lengths in tokens; the first and last three frames are dropped.
+        alignment: The `(origin, yaw)` from `encode_guides`.
+        neutral_joints: `(30, 3)` joint positions of the neutral pose, which give the bone offsets.
+        endpoints: The `endpoints` from `encode_guides`, to correct the endpoints with `correct_endpoints`.
+
+    Returns:
+        A list with one dict per motion: `posed_joints` `(frames, 30, 3)`, `global_rot_mats` `(frames, 30, 3, 3)`,
+        `local_rot_mats` `(frames, 30, 3, 3)`, `root_positions` `(frames, 3)` and `foot_contacts` `(frames, 4)`.
+    """
+    origin, yaw = alignment
+    rotation = yaw_matrix(yaw)
+    joints = len(SOMA30_JOINTS)
+    position_end = 5 + 3 * (joints - 1)
+    global_rots = rotation[:, None, None] @ rotation_from_6d(
+        motion[..., position_end : position_end + 6 * joints].unflatten(-1, (joints, 6))
+    )
+    roots = (rotation[:, None] @ motion[..., :3, None])[..., 0] + origin[:, None]
+    contacts = motion[..., -len(FOOT_CONTACT_JOINTS) :].clamp(0, 1)
+    local_rots = local_rotations(global_rots)
+    if endpoints is not None:
+        local_rots, roots = correct_endpoints(local_rots, roots, num_tokens, endpoints)
+        global_rots = global_rotations(local_rots)
+    positions = joint_positions(global_rots, roots, neutral_joints)
+    outputs = []
+    for i, count in enumerate((num_tokens * FRAMES_PER_TOKEN).tolist()):
+        frames = slice(BOUNDARY_FRAMES, count - BOUNDARY_FRAMES)
+        outputs.append(
+            {
+                "posed_joints": positions[i, frames],
+                "global_rot_mats": global_rots[i, frames],
+                "local_rot_mats": local_rots[i, frames],
+                "root_positions": roots[i, frames],
+                "foot_contacts": contacts[i, frames],
+            }
+        )
+    return outputs
+
+
+@functools.cache
+def skeleton(device: torch.device) -> tuple[torch.Tensor, list[torch.Tensor], torch.Tensor]:
+    """SOMA30 parent indices (the root its own parent), the non-root joints grouped by depth, and the `(30, 30)`
+    matrix whose row `j` marks the non-root joints on the path from `j` to the root."""
+    parents = torch.tensor(SOMA30_PARENTS).clamp(min=0)
+    depth = [0] * len(SOMA30_PARENTS)
+    ancestors = torch.zeros(len(SOMA30_PARENTS), len(SOMA30_PARENTS))
+    for joint in range(1, len(SOMA30_PARENTS)):
+        depth[joint] = depth[SOMA30_PARENTS[joint]] + 1
+        ancestors[joint] = ancestors[SOMA30_PARENTS[joint]]
+        ancestors[joint, joint] = 1
+    levels = [torch.tensor([j for j, d in enumerate(depth) if d == level]) for level in range(1, max(depth) + 1)]
+    return parents.to(device), [level.to(device) for level in levels], ancestors.to(device)
+
+
+def local_rotations(global_rots: torch.Tensor) -> torch.Tensor:
+    """Local rotations of `(..., 30, 3, 3)` global rotations; the root's local rotation is its global one."""
+    parents, _, _ = skeleton(global_rots.device)
+    local = global_rots[..., parents, :, :].transpose(-1, -2) @ global_rots
+    return torch.cat([global_rots[..., :1, :, :], local[..., 1:, :, :]], dim=-3)
+
+
+def joint_positions(global_rots: torch.Tensor, root_positions: torch.Tensor, neutral_joints: torch.Tensor):
+    """Joint positions from global rotations and a root trajectory, with the bone offsets of the neutral pose."""
+    parents, _, ancestors = skeleton(global_rots.device)
+    offsets = (neutral_joints - neutral_joints[parents]).to(global_rots.dtype)
+    bones = (global_rots[..., parents, :, :] @ offsets[:, :, None])[..., 0]
+    return root_positions[..., None, :] + torch.einsum("jk,...kc->...jc", ancestors.to(bones.dtype), bones)
+
+
+def global_rotations(local_rots: torch.Tensor) -> torch.Tensor:
+    """Global rotations of `(..., 30, 3, 3)` local rotations."""
+    parents, levels, _ = skeleton(local_rots.device)
+    global_rots = local_rots.clone()
+    for joints in levels:
+        global_rots[..., joints, :, :] = global_rots[..., parents[joints], :, :] @ local_rots[..., joints, :, :]
+    return global_rots
+
+
+def axis_angle(rotations: torch.Tensor) -> torch.Tensor:
+    """Axis-angle vectors of `(..., 3, 3)` rotation matrices."""
+    quaternion = matrix_to_quaternion(rotations)
+    quaternion = quaternion * torch.where(quaternion[..., :1] < 0, -1.0, 1.0)
+    vector = quaternion[..., 1:]
+    sin_half = vector.norm(dim=-1, keepdim=True)
+    angle = 2 * torch.atan2(sin_half, quaternion[..., :1])
+    scale = torch.where(sin_half > 1e-8, angle / sin_half.clamp(min=1e-8), torch.full_like(angle, 2.0))
+    return vector * scale
+
+
+def axis_angle_matrix(vector: torch.Tensor) -> torch.Tensor:
+    angle = vector.norm(dim=-1, keepdim=True)[..., None]
+    axis = F.normalize(vector, dim=-1, eps=1e-12)
+    x, y, z = axis.unbind(-1)
+    zero = torch.zeros_like(x)
+    skew = torch.stack([zero, -z, y, z, zero, -x, -y, x, zero], dim=-1).unflatten(-1, (3, 3))
+    eye = torch.eye(3, dtype=vector.dtype, device=vector.device).expand_as(skew)
+    return eye + torch.sin(angle) * skew + (1 - torch.cos(angle)) * (skew @ skew)
+
+
+def matrix_to_quaternion(rotations: torch.Tensor) -> torch.Tensor:
+    """`(w, x, y, z)` unit quaternions of `(..., 3, 3)` rotation matrices."""
+    m = rotations
+    trace = m[..., 0, 0] + m[..., 1, 1] + m[..., 2, 2]
+    candidates = torch.stack(
+        [
+            torch.stack(
+                [1 + trace, m[..., 2, 1] - m[..., 1, 2], m[..., 0, 2] - m[..., 2, 0], m[..., 1, 0] - m[..., 0, 1]], -1
+            ),
+            torch.stack(
+                [
+                    m[..., 2, 1] - m[..., 1, 2],
+                    1 + m[..., 0, 0] - m[..., 1, 1] - m[..., 2, 2],
+                    m[..., 0, 1] + m[..., 1, 0],
+                    m[..., 0, 2] + m[..., 2, 0],
+                ],
+                -1,
+            ),
+            torch.stack(
+                [
+                    m[..., 0, 2] - m[..., 2, 0],
+                    m[..., 0, 1] + m[..., 1, 0],
+                    1 - m[..., 0, 0] + m[..., 1, 1] - m[..., 2, 2],
+                    m[..., 1, 2] + m[..., 2, 1],
+                ],
+                -1,
+            ),
+            torch.stack(
+                [
+                    m[..., 1, 0] - m[..., 0, 1],
+                    m[..., 0, 2] + m[..., 2, 0],
+                    m[..., 1, 2] + m[..., 2, 1],
+                    1 - m[..., 0, 0] - m[..., 1, 1] + m[..., 2, 2],
+                ],
+                -1,
+            ),
+        ],
+        dim=-2,
+    )
+    best = torch.stack([trace, m[..., 0, 0], m[..., 1, 1], m[..., 2, 2]], -1).argmax(-1)
+    quaternion = torch.gather(candidates, -2, best[..., None, None].expand(*best.shape, 1, 4))[..., 0, :]
+    return F.normalize(quaternion, dim=-1)
+
+
+def scaled_rotation(rotations: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
+    """The rotations scaled by `weight` along their geodesic from the identity, as slerp(identity, q, weight)."""
+    return axis_angle_matrix(axis_angle(rotations) * weight[..., None])
+
+
+def correct_endpoints(
+    local_rots: torch.Tensor, roots: torch.Tensor, num_tokens: torch.Tensor, endpoints: dict
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Blend corrections over each motion so its first valid frame is the last context pose and its last valid frame
+    the first target pose.
+
+    Args:
+        local_rots: `(batch_size, num_frames, 30, 3, 3)` local rotations.
+        roots: `(batch_size, num_frames, 3)` root positions.
+        num_tokens: `(batch_size,)` motion lengths in tokens.
+        endpoints: `context_local` and `target_local` `(batch_size, 30, 3, 3)` local rotations of the last context
+            pose and the first target pose, `context_root` and `target_root` `(batch_size, 3)` their root positions.
+
+    Returns:
+        The corrected local rotations and root positions.
+    """
+    batch = torch.arange(len(local_rots), device=local_rots.device)
+    first = BOUNDARY_FRAMES
+    last = num_tokens * FRAMES_PER_TOKEN - BOUNDARY_FRAMES - 1
+    frames = torch.arange(local_rots.shape[1], dtype=local_rots.dtype, device=local_rots.device)
+    end_weight = ((frames - first) / (last - first).to(local_rots.dtype)[:, None]).clamp(0, 1)
+    start_weight = 1 - end_weight
+    ends = torch.stack([local_rots[batch, first], local_rots[batch, last]], dim=1)
+    deltas = ends.transpose(-1, -2) @ torch.stack([endpoints["context_local"], endpoints["target_local"]], dim=1)
+    weights = torch.stack([start_weight, end_weight], dim=1)[..., None]
+    start, end = scaled_rotation(deltas[:, :, None], weights).unbind(1)
+    local_rots = local_rots @ start @ end
+    roots = (
+        roots
+        + (endpoints["context_root"] - roots[batch, first])[:, None] * start_weight[..., None]
+        + (endpoints["target_root"] - roots[batch, last])[:, None] * end_weight[..., None]
+    )
+    return local_rots, roots