src/diffusion/model/mdm.py

from typing import Tuple, List
import numpy as np
import torch
import torch.nn as nn
from einops import rearrange
from transformers import CLIPProcessor, CLIPModel, T5Tokenizer, T5EncoderModel

from ...constant import TEXT_MODEL_DIMS
from ..metric.interaction import intra_tip_pairs, all_tips

class MotionDiffusionModel(nn.Module):
    def __init__(
        self, arch, latent_dim, num_heads, ff_size, dropout, activation, num_layers,
        njoints, nfeats,
        cond_mode, cond_mask_prob,
        treble_mask_prob=1.0,  # New parameter: probability to keep each treble branch (1.0 = no masking)
        *args, **kwargs
    ):
        super(MotionDiffusionModel, self).__init__()
        self.arch = arch
        self.latent_dim = latent_dim
        self.num_heads = num_heads
        self.ff_size = ff_size
        self.dropout = dropout
        self.activation = activation
        self.num_layers = num_layers
        self.njoints = njoints
        self.nfeats = nfeats

        self.cond_mode = cond_mode
        if self.cond_mode != 'no_cond':
            self.cond_mask_prob = cond_mask_prob
            if self.cond_mask_prob < 0 or self.cond_mask_prob > 1:
                raise ValueError(f"cond_mask_prob should be in [0, 1], but got {self.cond_mask_prob}")

        # Treble masking probability for trans_dec_treble_residual
        self.treble_mask_prob = treble_mask_prob
        if self.treble_mask_prob < 0 or self.treble_mask_prob > 1:
            raise ValueError(f"treble_mask_prob should be in [0, 1], but got {self.treble_mask_prob}")

        self.input_process = InputProcess(
            input_feats=njoints * nfeats,
            latent_dim=latent_dim
        )
        self.sequence_pos_encoder = PositionalEncoding(
            self.latent_dim,
            self.dropout
        )

        if self.cond_mode == 'text':
            self.text_model_name = kwargs['text_model_name']
            self.text_max_length = kwargs['text_max_length']
            if self.text_model_name.startswith("t5"):
                self.text_tokenizer = T5Tokenizer.from_pretrained(self.text_model_name)
                self._text_model = T5EncoderModel.from_pretrained(self.text_model_name)

                if kwargs.get("finetune_text_model", False):

                    if self.arch == 'trans_dec_treble_concat':
                        special_tokens_to_add = ['[LEFT]', '[RIGHT]', '[TWO_HANDS_RELATION]']
                        self.text_tokenizer.add_special_tokens({'additional_special_tokens': special_tokens_to_add})
                        self._text_model.resize_token_embeddings(len(self.text_tokenizer))

                    lora_config = LoraConfig(
                        r=8,
                        lora_alpha=32,
                        target_modules=["q", "v"],
                        lora_dropout=0.1,
                        bias="none",
                    )

                    self._text_model = get_peft_model(self._text_model, lora_config)
                    if self.arch == 'trans_dec_treble_concat':
                        for name, param in self._text_model.named_parameters():
                            # T5的词嵌入层名为 'shared'
                            if 'shared' in name:
                                param.requires_grad = True

                else:
                    for param in self._text_model.parameters():
                        param.requires_grad = False

            elif 'clip' in self.text_model_name:
                self.text_processor = CLIPProcessor.from_pretrained(self.text_model_name, local_files_only=True)
                self._text_model = CLIPModel.from_pretrained(self.text_model_name, use_safetensors=True, local_files_only=True)
                for param in self._text_model.parameters():
                    param.requires_grad = False
            else:
                raise NotImplementedError(f"Text model {self.text_model_name} is not implemented.")

            self.text_embedding_project = nn.Linear(
                TEXT_MODEL_DIMS[self.text_model_name], self.latent_dim
            )
        elif self.cond_mode == 'action':
            self.embed_action = EmbedAction(kwargs['num_actions'], self.latent_dim)


        if self.arch == 'trans_enc':
            print("Transformer Encoder initialize.")
            seq_trans_encoder_layer = nn.TransformerEncoderLayer(
                d_model=self.latent_dim,
                nhead=self.num_heads,
                dim_feedforward=self.ff_size,
                dropout=self.dropout,
                activation=self.activation,
                batch_first=True,
            )
            self.seq_trans_encoder = nn.TransformerEncoder(
                seq_trans_encoder_layer,
                num_layers=self.num_layers
            )
        elif self.arch.startswith('trans_dec'):
            print("Transformer Decoder initialize.")
            seq_trans_decoder_layer = nn.TransformerDecoderLayer(
                d_model=self.latent_dim,
                nhead=self.num_heads,
                dim_feedforward=self.ff_size,
                dropout=self.dropout,
                activation=self.activation,
                batch_first=True,
            )
            self.seq_trans_decoder = nn.TransformerDecoder(
                seq_trans_decoder_layer,
                num_layers=self.num_layers
            )
            self.null_text_embedding = nn.Parameter(torch.randn(1, 1, self.latent_dim)) # (1, 1, D)
            if self.arch.startswith('trans_dec_treble'):
                self.left_hand_cls_token = nn.Parameter(torch.randn(1, 1, self.latent_dim)) # (1, 1, D)
                self.right_hand_cls_token = nn.Parameter(torch.randn(1, 1, self.latent_dim)) # (1, 1, D)
                self.two_hands_relation_cls_token = nn.Parameter(torch.randn(1, 1, self.latent_dim)) # (1, 1, D)
        else:
            raise NotImplementedError(f"Architecture {self.arch} is not implemented.")

        self.contact_prediction = kwargs.get("contact_prediction", False)
        if self.contact_prediction:
            contact_predict_decoder_layer = nn.TransformerDecoderLayer(
                d_model=self.latent_dim,
                nhead=self.num_heads,
                dim_feedforward=self.ff_size,
                dropout=self.dropout,
                activation=self.activation,
                batch_first=True
            )
            self.contact_predict_decoder = nn.TransformerDecoder(
                contact_predict_decoder_layer,
                num_layers=self.num_layers
            )
            self.contact_predict_head = nn.Linear(self.latent_dim, (len(all_tips) + len(intra_tip_pairs)) * 2 + 1)

        self.embed_timestep = TimestepEmbedder(
            latent_dim=self.latent_dim,
            positional_encode=self.sequence_pos_encoder.pe.squeeze(0) # (max_len, D)
        )

        self.output_process = OutputProcess(
            latent_dim=self.latent_dim,
            njoints=self.njoints,
            nfeats=self.nfeats
        )

        self.apply(self._init_weights)

    def _init_weights(self, module:nn.Module):
        if isinstance(module, nn.Linear):
            std = 0.02
            torch.nn.init.normal_(module.weight, mean=0.0, std=std)
            if module.bias is not None:
                torch.nn.init.zeros_(module.bias)

    def get_cond_mask(self, batch_size, device):
        if self.training and self.cond_mask_prob > 0.:
            mask = torch.bernoulli(
                torch.ones(batch_size, device=device) * self.cond_mask_prob
            ) # (B,)
            return (1 - mask).bool() # (B,)
        else:
            return torch.ones(batch_size, device=device).bool() # (B,)

    def get_text_embeddings(self, texts : List[str]) -> torch.Tensor | Tuple[torch.Tensor, torch.Tensor]:
        device = next(self.parameters()).device
        if self.text_model_name.startswith("t5"):
            text_inputs = self.text_tokenizer(
                texts, padding=True, truncation=True, max_length=self.text_max_length,
                return_tensors='pt'
            )
            text_inputs = {key: value.to(device) for key, value in text_inputs.items()}
            outputs = self._text_model(**text_inputs)
            return self.text_embedding_project(outputs.last_hidden_state), text_inputs['attention_mask']
        elif 'clip' in self.text_model_name:
            inputs = self.text_processor(
                texts, return_tensors='pt', padding=True, truncation=True, max_length=self.text_max_length
            )
            inputs = {key: value.to(device) for key, value in inputs.items()}
            text_features = self._text_model.get_text_features(**inputs)
            return self.text_embedding_project(text_features)
        else:
            raise NotImplementedError(f"Text model {self.text_model_name} is not implemented.")

    def cross_attention_w_single_text(self, motion_embedding: torch.Tensor, single_text_list: List[str], cls_tokens: List[torch.Tensor], decoder: nn.TransformerDecoder, y_lengths: torch.Tensor) -> torch.Tensor:
        B, T, _ = motion_embedding.shape
        if not self.text_model_name.startswith("t5"):
            raise NotImplementedError(f"Text model {self.text_model_name} is not implemented for cross attention with single text.")
        text_embeddings, text_padding_mask = self.get_text_embeddings(single_text_list) # (B, L_Text, D), (B, L_Text)
        if self.training:
            condition_mask = self.get_cond_mask(B, motion_embedding.device) # (B,)
            text_embeddings = torch.where(
                condition_mask.view(B, 1, 1), text_embeddings, self.null_text_embedding
            )
            text_padding_mask = torch.where(
                condition_mask.view(B, 1), text_padding_mask, torch.ones((1, 1), device=motion_embedding.device)
            )

        text_padding_mask = (text_padding_mask == 0) # True means this position SHOULD be masked.

        xseq = torch.cat([*cls_tokens, motion_embedding], dim=1) # (B, len(cls_tokens)+T, D)
        xseq = self.sequence_pos_encoder(xseq) # (B, len(cls_tokens)+T, D)
        tgt_key_padding_mask = torch.arange(len(cls_tokens) + T, device=motion_embedding.device)[None, :] - len(cls_tokens) >= y_lengths[:, None] # (B, len(cls_tokens)+T)
        output = decoder(
            tgt=xseq, # (B, len(cls_tokens)+T, D)
            memory=text_embeddings, # (B, L_Text, D)
            tgt_key_padding_mask=tgt_key_padding_mask, # (B, len(cls_tokens)+T)
            memory_key_padding_mask=text_padding_mask, # (B, L_Text)
        )[:, len(cls_tokens):] # (B, T, D)
        return output

    def _forward(self, x, time_embedder_token, y:dict=None, decoder:nn.TransformerDecoder=None, *arg, **kwargs):
        '''
        x: (B, T, D)
        time_embedder_token: (B, 1, D)
        '''
        B, T, _ = x.shape


        if self.arch == 'trans_enc': # Deprecated
            raise DeprecationWarning("Transformer Encoder architecture is deprecated. Please use Transformer Decoder architectures.")
            if 'text' in self.cond_mode and not y.get("uncond", False):
                if 'clip' in self.text_model_name:
                    text_embeddings = self.get_text_embeddings(y['text']) # (B, D)
                    if self.training:
                        condition_mask = self.get_cond_mask(B, x.device) # (B,)
                        text_embeddings = text_embeddings * condition_mask[:, None] # (B, D)
                    time_embedder_token = time_embedder_token + text_embeddings[:, None, :] # (B, 1, D)
                else:
                    raise NotImplementedError(f"Text model {self.text_model_name} is not implemented for Transformer Encoder.")
            elif 'action' in self.cond_mode and not y.get("uncond", False):
                action_embeddings = self.embed_action(y['action']) # (B, D)
                if self.training:
                    condition_mask = self.get_cond_mask(B, x.device) # (B,)
                    action_embeddings = action_embeddings * condition_mask[:, None] # (B, D)
                time_embedder_token = time_embedder_token + action_embeddings[:, None, :] # (B, 1, D)

            xseq = torch.cat([time_embedder_token, x], dim=1) # (B, T+1, D)
            xseq = self.sequence_pos_encoder(xseq) # (B, T+1, D)
            src_key_padding_mask = torch.arange(T + 1, device=x.device)[None, :] - 1 >= y['lengths'][:, None] # (B, T+1)
            output = self.seq_trans_encoder(
                xseq, src_key_padding_mask=src_key_padding_mask
            )[:, 1:] # (B, T, D)

        elif self.arch == 'trans_dec':
            assert 'text' in self.cond_mode, "Transformer Decoder requires text condition."
            output = self.cross_attention_w_single_text(x, y['text'], [time_embedder_token], decoder=decoder, y_lengths=y['lengths']) # (B, T, D)

        elif self.arch.startswith('trans_dec_treble'):
            assert 'text' in self.cond_mode, "Transformer Decoder with treble residual requires text condition."
            assert 'text' in y and 'left' in y['text'] and 'right' in y['text'] and 'two_hands_relation' in y['text'], "Transformer Decoder with treble residual requires 'left', 'right' and 'two_hands_relation' text conditions."
            if self.arch == 'trans_dec_treble_residual':
                # Random masking: each branch has treble_mask_prob probability to be kept
                if self.training:
                    left_keep = torch.bernoulli(torch.ones(B, device=x.device) * self.treble_mask_prob).bool().cpu().tolist()  # (B,)
                    right_keep = torch.bernoulli(torch.ones(B, device=x.device) * self.treble_mask_prob).bool().cpu().tolist()  # (B,)
                    relation_keep = torch.bernoulli(torch.ones(B, device=x.device) * self.treble_mask_prob).bool().cpu().tolist()  # (B,)

                    masked_left_text = [text if keep else "" for text, keep in zip(y['text']['left'], left_keep)]
                    masked_right_text = [text if keep else "" for text, keep in zip(y['text']['right'], right_keep)]
                    masked_relation_text = [text if keep else "" for text, keep in zip(y['text']['two_hands_relation'], relation_keep)]
                else:
                    masked_left_text = y['text']['left']
                    masked_right_text = y['text']['right']
                    masked_relation_text = y['text']['two_hands_relation']

                left_hand_output = self.cross_attention_w_single_text(x, masked_left_text, [time_embedder_token, self.left_hand_cls_token.expand(B, -1, -1)], decoder=decoder, y_lengths=y['lengths']) # (B, T, D)
                right_hand_output = self.cross_attention_w_single_text(x, masked_right_text, [time_embedder_token, self.right_hand_cls_token.expand(B, -1, -1)], decoder=decoder, y_lengths=y['lengths']) # (B, T, D)
                two_hands_relation_output = self.cross_attention_w_single_text(x, masked_relation_text, [time_embedder_token, self.two_hands_relation_cls_token.expand(B, -1, -1)], decoder=decoder, y_lengths=y['lengths']) # (B, T, D)

                output = x + left_hand_output + right_hand_output + two_hands_relation_output # (B, T, D)

            elif self.arch == 'trans_dec_treble_concat':
                concated_texts = []
                for left_text, right_text, two_hands_relation_text in zip(y['text']['left'], y['text']['right'], y['text']['two_hands_relation']):
                    concated_texts.append('[LEFT] ' + left_text + ' [RIGHT] ' + right_text + ' [TWO_HANDS_RELATION] ' + two_hands_relation_text)

                output = self.cross_attention_w_single_text(x, concated_texts, [time_embedder_token], decoder=decoder, y_lengths=y['lengths']) # (B, T, D)


        else:
            raise NotImplementedError(f"Architecture {self.arch} is not implemented.")

        return output

    def forward(self, x, timesteps, y:dict=None, predict_contact=False, *arg, **kwargs):
        cls_token = self.embed_timestep(timesteps) # (B, 1, D)
        if not predict_contact:
            x = self.input_process(x) # (B, T, D)
            output = self._forward(x, cls_token, y=y, decoder=self.seq_trans_decoder, *arg, **kwargs)
            output = self.output_process(output) # (B, J, feats_per_joint, T)
            return output
        else:
            motion_embedding = self.input_process(x.clone().detach())
            output = self._forward(motion_embedding, cls_token.clone().detach(), y=y, decoder=self.contact_predict_decoder, *arg, **kwargs) # (B, T, D)
            contact_logits = self.contact_predict_head(output) # (B, T, P)
            return contact_logits


