Blame for src/models/transformers/anymo.py

import torch, math
import torch.nn.functional as F
from mmengine.registry import MODELS
from mmengine.model import BaseModel
from typing import Dict, List
from torch.distributions.categorical import Categorical
import numpy as np
from random import random
from ...utils.motion import recover_from_local_position_fast
from ..tokenizers.encdoc.encdoc_attn import length_to_mask

def cosine_schedule(t):
    return torch.cos(t * math.pi * 0.5)

def get_mask_subset_prob(mask, prob):
    subset_mask = torch.bernoulli(mask, p=prob) & mask
    return subset_mask

def uniform(shape, device=None):
    return torch.zeros(shape, device=device).float().uniform_(0, 1)

def top_k(logits, thres = 0.9, dim = 1):
    k = math.ceil((1 - thres) * logits.shape[dim])
    val, ind = logits.topk(k, dim = dim)
    probs = torch.full_like(logits, float('-inf'))
    probs.scatter_(dim, ind, val)
    return probs

@MODELS.register_module()
class AnyMo(BaseModel):
    def __init__(
        self, 
        vqvae_cfg: Dict,
        transformer_cfg: Dict,

        max_motion_len: int,
        motion_downsample_ratio: int, 
        audio_downsample_ratio: int, 

        text_encoder_cfg: Dict=None,
        traj_encoder_cfg: Dict=None, 
        audio_encoder_cfg: Dict=None,
        
        stage=1, 
        inference_timesteps=6, 
        cond_types = ['text'], 
        **kwargs
    ):
        super().__init__(**kwargs)

        self.transformer = MODELS.build(transformer_cfg)
        # import ipdb; ipdb.set_trace()
        if text_encoder_cfg is not None:
            self.text_encoder = MODELS.build(text_encoder_cfg)
        if traj_encoder_cfg is not None:
            self.traj_encoder = MODELS.build(traj_encoder_cfg)
        if audio_encoder_cfg is not None:
            self.audio_encoder = MODELS.build(audio_encoder_cfg)
        
        # 第一阶段tm训练，第二阶段只训练projector，第三阶段全量训练, 4阶段测评
        assert stage in [1, 2, 3, 4]
        self.stage = stage

        # build vqvae
        self.vqvae = MODELS.build(vqvae_cfg)
        self.vqvae.init_weights()
        for param in self.vqvae.parameters():
            param.requires_grad = False
        self.vqvae.eval()

        self.vocab_size = transformer_cfg.vocab_size # type: ignore
        self.n_embd = transformer_cfg.n_embd
        self.mask_idx = self.vocab_size - 1
        self.pad_idx = self.vocab_size - 2
        self.ignore_index = -100

        self.max_motion_len = max_motion_len
        self.motion_downsample_ratio = motion_downsample_ratio
        self.audio_downsample_ratio = audio_downsample_ratio
        self.topk_filter_thres = 0.9
        self.noise_schedule = cosine_schedule
        self.n_quant = len(self.vqvae.quantizer.quantizers)
        self.inference_timesteps = inference_timesteps
        self.cond_types = cond_types
        for cond_type in cond_types:
            assert cond_type in ['text', 'audio', 'traj']

        if self.stage == 2:
            self.init_weights()
            # only train audio and traj projector
            for param in self.transformer.parameters():
                param.requires_grad = False
            for param in self.text_encoder.parameters():
                param.requires_grad = False

    def forward_stage_1(self, motion, motion_length, text_data, mode='loss'):
        '''
        第一阶段只训t2m
        '''
        B = motion.shape[0]
        # 1. prepare text feature
        # text feat: B, text_token_cnt, C
        # text mask: B, text_token_cnt
        # import ipdb; ipdb.set_trace()
        text_feat, text_mask = self.text_encoder(text_data)
        audio_feat = torch.zeros(B, 0, self.n_embd).to(text_feat)
        audio_mask = torch.zeros(B, 0).to(text_feat)
        traj_feat = torch.zeros(B, 0, self.n_embd).to(text_feat)
        traj_mask = torch.zeros(B, 0).to(text_feat)
        # import ipdb; ipdb.set_trace()

        if mode == 'loss':
            tokens, label_tokens, non_pad_mask = self.prepare_train_input(motion, motion_length)
            loss =  self.forward_loss(tokens, label_tokens, text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask, non_pad_mask)
            return loss
        else:
            non_pad_mask = self.prepare_test_input(motion_length)
            return self.forward_predict(text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask, non_pad_mask)

    def forward_stage_2(self, motion, motion_length, audio_data, traj_data, audio_length, mode='loss', **kwargs):
        B = motion.shape[0]
        # 2. prepare audio and traj feature
        # audio feat: B, audio_token_cnt, C
        # traj feat: B, traj_token_cnt, C

        prob = random()
        # import ipdb; ipdb.set_trace()
        if prob < 0.7: # 百分之70的语音数据
            audio_feat = self.audio_encoder(audio_data.permute(0, 2, 1)).permute(0, 2, 1)
            audio_mask = length_to_mask(audio_length // self.audio_downsample_ratio)
            traj_feat = torch.zeros(B, 0, self.n_embd).to(audio_feat)
            traj_mask = torch.zeros(B, 0).to(audio_feat)
            dtype = audio_feat.dtype
            device = audio_feat.device
        elif prob >= 0.7 and prob < 0.9: # 百分之20的轨迹数据
            traj_feat = self.traj_encoder(traj_data.permute(0, 2, 1)).permute(0, 2, 1)
            traj_mask = length_to_mask(motion_length // self.motion_downsample_ratio)
            audio_feat = torch.zeros(B, 0, self.n_embd).to(traj_feat)
            audio_mask = torch.zeros(B, 0).to(traj_feat)
            dtype = traj_feat.dtype
            device = traj_feat.device
        else:
            audio_feat = self.audio_encoder(audio_data.permute(0, 2, 1)).permute(0, 2, 1)
            audio_mask = length_to_mask(audio_length // self.audio_downsample_ratio)
            traj_feat = self.traj_encoder(traj_data.permute(0, 2, 1)).permute(0, 2, 1)
            traj_mask = length_to_mask(motion_length // self.motion_downsample_ratio)
            dtype = audio_feat.dtype
            device = audio_feat.device

        text_feat = torch.zeros(B, 0, self.n_embd).to(dtype).to(device)
        text_mask = torch.zeros(B, 0).to(dtype).to(device)

        if mode == 'loss':
            tokens, label_tokens, non_pad_mask = self.prepare_train_input(motion, motion_length)
            loss =  self.forward_loss(
                tokens, label_tokens, text_feat, text_mask, audio_feat, 
                audio_mask, traj_feat, traj_mask, non_pad_mask, motion=motion, mean=kwargs['mean'], 
                std=kwargs['std']
            )
            return loss
        else:
            non_pad_mask = self.prepare_test_input(motion_length)
            return self.forward_predict(text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask, non_pad_mask, **kwargs)
    
    def forward_stage_3(self, motion, motion_length, text_data, audio_data, traj_data, audio_length, mode='loss', **kwargs):
        B = motion.shape[0]
        # import ipdb; ipdb.set_trace()
        has_traj = random() < 0.1 # 10%的几率加traj
        has_text_lst = (audio_length == 0) # 没 audio 的加 text
        has_text_lst[has_text_lst == False] = random() < 0.1 # audio 10% 的几率加 text

        text_feat, text_mask = self.text_encoder(text_data)

        traj_feat = self.traj_encoder(traj_data.permute(0, 2, 1)).permute(0, 2, 1)
        traj_mask = length_to_mask(motion_length // self.motion_downsample_ratio)
        traj_mask[:] = False if not has_traj else True

        text_mask = text_mask * has_text_lst[:, None]

        if not torch.all(audio_length == 0):
            audio_feat = self.audio_encoder(audio_data.permute(0, 2, 1)).permute(0, 2, 1)
            audio_mask = length_to_mask(audio_length // self.audio_downsample_ratio)
        else:
            # audio下采样8倍
            audio_data = torch.zeros((B, 8, 512)).to(motion)
            audio_feat = self.audio_encoder(audio_data.permute(0, 2, 1)).permute(0, 2, 1).mean() * 0.0
            audio_mask = torch.zeros(B, 0).to(text_feat.device)

            text_feat += audio_feat * 0.0
            traj_feat += audio_feat * 0.0
            audio_feat = torch.zeros(B, 0, self.n_embd).to(text_feat.device)

        if mode == 'loss':
            tokens, label_tokens, non_pad_mask = self.prepare_train_input(motion, motion_length)
            loss = self.forward_loss(
                tokens, label_tokens, text_feat, text_mask, audio_feat, 
                audio_mask, traj_feat, traj_mask, non_pad_mask, motion=motion, 
                mean=kwargs['mean'], std=kwargs['std']
            )
            return loss
        else:
            # 只测text
            text_feat, text_mask = self.text_encoder(text_data)
            B = text_feat.shape[0]
            audio_feat = torch.zeros(B, 0, self.n_embd).to(device)
            audio_mask = torch.zeros(B, 0).to(device)
            traj_feat = torch.zeros(B, 0, self.n_embd).to(device)
            traj_mask = torch.zeros(B, 0).to(device)

            non_pad_mask = self.prepare_test_input(motion_length)
            return self.forward_predict(text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask, non_pad_mask, **kwargs)
    
    def forward(
            self, 
            motion: torch.Tensor, 
            motion_length: torch.Tensor,

            text_data: List[str]=None,
            audio_data: torch.Tensor=None,
            traj_data: torch.Tensor=None,
            audio_length: torch.Tensor=None,
            
            mode = 'loss', 
            **kwargs
        ): 
        '''
        motion: tensor, B, T, C
        motion_length: tensor, B
        text_data: list[str], len B
        has_xxx: list[bool], len B
        music_data: tensor, B, T, C
        '''
        if self.stage == 1:
            return self.forward_stage_1(motion, motion_length, text_data, mode)
        elif self.stage == 2:
            return self.forward_stage_2(motion, motion_length, audio_data, traj_data, audio_length, mode, **kwargs)
        elif self.stage == 3:
            return self.forward_stage_3(motion, motion_length, text_data, audio_data, traj_data, audio_length, mode, **kwargs)
        elif self.stage == 4: # for evaluation
            B = motion.shape[0]
            device = motion.device
            text_feat = torch.zeros(B, 0, self.n_embd).to(device)
            audio_feat = torch.zeros(B, 0, self.n_embd).to(device)
            traj_feat = torch.zeros(B, 0, self.n_embd).to(device)
            text_mask = torch.zeros(B, 0).to(device)
            audio_mask = torch.zeros(B, 0).to(device)
            traj_mask = torch.zeros(B, 0).to(device)

            if 'text' in self.cond_types:
                text_feat, text_mask = self.text_encoder(text_data)
                dtype = text_feat.dtype
            if 'audio' in self.cond_types:
                audio_feat = self.audio_encoder(audio_data.permute(0, 2, 1)).permute(0, 2, 1)
                audio_mask = length_to_mask(audio_length // self.audio_downsample_ratio)
                dtype = audio_feat.dtype
            if 'traj' in self.cond_types:
                traj_feat = self.traj_encoder(traj_data.permute(0, 2, 1)).permute(0, 2, 1)
                traj_mask = length_to_mask(motion_length // self.motion_downsample_ratio)
                dtype = traj_feat.dtype

            text_feat, text_mask = text_feat.to(dtype), text_mask.to(device)
            audio_feat, audio_mask = audio_feat.to(dtype), audio_mask.to(dtype)
            traj_feat, traj_mask = traj_feat.to(dtype), traj_mask.to(dtype)

            non_pad_mask = self.prepare_test_input(motion_length)
            return self.forward_predict(text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask, non_pad_mask, **kwargs)
        else:
            raise NotImplementedError
    
    def prepare_train_input(self, motion, motion_length):
        if self.vqvae.training:
            self.vqvae.eval()

        B = motion.shape[0]
        device = motion.device
        m_token_lens = motion_length // self.motion_downsample_ratio
        max_token_len = m_token_lens.max()
        ids = self.vqvae.encode(motion, motion_length) # B, T, n_quant
        # True, True, ..., False
        non_pad_mask = torch.arange(max_token_len, device=device).expand(B, max_token_len) < m_token_lens.unsqueeze(1)

        all_x_ids = []
        all_labels = []
        mask = None
        for i in range(ids.shape[-1]): # n_quant
            # import ipdb; ipdb.set_trace()
            m_tokens_i = ids[..., i] # B, T

            # 1. prepare label token
            labels = torch.where(non_pad_mask, m_tokens_i, self.ignore_index)

            # 2. prepare input token
            m_tokens_i = torch.where(non_pad_mask, m_tokens_i, self.pad_idx) # [m1, m2, m3, ..., pad1, pad2, ...]
            
            # 2.1 prepare mask
            if mask is None:
                rand_time = uniform((B,), device=device)
                rand_mask_probs = self.noise_schedule(rand_time)
                num_token_masked = (max_token_len * rand_mask_probs).round().clamp(min=1)
                batch_randperm = torch.rand((B, max_token_len), device=device).argsort(dim=-1)
                # Positions to be masked are all true
                mask = batch_randperm < num_token_masked.unsqueeze(-1)
                # Positions to be MASKED must also be NON-PADDED
                mask &= non_pad_mask

            # 2.2 Further Apply Bert Masking Scheme
            # Step 1: 10% replace with an incorrect token
            x_ids = m_tokens_i.clone()
            mask_rid = get_mask_subset_prob(mask, 0.1)
            rand_id = torch.randint_like(x_ids, high=self.vocab_size - 2)
            x_ids = torch.where(mask_rid, rand_id, x_ids)
            # Step 2: 90% x 10% replace with correct token, and 90% x 88% replace with mask token
            mask_mid = get_mask_subset_prob(mask & ~mask_rid, 0.88)
            x_ids = torch.where(mask_mid, self.mask_idx, x_ids)
            all_x_ids.append(x_ids)
            all_labels.append(labels)

        # import ipdb; ipdb.set_trace()
        return all_x_ids, all_labels, non_pad_mask

    def forward_loss(self, 
                     tokens, label_tokens, 
                     text_feat, text_mask, 
                     audio_feat, audio_mask, 
                     traj_feat, traj_mask, 
                     non_pad_mask, mean=None, std=None, motion=None):
        all_logits = self.transformer(tokens, non_pad_mask, text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask)

        ce_loss = torch.tensor(0.).to(text_feat)
        acc_dict = dict()

        for i, (logits, labels) in enumerate(zip(all_logits, label_tokens)):
            logits = logits.flatten(0, 1)
            labels = labels.flatten(0, 1)
            ce_loss = ce_loss + F.cross_entropy(logits, labels.to(torch.int64), ignore_index=self.ignore_index)

            # calculate acc
            probs = torch.softmax(logits, dim=-1)
            cls_pred_index = torch.argmax(probs, dim=-1)
            token_mask = (labels != self.ignore_index).int()
            right_num = (cls_pred_index == labels).int() * token_mask
            acc_dict[f'ACC_{i}'] = (right_num.sum() / token_mask.sum()).float().detach()

        # add traj loss
        traj_loss = torch.tensor(0.0)
        if traj_mask.sum() > 0:
            all_pred_motion_tokens = []
            for i, (logits, labels) in enumerate(zip(all_logits, label_tokens)):
                # import ipdb; ipdb.set_trace()
                probs_gumbel_softmax = F.gumbel_softmax(logits, tau=1, eps=1e-10, hard=True, dim = -1)
                pred_motion_tokens = torch.topk(probs_gumbel_softmax, k=1, dim=-1)[1].reshape(labels.shape[0], -1) # B, T
                all_pred_motion_tokens.append(pred_motion_tokens)
            all_pred_motion_tokens = torch.stack(all_pred_motion_tokens).permute(1, 2, 0).to(logits.dtype) # nq, B, T

            pred_motion = self.vqvae.decode(all_pred_motion_tokens, motion_length=non_pad_mask.sum(1))

            pred_traj = recover_from_local_position_fast(pred_motion.float() * std + mean, njoint=22)[:, :, 0].to(logits.dtype)
            gt_traj = recover_from_local_position_fast(motion.float() * std + mean, njoint=22)[:, :, 0].to(logits.dtype)

            arange_t = torch.arange(motion.shape[1])[None, :].to(motion.device)  # (1, T)
            mask = arange_t < (non_pad_mask.sum(1)[:, None] * self.motion_downsample_ratio)       # (B, T)
            loss_raw = F.l1_loss(pred_traj, gt_traj, reduction='none') # (B, T, C)
            traj_loss = (loss_raw * mask[..., None]).sum() / (mask.sum() * loss_raw.shape[-1])
        return dict(
            ce_loss=ce_loss / self.n_quant, 
            traj_loss=traj_loss, 
            **acc_dict
        )
    
    @torch.no_grad()
    def forward_predict(self, text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask, non_pad_mask, return_ids=False, **kwargs):
        device = non_pad_mask.device
        # tokens: all pad mask seqence
        ids = torch.where(non_pad_mask, self.mask_idx, self.pad_idx)[None].repeat(self.n_quant, 1, 1)
        scores = torch.where(non_pad_mask, 0., 1e5)[None].repeat(self.n_quant, 1, 1)
        m_lens = non_pad_mask.sum(-1)

        for timestep in torch.linspace(0, 1, self.inference_timesteps, device=device):
            # import ipdb; ipdb.set_trace()
            rand_mask_prob = self.noise_schedule(timestep)
            # Maskout, and cope with variable length
            num_token_masked = torch.round(rand_mask_prob * m_lens).clamp(min=1)  # (b, )
            # select num_token_masked tokens with lowest scores to be masked
            sorted_indices = scores.argsort(dim=-1)  # (b, k), sorted_indices[i, j] = the index of j-th lowest element in scores on dim=1
            ranks = sorted_indices.argsort(dim=-1)  # (b, k), rank[i, j] = the rank (0: lowest) of scores[i, j] on dim=1
            is_mask = (ranks < num_token_masked.unsqueeze(-1).unsqueeze(0))
            ids = torch.where(is_mask, self.mask_idx, ids)
            # all_logits = self.transformer(ids, non_pad_mask, text_feat, text_mask)
            all_logits = self.transformer(ids, non_pad_mask, text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask)
            
            for i, logits in enumerate(all_logits):
                # import ipdb; ipdb.set_trace()
                filtered_logits = top_k(logits, self.topk_filter_thres, dim=-1)
                probs = F.softmax(filtered_logits, dim=-1)  # (b, seqlen, ntoken)
                pred_ids = Categorical(probs).sample()  # (b, seqlen)

                ids[i] = torch.where(is_mask[i], pred_ids, ids[i])

                # Updating scores
                probs_without_temperature = logits.softmax(dim=-1)  # (b, seqlen, ntoken)
                scores[i] = probs_without_temperature.gather(2, pred_ids.unsqueeze(dim=-1)).squeeze(-1)  # (b, seqlen)

                # We do not want to re-mask the previously kept tokens, or pad tokens
                scores[i] = scores[i].masked_fill(~is_mask[i], 1e5)

        ids = torch.where(non_pad_mask[None].repeat(self.n_quant, 1, 1), ids, -1)
        # import ipdb; ipdb.set_trace()
        pred_motions = []
        pred_motion_lengths = []
        
        # import ipdb; ipdb.set_trace()
        pred_motions = self.vqvae.decode(ids.permute(1, 2, 0), motion_length=non_pad_mask.sum(1))

        if return_ids:
            return pred_motions, pred_motion_lengths, ids
        else:
            return pred_motions

    def prepare_test_input(self, motion_length):
        B = motion_length.shape[0]
        device = motion_length.device
        m_lens = motion_length // self.motion_downsample_ratio
        max_len = motion_length.max() // self.motion_downsample_ratio
        non_pad_mask = torch.arange(max_len, device=device).expand(B, max_len) < m_lens.unsqueeze(1)
        return non_pad_mask

@MODELS.register_module()
class AnyMo_audio(BaseModel):
    def __init__(
            self, 
            vqvae_cfg: Dict,
            transformer_cfg: Dict,

            max_motion_len: int,
            motion_downsample_ratio: int, 
            audio_downsample_ratio: int, 

            text_encoder_cfg: Dict=None,
            traj_encoder_cfg: Dict=None, 
            audio_encoder_cfg: Dict=None,
            
            inference_timesteps=6, 
            cond_types = ['audio'], 
            **kwargs
    ):
        super().__init__(**kwargs)

        self.transformer = MODELS.build(transformer_cfg)
        self.text_encoder = MODELS.build(text_encoder_cfg)
        self.traj_encoder = MODELS.build(traj_encoder_cfg)
        self.audio_encoder = MODELS.build(audio_encoder_cfg)

        # build vqvae
        self.vqvae = MODELS.build(vqvae_cfg)
        self.vqvae.init_weights()
        for param in self.vqvae.parameters():
            param.requires_grad = False
        self.vqvae.eval()

        self.vocab_size = transformer_cfg.vocab_size # type: ignore
        self.n_embd = transformer_cfg.n_embd
        self.mask_idx = self.vocab_size - 1
        self.pad_idx = self.vocab_size - 2
        self.ignore_index = -100

        self.max_motion_len = max_motion_len
        self.motion_downsample_ratio = motion_downsample_ratio
        self.audio_downsample_ratio = audio_downsample_ratio
        self.topk_filter_thres = 0.9
        self.noise_schedule = cosine_schedule
        self.n_quant = len(self.vqvae.quantizer.quantizers)
        self.inference_timesteps = inference_timesteps
        self.cond_types = cond_types
        for cond_type in cond_types:
            assert cond_type in ['text', 'audio', 'traj']

    def forward_train(self, motion, motion_length, text_data, audio_data, traj_data, audio_length, **kwargs):
        audio_feat = self.audio_encoder(audio_data.permute(0, 2, 1)).permute(0, 2, 1)
        audio_mask = length_to_mask(audio_length // self.audio_downsample_ratio)

        B = motion.shape[0]
        device = motion.device
        text_feat = torch.zeros(B, 0, self.n_embd).to(device).to(audio_feat.dtype)
        traj_feat = torch.zeros(B, 0, self.n_embd).to(device).to(audio_feat.dtype)
        text_mask = torch.zeros(B, 0).to(device).to(audio_feat.dtype)
        traj_mask = torch.zeros(B, 0).to(device).to(audio_feat.dtype)

        tokens, label_tokens, non_pad_mask = self.prepare_train_input(motion, motion_length)
        loss =  self.forward_loss(
            tokens, label_tokens, text_feat, text_mask, audio_feat, 
            audio_mask, traj_feat, traj_mask, non_pad_mask, motion=motion, mean=kwargs['mean'], 
            std=kwargs['std']
        )
        return loss
    
    def forward(
            self, 
            motion: torch.Tensor, 
            motion_length: torch.Tensor,

            text_data: List[str]=None,
            audio_data: torch.Tensor=None,
            traj_data: torch.Tensor=None,
            audio_length: torch.Tensor=None,
            
            mode = 'loss', 
            **kwargs
        ): 
        '''
        motion: tensor, B, T, C
        motion_length: tensor, B
        text_data: list[str], len B
        has_xxx: list[bool], len B
        music_data: tensor, B, T, C
        '''
        if mode == 'loss':
            return self.forward_train(motion, motion_length, text_data, audio_data, traj_data, audio_length, **kwargs)
        else: # for evaluation
            B = motion.shape[0]
            device = motion.device
            text_feat = torch.zeros(B, 0, self.n_embd).to(device)
            audio_feat = torch.zeros(B, 0, self.n_embd).to(device)
            traj_feat = torch.zeros(B, 0, self.n_embd).to(device)
            text_mask = torch.zeros(B, 0).to(device)
            audio_mask = torch.zeros(B, 0).to(device)
            traj_mask = torch.zeros(B, 0).to(device)
            # import ipdb; ipdb.set_trace()
            if 'text' in self.cond_types:
                text_feat, text_mask = self.text_encoder(text_data)
                dtype = text_feat.dtype
            if 'audio' in self.cond_types:
                audio_feat = self.audio_encoder(audio_data.permute(0, 2, 1)).permute(0, 2, 1)
                audio_mask = length_to_mask(audio_length // self.audio_downsample_ratio)
                dtype = audio_feat.dtype
            if 'traj' in self.cond_types:
                traj_feat = self.traj_encoder(traj_data.permute(0, 2, 1)).permute(0, 2, 1)
                traj_mask = length_to_mask(motion_length // self.motion_downsample_ratio)
                dtype = traj_feat.dtype

            text_feat, text_mask = text_feat.to(dtype), text_mask.to(device)
            audio_feat, audio_mask = audio_feat.to(dtype), audio_mask.to(dtype)
            traj_feat, traj_mask = traj_feat.to(dtype), traj_mask.to(dtype)

            non_pad_mask = self.prepare_test_input(motion_length)
            return self.forward_predict(text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask, non_pad_mask, **kwargs)
    
    def prepare_train_input(self, motion, motion_length):
        if self.vqvae.training:
            self.vqvae.eval()

        B = motion.shape[0]
        device = motion.device
        m_token_lens = motion_length // self.motion_downsample_ratio
        max_token_len = m_token_lens.max()
        ids = self.vqvae.encode(motion, motion_length) # B, T, n_quant
        # True, True, ..., False
        non_pad_mask = torch.arange(max_token_len, device=device).expand(B, max_token_len) < m_token_lens.unsqueeze(1)

        all_x_ids = []
        all_labels = []
        mask = None
        for i in range(ids.shape[-1]): # n_quant
            # import ipdb; ipdb.set_trace()
            m_tokens_i = ids[..., i] # B, T

            # 1. prepare label token
            labels = torch.where(non_pad_mask, m_tokens_i, self.ignore_index)

            # 2. prepare input token
            m_tokens_i = torch.where(non_pad_mask, m_tokens_i, self.pad_idx) # [m1, m2, m3, ..., pad1, pad2, ...]
            
            # 2.1 prepare mask
            if mask is None:
                rand_time = uniform((B,), device=device)
                rand_mask_probs = self.noise_schedule(rand_time)
                num_token_masked = (max_token_len * rand_mask_probs).round().clamp(min=1)
                batch_randperm = torch.rand((B, max_token_len), device=device).argsort(dim=-1)
                # Positions to be masked are all true
                mask = batch_randperm < num_token_masked.unsqueeze(-1)
                # Positions to be MASKED must also be NON-PADDED
                mask &= non_pad_mask

            # 2.2 Further Apply Bert Masking Scheme
            # Step 1: 10% replace with an incorrect token
            x_ids = m_tokens_i.clone()
            mask_rid = get_mask_subset_prob(mask, 0.1)
            rand_id = torch.randint_like(x_ids, high=self.vocab_size - 2)
            x_ids = torch.where(mask_rid, rand_id, x_ids)
            # Step 2: 90% x 10% replace with correct token, and 90% x 88% replace with mask token
            mask_mid = get_mask_subset_prob(mask & ~mask_rid, 0.88)
            x_ids = torch.where(mask_mid, self.mask_idx, x_ids)
            all_x_ids.append(x_ids)
            all_labels.append(labels)

        # import ipdb; ipdb.set_trace()
        return all_x_ids, all_labels, non_pad_mask

    def forward_loss(self, 
                     tokens, label_tokens, 
                     text_feat, text_mask, 
                     audio_feat, audio_mask, 
                     traj_feat, traj_mask, 
                     non_pad_mask, mean=None, std=None, motion=None):
        all_logits = self.transformer(tokens, non_pad_mask, text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask)

        ce_loss = torch.tensor(0.).to(text_feat)
        acc_dict = dict()

        for i, (logits, labels) in enumerate(zip(all_logits, label_tokens)):
            logits = logits.flatten(0, 1)
            labels = labels.flatten(0, 1)
            ce_loss = ce_loss + F.cross_entropy(logits, labels.to(torch.int64), ignore_index=self.ignore_index)

            # calculate acc
            probs = torch.softmax(logits, dim=-1)
            cls_pred_index = torch.argmax(probs, dim=-1)
            token_mask = (labels != self.ignore_index).int()
            right_num = (cls_pred_index == labels).int() * token_mask
            acc_dict[f'ACC_{i}'] = (right_num.sum() / token_mask.sum()).float().detach()

        # add traj loss
        traj_loss = torch.tensor(0.0)
        if traj_mask.sum() > 0:
            all_pred_motion_tokens = []
            for i, (logits, labels) in enumerate(zip(all_logits, label_tokens)):
                # import ipdb; ipdb.set_trace()
                probs_gumbel_softmax = F.gumbel_softmax(logits, tau=1, eps=1e-10, hard=True, dim = -1)
                pred_motion_tokens = torch.topk(probs_gumbel_softmax, k=1, dim=-1)[1].reshape(labels.shape[0], -1) # B, T
                all_pred_motion_tokens.append(pred_motion_tokens)
            all_pred_motion_tokens = torch.stack(all_pred_motion_tokens) # nq, B, T
            pred_motion = self.vqvae.decode(all_pred_motion_tokens.permute(1, 2, 0), motion_length=non_pad_mask.sum(1))

            pred_traj = recover_from_local_position_fast(pred_motion * std + mean, njoint=22)[:, :, 0]
            gt_traj = recover_from_local_position_fast(motion * std + mean, njoint=22)[:, :, 0]

            arange_t = torch.arange(motion.shape[1])[None, :].to(motion.device)  # (1, T)
            mask = arange_t < (non_pad_mask.sum(1)[:, None] * self.motion_downsample_ratio)       # (B, T)
            loss_raw = F.l1_loss(pred_traj, gt_traj, reduction='none') # (B, T, C)
            traj_loss = (loss_raw * mask[..., None]).sum() / (mask.sum() * loss_raw.shape[-1])

        return dict(
            ce_loss=ce_loss / self.n_quant, 
            traj_loss=traj_loss, 
            **acc_dict
        )
    
    @torch.no_grad()
    def forward_predict(self, text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask, non_pad_mask, return_ids=False, **kwargs):
        device = non_pad_mask.device
        # tokens: all pad mask seqence
        ids = torch.where(non_pad_mask, self.mask_idx, self.pad_idx)[None].repeat(self.n_quant, 1, 1)
        scores = torch.where(non_pad_mask, 0., 1e5)[None].repeat(self.n_quant, 1, 1)
        m_lens = non_pad_mask.sum(-1)

        for timestep in torch.linspace(0, 1, self.inference_timesteps, device=device):
            # import ipdb; ipdb.set_trace()
            rand_mask_prob = self.noise_schedule(timestep)
            # Maskout, and cope with variable length
            num_token_masked = torch.round(rand_mask_prob * m_lens).clamp(min=1)  # (b, )
            # select num_token_masked tokens with lowest scores to be masked
            sorted_indices = scores.argsort(dim=-1)  # (b, k), sorted_indices[i, j] = the index of j-th lowest element in scores on dim=1
            ranks = sorted_indices.argsort(dim=-1)  # (b, k), rank[i, j] = the rank (0: lowest) of scores[i, j] on dim=1
            is_mask = (ranks < num_token_masked.unsqueeze(-1).unsqueeze(0))
            ids = torch.where(is_mask, self.mask_idx, ids)
            # all_logits = self.transformer(ids, non_pad_mask, text_feat, text_mask)
            all_logits = self.transformer(ids, non_pad_mask, text_feat, text_mask, audio_feat, audio_mask, traj_feat, traj_mask)
            
            for i, logits in enumerate(all_logits):
                # import ipdb; ipdb.set_trace()
                filtered_logits = top_k(logits, self.topk_filter_thres, dim=-1)
                probs = F.softmax(filtered_logits, dim=-1)  # (b, seqlen, ntoken)
                pred_ids = Categorical(probs).sample()  # (b, seqlen)

                ids[i] = torch.where(is_mask[i], pred_ids, ids[i])

                # Updating scores
                probs_without_temperature = logits.softmax(dim=-1)  # (b, seqlen, ntoken)
                scores[i] = probs_without_temperature.gather(2, pred_ids.unsqueeze(dim=-1)).squeeze(-1)  # (b, seqlen)

                # We do not want to re-mask the previously kept tokens, or pad tokens
                scores[i] = scores[i].masked_fill(~is_mask[i], 1e5)

        ids = torch.where(non_pad_mask[None].repeat(self.n_quant, 1, 1), ids, -1)
        # import ipdb; ipdb.set_trace()
        pred_motions = []
        pred_motion_lengths = []
        
        # import ipdb; ipdb.set_trace()
        pred_motions = self.vqvae.decode(ids.permute(1, 2, 0), motion_length=non_pad_mask.sum(1))

        if return_ids:
            return pred_motions, pred_motion_lengths, ids
        else:
            return pred_motions

    def prepare_test_input(self, motion_length):
        B = motion_length.shape[0]
        device = motion_length.device
        m_lens = motion_length // self.motion_downsample_ratio
        max_len = motion_length.max() // self.motion_downsample_ratio
        non_pad_mask = torch.arange(max_len, device=device).expand(B, max_len) < m_lens.unsqueeze(1)
        return non_pad_mask