|
import torch |
|
from torch import nn |
|
from typing import Optional |
|
|
|
from mogen.models.utils.misc import zero_module |
|
from mogen.models.utils.position_encoding import timestep_embedding |
|
from mogen.models.utils.stylization_block import StylizationBlock |
|
|
|
from ..builder import SUBMODULES, build_attention |
|
from .remodiffuse import ReMoDiffuseTransformer |
|
|
|
|
|
class FFN(nn.Module): |
|
""" |
|
A feed-forward network (FFN) with optional stylization block. |
|
|
|
Args: |
|
latent_dim (int): The dimension of the input and output latent space. |
|
ffn_dim (int): The dimension of the hidden feed-forward network. |
|
dropout (float): The dropout rate to apply after activation. |
|
time_embed_dim (int): The dimension of the time embedding. |
|
""" |
|
def __init__(self, latent_dim: int, ffn_dim: int, dropout: float, time_embed_dim: int): |
|
super().__init__() |
|
self.latent_dim = latent_dim |
|
self.linear1 = nn.Linear(latent_dim, ffn_dim) |
|
self.linear2 = zero_module(nn.Linear(ffn_dim, latent_dim)) |
|
self.activation = nn.GELU() |
|
self.dropout = nn.Dropout(dropout) |
|
self.proj_out = StylizationBlock(latent_dim, time_embed_dim, dropout) |
|
|
|
def forward(self, x: torch.Tensor, emb: torch.Tensor, **kwargs) -> torch.Tensor: |
|
""" |
|
Forward pass of the FFN layer. |
|
|
|
Args: |
|
x (torch.Tensor): Input tensor of shape (B, T, latent_dim*2). |
|
emb (torch.Tensor): Time embedding tensor. |
|
|
|
Returns: |
|
torch.Tensor: Output tensor after FFN and stylization block. |
|
""" |
|
x1 = x[:, :, :self.latent_dim].contiguous() |
|
x2 = x[:, :, self.latent_dim:].contiguous() |
|
y1 = self.linear2(self.dropout(self.activation(self.linear1(x1)))) |
|
y1 = x1 + self.proj_out(y1, emb) |
|
y2 = self.linear2(self.dropout(self.activation(self.linear1(x2)))) |
|
y2 = x2 + self.proj_out(y2, emb) |
|
y = torch.cat((y1, y2), dim=-1) |
|
return y |
|
|
|
|
|
class DecoderLayer(nn.Module): |
|
""" |
|
A single decoder layer consisting of a cross-attention block and a feed-forward network (FFN). |
|
|
|
Args: |
|
ca_block_cfg (Optional[dict]): Configuration for the cross-attention block. |
|
ffn_cfg (Optional[dict]): Configuration for the feed-forward network. |
|
""" |
|
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 = FFN(**ffn_cfg) |
|
|
|
def forward(self, **kwargs) -> torch.Tensor: |
|
""" |
|
Forward pass of the decoder layer. |
|
|
|
Args: |
|
**kwargs: Arguments passed to the cross-attention and FFN layers. |
|
|
|
Returns: |
|
torch.Tensor: Output tensor after passing through the 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 MoMatMoGenTransformer(ReMoDiffuseTransformer): |
|
""" |
|
MoMatMoGenTransformer is a motion generation transformer model, which uses ReMoDiffuse as a base. |
|
|
|
Args: |
|
ReMoDiffuseTransformer: Base transformer class. |
|
""" |
|
def build_temporal_blocks(self, sa_block_cfg: Optional[dict], ca_block_cfg: Optional[dict], ffn_cfg: Optional[dict]): |
|
""" |
|
Build temporal decoder blocks using the provided configurations. |
|
|
|
Args: |
|
sa_block_cfg (Optional[dict]): Self-attention block configuration. |
|
ca_block_cfg (Optional[dict]): Cross-attention block configuration. |
|
ffn_cfg (Optional[dict]): Feed-forward network configuration. |
|
""" |
|
self.temporal_decoder_blocks = nn.ModuleList() |
|
for i in range(self.num_layers): |
|
self.temporal_decoder_blocks.append( |
|
DecoderLayer(ca_block_cfg=ca_block_cfg, ffn_cfg=ffn_cfg)) |
|
|
|
def forward(self, |
|
motion: torch.Tensor, |
|
timesteps: torch.Tensor, |
|
motion_mask: Optional[torch.Tensor] = None, |
|
**kwargs) -> torch.Tensor: |
|
""" |
|
Forward pass for motion generation. |
|
|
|
Args: |
|
motion (torch.Tensor): Input motion tensor of shape (B, T, D). |
|
timesteps (torch.Tensor): Timestep embeddings. |
|
motion_mask (Optional[torch.Tensor]): Motion mask, if any. |
|
|
|
Returns: |
|
torch.Tensor: Output tensor after processing the motion data. |
|
""" |
|
T = motion.shape[1] |
|
conditions = self.get_precompute_condition(device=motion.device, |
|
**kwargs) |
|
if len(motion_mask.shape) == 2: |
|
src_mask = motion_mask.clone().unsqueeze(-1) |
|
else: |
|
src_mask = motion_mask.clone() |
|
|
|
if self.time_embedding_type == 'sinusoidal': |
|
emb = self.time_embed( |
|
timestep_embedding(timesteps, self.latent_dim)) |
|
else: |
|
emb = self.time_embed(self.time_tokens(timesteps)) |
|
|
|
if self.use_text_proj: |
|
emb = emb + conditions['xf_proj'] |
|
|
|
motion1 = motion[:, :, :self.input_feats].contiguous() |
|
motion2 = motion[:, :, self.input_feats:].contiguous() |
|
h1 = self.joint_embed(motion1) |
|
h2 = self.joint_embed(motion2) |
|
if self.use_pos_embedding: |
|
h1 = h1 + self.sequence_embedding.unsqueeze(0)[:, :T, :] |
|
h2 = h2 + self.sequence_embedding.unsqueeze(0)[:, :T, :] |
|
h = torch.cat((h1, h2), dim=-1) |
|
|
|
if self.training: |
|
output = self.forward_train(h=h, |
|
src_mask=src_mask, |
|
emb=emb, |
|
timesteps=timesteps, |
|
**conditions) |
|
else: |
|
output = self.forward_test(h=h, |
|
src_mask=src_mask, |
|
emb=emb, |
|
timesteps=timesteps, |
|
**conditions) |
|
if self.use_residual_connection: |
|
output = motion + output |
|
return output |
|
|
|
def forward_train(self, |
|
h: Optional[torch.Tensor] = None, |
|
src_mask: Optional[torch.Tensor] = None, |
|
emb: Optional[torch.Tensor] = None, |
|
xf_out: Optional[torch.Tensor] = None, |
|
re_dict: Optional[dict] = None, |
|
**kwargs) -> torch.Tensor: |
|
""" |
|
Training forward pass for the motion generation transformer. |
|
|
|
Args: |
|
h (Optional[torch.Tensor]): Input tensor. |
|
src_mask (Optional[torch.Tensor]): Source mask. |
|
emb (Optional[torch.Tensor]): Embedding tensor. |
|
xf_out (Optional[torch.Tensor]): Output of the cross-attention block. |
|
re_dict (Optional[dict]): Dictionary for recurrent features. |
|
|
|
Returns: |
|
torch.Tensor: Output tensor after processing. |
|
""" |
|
B, T = h.shape[0], h.shape[1] |
|
cond_type = 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, |
|
re_dict=re_dict) |
|
|
|
out1 = self.out(h[:, :, :self.latent_dim].contiguous()) |
|
out1 = out1.view(B, T, -1).contiguous() |
|
out2 = self.out(h[:, :, self.latent_dim:].contiguous()) |
|
out2 = out2.view(B, T, -1).contiguous() |
|
output = torch.cat((out1, out2), dim=-1) |
|
return output |
|
|
|
def forward_test(self, |
|
h: Optional[torch.Tensor] = None, |
|
src_mask: Optional[torch.Tensor] = None, |
|
emb: Optional[torch.Tensor] = None, |
|
xf_out: Optional[torch.Tensor] = None, |
|
re_dict: Optional[dict] = None, |
|
timesteps: Optional[torch.Tensor] = None, |
|
**kwargs) -> torch.Tensor: |
|
""" |
|
Testing forward pass for the motion generation transformer. |
|
|
|
Args: |
|
h (Optional[torch.Tensor]): Input tensor. |
|
src_mask (Optional[torch.Tensor]): Source mask. |
|
emb (Optional[torch.Tensor]): Embedding tensor. |
|
xf_out (Optional[torch.Tensor]): Output of the cross-attention block. |
|
re_dict (Optional[dict]): Dictionary for recurrent features. |
|
timesteps (Optional[torch.Tensor]): Timestep embeddings. |
|
|
|
Returns: |
|
torch.Tensor: Output tensor after processing. |
|
""" |
|
B, T = h.shape[0], h.shape[1] |
|
both_cond_type = torch.zeros(B, 1, 1).to(h.device) + 99 |
|
text_cond_type = torch.zeros(B, 1, 1).to(h.device) + 1 |
|
retr_cond_type = torch.zeros(B, 1, 1).to(h.device) + 10 |
|
none_cond_type = torch.zeros(B, 1, 1).to(h.device) |
|
|
|
all_cond_type = torch.cat( |
|
(both_cond_type, text_cond_type, retr_cond_type, none_cond_type), |
|
dim=0) |
|
h = h.repeat(4, 1, 1) |
|
xf_out = xf_out.repeat(4, 1, 1) |
|
emb = emb.repeat(4, 1) |
|
src_mask = src_mask.repeat(4, 1, 1) |
|
if re_dict['re_motion'].shape[0] != h.shape[0]: |
|
re_dict['re_motion'] = re_dict['re_motion'].repeat(4, 1, 1, 1) |
|
re_dict['re_text'] = re_dict['re_text'].repeat(4, 1, 1, 1) |
|
re_dict['re_mask'] = re_dict['re_mask'].repeat(4, 1, 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, |
|
re_dict=re_dict) |
|
|
|
out1 = self.out(h[:, :, :self.latent_dim].contiguous()) |
|
out1 = out1.view(4 * B, T, -1).contiguous() |
|
out2 = self.out(h[:, :, self.latent_dim:].contiguous()) |
|
out2 = out2.view(4 * B, T, -1).contiguous() |
|
out = torch.cat((out1, out2), dim=-1) |
|
out_both = out[:B].contiguous() |
|
out_text = out[B:2 * B].contiguous() |
|
out_retr = out[2 * B:3 * B].contiguous() |
|
out_none = out[3 * B:].contiguous() |
|
|
|
coef_cfg = self.scale_func(int(timesteps[0])) |
|
both_coef = coef_cfg['both_coef'] |
|
text_coef = coef_cfg['text_coef'] |
|
retr_coef = coef_cfg['retr_coef'] |
|
none_coef = coef_cfg['none_coef'] |
|
output = out_both * both_coef |
|
output += out_text * text_coef |
|
output += out_retr * retr_coef |
|
output += out_none * none_coef |
|
return output |
|
|