class InputProcess(nn.Module):
    def __init__(self, input_feats, latent_dim):
        '''
        input_feats: num_joints * feat_dim_per_joint
        '''
        super(InputProcess, self).__init__()
        self.input_feats = input_feats
        self.latent_dim = latent_dim
        self.pose_embedding = nn.Linear(input_feats, latent_dim)

    def forward(self, x):
        '''
        x: (B, J, D, T)
        '''
        B, J, D, T = x.shape
        x = rearrange(x, 'b j d t -> b t (j d)') # (B, T, J*D)
        x = self.pose_embedding(x) # (B, T, D)
        return x

class OutputProcess(nn.Module):
    def __init__(self, latent_dim, njoints, nfeats):
        '''
        nfeats: feature dimension per joint
        '''
        super(OutputProcess, self).__init__()
        self.latent_dim = latent_dim
        self.njoints = njoints
        self.nfeats = nfeats
        self.pose_head = nn.Linear(self.latent_dim, self.njoints * self.nfeats)

    def forward(self, x):
        '''
        x: (B, T, D)
        '''
        B, T, D = x.shape
        x = self.pose_head(x)
        ret = rearrange(x, 'b t (j d) -> b j d t', j=self.njoints, d=self.nfeats) # (B, J, D, T)
        return ret

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0) # (1, max_len, d_model)

        self.register_buffer('pe', pe) # 不会计算梯度，但是会保存到state_dict中，移动设备时也会随模型一起移动

    def forward(self, x):
        '''
        x: (B, T, D)
        '''
        x = x + self.pe[:, :x.shape[1]]
        return self.dropout(x)

class TimestepEmbedder(nn.Module):
    def __init__(self, latent_dim, positional_encode:torch.Tensor):
        '''
        positional_encode: (max_len, D)
        '''
        super(TimestepEmbedder, self).__init__()
        self.latent_dim = latent_dim

        self.time_embedding = nn.Sequential(
            nn.Linear(self.latent_dim, self.latent_dim),
            nn.SiLU(),
            nn.Linear(self.latent_dim, self.latent_dim)
        )

        self.register_buffer(
            'positional_encode',
            positional_encode
        )

    def forward(self, timesteps) -> torch.Tensor:
        '''
        timesteps: (B,)
        '''
        return self.time_embedding(
            self.positional_encode[timesteps] # (B, D)
        ).unsqueeze(1) # (B, 1, D)


class EmbedAction(nn.Module):
    def __init__(self, num_actions, latent_dim):
        super(EmbedAction, self).__init__()
        self.action_embedding = nn.Parameter(torch.randn(num_actions, latent_dim))

    def forward(self, input:torch.Tensor):
        '''
        input: (B,)
        '''
        idx = input.long()
        output = self.action_embedding[idx]
        return output