Upload 13 files

LICENSE

@@ -0,0 +1,59 @@
+NVIDIA Open Model Agreement
+
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+Release Date: April 2, 2026
+
+NVIDIA Works released under this License are intended to be used permissively and enable the further development of AI technologies. Subject to the terms of this License, NVIDIA confirms that:
+
+• Works are commercially usable.
+
+• You are free to create and distribute Derivative Works.
+
+• NVIDIA does not claim ownership to any outputs generated using the Works or Derivative Works.
+
+By using, reproducing, modifying, distributing, performing or displaying any portion or element of the Works or Derivative Works, or otherwise accepting the terms of this License, you agree to be bound by this License.
+
+1. Definitions.
+
+“License” shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 10 of this document.
+
+“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.
+
+“Work” shall mean the work of authorship, including machine learning model, software, checkpoints, learnt weights, algorithms, parameters, configuration files and documentation, 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.
+
+2. Grant of License. Subject to the terms and conditions of this License, NVIDIA hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable 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.
+
+If You institute patent or copyright litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or an output from the Work constitutes direct or contributory patent or copyright infringement, then any licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed.
+
+3. 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 a copy of this License; and
+
+   b. 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
+
+   c. 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 following attribution notice within a “Notice” text file with such copies and the following statement: “Licensed by NVIDIA Corporation under the NVIDIA Open Model Agreement.”
+
+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.
+
+4. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of NVIDIA, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file.
+
+5. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, NVIDIA provides the Work 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.
+
+6. 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 NVIDIA 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, a Derivative Work or an output from the Work or Derivative 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 NVIDIA has been advised of the possibility of such damages.
+
+7. 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 NVIDIA, and only if You agree to indemnify, defend, and hold NVIDIA harmless for any liability incurred by, or claims asserted against, NVIDIA by reason of your accepting any such warranty or additional liability. You will indemnify and hold harmless NVIDIA from and against any claim by any third party arising out of or related to your use or distribution of the Works, Derivative Works thereof, or output from the Works or Derivative Works.
+
+8. Feedback. NVIDIA appreciates your feedback, and You agree that NVIDIA may use it without restriction or compensation to You.
+
+9. Governing Law. This Agreement will be governed in all respects by the laws of the United States and the laws of the State of Delaware, without regard to conflict of laws principles or the United Nations Convention on Contracts for the International Sale of Goods. The state and federal courts residing in Santa Clara County, California will have exclusive jurisdiction over any dispute or claim arising out of or related to this Agreement, and the parties irrevocably consent to personal jurisdiction and venue in those courts; except that, either party may apply for injunctive remedies or an equivalent type of urgent legal relief in any jurisdiction.
+
+10. Trade and Compliance. You agree to comply with all applicable export, import, trade and economic sanctions laws and regulations, as amended, including without limitation U.S. Export Administration Regulations and Office of Foreign Assets Control regulations. These laws include restrictions on destinations, end-users and end-use.
+
+(v. March 9, 2026)
+
+END OF TERMS AND CONDITIONS

NOTICE

@@ -0,0 +1,14 @@
+Licensed by NVIDIA Corporation under the NVIDIA Open Model Agreement.
+Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+The agreement is included as LICENSE and published at
+https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-agreement/
+
+This repository repacks the SOMA checkpoints of nvidia/MotionBricks (soma/soma-v6-nova-4n8g-ea43dd8-fp32-v1) for
+Diffusers. The float32 tensors of the root model, the pose model and the decoder of the pose VQ-VAE are unchanged,
+renamed for the Diffusers model code next to them; the fused attention input projections were split into query, key
+and value. The VQ-VAE encoder and the training-only codebook statistics were left out. The normalization statistics and
+the neutral joint positions of the skeleton come from the same package and are stored as buffers, with the standard
+deviations as sqrt(std^2 + 1e-5), as in NVIDIA's exported ONNX model.
+
+The model code is new, Apache-2.0, and written by The HuggingFace Team.

README.md

