Spaces:
Running on Zero
Running on Zero
File size: 2,498 Bytes
4cc0d6c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | import torch
import torch.nn.functional as F
from mmengine.registry import MODELS
from mmengine.model import BaseModel
from typing import Dict
import torch.nn as nn
@MODELS.register_module()
class RetargetTransformerPredMotion_no_smplvq(BaseModel):
def __init__(self,
transformer_cfg: Dict,
smplx_vqvae_cfg: Dict,
n_embd: int = 512,
**kwargs):
super().__init__(**kwargs)
self.transformer = MODELS.build(transformer_cfg)
self.motion_encoder = nn.Sequential(
nn.Linear(512, n_embd),
nn.SiLU(),
nn.Linear(n_embd, n_embd),
)
self.smplx_vqvae = MODELS.build(smplx_vqvae_cfg)
self.upsample = nn.Sequential(
nn.Upsample(scale_factor=2, mode='nearest'),
nn.Conv1d(n_embd, n_embd, 3, 1, 1),
nn.Upsample(scale_factor=2, mode='nearest'),
nn.Conv1d(n_embd, n_embd, 3, 1, 1),
)
self.projector = nn.Linear(n_embd, 217)
def forward_predict(self, cond_embd: torch.Tensor, cond_mask: torch.Tensor):
feat = self.transformer(cond_embd, cond_mask)
up_sample_feat = self.upsample(feat.permute(0, 2, 1)).permute(0, 2, 1)
pred_motions = self.projector(up_sample_feat)
new_mask = F.interpolate(cond_mask[None].float(), scale_factor=2, mode='nearest')[0]
pred_motion_lengths = new_mask.sum(dim=1)
_pred_motions = []
for pred_motion, length in zip(pred_motions, pred_motion_lengths):
_pred_motions.append(pred_motion[:length.int()])
return pred_motions, pred_motion_lengths
def forward(self,
motion: torch.Tensor = None,
smplx_motion: torch.Tensor = None,
motion_length: torch.Tensor = None,
mode='loss', **kwargs):
B, T, _ = smplx_motion.shape
smplx_motion_in = self.smplx_vqvae.preprocess(smplx_motion)
smplx_motion_embd = self.smplx_vqvae.encoder(smplx_motion_in, motion_length=motion_length)
smplx_motion_embd = smplx_motion_embd.permute(0, 2, 1)
smplx_motion_embd = self.motion_encoder(smplx_motion_embd)
masks = torch.arange(T // 2, device=smplx_motion_embd.device)[None].repeat(B, 1) < motion_length[:, None] // 2
if mode == 'loss':
return self.forward_loss(smplx_motion_embd, masks, motion)
else:
return self.forward_predict(smplx_motion_embd, masks)
|