import random import clip import numpy as np import torch import torch.nn.functional as F from torch import nn from torch import Tensor from typing import List, Dict, Optional, Union from mogen.models.utils.misc import zero_module from ..builder import SUBMODULES, build_attention from .motion_transformer import MotionTransformer class FFN(nn.Module): """ Feed-forward network (FFN) used in the transformer layers. It consists of two linear layers with a GELU activation in between. Args: latent_dim (int): Input dimension of the FFN. ffn_dim (int): Hidden dimension of the FFN. dropout (float): Dropout rate applied after activation. """ def __init__(self, latent_dim: int, ffn_dim: int, dropout: float): super().__init__() 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) def forward(self, x: Tensor, **kwargs) -> Tensor: """ Forward pass for the FFN. Args: x (Tensor): Input tensor of shape (B, T, D). Returns: Tensor: Output tensor after the FFN, of shape (B, T, D). """ y = self.linear2(self.dropout(self.activation(self.linear1(x)))) y = x + y return y class EncoderLayer(nn.Module): """ Encoder layer consisting of self-attention and feed-forward network. Args: sa_block_cfg (Optional[dict]): Configuration for the self-attention block. ca_block_cfg (Optional[dict]): Configuration for the cross-attention block (if applicable). ffn_cfg (dict): Configuration for the feed-forward network. """ def __init__(self, sa_block_cfg: Optional[dict] = None, ca_block_cfg: Optional[dict] = None, ffn_cfg: dict = None): super().__init__() self.sa_block = build_attention(sa_block_cfg) self.ffn = FFN(**ffn_cfg) def forward(self, **kwargs) -> Tensor: """ Forward pass for the encoder layer. Args: kwargs: Dictionary containing the input tensor (x) and other related parameters. Returns: Tensor: Output tensor after the encoder layer. """ if self.sa_block is not None: x = self.sa_block(**kwargs) kwargs.update({'x': x}) if self.ffn is not None: x = self.ffn(**kwargs) return x class RetrievalDatabase(nn.Module): """ Retrieval database for retrieving motions and text features based on given captions. Args: num_retrieval (int): Number of retrievals for each caption. topk (int): Number of top results to consider. retrieval_file (str): Path to the retrieval file containing text, motion, and length data. latent_dim (Optional[int]): Dimension of the latent space. output_dim (Optional[int]): Output dimension of the retrieved features. num_layers (Optional[int]): Number of layers in the text encoder. num_motion_layers (Optional[int]): Number of layers in the motion encoder. kinematic_coef (Optional[float]): Coefficient for scaling kinematic similarity. max_seq_len (Optional[int]): Maximum sequence length. num_heads (Optional[int]): Number of attention heads. ff_size (Optional[int]): Feed-forward size for the transformer layers. stride (Optional[int]): Stride for downsampling motion data. sa_block_cfg (Optional[dict]): Configuration for the self-attention block. ffn_cfg (Optional[dict]): Configuration for the feed-forward network. dropout (Optional[float]): Dropout rate. """ def __init__(self, num_retrieval: int, topk: int, retrieval_file: str, latent_dim: Optional[int] = 512, output_dim: Optional[int] = 512, num_layers: Optional[int] = 2, num_motion_layers: Optional[int] = 4, kinematic_coef: Optional[float] = 0.1, max_seq_len: Optional[int] = 196, num_heads: Optional[int] = 8, ff_size: Optional[int] = 1024, stride: Optional[int] = 4, sa_block_cfg: Optional[dict] = None, ffn_cfg: Optional[dict] = None, dropout: Optional[float] = 0): super().__init__() self.num_retrieval = num_retrieval self.topk = topk self.latent_dim = latent_dim self.stride = stride self.kinematic_coef = kinematic_coef self.num_layers = num_layers self.num_motion_layers = num_motion_layers self.max_seq_len = max_seq_len # Load data from the retrieval file data = np.load(retrieval_file) self.text_features = torch.Tensor(data['text_features']) self.captions = data['captions'] self.motions = data['motions'] self.m_lengths = data['m_lengths'] self.clip_seq_features = data['clip_seq_features'] self.train_indexes = data.get('train_indexes', None) self.test_indexes = data.get('test_indexes', None) self.latent_dim = latent_dim self.output_dim = output_dim self.motion_proj = nn.Linear(self.motions.shape[-1], self.latent_dim) self.motion_pos_embedding = nn.Parameter( torch.randn(max_seq_len, self.latent_dim)) self.motion_encoder_blocks = nn.ModuleList() # Build motion encoder blocks for i in range(num_motion_layers): self.motion_encoder_blocks.append( EncoderLayer(sa_block_cfg=sa_block_cfg, ffn_cfg=ffn_cfg)) # Transformer for encoding text TransEncoderLayer = nn.TransformerEncoderLayer(d_model=self.latent_dim, nhead=num_heads, dim_feedforward=ff_size, dropout=dropout, activation="gelu") self.text_encoder = nn.TransformerEncoder(TransEncoderLayer, num_layers=num_layers) self.results = {} def extract_text_feature(self, text: str, clip_model: nn.Module, device: torch.device) -> Tensor: """ Extract text features from CLIP model. Args: text (str): Input text caption. clip_model (nn.Module): CLIP model for encoding the text. device (torch.device): Device for computation. Returns: Tensor: Extracted text features of shape (1, 512). """ text = clip.tokenize([text], truncate=True).to(device) with torch.no_grad(): text_features = clip_model.encode_text(text) return text_features def encode_text(self, text: List[str], device: torch.device) -> Tensor: """ Encode text using the CLIP model's text encoder. Args: text (List[str]): List of input text captions. device (torch.device): Device for computation. Returns: Tensor: Encoded text features of shape (B, T, D). """ with torch.no_grad(): text = clip.tokenize(text, truncate=True).to(device) x = self.clip.token_embedding(text).type(self.clip.dtype) x = x + self.clip.positional_embedding.type(self.clip.dtype) x = x.permute(1, 0, 2) # NLD -> LND x = self.clip.transformer(x) x = self.clip.ln_final(x).type(self.clip.dtype) # B, T, D xf_out = x.permute(1, 0, 2) return xf_out def retrieve(self, caption: str, length: int, clip_model: nn.Module, device: torch.device, idx: Optional[int] = None) -> List[int]: """ Retrieve motions and text features based on a given caption. Args: caption (str): Input text caption. length (int): Length of the corresponding motion sequence. clip_model (nn.Module): CLIP model for encoding the text. device (torch.device): Device for computation. idx (Optional[int]): Index for retrieval (if provided). Returns: List[int]: List of indexes for the retrieved motions. """ value = hash(caption) if value in self.results: return self.results[value] text_feature = self.extract_text_feature(caption, clip_model, device) rel_length = torch.LongTensor(self.m_lengths).to(device) rel_length = torch.abs(rel_length - length) rel_length = rel_length / torch.clamp(rel_length, min=length) semantic_score = F.cosine_similarity(self.text_features.to(device), text_feature) kinematic_score = torch.exp(-rel_length * self.kinematic_coef) score = semantic_score * kinematic_score indexes = torch.argsort(score, descending=True) data = [] cnt = 0 for idx in indexes: caption, m_length = self.captions[idx], self.m_lengths[idx] if not self.training or m_length != length: cnt += 1 data.append(idx.item()) if cnt == self.num_retrieval: self.results[value] = data return data assert False def generate_src_mask(self, T: int, length: List[int]) -> Tensor: """ Generate source mask for the motion sequences based on the motion lengths. Args: T (int): Maximum sequence length. length (List[int]): List of motion lengths for each sample. Returns: Tensor: A binary mask tensor of shape (B, T), where `B` is the batch size, and `T` is the maximum sequence length. Mask values are 1 for valid positions and 0 for padded positions. """ B = len(length) src_mask = torch.ones(B, T) for i in range(B): for j in range(length[i], T): src_mask[i, j] = 0 return src_mask def forward(self, captions: List[str], lengths: List[int], clip_model: nn.Module, device: torch.device, idx: Optional[List[int]] = None) -> Dict[str, Tensor]: """ Forward pass for retrieving motion sequences and text features. Args: captions (List[str]): List of input text captions. lengths (List[int]): List of corresponding motion lengths. clip_model (nn.Module): CLIP model for encoding the text. device (torch.device): Device for computation. idx (Optional[List[int]]): Optional list of indices for retrieval. Returns: Dict[str, Tensor]: Dictionary containing retrieved text and motion features. - re_text: Retrieved text features of shape (B, num_retrieval, T, D). - re_motion: Retrieved motion features of shape (B, num_retrieval, T, D). - re_mask: Source mask for the retrieved motion of shape (B, num_retrieval, T). - raw_motion: Raw motion features of shape (B, T, motion_dim). - raw_motion_length: Motion sequence lengths (before any stride). - raw_motion_mask: Raw binary mask for valid motion positions of shape (B, T). """ B = len(captions) all_indexes = [] for b_ix in range(B): length = int(lengths[b_ix]) if idx is None: batch_indexes = self.retrieve(captions[b_ix], length, clip_model, device) else: batch_indexes = self.retrieve(captions[b_ix], length, clip_model, device, idx[b_ix]) all_indexes.extend(batch_indexes) all_indexes = np.array(all_indexes) all_motions = torch.Tensor(self.motions[all_indexes]).to(device) all_m_lengths = torch.Tensor(self.m_lengths[all_indexes]).long() # Generate masks and positional encodings T = all_motions.shape[1] src_mask = self.generate_src_mask(T, all_m_lengths).to(device) raw_src_mask = src_mask.clone() re_motion = self.motion_proj(all_motions) + self.motion_pos_embedding.unsqueeze(0) for module in self.motion_encoder_blocks: re_motion = module(x=re_motion, src_mask=src_mask.unsqueeze(-1)) re_motion = re_motion.view(B, self.num_retrieval, T, -1).contiguous() re_motion = re_motion[:, :, ::self.stride, :].contiguous() # Apply stride src_mask = src_mask[:, ::self.stride].contiguous() src_mask = src_mask.view(B, self.num_retrieval, -1).contiguous() # Process text sequences T = 77 # CLIP's max token length all_text_seq_features = torch.Tensor(self.clip_seq_features[all_indexes]).to(device) all_text_seq_features = all_text_seq_features.permute(1, 0, 2) re_text = self.text_encoder(all_text_seq_features) re_text = re_text.permute(1, 0, 2) re_text = re_text.view(B, self.num_retrieval, T, -1).contiguous() re_text = re_text[:, :, -1:, :].contiguous() # Use the last token only for each sequence re_dict = { 're_text': re_text, 're_motion': re_motion, 're_mask': src_mask, 'raw_motion': all_motions, 'raw_motion_length': all_m_lengths, 'raw_motion_mask': raw_src_mask } return re_dict @SUBMODULES.register_module() class ReMoDiffuseTransformer(MotionTransformer): """ Transformer model for motion retrieval and diffusion. Args: retrieval_cfg (dict): Configuration for the retrieval database. scale_func_cfg (dict): Configuration for scaling functions. kwargs: Additional arguments for the base DiffusionTransformer. """ def __init__(self, retrieval_cfg: dict, scale_func_cfg: dict, **kwargs): super().__init__(**kwargs) self.database = RetrievalDatabase(**retrieval_cfg) self.scale_func_cfg = scale_func_cfg def scale_func(self, timestep: int) -> Dict[str, float]: """ Scale function for adjusting the guidance between text and retrieval. Args: timestep (int): Current diffusion timestep. Returns: Dict[str, float]: Scaling coefficients for different guidance types. - both_coef: Coefficient for both text and retrieval guidance. - text_coef: Coefficient for text-only guidance. - retr_coef: Coefficient for retrieval-only guidance. - none_coef: Coefficient for no guidance. """ coarse_scale = self.scale_func_cfg['coarse_scale'] w = (1 - (1000 - timestep) / 1000) * coarse_scale + 1 if timestep > 100: if random.randint(0, 1) == 0: output = { 'both_coef': w, 'text_coef': 0, 'retr_coef': 1 - w, 'none_coef': 0 } else: output = { 'both_coef': 0, 'text_coef': w, 'retr_coef': 0, 'none_coef': 1 - w } else: both_coef = self.scale_func_cfg['both_coef'] text_coef = self.scale_func_cfg['text_coef'] retr_coef = self.scale_func_cfg['retr_coef'] none_coef = 1 - both_coef - text_coef - retr_coef output = { 'both_coef': both_coef, 'text_coef': text_coef, 'retr_coef': retr_coef, 'none_coef': none_coef } return output def get_precompute_condition(self, text: Optional[str] = None, motion_length: Optional[Tensor] = None, xf_out: Optional[Tensor] = None, re_dict: Optional[Dict] = None, device: Optional[torch.device] = None, sample_idx: Optional[Tensor] = None, clip_feat: Optional[Tensor] = None, **kwargs) -> Dict[str, Union[Tensor, Dict]]: """ Precompute conditions for both text and retrieval-guided diffusion. Args: text (Optional[str]): Input text string for guidance. motion_length (Optional[Tensor]): Lengths of the motion sequences. xf_out (Optional[Tensor]): Encoded text feature (if precomputed). re_dict (Optional[Dict]): Dictionary of retrieval results (if precomputed). device (Optional[torch.device]): Device to perform computation on. sample_idx (Optional[Tensor]): Sample indices for retrieval. clip_feat (Optional[Tensor]): Clip features (if used). Returns: Dict[str, Union[Tensor, Dict]]: Dictionary containing encoded features and retrieval results. """ if xf_out is None: xf_out = self.encode_text(text, clip_feat, device) output = {'xf_out': xf_out} if re_dict is None: re_dict = self.database(text, motion_length, self.clip, device, idx=sample_idx) output['re_dict'] = re_dict return output def post_process(self, motion: Tensor) -> Tensor: """ Post-process the generated motion by normalizing or un-normalizing it. Args: motion (Tensor): Generated motion data. Returns: Tensor: Post-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: Tensor, src_mask: Tensor, emb: Tensor, xf_out: Optional[Tensor] = None, re_dict: Optional[Dict] = None, **kwargs) -> Tensor: """ Forward training pass for motion retrieval and diffusion model. Args: h (Tensor): Input motion features of shape (B, T, D). src_mask (Tensor): Mask for the motion data of shape (B, T, 1). emb (Tensor): Embedding tensor for timesteps. xf_out (Optional[Tensor]): Precomputed text features. re_dict (Optional[Dict]): Dictionary of retrieval features. Returns: Tensor: Output motion data of shape (B, T, D). """ 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) output = self.out(h).view(B, T, -1).contiguous() return output def forward_test(self, h: Tensor, src_mask: Tensor, emb: Tensor, xf_out: Optional[Tensor] = None, re_dict: Optional[Dict] = None, timesteps: Optional[Tensor] = None, **kwargs) -> Tensor: """ Forward testing pass for motion retrieval and diffusion model. This method handles multiple conditional types such as both text and retrieval-based guidance. Args: h (Tensor): Input motion features of shape (B, T, D). src_mask (Tensor): Mask for the motion data of shape (B, T, 1). emb (Tensor): Embedding tensor for timesteps. xf_out (Optional[Tensor]): Precomputed text features. re_dict (Optional[Dict]): Dictionary of retrieval features. timesteps (Optional[Tensor]): Tensor containing current timesteps in the diffusion process. Returns: Tensor: Output motion data after applying multiple conditional types, of shape (B, T, D). """ B, T = h.shape[0], h.shape[1] # Define condition types for different guidance types 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) # Concatenate all conditional types and repeat inputs for different guidance modes 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) # Repeat retrieval features if necessary 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) # Pass through the temporal decoder blocks 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) # Retrieve output features and handle different guidance coefficients out = self.out(h).view(4 * B, T, -1).contiguous() 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() # Apply scaling coefficients based on the timestep 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'] # Compute the final output by blending the different guidance outputs output = out_both * both_coef output += out_text * text_coef output += out_retr * retr_coef output += out_none * none_coef return output