@@ -0,0 +1,71 @@
+---
+language:
+- en
+license: other
+license_name: nvidia-open-model-agreement
+license_link: LICENSE
+base_model:
+- nvidia/MotionBricks
+library_name: diffusers
+tags:
+- motionbricks
+- modular-diffusers
+- motion-generation
+- motion-inbetweening
+- soma
+---
+
+# MotionBricks SOMA
+
+<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.*
+
+This is the SOMA model of [nvidia/MotionBricks](https://huggingface.co/nvidia/MotionBricks), repacked into the Diffusers format in float32. It runs as a Modular Diffusers pipeline with the [MotionBricks SOMA custom blocks](https://huggingface.co/OzzyGT/motionbricks_soma_custom_blocks) and generates the transition from the last poses of one motion to target poses in one forward pass, choosing the length of the transition itself.
+
+The three components work together:
+
+- `root_model` (40M parameters) chooses the length and predicts the root trajectory.
+- `pose_model` (137M parameters) predicts the pose codes of every token of four frames.
+- `pose_decoder` (12M parameters) decodes the codes to frames.
+
+Note: You need `diffusers>=0.41.0`.
+
+## Usage
+
+The custom blocks load the model code shipped in this repository, hence `trust_remote_code` and `trust_components_code`.
+
+**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.
+
+```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()})
+```
+
+Both clips come from [Kimodo SOMA Modular](https://huggingface.co/OzzyGT/Kimodo_SOMA_Modular): the start of a run, and a standing pose that `target_offset` places 1.4 m ahead of it. The model accepts SOMA30 or SOMA77 poses at 30 FPS and returns a SOMA30 motion; here it brings the run to a stop in 30 frames. See the [custom blocks](https://huggingface.co/OzzyGT/motionbricks_soma_custom_blocks) for the options.
+
+Float32 reproduces NVIDIA's ONNX graph with PyTorch's default settings; float16 halves the memory with small drift, while bfloat16 changes the motion noticeably.
+
+## Licenses
+
+The model weights are licensed under the [NVIDIA Open Model Agreement](LICENSE), the same as the original. The model code in this repository is Apache-2.0. See [NOTICE](NOTICE) for what was changed.

modular_model_index.json

@@ -0,0 +1,46 @@
+{
+  "_class_name": "ModularPipeline",
+  "_blocks_class_name": "MotionBricksBlocks",
+  "root_model": [
+    "root_motionbricks",
+    "MotionBricksRootModel",
+    {
+      "pretrained_model_name_or_path": null,
+      "revision": null,
+      "subfolder": "root_model",
+      "type_hint": [
+        "root_motionbricks",
+        "MotionBricksRootModel"
+      ],
+      "variant": null
+    }
+  ],
+  "pose_model": [
+    "transformer_motionbricks",
+    "MotionBricksPoseTransformer",
+    {
+      "pretrained_model_name_or_path": null,
+      "revision": null,
+      "subfolder": "pose_model",
+      "type_hint": [
+        "transformer_motionbricks",
+        "MotionBricksPoseTransformer"
+      ],
+      "variant": null
+    }
+  ],
+  "pose_decoder": [
+    "decoder_motionbricks",
+    "MotionBricksPoseDecoder",
+    {
+      "pretrained_model_name_or_path": null,
+      "revision": null,
+      "subfolder": "pose_decoder",
+      "type_hint": [
+        "decoder_motionbricks",
+        "MotionBricksPoseDecoder"
+      ],
+      "variant": null
+    }
+  ]
+}

pose_decoder/config.json

@@ -0,0 +1,19 @@
+{
+  "_class_name": "MotionBricksPoseDecoder",
+  "_diffusers_version": "0.41.0",
+  "auto_map": {
+    "AutoModel": "decoder_motionbricks.MotionBricksPoseDecoder"
+  },
+  "code_embed_dim": 32,
+  "codebook_size": 16,
+  "depth": 4,
+  "dilation_growth_rate": 3,
+  "feature_dim": 365,
+  "frames_per_token": 4,
+  "max_tokens": 16,
+  "num_heads_codebook": 5,
+  "num_joints": 30,
+  "pose_dim": 268,
+  "velocity_dim": 2,
+  "width": 512
+}

pose_decoder/decoder_motionbricks.py

@@ -0,0 +1,230 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# 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.
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import BaseOutput
+
+
+class MotionBricksConv1d(nn.Conv1d):
+    """Stride-1 `nn.Conv1d` computed as one matrix product over the shifted inputs, faster than cuDNN at these
+    sizes."""
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        kernel, dilation, padding = self.kernel_size[0], self.dilation[0], self.padding[0]
+        length = hidden_states.shape[-1]
+        if kernel > 1:
+            padded = F.pad(hidden_states, (padding, padding))
+            shifted = [padded[..., i * dilation : i * dilation + length] for i in range(kernel)]
+            hidden_states = torch.stack(shifted, dim=2).flatten(1, 2)
+        weight = self.weight.reshape(self.out_channels, -1)
+        return torch.matmul(weight, hidden_states) + self.bias[:, None]
+
+
+class MotionBricksResnetBlock(nn.Module):
+    def __init__(self, channels: int, dilation: int):
+        super().__init__()
+        self.conv1 = MotionBricksConv1d(channels, channels, 3, padding=dilation, dilation=dilation)
+        self.conv2 = MotionBricksConv1d(channels, channels, 1)
+
+    def forward(self, hidden_states: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        residual = hidden_states * mask
+        hidden_states = self.conv1(F.relu(residual) * mask)
+        hidden_states = self.conv2(F.relu(hidden_states) * mask)
+        return (hidden_states + residual) * mask
+
+
+class MotionBricksUpBlock(nn.Module):
+    def __init__(self, channels: int, depth: int, dilation_growth_rate: int):
+        super().__init__()
+        dilations = [dilation_growth_rate**i for i in reversed(range(depth))]
+        self.resnets = nn.ModuleList([MotionBricksResnetBlock(channels, dilation) for dilation in dilations])
+        self.conv = MotionBricksConv1d(channels, channels, 3, padding=1)
+
+
+class MotionBricksConditionedDecoder(nn.Module):
+    """Dilated 1D convolution decoder that upsamples tokens to frames, conditioned on per-frame targets and
+    external features."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        target_channels: int,
+        external_channels: int,
+        width: int = 512,
+        depth: int = 4,
+        dilation_growth_rate: int = 3,
+        num_tokens: int = 16,
+        num_frames: int = 64,
+    ):
+        super().__init__()
+        self.width = width
+        self.num_frames = num_frames
+        lengths = [num_tokens, num_tokens * 2]
+        self.conv_in = MotionBricksConv1d(in_channels, width, 3, padding=1)
+        self.target_projections = nn.ModuleList(
+            [nn.Linear(target_channels, width * length // num_frames) for length in lengths]
+        )
+        self.external_projections = nn.ModuleList(
+            [nn.Linear(width + external_channels * num_frames // length, width) for length in lengths]
+        )
+        self.up_blocks = nn.ModuleList([MotionBricksUpBlock(width, depth, dilation_growth_rate) for _ in lengths])
+        self.conv_mid = MotionBricksConv1d(width, width, 3, padding=1)
+        self.conv_out = MotionBricksConv1d(width, out_channels, 3, padding=1)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        token_mask: torch.Tensor,
+        target: torch.Tensor,
+        target_mask: torch.Tensor,
+        external: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Args:
+            hidden_states: `(batch_size, in_channels, num_tokens)` token features.
+            token_mask: `(batch_size, num_tokens)` float mask of the valid tokens.
+            target: `(batch_size, num_frames, target_channels)` per-frame targets.
+            target_mask: `(batch_size, num_frames)` boolean mask of the frames with a target.
+            external: `(batch_size, num_frames, external_channels)` per-frame conditioning.
+
+        Returns:
+            `(batch_size, out_channels, num_frames)` frame features.
+        """
+        batch_size = hidden_states.shape[0]
+        mask = token_mask[:, None]
+        hidden_states = F.relu(self.conv_in(hidden_states * mask))
+        for block, target_projection, external_projection in zip(
+            self.up_blocks, self.target_projections, self.external_projections
+        ):
+            length = hidden_states.shape[-1]
+            projected = F.relu(target_projection(target))
+            grouped = hidden_states.transpose(1, 2).reshape(batch_size, self.num_frames, -1)
+            grouped = torch.where(target_mask[..., None], projected, grouped)
+            grouped = grouped.reshape(batch_size, length, self.width)
+            conditioned = torch.cat([grouped, external.reshape(batch_size, length, -1)], dim=-1)
+            hidden_states = F.relu(external_projection(conditioned).transpose(1, 2))
+            for resnet in block.resnets:
+                hidden_states = resnet(hidden_states, mask)
+            hidden_states = F.interpolate(hidden_states, scale_factor=2, mode="nearest")
+            mask = mask.repeat_interleave(2, dim=-1)
+            hidden_states = block.conv(hidden_states * mask)
+        hidden_states = F.relu(self.conv_mid(hidden_states * mask)) * mask
+        return self.conv_out(hidden_states) * mask
+
+
+@dataclass
+class MotionBricksPoseDecoderOutput(BaseOutput):
+    """
+    Args:
+        features (`torch.Tensor` of shape `(batch_size, num_frames, feature_dim)`):
+            The normalized features of every frame.
+    """
+
+    features: torch.Tensor
+
+
+class MotionBricksPoseDecoder(ModelMixin, ConfigMixin):
+    """Decoder of the pose VQ-VAE: looks up the codes of every token and decodes them to normalized frame features,
+    conditioned on the given poses and the root velocity.
+
+    The features of a frame are the local root (heading rate, ground velocity, height) followed by the body: joint
+    positions relative to the root, 6D global rotations, joint velocities and foot contacts. Their normalization
+    statistics and the neutral joint positions of the skeleton are stored as buffers.
+    """
+
+    _no_split_modules = ["MotionBricksConditionedDecoder"]
+    _repeated_blocks = ["MotionBricksResnetBlock"]
+
+    @register_to_config
+    def __init__(
+        self,
+        num_heads_codebook: int = 5,
+        codebook_size: int = 16,
+        code_embed_dim: int = 32,
+        pose_dim: int = 268,
+        velocity_dim: int = 2,
+        feature_dim: int = 365,
+        num_joints: int = 30,
+        width: int = 512,
+        depth: int = 4,
+        dilation_growth_rate: int = 3,
+        frames_per_token: int = 4,
+        max_tokens: int = 16,
+    ):
+        super().__init__()
+        self.codebook = nn.Parameter(torch.zeros(num_heads_codebook, codebook_size, code_embed_dim))
+        self.decoder = MotionBricksConditionedDecoder(
+            num_heads_codebook * code_embed_dim,
+            feature_dim,
+            target_channels=pose_dim,
+            external_channels=velocity_dim,
+            width=width,
+            depth=depth,
+            dilation_growth_rate=dilation_growth_rate,
+            num_tokens=max_tokens,
+            num_frames=max_tokens * frames_per_token,
+        )
+        self.register_buffer("feature_mean", torch.zeros(feature_dim))
+        self.register_buffer("feature_std", torch.ones(feature_dim))
+        self.register_buffer("neutral_joints", torch.zeros(num_joints, 3))
+
+    def denormalize(self, features: torch.Tensor) -> torch.Tensor:
+        # The buffers can sit on another device while the model is offloaded.
+        std = self.feature_std.to(features.device, features.dtype)
+        return features * std + self.feature_mean.to(features.device, features.dtype)
+
+    def forward(
+        self,
+        codes: torch.LongTensor,
+        num_tokens: torch.Tensor,
+        pose: torch.Tensor,
+        has_pose: torch.Tensor,
+        velocity: torch.Tensor,
+        return_dict: bool = True,
+    ) -> MotionBricksPoseDecoderOutput | tuple:
+        """
+        Args:
+            codes (`torch.LongTensor` of shape `(batch_size, max_tokens, num_heads_codebook)`):
+                The pose code of every token and head.
+            num_tokens (`torch.LongTensor` of shape `(batch_size,)`):
+                The motion length in tokens.
+            pose (`torch.Tensor` of shape `(batch_size, num_frames, pose_dim)`):
+                Normalized pose features on the frame timeline.
+            has_pose (`torch.BoolTensor` of shape `(batch_size, num_frames)`):
+                The frames whose pose is given.
+            velocity (`torch.Tensor` of shape `(batch_size, num_frames, velocity_dim)`):
+                Normalized ground velocity of the root.
+            return_dict (`bool`, defaults to `True`):
+                Whether to return a `MotionBricksPoseDecoderOutput` instead of a tuple.
+
+        Returns:
+            `MotionBricksPoseDecoderOutput` or `tuple`: The normalized frame features; a one-element tuple when
+            `return_dict` is `False`.
+        """
+        heads = torch.arange(self.config.num_heads_codebook, device=codes.device)
+        embeds = self.codebook[heads, codes].flatten(2)
+        token_mask = (torch.arange(codes.shape[1], device=codes.device)[None] < num_tokens[:, None]).to(embeds.dtype)
+        features = self.decoder(embeds.transpose(1, 2), token_mask, pose, has_pose, velocity).transpose(1, 2)
+
+        if not return_dict:
+            return (features,)
+        return MotionBricksPoseDecoderOutput(features=features)

pose_decoder/diffusion_pytorch_model.safetensors

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

pose_model/config.json

@@ -0,0 +1,22 @@
+{
+  "_class_name": "MotionBricksPoseTransformer",
+  "_diffusers_version": "0.41.0",
+  "auto_map": {
+    "AutoModel": "transformer_motionbricks.MotionBricksPoseTransformer"
+  },
+  "code_embed_dim": 32,
+  "codebook_size": 16,
+  "ff_dim": 2048,
+  "frames_per_token": 4,
+  "hidden_dim": 1024,
+  "length_embed_dim": 128,
+  "max_tokens": 16,
+  "min_tokens": 6,
+  "num_heads": 16,
+  "num_heads_codebook": 5,
+  "num_layers": 16,
+  "pose_dim": 268,
+  "root_dim": 4,
+  "root_embed_dim": 256,
+  "token_dim": 640
+}

pose_model/diffusion_pytorch_model.safetensors

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

pose_model/transformer_motionbricks.py

@@ -0,0 +1,215 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# 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.
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.attention import AttentionMixin, AttentionModuleMixin
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import BaseOutput
+
+
+class MotionBricksMLP(nn.Module):
+    """Two linear layers, each followed by LeakyReLU(0.01), then a linear projection."""
+
+    def __init__(self, in_features: int, hidden_features: int, out_features: int):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [nn.Linear(in_features, hidden_features), nn.Linear(hidden_features, hidden_features)]
+        )
+        self.proj_out = nn.Linear(hidden_features, out_features)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for layer in self.layers:
+            hidden_states = F.leaky_relu(layer(hidden_states), 0.01)
+        return self.proj_out(hidden_states)
+
+
+class MotionBricksAttnProcessor:
+    _attention_backend = None
+    _parallel_config = None
+
+    def __call__(
+        self,
+        attn: "MotionBricksAttention",
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = hidden_states.shape
+        shape = (batch_size, seq_len, attn.heads, -1)
+        query = attn.to_q(hidden_states).view(shape)
+        key = attn.to_k(hidden_states).view(shape)
+        value = attn.to_v(hidden_states).view(shape)
+        hidden_states = dispatch_attention_fn(
+            query,
+            key,
+            value,
+            attn_mask=attention_mask,
+            backend=self._attention_backend,
+            parallel_config=self._parallel_config,
+        )
+        hidden_states = hidden_states.flatten(2, 3).to(query.dtype)
+        return attn.to_out[0](hidden_states)
+
+
+class MotionBricksAttention(nn.Module, AttentionModuleMixin):
+    """Multi-head self-attention. `attention_mask` is a boolean `(batch_size, 1, 1, seq_len)` mask of the keys to
+    attend to."""
+
+    _default_processor_cls = MotionBricksAttnProcessor
+    _available_processors = [MotionBricksAttnProcessor]
+
+    def __init__(self, dim: int, heads: int):
+        super().__init__()
+        self.heads = heads
+        self.to_q = nn.Linear(dim, dim)
+        self.to_k = nn.Linear(dim, dim)
+        self.to_v = nn.Linear(dim, dim)
+        self.to_out = nn.ModuleList([nn.Linear(dim, dim), nn.Dropout(0.0)])
+        self.set_processor(MotionBricksAttnProcessor())
+
+    def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
+        return self.processor(self, hidden_states, attention_mask)
+
+
+class MotionBricksTransformerLayer(nn.Module):
+    """Post-norm transformer encoder layer with a ReLU feed-forward."""
+
+    def __init__(self, dim: int, heads: int, ff_dim: int):
+        super().__init__()
+        self.attn = MotionBricksAttention(dim, heads)
+        self.norm1 = nn.LayerNorm(dim)
+        self.ff_in = nn.Linear(dim, ff_dim)
+        self.ff_out = nn.Linear(ff_dim, dim)
+        self.norm2 = nn.LayerNorm(dim)
+
+    def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
+        hidden_states = self.norm1(hidden_states + self.attn(hidden_states, attention_mask))
+        return self.norm2(hidden_states + self.ff_out(F.relu(self.ff_in(hidden_states))))
+
+
+def key_mask(num_tokens: torch.Tensor, max_tokens: int, prefix: int = 0) -> torch.Tensor:
+    """Boolean `(batch_size, 1, 1, prefix + max_tokens)` attention mask: the `prefix` positions, then the first
+    `num_tokens` of `max_tokens` token positions."""
+    positions = torch.arange(max_tokens, device=num_tokens.device)
+    valid = positions[None] < num_tokens[:, None]
+    valid = F.pad(valid, (prefix, 0), value=True)
+    return valid[:, None, None]
+
+
+@dataclass
+class MotionBricksPoseTransformerOutput(BaseOutput):
+    """
+    Args:
+        logits (`torch.Tensor` of shape `(batch_size, max_tokens, num_heads_codebook, codebook_size)`):
+            The code logits of every token and head.
+    """
+
+    logits: torch.Tensor
+
+
+class MotionBricksPoseTransformer(ModelMixin, ConfigMixin, AttentionMixin):
+    """Masked token transformer that predicts the pose codes of every token in one pass, from the given poses, the
+    root of every frame and the length of the motion. Each token covers `frames_per_token` frames and has one code
+    per head."""
+
+    _no_split_modules = ["MotionBricksTransformerLayer"]
+    _repeated_blocks = ["MotionBricksTransformerLayer"]
+    _skip_layerwise_casting_patterns = ["norm", "embed"]
+
+    @register_to_config
+    def __init__(
+        self,
+        pose_dim: int = 268,
+        root_dim: int = 4,
+        num_heads_codebook: int = 5,
+        codebook_size: int = 16,
+        code_embed_dim: int = 32,
+        token_dim: int = 640,
+        root_embed_dim: int = 256,
+        length_embed_dim: int = 128,
+        hidden_dim: int = 1024,
+        num_heads: int = 16,
+        ff_dim: int = 2048,
+        num_layers: int = 16,
+        frames_per_token: int = 4,
+        min_tokens: int = 6,
+        max_tokens: int = 16,
+    ):
+        super().__init__()
+        code_dim = num_heads_codebook * code_embed_dim
+        # Each head has codebook_size codes and a mask code.
+        self.code_embed = nn.Embedding(num_heads_codebook * (codebook_size + 1), code_embed_dim)
+        self.code_proj = MotionBricksMLP(code_dim, token_dim, token_dim)
+        self.pose_proj = nn.Linear(pose_dim, token_dim // frames_per_token)
+        self.root_proj = nn.Linear(root_dim * frames_per_token, root_embed_dim)
+        self.length_embed = nn.Embedding(max_tokens - min_tokens + 1, length_embed_dim)
+        self.proj_in = nn.Linear(token_dim + root_embed_dim + length_embed_dim, hidden_dim)
+        self.position_embed = nn.Parameter(torch.zeros(max_tokens, hidden_dim))
+        self.layers = nn.ModuleList(
+            [MotionBricksTransformerLayer(hidden_dim, num_heads, ff_dim) for _ in range(num_layers)]
+        )
+        self.proj_out = nn.Linear(hidden_dim, num_heads_codebook * codebook_size)
+
+    def forward(
+        self,
+        pose: torch.Tensor,
+        has_pose: torch.Tensor,
+        root: torch.Tensor,
+        num_tokens: torch.Tensor,
+        return_dict: bool = True,
+    ) -> MotionBricksPoseTransformerOutput | tuple:
+        """
+        Args:
+            pose (`torch.Tensor` of shape `(batch_size, num_frames, pose_dim)`):
+                Normalized pose features on the frame timeline.
+            has_pose (`torch.BoolTensor` of shape `(batch_size, num_frames)`):
+                The frames whose pose is given.
+            root (`torch.Tensor` of shape `(batch_size, num_frames, root_dim)`):
+                Normalized root of every frame: ground position and the cosine and sine of the heading.
+            num_tokens (`torch.LongTensor` of shape `(batch_size,)`):
+                The motion length in tokens.
+            return_dict (`bool`, defaults to `True`):
+                Whether to return a `MotionBricksPoseTransformerOutput` instead of a tuple.
+
+        Returns:
+            `MotionBricksPoseTransformerOutput` or `tuple`: The code logits; a one-element tuple when `return_dict` is
+            `False`.
+        """
+        config = self.config
+        batch_size = pose.shape[0]
+        heads, size = config.num_heads_codebook, config.codebook_size
+        mask_codes = size + torch.arange(heads, device=pose.device) * (size + 1)
+        codes = self.code_embed(mask_codes).flatten()
+        codes = self.code_proj(codes).expand(batch_size, config.max_tokens, -1)
+        codes = codes.reshape(batch_size, -1, self.pose_proj.out_features)
+        has_pose = has_pose[..., None].to(pose.dtype)
+        frames = self.pose_proj(pose) * has_pose + codes * (1 - has_pose)
+        frames = frames.reshape(batch_size, config.max_tokens, -1)
+        root = self.root_proj(root.reshape(batch_size, config.max_tokens, -1))
+        length = self.length_embed(num_tokens - config.min_tokens)[:, None].expand(-1, config.max_tokens, -1)
+        hidden_states = F.relu(self.proj_in(torch.cat([frames, root, length], dim=-1))) + self.position_embed
+        mask = key_mask(num_tokens, config.max_tokens)
+        for layer in self.layers:
+            hidden_states = layer(hidden_states, mask)
+        logits = self.proj_out(hidden_states).unflatten(-1, (heads, size))
+
+        if not return_dict:
+            return (logits,)
+        return MotionBricksPoseTransformerOutput(logits=logits)

root_model/config.json

@@ -0,0 +1,24 @@
+{
+  "_class_name": "MotionBricksRootModel",
+  "_diffusers_version": "0.41.0",
+  "auto_map": {
+    "AutoModel": "root_motionbricks.MotionBricksRootModel"
+  },
+  "context_frames": 4,
+  "decoder_depth": 4,
+  "decoder_dilation_growth_rate": 3,
+  "ff_dim": 2048,
+  "frames_per_token": 4,
+  "global_root_dim": 5,
+  "hidden_dim": 512,
+  "local_root_dim": 4,
+  "max_tokens": 16,
+  "min_tokens": 6,
+  "num_heads": 16,
+  "num_shared_layers": 6,
+  "num_token_layers": 2,
+  "pose_dim": 268,
+  "pose_embed_dim": 256,
+  "root_embed_dim": 64,
+  "target_frames": 4
+}

root_model/diffusion_pytorch_model.safetensors

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

root_model/root_motionbricks.py

@@ -0,0 +1,430 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# 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.
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.attention import AttentionMixin, AttentionModuleMixin
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import BaseOutput
+
+
+class MotionBricksMLP(nn.Module):
+    """Two linear layers, each followed by LeakyReLU(0.01), then a linear projection."""
+
+    def __init__(self, in_features: int, hidden_features: int, out_features: int):
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [nn.Linear(in_features, hidden_features), nn.Linear(hidden_features, hidden_features)]
+        )
+        self.proj_out = nn.Linear(hidden_features, out_features)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for layer in self.layers:
+            hidden_states = F.leaky_relu(layer(hidden_states), 0.01)
+        return self.proj_out(hidden_states)
+
+
+class MotionBricksAttnProcessor:
+    _attention_backend = None
+    _parallel_config = None
+
+    def __call__(
+        self,
+        attn: "MotionBricksAttention",
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = hidden_states.shape
+        shape = (batch_size, seq_len, attn.heads, -1)
+        query = attn.to_q(hidden_states).view(shape)
+        key = attn.to_k(hidden_states).view(shape)
+        value = attn.to_v(hidden_states).view(shape)
+        hidden_states = dispatch_attention_fn(
+            query,
+            key,
+            value,
+            attn_mask=attention_mask,
+            backend=self._attention_backend,
+            parallel_config=self._parallel_config,
+        )
+        hidden_states = hidden_states.flatten(2, 3).to(query.dtype)
+        return attn.to_out[0](hidden_states)
+
+
+class MotionBricksAttention(nn.Module, AttentionModuleMixin):
+    """Multi-head self-attention. `attention_mask` is a boolean `(batch_size, 1, 1, seq_len)` mask of the keys to
+    attend to."""
+
+    _default_processor_cls = MotionBricksAttnProcessor
+    _available_processors = [MotionBricksAttnProcessor]
+
+    def __init__(self, dim: int, heads: int):
+        super().__init__()
+        self.heads = heads
+        self.to_q = nn.Linear(dim, dim)
+        self.to_k = nn.Linear(dim, dim)
+        self.to_v = nn.Linear(dim, dim)
+        self.to_out = nn.ModuleList([nn.Linear(dim, dim), nn.Dropout(0.0)])
+        self.set_processor(MotionBricksAttnProcessor())
+
+    def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
+        return self.processor(self, hidden_states, attention_mask)
+
+
+class MotionBricksTransformerLayer(nn.Module):
+    """Post-norm transformer encoder layer with a ReLU feed-forward."""
+
+    def __init__(self, dim: int, heads: int, ff_dim: int):
+        super().__init__()
+        self.attn = MotionBricksAttention(dim, heads)
+        self.norm1 = nn.LayerNorm(dim)
+        self.ff_in = nn.Linear(dim, ff_dim)
+        self.ff_out = nn.Linear(ff_dim, dim)
+        self.norm2 = nn.LayerNorm(dim)
+
+    def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
+        hidden_states = self.norm1(hidden_states + self.attn(hidden_states, attention_mask))
+        return self.norm2(hidden_states + self.ff_out(F.relu(self.ff_in(hidden_states))))
+
+
+def key_mask(num_tokens: torch.Tensor, max_tokens: int, prefix: int = 0) -> torch.Tensor:
+    """Boolean `(batch_size, 1, 1, prefix + max_tokens)` attention mask: the `prefix` positions, then the first
+    `num_tokens` of `max_tokens` token positions."""
+    positions = torch.arange(max_tokens, device=num_tokens.device)
+    valid = positions[None] < num_tokens[:, None]
+    valid = F.pad(valid, (prefix, 0), value=True)
+    return valid[:, None, None]
+
+
+class MotionBricksConv1d(nn.Conv1d):
+    """Stride-1 `nn.Conv1d` computed as one matrix product over the shifted inputs, faster than cuDNN at these
+    sizes."""
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        kernel, dilation, padding = self.kernel_size[0], self.dilation[0], self.padding[0]
+        length = hidden_states.shape[-1]
+        if kernel > 1:
+            padded = F.pad(hidden_states, (padding, padding))
+            shifted = [padded[..., i * dilation : i * dilation + length] for i in range(kernel)]
+            hidden_states = torch.stack(shifted, dim=2).flatten(1, 2)
+        weight = self.weight.reshape(self.out_channels, -1)
+        return torch.matmul(weight, hidden_states) + self.bias[:, None]
+
+
+class MotionBricksResnetBlock(nn.Module):
+    def __init__(self, channels: int, dilation: int):
+        super().__init__()
+        self.conv1 = MotionBricksConv1d(channels, channels, 3, padding=dilation, dilation=dilation)
+        self.conv2 = MotionBricksConv1d(channels, channels, 1)
+
+    def forward(self, hidden_states: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        residual = hidden_states * mask
+        hidden_states = self.conv1(F.relu(residual) * mask)
+        hidden_states = self.conv2(F.relu(hidden_states) * mask)
+        return (hidden_states + residual) * mask
+
+
+class MotionBricksUpBlock(nn.Module):
+    def __init__(self, channels: int, depth: int, dilation_growth_rate: int):
+        super().__init__()
+        dilations = [dilation_growth_rate**i for i in reversed(range(depth))]
+        self.resnets = nn.ModuleList([MotionBricksResnetBlock(channels, dilation) for dilation in dilations])
+        self.conv = MotionBricksConv1d(channels, channels, 3, padding=1)
+
+
+class MotionBricksConditionedDecoder(nn.Module):
+    """Dilated 1D convolution decoder that upsamples tokens to frames, conditioned on per-frame targets and
+    external features."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        target_channels: int,
+        external_channels: int,
+        width: int = 512,
+        depth: int = 4,
+        dilation_growth_rate: int = 3,
+        num_tokens: int = 16,
+        num_frames: int = 64,
+    ):
+        super().__init__()
+        self.width = width
+        self.num_frames = num_frames
+        lengths = [num_tokens, num_tokens * 2]
+        self.conv_in = MotionBricksConv1d(in_channels, width, 3, padding=1)
+        self.target_projections = nn.ModuleList(
+            [nn.Linear(target_channels, width * length // num_frames) for length in lengths]
+        )
+        self.external_projections = nn.ModuleList(
+            [nn.Linear(width + external_channels * num_frames // length, width) for length in lengths]
+        )
+        self.up_blocks = nn.ModuleList([MotionBricksUpBlock(width, depth, dilation_growth_rate) for _ in lengths])
+        self.conv_mid = MotionBricksConv1d(width, width, 3, padding=1)
+        self.conv_out = MotionBricksConv1d(width, out_channels, 3, padding=1)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        token_mask: torch.Tensor,
+        target: torch.Tensor,
+        target_mask: torch.Tensor,
+        external: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Args:
+            hidden_states: `(batch_size, in_channels, num_tokens)` token features.
+            token_mask: `(batch_size, num_tokens)` float mask of the valid tokens.
+            target: `(batch_size, num_frames, target_channels)` per-frame targets.
+            target_mask: `(batch_size, num_frames)` boolean mask of the frames with a target.
+            external: `(batch_size, num_frames, external_channels)` per-frame conditioning.
+
+        Returns:
+            `(batch_size, out_channels, num_frames)` frame features.
+        """
+        batch_size = hidden_states.shape[0]
+        mask = token_mask[:, None]
+        hidden_states = F.relu(self.conv_in(hidden_states * mask))
+        for block, target_projection, external_projection in zip(
+            self.up_blocks, self.target_projections, self.external_projections
+        ):
+            length = hidden_states.shape[-1]
+            projected = F.relu(target_projection(target))
+            grouped = hidden_states.transpose(1, 2).reshape(batch_size, self.num_frames, -1)
+            grouped = torch.where(target_mask[..., None], projected, grouped)
+            grouped = grouped.reshape(batch_size, length, self.width)
+            conditioned = torch.cat([grouped, external.reshape(batch_size, length, -1)], dim=-1)
+            hidden_states = F.relu(external_projection(conditioned).transpose(1, 2))
+            for resnet in block.resnets:
+                hidden_states = resnet(hidden_states, mask)
+            hidden_states = F.interpolate(hidden_states, scale_factor=2, mode="nearest")
+            mask = mask.repeat_interleave(2, dim=-1)
+            hidden_states = block.conv(hidden_states * mask)
+        hidden_states = F.relu(self.conv_mid(hidden_states * mask)) * mask
+        return self.conv_out(hidden_states) * mask
+
+
+@dataclass
+class MotionBricksRootOutput(BaseOutput):
+    """
+    Args:
+        num_tokens (`torch.LongTensor` of shape `(batch_size,)`):
+            The chosen motion length in tokens; the motion spans `num_tokens * frames_per_token` frames.
+        root (`torch.Tensor` of shape `(batch_size, num_frames, 5)`):
+            The normalized global root of every frame: position, then cosine and sine of the heading.
+    """
+
+    num_tokens: torch.LongTensor
+    root: torch.Tensor
+
+
+class MotionBricksRootModel(ModelMixin, ConfigMixin, AttentionMixin):
+    """Chooses the length of a transition and predicts its root trajectory from `context_frames` context poses and
+    `target_frames` target poses. Stores the normalization statistics of the guide features as buffers."""
+
+    _no_split_modules = ["MotionBricksTransformerLayer", "MotionBricksConditionedDecoder"]
+    _repeated_blocks = ["MotionBricksTransformerLayer"]
+    _skip_layerwise_casting_patterns = ["norm", "embed"]
+
+    @register_to_config
+    def __init__(
+        self,
+        pose_dim: int = 268,
+        global_root_dim: int = 5,
+        local_root_dim: int = 4,
+        pose_embed_dim: int = 256,
+        root_embed_dim: int = 64,
+        hidden_dim: int = 512,
+        num_heads: int = 16,
+        ff_dim: int = 2048,
+        num_shared_layers: int = 6,
+        num_token_layers: int = 2,
+        context_frames: int = 4,
+        target_frames: int = 4,
+        frames_per_token: int = 4,
+        min_tokens: int = 6,
+        max_tokens: int = 16,
+        decoder_depth: int = 4,
+        decoder_dilation_growth_rate: int = 3,
+    ):
+        super().__init__()
+        guide_frames = context_frames + target_frames
+        num_lengths = max_tokens - min_tokens + 1
+        guide_dim = pose_embed_dim + 2 * root_embed_dim
+
+        self.pose_proj = nn.Linear(pose_dim, pose_embed_dim)
+        self.local_root_proj = nn.Linear(local_root_dim, root_embed_dim)
+        self.global_root_proj = nn.Linear(global_root_dim, root_embed_dim)
+        self.missing_pose_embed = nn.Parameter(torch.zeros(pose_embed_dim))
+        self.missing_local_root_embed = nn.Parameter(torch.zeros(root_embed_dim))
+        self.missing_global_root_embed = nn.Parameter(torch.zeros(root_embed_dim))
+        self.context_proj = MotionBricksMLP(guide_dim, hidden_dim, hidden_dim)
+        self.target_proj = MotionBricksMLP(guide_dim, hidden_dim, hidden_dim)
+        self.guide_position_embed = nn.Parameter(torch.zeros(guide_frames, hidden_dim))
+
+        # Index num_lengths + 1 asks for a prediction; the length classes are min_tokens..max_tokens and one unused.
+        self.length_query_embed = nn.Embedding(num_lengths + 2, hidden_dim)
+        self.shared_layers = nn.ModuleList(
+            [MotionBricksTransformerLayer(hidden_dim, num_heads, ff_dim) for _ in range(num_shared_layers)]
+        )
+        self.length_head = nn.Linear(hidden_dim, num_lengths + 1)
+
+        self.length_embed = nn.Embedding(num_lengths + 1, hidden_dim)
+        self.token_position_embed = nn.Parameter(torch.zeros(max_tokens, hidden_dim))
+        self.token_layers = nn.ModuleList(
+            [MotionBricksTransformerLayer(hidden_dim, num_heads, ff_dim) for _ in range(num_token_layers)]
+        )
+
+        self.missing_frame_embed = nn.Parameter(torch.zeros(hidden_dim))
+        self.decoder = MotionBricksConditionedDecoder(
+            hidden_dim,
+            global_root_dim,
+            target_channels=global_root_dim,
+            external_channels=hidden_dim,
+            width=hidden_dim,
+            depth=decoder_depth,
+            dilation_growth_rate=decoder_dilation_growth_rate,
+            num_tokens=max_tokens,
+            num_frames=max_tokens * frames_per_token,
+        )
+
+        self.register_buffer("global_root_mean", torch.zeros(global_root_dim))
+        self.register_buffer("global_root_std", torch.ones(global_root_dim))
+        self.register_buffer("local_root_mean", torch.zeros(local_root_dim))
+        self.register_buffer("local_root_std", torch.ones(local_root_dim))
+        self.register_buffer("pose_mean", torch.zeros(pose_dim))
+        self.register_buffer("pose_std", torch.ones(pose_dim))
+
+    def _statistics(self, value: torch.Tensor, name: str) -> tuple[torch.Tensor, torch.Tensor]:
+        # The buffers can sit on another device while the model is offloaded.
+        mean, std = getattr(self, f"{name}_mean"), getattr(self, f"{name}_std")
+        return mean.to(value.device, value.dtype), std.to(value.device, value.dtype)
+
+    def normalize(self, value: torch.Tensor, name: str) -> torch.Tensor:
+        """Normalize `global_root`, `local_root` or `pose` features with this model's statistics."""
+        mean, std = self._statistics(value, name)
+        return (value - mean) / std
+
+    def denormalize(self, value: torch.Tensor, name: str) -> torch.Tensor:
+        mean, std = self._statistics(value, name)
+        return value * std + mean
+
+    def place_guides(self, guides: torch.Tensor, num_tokens: torch.Tensor, fill) -> torch.Tensor:
+        """Place `(batch_size, guide_frames, ...)` guide values on the frame timeline: the context poses at the start
+        and the target poses at the last frames of the chosen length. Other frames hold `fill`."""
+        config = self.config
+        shape = (guides.shape[0], config.max_tokens * config.frames_per_token, *guides.shape[2:])
+        if isinstance(fill, torch.Tensor):
+            timeline = fill.to(guides.dtype).expand(shape).clone()
+        else:
+            timeline = torch.full(shape, fill, dtype=guides.dtype, device=guides.device)
+        timeline[:, : config.context_frames] = guides[:, : config.context_frames]
+        end = num_tokens * config.frames_per_token
+        frames = end[:, None] - config.target_frames + torch.arange(config.target_frames, device=end.device)
+        batch = torch.arange(guides.shape[0], device=guides.device)[:, None]
+        timeline[batch, frames] = guides[:, config.context_frames :]
+        return timeline
+
+    def forward(
+        self,
+        global_root: torch.Tensor,
+        local_root: torch.Tensor,
+        pose: torch.Tensor,
+        has_global_root: torch.Tensor,
+        has_local_root: torch.Tensor,
+        has_pose: torch.Tensor,
+        allowed_num_tokens: torch.Tensor,
+        return_dict: bool = True,
+    ) -> MotionBricksRootOutput | tuple:
+        """
+        Args:
+            global_root (`torch.Tensor` of shape `(batch_size, guide_frames, global_root_dim)`):
+                Normalized global root of the guide poses.
+            local_root (`torch.Tensor` of shape `(batch_size, guide_frames, local_root_dim)`):
+                Normalized local root of the guide poses.
+            pose (`torch.Tensor` of shape `(batch_size, guide_frames, pose_dim)`):
+                Normalized pose features of the guide poses.
+            has_global_root (`torch.BoolTensor` of shape `(batch_size, guide_frames)`):
+                Which guide poses have a global root.
+            has_local_root (`torch.BoolTensor` of shape `(batch_size, guide_frames)`):
+                Which guide poses have a local root.
+            has_pose (`torch.BoolTensor` of shape `(batch_size, guide_frames)`):
+                Which guide poses have pose features.
+            allowed_num_tokens (`torch.BoolTensor` of shape `(batch_size, max_tokens - min_tokens + 1)`):
+                The lengths the model may choose, from `min_tokens`.
+            return_dict (`bool`, defaults to `True`):
+                Whether to return a `MotionBricksRootOutput` instead of a tuple.
+
+        Returns:
+            `MotionBricksRootOutput` or `tuple`: The chosen length and the normalized root; a tuple of the two when
+            `return_dict` is `False`.
+        """
+        config = self.config
+        has_pose_f = has_pose[..., None].to(pose.dtype)
+        has_local_f = has_local_root[..., None].to(pose.dtype)
+        has_global_f = has_global_root[..., None].to(pose.dtype)
+        pose_embeds = self.pose_proj(pose) * has_pose_f + self.missing_pose_embed * (1 - has_pose_f)
+        local_embeds = self.local_root_proj(local_root) * has_local_f + self.missing_local_root_embed * (
+            1 - has_local_f
+        )
+        global_embeds = self.global_root_proj(global_root) * has_global_f + self.missing_global_root_embed * (
+            1 - has_global_f
+        )
+        guides = torch.cat([pose_embeds, local_embeds, global_embeds], dim=-1)
+        guides = torch.cat(
+            [
+                self.context_proj(guides[:, : config.context_frames]),
+                self.target_proj(guides[:, config.context_frames :]),
+            ],
+            dim=1,
+        )
+        positioned = guides + self.guide_position_embed
+
+        batch_size = guides.shape[0]
+        num_lengths = config.max_tokens - config.min_tokens + 1
+        query = self.length_query_embed(torch.full((batch_size, 1), num_lengths + 1, device=guides.device))
+        hidden_states = torch.cat([query, positioned], dim=1)
+        for layer in self.shared_layers:
+            hidden_states = layer(hidden_states)
+        logits = self.length_head(hidden_states[:, 0])
+        allowed = F.pad(allowed_num_tokens, (0, 1), value=False)
+        length = logits.masked_fill(~allowed, float("-inf")).argmax(dim=-1)
+        num_tokens = length + config.min_tokens
+
+        tokens = self.token_position_embed.expand(batch_size, -1, -1)
+        hidden_states = torch.cat([self.length_embed(length)[:, None], positioned, tokens], dim=1)
+        mask = key_mask(num_tokens, config.max_tokens, prefix=1 + guides.shape[1])
+        for layer in self.token_layers:
+            hidden_states = layer(hidden_states, mask)
+        tokens = hidden_states[:, -config.max_tokens :]
+
+        has_guide = has_pose | has_local_root | has_global_root
+        frame_guides = torch.where(has_guide[..., None], guides, self.missing_frame_embed)
+        external = self.place_guides(frame_guides, num_tokens, fill=self.missing_frame_embed)
+        target = self.place_guides(global_root, num_tokens, fill=0)
+        target_mask = self.place_guides(has_global_root, num_tokens, fill=False)
+        token_mask = (torch.arange(config.max_tokens, device=num_tokens.device)[None] < num_tokens[:, None]).to(
+            tokens.dtype
+        )
+        root = self.decoder(tokens.transpose(1, 2), token_mask, target, target_mask, external).transpose(1, 2)
+
+        if not return_dict:
+            return (num_tokens, root)
+        return MotionBricksRootOutput(num_tokens=num_tokens, root=root)