|
import numpy as np |
|
import torch |
|
from torch import nn |
|
|
|
from typing import Optional, Dict, List |
|
|
|
from mogen.models.utils.misc import zero_module |
|
|
|
from ..builder import SUBMODULES, build_attention |
|
from ..utils.stylization_block import StylizationBlock |
|
from .motion_transformer import MotionTransformer |
|
|
|
|
|
def get_kit_slice(idx: int) -> List[int]: |
|
""" |
|
Get the slice indices for the KIT skeleton. |
|
|
|
Args: |
|
idx (int): The index of the skeleton part. |
|
|
|
Returns: |
|
List[int]: Slice indices for the specified skeleton part. |
|
""" |
|
if idx == 0: |
|
return [0, 1, 2, 3, 184, 185, 186, 247, 248, 249, 250] |
|
return [ |
|
4 + (idx - 1) * 3, |
|
4 + (idx - 1) * 3 + 1, |
|
4 + (idx - 1) * 3 + 2, |
|
64 + (idx - 1) * 6, |
|
64 + (idx - 1) * 6 + 1, |
|
64 + (idx - 1) * 6 + 2, |
|
64 + (idx - 1) * 6 + 3, |
|
64 + (idx - 1) * 6 + 4, |
|
64 + (idx - 1) * 6 + 5, |
|
184 + idx * 3, |
|
184 + idx * 3 + 1, |
|
184 + idx * 3 + 2, |
|
] |
|
|
|
|
|
def get_t2m_slice(idx: int) -> List[int]: |
|
""" |
|
Get the slice indices for the T2M skeleton. |
|
|
|
Args: |
|
idx (int): The index of the skeleton part. |
|
|
|
Returns: |
|
List[int]: Slice indices for the specified skeleton part. |
|
""" |
|
if idx == 0: |
|
return [0, 1, 2, 3, 193, 194, 195, 259, 260, 261, 262] |
|
return [ |
|
4 + (idx - 1) * 3, |
|
4 + (idx - 1) * 3 + 1, |
|
4 + (idx - 1) * 3 + 2, |
|
67 + (idx - 1) * 6, |
|
67 + (idx - 1) * 6 + 1, |
|
67 + (idx - 1) * 6 + 2, |
|
67 + (idx - 1) * 6 + 3, |
|
67 + (idx - 1) * 6 + 4, |
|
67 + (idx - 1) * 6 + 5, |
|
193 + idx * 3, |
|
193 + idx * 3 + 1, |
|
193 + idx * 3 + 2, |
|
] |
|
|
|
|
|
def get_part_slice(idx_list: List[int], func) -> List[int]: |
|
""" |
|
Get the slice indices for a list of indices. |
|
|
|
Args: |
|
idx_list (List[int]): List of part indices. |
|
func (Callable): Function to get slice indices for each part. |
|
|
|
Returns: |
|
List[int]: Concatenated list of slice indices for the parts. |
|
""" |
|
result = [] |
|
for idx in idx_list: |
|
result.extend(func(idx)) |
|
return result |
|
|
|
|
|
class PoseEncoder(nn.Module): |
|
""" |
|
Pose Encoder to process motion data and encode body parts into latent representations. |
|
""" |
|
|
|
def __init__(self, |
|
dataset_name: str = "human_ml3d", |
|
latent_dim: int = 64, |
|
input_dim: int = 263): |
|
super().__init__() |
|
self.dataset_name = dataset_name |
|
if dataset_name == "human_ml3d": |
|
func = get_t2m_slice |
|
self.head_slice = get_part_slice([12, 15], func) |
|
self.stem_slice = get_part_slice([3, 6, 9], func) |
|
self.larm_slice = get_part_slice([14, 17, 19, 21], func) |
|
self.rarm_slice = get_part_slice([13, 16, 18, 20], func) |
|
self.lleg_slice = get_part_slice([2, 5, 8, 11], func) |
|
self.rleg_slice = get_part_slice([1, 4, 7, 10], func) |
|
self.root_slice = get_part_slice([0], func) |
|
self.body_slice = get_part_slice([_ for _ in range(22)], func) |
|
elif dataset_name == "kit_ml": |
|
func = get_kit_slice |
|
self.head_slice = get_part_slice([4], func) |
|
self.stem_slice = get_part_slice([1, 2, 3], func) |
|
self.larm_slice = get_part_slice([8, 9, 10], func) |
|
self.rarm_slice = get_part_slice([5, 6, 7], func) |
|
self.lleg_slice = get_part_slice([16, 17, 18, 19, 20], func) |
|
self.rleg_slice = get_part_slice([11, 12, 13, 14, 15], func) |
|
self.root_slice = get_part_slice([0], func) |
|
self.body_slice = get_part_slice([_ for _ in range(21)], func) |
|
else: |
|
raise ValueError() |
|
|
|
self.head_embed = nn.Linear(len(self.head_slice), latent_dim) |
|
self.stem_embed = nn.Linear(len(self.stem_slice), latent_dim) |
|
self.larm_embed = nn.Linear(len(self.larm_slice), latent_dim) |
|
self.rarm_embed = nn.Linear(len(self.rarm_slice), latent_dim) |
|
self.lleg_embed = nn.Linear(len(self.lleg_slice), latent_dim) |
|
self.rleg_embed = nn.Linear(len(self.rleg_slice), latent_dim) |
|
self.root_embed = nn.Linear(len(self.root_slice), latent_dim) |
|
self.body_embed = nn.Linear(len(self.body_slice), latent_dim) |
|
|
|
assert len(set(self.body_slice)) == input_dim |
|
|
|
def forward(self, motion: torch.Tensor) -> torch.Tensor: |
|
""" |
|
Forward pass for encoding the motion into body part embeddings. |
|
|
|
Args: |
|
motion (torch.Tensor): Input motion tensor of shape (B, T, D). |
|
|
|
Returns: |
|
torch.Tensor: Concatenated latent representations of body parts. |
|
""" |
|
head_feat = self.head_embed(motion[:, :, self.head_slice].contiguous()) |
|
stem_feat = self.stem_embed(motion[:, :, self.stem_slice].contiguous()) |
|
larm_feat = self.larm_embed(motion[:, :, self.larm_slice].contiguous()) |
|
rarm_feat = self.rarm_embed(motion[:, :, self.rarm_slice].contiguous()) |
|
lleg_feat = self.lleg_embed(motion[:, :, self.lleg_slice].contiguous()) |
|
rleg_feat = self.rleg_embed(motion[:, :, self.rleg_slice].contiguous()) |
|
root_feat = self.root_embed(motion[:, :, self.root_slice].contiguous()) |
|
body_feat = self.body_embed(motion[:, :, self.body_slice].contiguous()) |
|
feat = torch.cat((head_feat, stem_feat, larm_feat, rarm_feat, |
|
lleg_feat, rleg_feat, root_feat, body_feat), |
|
dim=-1) |
|
return feat |
|
|
|
|
|
class PoseDecoder(nn.Module): |
|
""" |
|
Pose Decoder to decode the latent representations of body parts back into motion. |
|
""" |
|
|
|
def __init__(self, |
|
dataset_name: str = "human_ml3d", |
|
latent_dim: int = 64, |
|
output_dim: int = 263): |
|
super().__init__() |
|
self.dataset_name = dataset_name |
|
self.latent_dim = latent_dim |
|
self.output_dim = output_dim |
|
if dataset_name == "human_ml3d": |
|
func = get_t2m_slice |
|
self.head_slice = get_part_slice([12, 15], func) |
|
self.stem_slice = get_part_slice([3, 6, 9], func) |
|
self.larm_slice = get_part_slice([14, 17, 19, 21], func) |
|
self.rarm_slice = get_part_slice([13, 16, 18, 20], func) |
|
self.lleg_slice = get_part_slice([2, 5, 8, 11], func) |
|
self.rleg_slice = get_part_slice([1, 4, 7, 10], func) |
|
self.root_slice = get_part_slice([0], func) |
|
self.body_slice = get_part_slice([_ for _ in range(22)], func) |
|
elif dataset_name == "kit_ml": |
|
func = get_kit_slice |
|
self.head_slice = get_part_slice([4], func) |
|
self.stem_slice = get_part_slice([1, 2, 3], func) |
|
self.larm_slice = get_part_slice([8, 9, 10], func) |
|
self.rarm_slice = get_part_slice([5, 6, 7], func) |
|
self.lleg_slice = get_part_slice([16, 17, 18, 19, 20], func) |
|
self.rleg_slice = get_part_slice([11, 12, 13, 14, 15], func) |
|
self.root_slice = get_part_slice([0], func) |
|
self.body_slice = get_part_slice([_ for _ in range(21)], func) |
|
else: |
|
raise ValueError() |
|
|
|
self.head_out = nn.Linear(latent_dim, len(self.head_slice)) |
|
self.stem_out = nn.Linear(latent_dim, len(self.stem_slice)) |
|
self.larm_out = nn.Linear(latent_dim, len(self.larm_slice)) |
|
self.rarm_out = nn.Linear(latent_dim, len(self.rarm_slice)) |
|
self.lleg_out = nn.Linear(latent_dim, len(self.lleg_slice)) |
|
self.rleg_out = nn.Linear(latent_dim, len(self.rleg_slice)) |
|
self.root_out = nn.Linear(latent_dim, len(self.root_slice)) |
|
self.body_out = nn.Linear(latent_dim, len(self.body_slice)) |
|
|
|
def forward(self, motion: torch.Tensor) -> torch.Tensor: |
|
""" |
|
Forward pass to decode the latent body part features back to motion. |
|
|
|
Args: |
|
motion (torch.Tensor): Input tensor of shape (B, T, D). |
|
|
|
Returns: |
|
torch.Tensor: Output motion tensor of shape (B, T, output_dim). |
|
""" |
|
B, T = motion.shape[:2] |
|
D = self.latent_dim |
|
head_feat = self.head_out(motion[:, :, :D].contiguous()) |
|
stem_feat = self.stem_out(motion[:, :, D:2 * D].contiguous()) |
|
larm_feat = self.larm_out(motion[:, :, 2 * D:3 * D].contiguous()) |
|
rarm_feat = self.rarm_out(motion[:, :, 3 * D:4 * D].contiguous()) |
|
lleg_feat = self.lleg_out(motion[:, :, 4 * D:5 * D].contiguous()) |
|
rleg_feat = self.rleg_out(motion[:, :, 5 * D:6 * D].contiguous()) |
|
root_feat = self.root_out(motion[:, :, 6 * D:7 * D].contiguous()) |
|
body_feat = self.body_out(motion[:, :, 7 * D:].contiguous()) |
|
output = torch.zeros(B, T, self.output_dim).type_as(motion) |
|
output[:, :, self.head_slice] = head_feat |
|
output[:, :, self.stem_slice] = stem_feat |
|
output[:, :, self.larm_slice] = larm_feat |
|
output[:, :, self.rarm_slice] = rarm_feat |
|
output[:, :, self.lleg_slice] = lleg_feat |
|
output[:, :, self.rleg_slice] = rleg_feat |
|
output[:, :, self.root_slice] = root_feat |
|
output = (output + body_feat) / 2.0 |
|
return output |
|
|
|
|
|
class SFFN(nn.Module): |
|
""" |
|
A Stylized Feed-Forward Network (SFFN) module for transformer layers. |
|
|
|
Args: |
|
latent_dim (int): Dimensionality of the input. |
|
ffn_dim (int): Dimensionality of the feed-forward layer. |
|
dropout (float): Dropout probability. |
|
time_embed_dim (int): Dimensionality of the time embedding. |
|
norm (str): Normalization type ('None'). |
|
activation (str): Activation function ('GELU'). |
|
""" |
|
|
|
def __init__(self, |
|
latent_dim: int, |
|
ffn_dim: int, |
|
dropout: float, |
|
time_embed_dim: int, |
|
norm: str = "None", |
|
activation: str = "GELU", |
|
**kwargs): |
|
super().__init__() |
|
self.linear1_list = nn.ModuleList() |
|
self.linear2_list = nn.ModuleList() |
|
|
|
channel_mul = 1 |
|
if activation == "GELU": |
|
self.activation = nn.GELU() |
|
|
|
for i in range(8): |
|
self.linear1_list.append(nn.Linear(latent_dim, ffn_dim * channel_mul)) |
|
self.linear2_list.append(nn.Linear(ffn_dim, latent_dim)) |
|
|
|
self.dropout = nn.Dropout(dropout) |
|
self.proj_out = StylizationBlock(latent_dim * 8, time_embed_dim, dropout) |
|
|
|
if norm == "None": |
|
self.norm = nn.Identity() |
|
|
|
def forward(self, x: torch.Tensor, emb: torch.Tensor, **kwargs) -> torch.Tensor: |
|
""" |
|
Forward pass of the SFFN layer. |
|
|
|
Args: |
|
x (torch.Tensor): Input tensor of shape (B, T, D). |
|
emb (torch.Tensor): Embedding tensor for time step modulation. |
|
|
|
Returns: |
|
torch.Tensor: Output tensor of shape (B, T, D). |
|
""" |
|
B, T, D = x.shape |
|
x = self.norm(x) |
|
x = x.reshape(B, T, 8, -1) |
|
output = [] |
|
for i in range(8): |
|
feat = x[:, :, i].contiguous() |
|
feat = self.dropout(self.activation(self.linear1_list[i](feat))) |
|
feat = self.linear2_list[i](feat) |
|
output.append(feat) |
|
y = torch.cat(output, dim=-1) |
|
y = x.reshape(B, T, D) + self.proj_out(y, emb) |
|
return y |
|
|
|
|
|
class DecoderLayer(nn.Module): |
|
""" |
|
A transformer decoder layer with cross-attention and feed-forward network (SFFN). |
|
|
|
Args: |
|
ca_block_cfg (Optional[Dict]): Configuration for the cross-attention block. |
|
ffn_cfg (Optional[Dict]): Configuration for the feed-forward network (SFFN). |
|
""" |
|
|
|
def __init__(self, ca_block_cfg: Optional[Dict] = None, ffn_cfg: Optional[Dict] = None): |
|
super().__init__() |
|
self.ca_block = build_attention(ca_block_cfg) |
|
self.ffn = SFFN(**ffn_cfg) |
|
|
|
def forward(self, **kwargs) -> torch.Tensor: |
|
""" |
|
Forward pass of the decoder layer. |
|
|
|
Args: |
|
kwargs: Keyword arguments for attention and feed-forward layers. |
|
|
|
Returns: |
|
torch.Tensor: Output of the decoder layer. |
|
""" |
|
if self.ca_block is not None: |
|
x = self.ca_block(**kwargs) |
|
kwargs.update({'x': x}) |
|
if self.ffn is not None: |
|
x = self.ffn(**kwargs) |
|
return x |
|
|
|
|
|
@SUBMODULES.register_module() |
|
class FineMoGenTransformer(MotionTransformer): |
|
""" |
|
A transformer model for motion generation using fine-grained control with Diffusion. |
|
|
|
Args: |
|
scale_func_cfg (Optional[Dict]): Configuration for scaling function. |
|
pose_encoder_cfg (Optional[Dict]): Configuration for the PoseEncoder. |
|
pose_decoder_cfg (Optional[Dict]): Configuration for the PoseDecoder. |
|
moe_route_loss_weight (float): Weight for the Mixture of Experts (MoE) routing loss. |
|
template_kl_loss_weight (float): Weight for the KL loss in template generation. |
|
fine_mode (bool): Whether to enable fine mode for control over body parts. |
|
""" |
|
|
|
def __init__(self, |
|
scale_func_cfg: Optional[Dict] = None, |
|
pose_encoder_cfg: Optional[Dict] = None, |
|
pose_decoder_cfg: Optional[Dict] = None, |
|
moe_route_loss_weight: float = 1.0, |
|
template_kl_loss_weight: float = 0.0001, |
|
fine_mode: bool = False, |
|
**kwargs): |
|
super().__init__(**kwargs) |
|
self.scale_func_cfg = scale_func_cfg |
|
self.joint_embed = PoseEncoder(**pose_encoder_cfg) |
|
self.out = zero_module(PoseDecoder(**pose_decoder_cfg)) |
|
self.moe_route_loss_weight = moe_route_loss_weight |
|
self.template_kl_loss_weight = template_kl_loss_weight |
|
self.mean = np.load("data/datasets/kit_ml/mean.npy") |
|
self.std = np.load("data/datasets/kit_ml/std.npy") |
|
self.fine_mode = fine_mode |
|
|
|
def build_temporal_blocks(self, sa_block_cfg: Optional[Dict], ca_block_cfg: Optional[Dict], ffn_cfg: Optional[Dict]): |
|
""" |
|
Build temporal decoder blocks for the model. |
|
|
|
Args: |
|
sa_block_cfg (Optional[Dict]): Configuration for self-attention blocks. |
|
ca_block_cfg (Optional[Dict]): Configuration for cross-attention blocks. |
|
ffn_cfg (Optional[Dict]): Configuration for feed-forward networks. |
|
""" |
|
self.temporal_decoder_blocks = nn.ModuleList() |
|
for i in range(self.num_layers): |
|
if isinstance(ffn_cfg, list): |
|
ffn_cfg_block = ffn_cfg[i] |
|
else: |
|
ffn_cfg_block = ffn_cfg |
|
self.temporal_decoder_blocks.append(DecoderLayer(ca_block_cfg=ca_block_cfg, ffn_cfg=ffn_cfg_block)) |
|
|
|
def scale_func(self, timestep: int) -> Dict[str, float]: |
|
""" |
|
Scaling function for text and none coefficient based on timestep. |
|
|
|
Args: |
|
timestep (int): Current diffusion timestep. |
|
|
|
Returns: |
|
Dict[str, float]: Scaling factors for text and non-text conditioning. |
|
""" |
|
scale = self.scale_func_cfg['scale'] |
|
w = (1 - (1000 - timestep) / 1000) * scale + 1 |
|
return {'text_coef': w, 'none_coef': 1 - w} |
|
|
|
def aux_loss(self) -> Dict[str, torch.Tensor]: |
|
""" |
|
Auxiliary loss computation for MoE routing and KL loss. |
|
|
|
Returns: |
|
Dict[str, torch.Tensor]: Computed auxiliary losses. |
|
""" |
|
aux_loss = 0 |
|
kl_loss = 0 |
|
for module in self.temporal_decoder_blocks: |
|
if hasattr(module.ca_block, 'aux_loss'): |
|
aux_loss = aux_loss + module.ca_block.aux_loss |
|
if hasattr(module.ca_block, 'kl_loss'): |
|
kl_loss = kl_loss + module.ca_block.kl_loss |
|
losses = {} |
|
if aux_loss > 0: |
|
losses['moe_route_loss'] = aux_loss * self.moe_route_loss_weight |
|
if kl_loss > 0: |
|
losses['template_kl_loss'] = kl_loss * self.template_kl_loss_weight |
|
return losses |
|
|
|
def get_precompute_condition(self, |
|
text: Optional[str] = None, |
|
motion_length: Optional[torch.Tensor] = None, |
|
xf_out: Optional[torch.Tensor] = None, |
|
re_dict: Optional[Dict] = None, |
|
device: Optional[torch.device] = None, |
|
sample_idx: Optional[int] = None, |
|
clip_feat: Optional[torch.Tensor] = None, |
|
**kwargs) -> Dict[str, torch.Tensor]: |
|
""" |
|
Precompute conditioning features for text or other modalities. |
|
|
|
Args: |
|
text (Optional[str]): Text input for conditioning. |
|
motion_length (Optional[torch.Tensor]): Length of the motion sequence. |
|
xf_out (Optional[torch.Tensor]): Precomputed text features. |
|
re_dict (Optional[Dict]): Additional features dictionary. |
|
device (Optional[torch.device]): Target device for the model. |
|
sample_idx (Optional[int]): Sample index for specific conditioning. |
|
clip_feat (Optional[torch.Tensor]): Precomputed CLIP features. |
|
|
|
Returns: |
|
Dict[str, torch.Tensor]: Precomputed conditioning features. |
|
""" |
|
if xf_out is None: |
|
xf_out = self.encode_text(text, clip_feat, device) |
|
output = {'xf_out': xf_out} |
|
return output |
|
|
|
def post_process(self, motion: torch.Tensor) -> torch.Tensor: |
|
""" |
|
Post-process motion data by unnormalizing if necessary. |
|
|
|
Args: |
|
motion (torch.Tensor): Input motion data. |
|
|
|
Returns: |
|
torch.Tensor: Processed motion data. |
|
""" |
|
if self.post_process_cfg is not None: |
|
if self.post_process_cfg.get("unnormalized_infer", False): |
|
mean = torch.from_numpy(np.load(self.post_process_cfg['mean_path'])).type_as(motion) |
|
std = torch.from_numpy(np.load(self.post_process_cfg['std_path'])).type_as(motion) |
|
motion = motion * std + mean |
|
return motion |
|
|
|
def forward_train(self, |
|
h: torch.Tensor, |
|
src_mask: Optional[torch.Tensor] = None, |
|
emb: Optional[torch.Tensor] = None, |
|
xf_out: Optional[torch.Tensor] = None, |
|
motion_length: Optional[torch.Tensor] = None, |
|
num_intervals: int = 1, |
|
**kwargs) -> torch.Tensor: |
|
""" |
|
Forward pass during training. |
|
|
|
Args: |
|
h (torch.Tensor): Input tensor of shape (B, T, D). |
|
src_mask (Optional[torch.Tensor]): Source mask tensor. |
|
emb (Optional[torch.Tensor]): Time embedding tensor. |
|
xf_out (Optional[torch.Tensor]): Precomputed text features. |
|
motion_length (Optional[torch.Tensor]): Lengths of motion sequences. |
|
num_intervals (int): Number of intervals for processing. |
|
|
|
Returns: |
|
torch.Tensor: Output tensor of shape (B, T, D). |
|
""" |
|
B, T = h.shape[0], h.shape[1] |
|
cond_type = torch.randint(0, 100, size=(B, 1, 1)).repeat(1, 8, 1).to(h.device) if self.fine_mode else torch.randint(0, 100, size=(B, 1, 1)).to(h.device) |
|
for module in self.temporal_decoder_blocks: |
|
h = module(x=h, |
|
xf=xf_out, |
|
emb=emb, |
|
src_mask=src_mask, |
|
cond_type=cond_type, |
|
motion_length=motion_length, |
|
num_intervals=num_intervals) |
|
|
|
output = self.out(h).view(B, T, -1).contiguous() |
|
return output |
|
|
|
def forward_test(self, |
|
h: torch.Tensor, |
|
src_mask: Optional[torch.Tensor] = None, |
|
emb: Optional[torch.Tensor] = None, |
|
xf_out: Optional[torch.Tensor] = None, |
|
timesteps: Optional[torch.Tensor] = None, |
|
motion_length: Optional[torch.Tensor] = None, |
|
num_intervals: int = 1, |
|
**kwargs) -> torch.Tensor: |
|
""" |
|
Forward pass during inference. |
|
|
|
Args: |
|
h (torch.Tensor): Input tensor of shape (B, T, D). |
|
src_mask (Optional[torch.Tensor]): Source mask tensor. |
|
emb (Optional[torch.Tensor]): Time embedding tensor. |
|
xf_out (Optional[torch.Tensor]): Precomputed text features. |
|
timesteps (Optional[torch.Tensor]): Diffusion timesteps. |
|
motion_length (Optional[torch.Tensor]): Lengths of motion sequences. |
|
num_intervals (int): Number of intervals for processing. |
|
|
|
Returns: |
|
torch.Tensor: Output tensor of shape (B, T, D). |
|
""" |
|
B, T = h.shape[0], h.shape[1] |
|
text_cond_type = torch.zeros(B, 1, 1).to(h.device) + 1 |
|
none_cond_type = torch.zeros(B, 1, 1).to(h.device) |
|
|
|
all_cond_type = torch.cat((text_cond_type, none_cond_type), dim=0) |
|
h = h.repeat(2, 1, 1) |
|
xf_out = xf_out.repeat(2, 1, 1) |
|
emb = emb.repeat(2, 1) |
|
src_mask = src_mask.repeat(2, 1, 1) |
|
motion_length = motion_length.repeat(2, 1) |
|
for module in self.temporal_decoder_blocks: |
|
h = module(x=h, |
|
xf=xf_out, |
|
emb=emb, |
|
src_mask=src_mask, |
|
cond_type=all_cond_type, |
|
motion_length=motion_length, |
|
num_intervals=num_intervals) |
|
out = self.out(h).view(2 * B, T, -1).contiguous() |
|
out_text = out[:B].contiguous() |
|
out_none = out[B:].contiguous() |
|
|
|
coef_cfg = self.scale_func(int(timesteps[0])) |
|
text_coef = coef_cfg['text_coef'] |
|
none_coef = coef_cfg['none_coef'] |
|
output = out_text * text_coef + out_none * none_coef |
|
return output |
|
|
|
|
|
|
