OzzyGT/motionbricks_soma_custom_blocks

MotionBricks SOMA Modular Diffusers

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 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.

MotionBricks has no text input. It fills the gap between poses you provide, for example two motions from 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, adapted from the MotionBricks runtime. See NOTICE. This repository contains no model weights; they use the NVIDIA Open Model Agreement, see the model repository.

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: the start of a run, and a standing pose that the pipeline places 1.4 m ahead of it:

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, 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.

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:

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.