motionbricks_soma_diffusers / pose_model /transformer_motionbricks.py

OzzyGT HF Staff

Upload 13 files

338c3c5 verified 2 days ago

9.29 kB

# 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)

	# 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)