initial commit

.gitattributes

@@ -29,3 +29,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+checkpoints/ filter=lfs diff=lfs merge=lfs -text
+checkpoints/t2m/t2m_motiondiffuse/model/latest.tar filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,72 @@
+import os
+import sys
+import gradio as gr
+try:
+    os.system("pip install -r requirements.txt")
+except Exception as e:
+    print(e)
+
+sys.path.insert(0, '.')
+
+
+from utils.get_opt import get_opt
+from os.path import join as pjoin
+import numpy as np
+from trainers import DDPMTrainer
+from models import MotionTransformer
+
+device = 'cpu'
+opt = get_opt("checkpoints/t2m/t2m_motiondiffuse/opt.txt", device)
+opt.do_denoise = True
+
+assert opt.dataset_name == "t2m"
+opt.data_root = './dataset/HumanML3D'
+opt.motion_dir = pjoin(opt.data_root, 'new_joint_vecs')
+opt.text_dir = pjoin(opt.data_root, 'texts')
+opt.joints_num = 22
+opt.dim_pose = 263
+
+mean = np.load(pjoin(opt.meta_dir, 'mean.npy'))
+std = np.load(pjoin(opt.meta_dir, 'std.npy'))
+
+
+def build_models(opt):
+    encoder = MotionTransformer(
+        input_feats=opt.dim_pose,
+        num_frames=opt.max_motion_length,
+        num_layers=opt.num_layers,
+        latent_dim=opt.latent_dim,
+        no_clip=opt.no_clip,
+        no_eff=opt.no_eff)
+    return encoder
+
+
+encoder = build_models(opt).to(device)
+trainer = DDPMTrainer(opt, encoder)
+trainer.load(pjoin(opt.model_dir, opt.which_epoch + '.tar'))
+
+trainer.eval_mode()
+trainer.to(opt.device)
+
+def generate(prompt, length):
+    from tools.visualization import process
+    result_path = "outputs/" + str(hash(prompt)) + ".mp4"
+    process(trainer, opt, device, mean, std, prompt, int(length), result_path)
+    return result_path
+
+demo = gr.Interface(
+    fn=generate,
+    inputs=["text", gr.Slider(20, 196, value=60)],
+    examples=[
+        ["the man throws a punch with each hand.", 58],
+        ["a person jogs clockwise in a circle.", 178],
+        ["a person spins quickly and takes off running.", 29],
+        ["a person is walking slowly forward.", 142],
+        ["a person quickly waves with their right hand", 46],
+        ["a person performing a slight bow", 89],
+    ],
+    outputs="video",
+    title="MotionDiffuse: Text-Driven Human Motion Generation with Diffusion Model",
+    description="This is an interactive demo for MotionDiffuse. For more information, feel free to visit our project page(https://mingyuan-zhang.github.io/projects/MotionDiffuse.html).")
+
+demo.launch(share=True)

checkpoints/t2m/t2m_motiondiffuse/meta/mean.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ebc7c543b4e27e886dba7e1bcde8fc0149f12a981586a548df98977b4b7c1a6a
+size 1180

checkpoints/t2m/t2m_motiondiffuse/meta/std.npy

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9a91111b5b44b2b6785cedf426ede6f8bc412ecb995914563e2810738de179e4
+size 1180

checkpoints/t2m/t2m_motiondiffuse/model/latest.tar

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:60a8b95aa5d190a95a6b3c20515301f72d70646b7a0d7f927c7167b968fa1222
+size 953997124

checkpoints/t2m/t2m_motiondiffuse/opt.txt

@@ -0,0 +1,38 @@
+------------ Options -------------
+batch_size: 128
+checkpoints_dir: ./checkpoints
+dataset_name: t2m
+decomp_name: Decomp_SP001_SM001_H512
+dim_att_vec: 512
+dim_dec_hidden: 1024
+dim_movement_dec_hidden: 512
+dim_movement_enc_hidden: 512
+dim_movement_latent: 512
+dim_pos_hidden: 1024
+dim_pri_hidden: 1024
+dim_text_hidden: 512
+dim_z: 128
+early_stop_count: 3
+estimator_mod: bigru
+eval_every_e: 5
+feat_bias: 5
+gpu_id: -1
+is_continue: False
+is_train: True
+lambda_kld: 0.005
+lambda_rec_mot: 1
+lambda_rec_mov: 1
+log_every: 50
+lr: 0.0002
+max_sub_epoch: 50
+max_text_len: 20
+n_layers_dec: 1
+n_layers_pos: 1
+n_layers_pri: 1
+name: t2m_motiondiffuse
+save_every_e: 10
+save_latest: 500
+text_enc_mod: bigru
+times: 50
+unit_length: 4
+-------------- End ----------------

datasets/__init__.py

@@ -0,0 +1,11 @@
+from .dataset import Text2MotionDataset
+from .evaluator import (
+    EvaluationDataset,
+    get_dataset_motion_loader,
+    get_motion_loader,
+    EvaluatorModelWrapper)
+from .dataloader import build_dataloader
+
+__all__ = [
+    'Text2MotionDataset', 'EvaluationDataset', 'build_dataloader',
+    'get_dataset_motion_loader', 'get_motion_loader']

datasets/dataloader.py

@@ -0,0 +1,130 @@
+import platform
+import random
+from functools import partial
+from typing import Optional, Union
+
+import numpy as np
+from mmcv.runner import get_dist_info
+from mmcv.utils import Registry, build_from_cfg
+from torch.utils.data import DataLoader
+from torch.utils.data.dataset import Dataset
+
+import torch
+from torch.utils.data import DistributedSampler as _DistributedSampler
+
+
+class DistributedSampler(_DistributedSampler):
+
+    def __init__(self,
+                 dataset,
+                 num_replicas=None,
+                 rank=None,
+                 shuffle=True,
+                 round_up=True):
+        super().__init__(dataset, num_replicas=num_replicas, rank=rank)
+        self.shuffle = shuffle
+        self.round_up = round_up
+        if self.round_up:
+            self.total_size = self.num_samples * self.num_replicas
+        else:
+            self.total_size = len(self.dataset)
+
+    def __iter__(self):
+        # deterministically shuffle based on epoch
+        if self.shuffle:
+            g = torch.Generator()
+            g.manual_seed(self.epoch)
+            indices = torch.randperm(len(self.dataset), generator=g).tolist()
+        else:
+            indices = torch.arange(len(self.dataset)).tolist()
+
+        # add extra samples to make it evenly divisible
+        if self.round_up:
+            indices = (
+                indices *
+                int(self.total_size / len(indices) + 1))[:self.total_size]
+        assert len(indices) == self.total_size
+
+        # subsample
+        indices = indices[self.rank:self.total_size:self.num_replicas]
+        if self.round_up:
+            assert len(indices) == self.num_samples
+
+        return iter(indices)
+
+
+def build_dataloader(dataset: Dataset,
+                     samples_per_gpu: int,
+                     workers_per_gpu: int,
+                     num_gpus: Optional[int] = 1,
+                     dist: Optional[bool] = True,
+                     shuffle: Optional[bool] = True,
+                     round_up: Optional[bool] = True,
+                     seed: Optional[Union[int, None]] = None,
+                     persistent_workers: Optional[bool] = True,
+                     **kwargs):
+    """Build PyTorch DataLoader.
+
+    In distributed training, each GPU/process has a dataloader.
+    In non-distributed training, there is only one dataloader for all GPUs.
+
+    Args:
+        dataset (:obj:`Dataset`): A PyTorch dataset.
+        samples_per_gpu (int): Number of training samples on each GPU, i.e.,
+            batch size of each GPU.
+        workers_per_gpu (int): How many subprocesses to use for data loading
+            for each GPU.
+        num_gpus (int, optional): Number of GPUs. Only used in non-distributed
+            training.
+        dist (bool, optional): Distributed training/test or not. Default: True.
+        shuffle (bool, optional): Whether to shuffle the data at every epoch.
+            Default: True.
+        round_up (bool, optional): Whether to round up the length of dataset by
+            adding extra samples to make it evenly divisible. Default: True.
+        persistent_workers (bool): If True, the data loader will not shutdown
+            the worker processes after a dataset has been consumed once.
+            This allows to maintain the workers Dataset instances alive.
+            The argument also has effect in PyTorch>=1.7.0.
+            Default: True
+        kwargs: any keyword argument to be used to initialize DataLoader
+
+    Returns:
+        DataLoader: A PyTorch dataloader.
+    """
+    rank, world_size = get_dist_info()
+    if dist:
+        sampler = DistributedSampler(
+            dataset, world_size, rank, shuffle=shuffle, round_up=round_up)
+        shuffle = False
+        batch_size = samples_per_gpu
+        num_workers = workers_per_gpu
+    else:
+        sampler = None
+        batch_size = num_gpus * samples_per_gpu
+        num_workers = num_gpus * workers_per_gpu
+
+    init_fn = partial(
+        worker_init_fn, num_workers=num_workers, rank=rank,
+        seed=seed) if seed is not None else None
+
+    data_loader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        sampler=sampler,
+        num_workers=num_workers,
+        pin_memory=False,
+        shuffle=shuffle,
+        worker_init_fn=init_fn,
+        persistent_workers=persistent_workers,
+        **kwargs)
+
+    return data_loader
+
+
+def worker_init_fn(worker_id: int, num_workers: int, rank: int, seed: int):
+    """Init random seed for each worker."""
+    # The seed of each worker equals to
+    # num_worker * rank + worker_id + user_seed
+    worker_seed = num_workers * rank + worker_id + seed
+    np.random.seed(worker_seed)
+    random.seed(worker_seed)

datasets/dataset.py

@@ -0,0 +1,164 @@
+import torch
+from torch.utils import data
+import numpy as np
+import os
+from os.path import join as pjoin
+import random
+import codecs as cs
+from tqdm import tqdm
+
+
+class Text2MotionDataset(data.Dataset):
+    """Dataset for Text2Motion generation task.
+
+    """
+    def __init__(self, opt, mean, std, split_file, times=1, w_vectorizer=None, eval_mode=False):
+        self.opt = opt
+        self.max_length = 20
+        self.times = times
+        self.w_vectorizer = w_vectorizer
+        self.eval_mode = eval_mode
+        min_motion_len = 40 if self.opt.dataset_name =='t2m' else 24
+
+        joints_num = opt.joints_num
+
+        data_dict = {}
+        id_list = []
+        with cs.open(split_file, 'r') as f:
+            for line in f.readlines():
+                id_list.append(line.strip())
+
+        new_name_list = []
+        length_list = []
+        for name in tqdm(id_list):
+            try:
+                motion = np.load(pjoin(opt.motion_dir, name + '.npy'))
+                if (len(motion)) < min_motion_len or (len(motion) >= 200):
+                    continue
+                text_data = []
+                flag = False
+                with cs.open(pjoin(opt.text_dir, name + '.txt')) as f:
+                    for line in f.readlines():
+                        text_dict = {}
+                        line_split = line.strip().split('#')
+                        caption = line_split[0]
+                        tokens = line_split[1].split(' ')
+                        f_tag = float(line_split[2])
+                        to_tag = float(line_split[3])
+                        f_tag = 0.0 if np.isnan(f_tag) else f_tag
+                        to_tag = 0.0 if np.isnan(to_tag) else to_tag
+
+                        text_dict['caption'] = caption
+                        text_dict['tokens'] = tokens
+                        if f_tag == 0.0 and to_tag == 0.0:
+                            flag = True
+                            text_data.append(text_dict)
+                        else:
+                            n_motion = motion[int(f_tag*20) : int(to_tag*20)]
+                            if (len(n_motion)) < min_motion_len or (len(n_motion) >= 200):
+                                continue
+                            new_name = random.choice('ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name
+                            while new_name in data_dict:
+                                new_name = random.choice('ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name
+                            data_dict[new_name] = {'motion': n_motion,
+                                                    'length': len(n_motion),
+                                                    'text':[text_dict]}
+                            new_name_list.append(new_name)
+                            length_list.append(len(n_motion))
+
+                if flag:
+                    data_dict[name] = {'motion': motion,
+                                       'length': len(motion),
+                                       'text':text_data}
+                    new_name_list.append(name)
+                    length_list.append(len(motion))
+            except:
+                # Some motion may not exist in KIT dataset
+                pass
+
+
+        name_list, length_list = zip(*sorted(zip(new_name_list, length_list), key=lambda x: x[1]))
+
+        if opt.is_train:
+            # root_rot_velocity (B, seq_len, 1)
+            std[0:1] = std[0:1] / opt.feat_bias
+            # root_linear_velocity (B, seq_len, 2)
+            std[1:3] = std[1:3] / opt.feat_bias
+            # root_y (B, seq_len, 1)
+            std[3:4] = std[3:4] / opt.feat_bias
+            # ric_data (B, seq_len, (joint_num - 1)*3)
+            std[4: 4 + (joints_num - 1) * 3] = std[4: 4 + (joints_num - 1) * 3] / 1.0
+            # rot_data (B, seq_len, (joint_num - 1)*6)
+            std[4 + (joints_num - 1) * 3: 4 + (joints_num - 1) * 9] = std[4 + (joints_num - 1) * 3: 4 + (
+                        joints_num - 1) * 9] / 1.0
+            # local_velocity (B, seq_len, joint_num*3)
+            std[4 + (joints_num - 1) * 9: 4 + (joints_num - 1) * 9 + joints_num * 3] = std[
+                                                                                       4 + (joints_num - 1) * 9: 4 + (
+                                                                                                   joints_num - 1) * 9 + joints_num * 3] / 1.0
+            # foot contact (B, seq_len, 4)
+            std[4 + (joints_num - 1) * 9 + joints_num * 3:] = std[
+                                                              4 + (joints_num - 1) * 9 + joints_num * 3:] / opt.feat_bias
+
+            assert 4 + (joints_num - 1) * 9 + joints_num * 3 + 4 == mean.shape[-1]
+            np.save(pjoin(opt.meta_dir, 'mean.npy'), mean)
+            np.save(pjoin(opt.meta_dir, 'std.npy'), std)
+
+        self.mean = mean
+        self.std = std
+        self.length_arr = np.array(length_list)
+        self.data_dict = data_dict
+        self.name_list = name_list
+
+    def inv_transform(self, data):
+        return data * self.std + self.mean
+
+    def real_len(self):
+        return len(self.data_dict)
+
+    def __len__(self):
+        return self.real_len() * self.times
+
+    def __getitem__(self, item):
+        idx = item % self.real_len()
+        data = self.data_dict[self.name_list[idx]]
+        motion, m_length, text_list = data['motion'], data['length'], data['text']
+        # Randomly select a caption
+        text_data = random.choice(text_list)
+        caption = text_data['caption']
+
+        max_motion_length = self.opt.max_motion_length
+        if m_length >= self.opt.max_motion_length:
+            idx = random.randint(0, len(motion) - max_motion_length)
+            motion = motion[idx: idx + max_motion_length]
+        else:
+            padding_len = max_motion_length - m_length
+            D = motion.shape[1]
+            padding_zeros = np.zeros((padding_len, D))
+            motion = np.concatenate((motion, padding_zeros), axis=0)
+
+        assert len(motion) == max_motion_length
+        "Z Normalization"
+        motion = (motion - self.mean) / self.std
+
+        if self.eval_mode:
+            tokens = text_data['tokens']
+            if len(tokens) < self.opt.max_text_len:
+                # pad with "unk"
+                tokens = ['sos/OTHER'] + tokens + ['eos/OTHER']
+                sent_len = len(tokens)
+                tokens = tokens + ['unk/OTHER'] * (self.opt.max_text_len + 2 - sent_len)
+            else:
+                # crop
+                tokens = tokens[:self.opt.max_text_len]
+                tokens = ['sos/OTHER'] + tokens + ['eos/OTHER']
+                sent_len = len(tokens)
+            pos_one_hots = []
+            word_embeddings = []
+            for token in tokens:
+                word_emb, pos_oh = self.w_vectorizer[token]
+                pos_one_hots.append(pos_oh[None, :])
+                word_embeddings.append(word_emb[None, :])
+            pos_one_hots = np.concatenate(pos_one_hots, axis=0)
+            word_embeddings = np.concatenate(word_embeddings, axis=0)
+            return word_embeddings, pos_one_hots, caption, sent_len, motion, m_length
+        return caption, motion, m_length

datasets/evaluator.py

@@ -0,0 +1,441 @@
+import torch
+from utils.word_vectorizer import WordVectorizer, POS_enumerator
+from utils.get_opt import get_opt
+from models import MotionTransformer
+from torch.utils.data import Dataset, DataLoader
+from os.path import join as pjoin
+from tqdm import tqdm
+import numpy as np
+from .evaluator_models import *
+import os
+import codecs as cs
+import random
+from torch.utils.data._utils.collate import default_collate
+
+
+class EvaluationDataset(Dataset):
+
+    def __init__(self, opt, trainer, dataset, w_vectorizer, mm_num_samples, mm_num_repeats):
+        assert mm_num_samples < len(dataset)
+        print(opt.model_dir)
+
+        dataloader = DataLoader(dataset, batch_size=1, num_workers=1, shuffle=True)
+        epoch, it = trainer.load(pjoin(opt.model_dir, opt.which_epoch + '.tar'))
+
+        generated_motion = []
+        min_mov_length = 10 if opt.dataset_name == 't2m' else 6
+
+        trainer.eval_mode()
+        trainer.to(opt.device)
+
+        # Pre-process all target captions
+        mm_generated_motions = []
+        mm_idxs = np.random.choice(len(dataset), mm_num_samples, replace=False)
+        mm_idxs = np.sort(mm_idxs)
+        all_caption = []
+        all_m_lens = []
+        all_data = []
+        with torch.no_grad():
+            for i, data in tqdm(enumerate(dataloader)):
+                word_emb, pos_ohot, caption, cap_lens, motions, m_lens, tokens = data
+                all_data.append(data)
+                tokens = tokens[0].split('_')
+                mm_num_now = len(mm_generated_motions)
+                is_mm = True if ((mm_num_now < mm_num_samples) and (i == mm_idxs[mm_num_now])) else False
+                repeat_times = mm_num_repeats if is_mm else 1
+                m_lens = max(m_lens // opt.unit_length * opt.unit_length, min_mov_length * opt.unit_length)
+                m_lens = min(m_lens, opt.max_motion_length)
+                if isinstance(m_lens, int):
+                    m_lens = torch.LongTensor([m_lens]).to(opt.device)
+                else:
+                    m_lens = m_lens.to(opt.device)
+                for t in range(repeat_times):
+                    all_m_lens.append(m_lens)
+                    all_caption.extend(caption)
+                if is_mm:
+                    mm_generated_motions.append(0)
+        all_m_lens = torch.stack(all_m_lens)
+        
+        # Generate all sequences
+        with torch.no_grad():
+            all_pred_motions = trainer.generate(all_caption, all_m_lens, opt.dim_pose)
+        
+        cur_idx = 0
+        mm_generated_motions = []
+        with torch.no_grad():
+            for i, data_dummy in tqdm(enumerate(dataloader)):
+                data = all_data[i]
+                word_emb, pos_ohot, caption, cap_lens, motions, m_lens, tokens = data
+                tokens = tokens[0].split('_')
+                mm_num_now = len(mm_generated_motions)
+                is_mm = True if ((mm_num_now < mm_num_samples) and (i == mm_idxs[mm_num_now])) else False
+                repeat_times = mm_num_repeats if is_mm else 1
+                mm_motions = []
+                m_lens = max(m_lens // opt.unit_length * opt.unit_length, min_mov_length * opt.unit_length)
+                m_lens = min(m_lens, opt.max_motion_length)
+                if isinstance(m_lens, int):
+                    m_lens = torch.LongTensor([m_lens]).to(opt.device)
+                else:
+                    m_lens = m_lens.to(opt.device)
+                for t in range(repeat_times):
+                    m_len = m_lens[0].item()
+                    pred_motions = all_pred_motions[cur_idx][:m_lens[0].item()]
+                    assert pred_motions.shape[0] == m_lens[0].item()
+                    cur_idx += 1
+                    if t == 0:
+                        sub_dict = {'motion': pred_motions.cpu().numpy(),
+                                    'length': pred_motions.shape[0],
+                                    'caption': caption[0],
+                                    'cap_len': cap_lens[0].item(),
+                                    'tokens': tokens}
+                        generated_motion.append(sub_dict)
+
+                    if is_mm:
+                        mm_motions.append({
+                            'motion': pred_motions.cpu().numpy(),
+                            'length': m_lens[0].item()
+                        })
+                if is_mm:
+                    mm_generated_motions.append({'caption': caption[0],
+                                                 'tokens': tokens,
+                                                 'cap_len': cap_lens[0].item(),
+                                                 'mm_motions': mm_motions})
+        self.generated_motion = generated_motion
+        self.mm_generated_motion = mm_generated_motions
+        self.opt = opt
+        self.w_vectorizer = w_vectorizer
+
+
+    def __len__(self):
+        return len(self.generated_motion)
+
+
+    def __getitem__(self, item):
+        data = self.generated_motion[item]
+        motion, m_length, caption, tokens = data['motion'], data['length'], data['caption'], data['tokens']
+        sent_len = data['cap_len']
+        pos_one_hots = []
+        word_embeddings = []
+        for token in tokens:
+            word_emb, pos_oh = self.w_vectorizer[token]
+            pos_one_hots.append(pos_oh[None, :])
+            word_embeddings.append(word_emb[None, :])
+        pos_one_hots = np.concatenate(pos_one_hots, axis=0)
+        word_embeddings = np.concatenate(word_embeddings, axis=0)
+
+        if m_length < self.opt.max_motion_length:
+            motion = np.concatenate([motion,
+                                     np.zeros((self.opt.max_motion_length - m_length, motion.shape[1]))
+                                     ], axis=0)
+        return word_embeddings, pos_one_hots, caption, sent_len, motion, m_length, '_'.join(tokens)
+
+
+def collate_fn(batch):
+    batch.sort(key=lambda x: x[3], reverse=True)
+    return default_collate(batch)
+
+
+'''For use of training text motion matching model, and evaluations'''
+class Text2MotionDatasetV2(Dataset):
+    def __init__(self, opt, mean, std, split_file, w_vectorizer):
+        self.opt = opt
+        self.w_vectorizer = w_vectorizer
+        self.max_length = 20
+        self.pointer = 0
+        self.max_motion_length = opt.max_motion_length
+        min_motion_len = 40 if self.opt.dataset_name =='t2m' else 24
+
+        data_dict = {}
+        id_list = []
+        with cs.open(split_file, 'r') as f:
+            for line in f.readlines():
+                id_list.append(line.strip())
+
+        new_name_list = []
+        length_list = []
+        for name in tqdm(id_list):
+            try:
+                motion = np.load(pjoin(opt.motion_dir, name + '.npy'))
+                if (len(motion)) < min_motion_len or (len(motion) >= 200):
+                    continue
+                text_data = []
+                flag = False
+                with cs.open(pjoin(opt.text_dir, name + '.txt')) as f:
+                    for line in f.readlines():
+                        text_dict = {}
+                        line_split = line.strip().split('#')
+                        caption = line_split[0]
+                        tokens = line_split[1].split(' ')
+                        f_tag = float(line_split[2])
+                        to_tag = float(line_split[3])
+                        f_tag = 0.0 if np.isnan(f_tag) else f_tag
+                        to_tag = 0.0 if np.isnan(to_tag) else to_tag
+
+                        text_dict['caption'] = caption
+                        text_dict['tokens'] = tokens
+                        if f_tag == 0.0 and to_tag == 0.0:
+                            flag = True
+                            text_data.append(text_dict)
+                        else:
+                            try:
+                                n_motion = motion[int(f_tag*20) : int(to_tag*20)]
+                                if (len(n_motion)) < min_motion_len or (len(n_motion) >= 200):
+                                    continue
+                                new_name = random.choice('ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name
+                                while new_name in data_dict:
+                                    new_name = random.choice('ABCDEFGHIJKLMNOPQRSTUVW') + '_' + name
+                                data_dict[new_name] = {'motion': n_motion,
+                                                       'length': len(n_motion),
+                                                       'text':[text_dict]}
+                                new_name_list.append(new_name)
+                                length_list.append(len(n_motion))
+                            except:
+                                print(line_split)
+                                print(line_split[2], line_split[3], f_tag, to_tag, name)
+                                # break
+
+                if flag:
+                    data_dict[name] = {'motion': motion,
+                                       'length': len(motion),
+                                       'text': text_data}
+                    new_name_list.append(name)
+                    length_list.append(len(motion))
+            except:
+                pass
+
+        name_list, length_list = zip(*sorted(zip(new_name_list, length_list), key=lambda x: x[1]))
+
+        self.mean = mean
+        self.std = std
+        self.length_arr = np.array(length_list)
+        self.data_dict = data_dict
+        self.name_list = name_list
+        self.reset_max_len(self.max_length)
+
+    def reset_max_len(self, length):
+        assert length <= self.max_motion_length
+        self.pointer = np.searchsorted(self.length_arr, length)
+        print("Pointer Pointing at %d"%self.pointer)
+        self.max_length = length
+
+    def inv_transform(self, data):
+        return data * self.std + self.mean
+
+    def __len__(self):
+        return len(self.data_dict) - self.pointer
+
+    def __getitem__(self, item):
+        idx = self.pointer + item
+        data = self.data_dict[self.name_list[idx]]
+        motion, m_length, text_list = data['motion'], data['length'], data['text']
+        # Randomly select a caption
+        text_data = random.choice(text_list)
+        caption, tokens = text_data['caption'], text_data['tokens']
+
+        if len(tokens) < self.opt.max_text_len:
+            # pad with "unk"
+            tokens = ['sos/OTHER'] + tokens + ['eos/OTHER']
+            sent_len = len(tokens)
+            tokens = tokens + ['unk/OTHER'] * (self.opt.max_text_len + 2 - sent_len)
+        else:
+            # crop
+            tokens = tokens[:self.opt.max_text_len]
+            tokens = ['sos/OTHER'] + tokens + ['eos/OTHER']
+            sent_len = len(tokens)
+        pos_one_hots = []
+        word_embeddings = []
+        for token in tokens:
+            word_emb, pos_oh = self.w_vectorizer[token]
+            pos_one_hots.append(pos_oh[None, :])
+            word_embeddings.append(word_emb[None, :])
+        pos_one_hots = np.concatenate(pos_one_hots, axis=0)
+        word_embeddings = np.concatenate(word_embeddings, axis=0)
+
+        # Crop the motions in to times of 4, and introduce small variations
+        if self.opt.unit_length < 10:
+            coin2 = np.random.choice(['single', 'single', 'double'])
+        else:
+            coin2 = 'single'
+
+        if coin2 == 'double':
+            m_length = (m_length // self.opt.unit_length - 1) * self.opt.unit_length
+        elif coin2 == 'single':
+            m_length = (m_length // self.opt.unit_length) * self.opt.unit_length
+        idx = random.randint(0, len(motion) - m_length)
+        motion = motion[idx:idx+m_length]
+
+        "Z Normalization"
+        motion = (motion - self.mean) / self.std
+
+        if m_length < self.max_motion_length:
+            motion = np.concatenate([motion,
+                                     np.zeros((self.max_motion_length - m_length, motion.shape[1]))
+                                     ], axis=0)
+        return word_embeddings, pos_one_hots, caption, sent_len, motion, m_length, '_'.join(tokens)
+
+
+def get_dataset_motion_loader(opt_path, batch_size, device):
+    opt = get_opt(opt_path, device)
+
+    # Configurations of T2M dataset and KIT dataset is almost the same
+    if opt.dataset_name == 't2m' or opt.dataset_name == 'kit':
+        print('Loading dataset %s ...' % opt.dataset_name)
+
+        mean = np.load(pjoin(opt.meta_dir, 'mean.npy'))
+        std = np.load(pjoin(opt.meta_dir, 'std.npy'))
+
+        w_vectorizer = WordVectorizer('./data/glove', 'our_vab')
+        split_file = pjoin(opt.data_root, 'test.txt')
+        dataset = Text2MotionDatasetV2(opt, mean, std, split_file, w_vectorizer)
+        dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=4, drop_last=True,
+                                collate_fn=collate_fn, shuffle=True)
+    else:
+        raise KeyError('Dataset not Recognized !!')
+
+    print('Ground Truth Dataset Loading Completed!!!')
+    return dataloader, dataset
+
+
+class MMGeneratedDataset(Dataset):
+    def __init__(self, opt, motion_dataset, w_vectorizer):
+        self.opt = opt
+        self.dataset = motion_dataset.mm_generated_motion
+        self.w_vectorizer = w_vectorizer
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, item):
+        data = self.dataset[item]
+        mm_motions = data['mm_motions']
+        m_lens = []
+        motions = []
+        for mm_motion in mm_motions:
+            m_lens.append(mm_motion['length'])
+            motion = mm_motion['motion']
+            if len(motion) < self.opt.max_motion_length:
+                motion = np.concatenate([motion,
+                                         np.zeros((self.opt.max_motion_length - len(motion), motion.shape[1]))
+                                         ], axis=0)
+            motion = motion[None, :]
+            motions.append(motion)
+        m_lens = np.array(m_lens, dtype=np.int)
+        motions = np.concatenate(motions, axis=0)
+        sort_indx = np.argsort(m_lens)[::-1].copy()
+        # print(m_lens)
+        # print(sort_indx)
+        # print(m_lens[sort_indx])
+        m_lens = m_lens[sort_indx]
+        motions = motions[sort_indx]
+        return motions, m_lens
+
+
+
+def get_motion_loader(opt, batch_size, trainer, ground_truth_dataset, mm_num_samples, mm_num_repeats):
+
+    # Currently the configurations of two datasets are almost the same
+    if opt.dataset_name == 't2m' or opt.dataset_name == 'kit':
+        w_vectorizer = WordVectorizer('./data/glove', 'our_vab')
+    else:
+        raise KeyError('Dataset not recognized!!')
+    print('Generating %s ...' % opt.name)
+
+    dataset = EvaluationDataset(opt, trainer, ground_truth_dataset, w_vectorizer, mm_num_samples, mm_num_repeats)
+    mm_dataset = MMGeneratedDataset(opt, dataset, w_vectorizer)
+
+    motion_loader = DataLoader(dataset, batch_size=batch_size, collate_fn=collate_fn, drop_last=True, num_workers=4)
+    mm_motion_loader = DataLoader(mm_dataset, batch_size=1, num_workers=1)
+
+    print('Generated Dataset Loading Completed!!!')
+
+    return motion_loader, mm_motion_loader
+
+
+def build_models(opt):
+    movement_enc = MovementConvEncoder(opt.dim_pose-4, opt.dim_movement_enc_hidden, opt.dim_movement_latent)
+    text_enc = TextEncoderBiGRUCo(word_size=opt.dim_word,
+                                  pos_size=opt.dim_pos_ohot,
+                                  hidden_size=opt.dim_text_hidden,
+                                  output_size=opt.dim_coemb_hidden,
+                                  device=opt.device)
+
+    motion_enc = MotionEncoderBiGRUCo(input_size=opt.dim_movement_latent,
+                                      hidden_size=opt.dim_motion_hidden,
+                                      output_size=opt.dim_coemb_hidden,
+                                      device=opt.device)
+
+    checkpoint = torch.load(pjoin('data/pretrained_models', opt.dataset_name, 'text_mot_match', 'model', 'finest.tar'),
+                            map_location=opt.device)
+    movement_enc.load_state_dict(checkpoint['movement_encoder'])
+    text_enc.load_state_dict(checkpoint['text_encoder'])
+    motion_enc.load_state_dict(checkpoint['motion_encoder'])
+    print('Loading Evaluation Model Wrapper (Epoch %d) Completed!!' % (checkpoint['epoch']))
+    return text_enc, motion_enc, movement_enc
+
+
+class EvaluatorModelWrapper(object):
+
+    def __init__(self, opt):
+
+        if opt.dataset_name == 't2m':
+            opt.dim_pose = 263
+        elif opt.dataset_name == 'kit':
+            opt.dim_pose = 251
+        else:
+            raise KeyError('Dataset not Recognized!!!')
+
+        opt.dim_word = 300
+        opt.max_motion_length = 196
+        opt.dim_pos_ohot = len(POS_enumerator)
+        opt.dim_motion_hidden = 1024
+        opt.max_text_len = 20
+        opt.dim_text_hidden = 512
+        opt.dim_coemb_hidden = 512
+
+        self.text_encoder, self.motion_encoder, self.movement_encoder = build_models(opt)
+        self.opt = opt
+        self.device = opt.device
+
+        self.text_encoder.to(opt.device)
+        self.motion_encoder.to(opt.device)
+        self.movement_encoder.to(opt.device)
+
+        self.text_encoder.eval()
+        self.motion_encoder.eval()
+        self.movement_encoder.eval()
+
+    # Please note that the results does not following the order of inputs
+    def get_co_embeddings(self, word_embs, pos_ohot, cap_lens, motions, m_lens):
+        with torch.no_grad():
+            word_embs = word_embs.detach().to(self.device).float()
+            pos_ohot = pos_ohot.detach().to(self.device).float()
+            motions = motions.detach().to(self.device).float()
+
+            align_idx = np.argsort(m_lens.data.tolist())[::-1].copy()
+            motions = motions[align_idx]
+            m_lens = m_lens[align_idx]
+
+            '''Movement Encoding'''
+            movements = self.movement_encoder(motions[..., :-4]).detach()
+            m_lens = m_lens // self.opt.unit_length
+            motion_embedding = self.motion_encoder(movements, m_lens)
+
+            '''Text Encoding'''
+            text_embedding = self.text_encoder(word_embs, pos_ohot, cap_lens)
+            text_embedding = text_embedding[align_idx]
+        return text_embedding, motion_embedding
+
+    # Please note that the results does not following the order of inputs
+    def get_motion_embeddings(self, motions, m_lens):
+        with torch.no_grad():
+            motions = motions.detach().to(self.device).float()
+
+            align_idx = np.argsort(m_lens.data.tolist())[::-1].copy()
+            motions = motions[align_idx]
+            m_lens = m_lens[align_idx]
+
+            '''Movement Encoding'''
+            movements = self.movement_encoder(motions[..., :-4]).detach()
+            m_lens = m_lens // self.opt.unit_length
+            motion_embedding = self.motion_encoder(movements, m_lens)
+        return motion_embedding

datasets/evaluator_models.py

@@ -0,0 +1,438 @@
+import torch
+import torch.nn as nn
+import numpy as np
+import time
+import math
+from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
+# from networks.layers import *
+import torch.nn.functional as F
+
+
+class ContrastiveLoss(torch.nn.Module):
+    """
+    Contrastive loss function.
+    Based on: http://yann.lecun.com/exdb/publis/pdf/hadsell-chopra-lecun-06.pdf
+    """
+    def __init__(self, margin=3.0):
+        super(ContrastiveLoss, self).__init__()
+        self.margin = margin
+
+    def forward(self, output1, output2, label):
+        euclidean_distance = F.pairwise_distance(output1, output2, keepdim=True)
+        loss_contrastive = torch.mean((1-label) * torch.pow(euclidean_distance, 2) +
+                                      (label) * torch.pow(torch.clamp(self.margin - euclidean_distance, min=0.0), 2))
+        return loss_contrastive
+
+
+def init_weight(m):
+    if isinstance(m, nn.Conv1d) or isinstance(m, nn.Linear) or isinstance(m, nn.ConvTranspose1d):
+        nn.init.xavier_normal_(m.weight)
+        # m.bias.data.fill_(0.01)
+        if m.bias is not None:
+            nn.init.constant_(m.bias, 0)
+
+
+def reparameterize(mu, logvar):
+    s_var = logvar.mul(0.5).exp_()
+    eps = s_var.data.new(s_var.size()).normal_()
+    return eps.mul(s_var).add_(mu)
+
+
+# batch_size, dimension and position
+# output: (batch_size, dim)
+def positional_encoding(batch_size, dim, pos):
+    assert batch_size == pos.shape[0]
+    positions_enc = np.array([
+        [pos[j] / np.power(10000, (i-i%2)/dim) for i in range(dim)]
+        for j in range(batch_size)
+    ], dtype=np.float32)
+    positions_enc[:, 0::2] = np.sin(positions_enc[:, 0::2])
+    positions_enc[:, 1::2] = np.cos(positions_enc[:, 1::2])
+    return torch.from_numpy(positions_enc).float()
+
+
+def get_padding_mask(batch_size, seq_len, cap_lens):
+    cap_lens = cap_lens.data.tolist()
+    mask_2d = torch.ones((batch_size, seq_len, seq_len), dtype=torch.float32)
+    for i, cap_len in enumerate(cap_lens):
+        mask_2d[i, :, :cap_len] = 0
+    return mask_2d.bool(), 1 - mask_2d[:, :, 0].clone()
+
+
+class PositionalEncoding(nn.Module):
+
+    def __init__(self, d_model, max_len=300):
+        super(PositionalEncoding, self).__init__()
+
+        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() * (-math.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).transpose(0, 1)
+        self.register_buffer('pe', pe)
+
+    def forward(self, pos):
+        return self.pe[pos]
+
+
+class MovementConvEncoder(nn.Module):
+    def __init__(self, input_size, hidden_size, output_size):
+        super(MovementConvEncoder, self).__init__()
+        self.main = nn.Sequential(
+            nn.Conv1d(input_size, hidden_size, 4, 2, 1),
+            nn.Dropout(0.2, inplace=True),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Conv1d(hidden_size, output_size, 4, 2, 1),
+            nn.Dropout(0.2, inplace=True),
+            nn.LeakyReLU(0.2, inplace=True),
+        )
+        self.out_net = nn.Linear(output_size, output_size)
+        self.main.apply(init_weight)
+        self.out_net.apply(init_weight)
+
+    def forward(self, inputs):
+        inputs = inputs.permute(0, 2, 1)
+        outputs = self.main(inputs).permute(0, 2, 1)
+        # print(outputs.shape)
+        return self.out_net(outputs)
+
+
+class MovementConvDecoder(nn.Module):
+    def __init__(self, input_size, hidden_size, output_size):
+        super(MovementConvDecoder, self).__init__()
+        self.main = nn.Sequential(
+            nn.ConvTranspose1d(input_size, hidden_size, 4, 2, 1),
+            # nn.Dropout(0.2, inplace=True),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.ConvTranspose1d(hidden_size, output_size, 4, 2, 1),
+            # nn.Dropout(0.2, inplace=True),
+            nn.LeakyReLU(0.2, inplace=True),
+        )
+        self.out_net = nn.Linear(output_size, output_size)
+
+        self.main.apply(init_weight)
+        self.out_net.apply(init_weight)
+
+    def forward(self, inputs):
+        inputs = inputs.permute(0, 2, 1)
+        outputs = self.main(inputs).permute(0, 2, 1)
+        return self.out_net(outputs)
+
+
+class TextVAEDecoder(nn.Module):
+    def __init__(self, text_size, input_size, output_size, hidden_size, n_layers):
+        super(TextVAEDecoder, self).__init__()
+        self.input_size = input_size
+        self.output_size = output_size
+        self.hidden_size = hidden_size
+        self.n_layers = n_layers
+        self.emb = nn.Sequential(
+            nn.Linear(input_size, hidden_size),
+            nn.LayerNorm(hidden_size),
+            nn.LeakyReLU(0.2, inplace=True))
+
+        self.z2init = nn.Linear(text_size, hidden_size * n_layers)
+        self.gru = nn.ModuleList([nn.GRUCell(hidden_size, hidden_size) for i in range(self.n_layers)])
+        self.positional_encoder = PositionalEncoding(hidden_size)
+
+
+        self.output = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size),
+            nn.LayerNorm(hidden_size),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Linear(hidden_size, output_size)
+        )
+
+        #
+        # self.output = nn.Sequential(
+        #     nn.Linear(hidden_size, hidden_size),
+        #     nn.LayerNorm(hidden_size),
+        #     nn.LeakyReLU(0.2, inplace=True),
+        #     nn.Linear(hidden_size, output_size-4)
+        # )
+
+        # self.contact_net = nn.Sequential(
+        #     nn.Linear(output_size-4, 64),
+        #     nn.LayerNorm(64),
+        #     nn.LeakyReLU(0.2, inplace=True),
+        #     nn.Linear(64, 4)
+        # )
+
+        self.output.apply(init_weight)
+        self.emb.apply(init_weight)
+        self.z2init.apply(init_weight)
+        # self.contact_net.apply(init_weight)
+
+    def get_init_hidden(self, latent):
+        hidden = self.z2init(latent)
+        hidden = torch.split(hidden, self.hidden_size, dim=-1)
+        return list(hidden)
+
+    def forward(self, inputs, last_pred, hidden, p):
+        h_in = self.emb(inputs)
+        pos_enc = self.positional_encoder(p).to(inputs.device).detach()
+        h_in = h_in + pos_enc
+        for i in range(self.n_layers):
+            # print(h_in.shape)
+            hidden[i] = self.gru[i](h_in, hidden[i])
+            h_in = hidden[i]
+        pose_pred = self.output(h_in)
+        # pose_pred = self.output(h_in) + last_pred.detach()
+        # contact = self.contact_net(pose_pred)
+        # return torch.cat([pose_pred, contact], dim=-1), hidden
+        return pose_pred, hidden
+
+
+class TextDecoder(nn.Module):
+    def __init__(self, text_size, input_size, output_size, hidden_size, n_layers):
+        super(TextDecoder, self).__init__()
+        self.input_size = input_size
+        self.output_size = output_size
+        self.hidden_size = hidden_size
+        self.n_layers = n_layers
+        self.emb = nn.Sequential(
+            nn.Linear(input_size, hidden_size),
+            nn.LayerNorm(hidden_size),
+            nn.LeakyReLU(0.2, inplace=True))
+
+        self.gru = nn.ModuleList([nn.GRUCell(hidden_size, hidden_size) for i in range(self.n_layers)])
+        self.z2init = nn.Linear(text_size, hidden_size * n_layers)
+        self.positional_encoder = PositionalEncoding(hidden_size)
+
+        self.mu_net = nn.Linear(hidden_size, output_size)
+        self.logvar_net = nn.Linear(hidden_size, output_size)
+
+        self.emb.apply(init_weight)
+        self.z2init.apply(init_weight)
+        self.mu_net.apply(init_weight)
+        self.logvar_net.apply(init_weight)
+
+    def get_init_hidden(self, latent):
+
+        hidden = self.z2init(latent)
+        hidden = torch.split(hidden, self.hidden_size, dim=-1)
+
+        return list(hidden)
+
+    def forward(self, inputs, hidden, p):
+        # print(inputs.shape)
+        x_in = self.emb(inputs)
+        pos_enc = self.positional_encoder(p).to(inputs.device).detach()
+        x_in = x_in + pos_enc
+
+        for i in range(self.n_layers):
+            hidden[i] = self.gru[i](x_in, hidden[i])
+            h_in = hidden[i]
+        mu = self.mu_net(h_in)
+        logvar = self.logvar_net(h_in)
+        z = reparameterize(mu, logvar)
+        return z, mu, logvar, hidden
+
+class AttLayer(nn.Module):
+    def __init__(self, query_dim, key_dim, value_dim):
+        super(AttLayer, self).__init__()
+        self.W_q = nn.Linear(query_dim, value_dim)
+        self.W_k = nn.Linear(key_dim, value_dim, bias=False)
+        self.W_v = nn.Linear(key_dim, value_dim)
+
+        self.softmax = nn.Softmax(dim=1)
+        self.dim = value_dim
+
+        self.W_q.apply(init_weight)
+        self.W_k.apply(init_weight)
+        self.W_v.apply(init_weight)
+
+    def forward(self, query, key_mat):
+        '''
+        query (batch, query_dim)
+        key (batch, seq_len, key_dim)
+        '''
+        # print(query.shape)
+        query_vec = self.W_q(query).unsqueeze(-1)       # (batch, value_dim, 1)
+        val_set = self.W_v(key_mat)                     # (batch, seq_len, value_dim)
+        key_set = self.W_k(key_mat)                     # (batch, seq_len, value_dim)
+
+        weights = torch.matmul(key_set, query_vec) / np.sqrt(self.dim)
+
+        co_weights = self.softmax(weights)              # (batch, seq_len, 1)
+        values = val_set * co_weights                   # (batch, seq_len, value_dim)
+        pred = values.sum(dim=1)                        # (batch, value_dim)
+        return pred, co_weights
+
+    def short_cut(self, querys, keys):
+        return self.W_q(querys), self.W_k(keys)
+
+
+class TextEncoderBiGRU(nn.Module):
+    def __init__(self, word_size, pos_size, hidden_size, device):
+        super(TextEncoderBiGRU, self).__init__()
+        self.device = device
+
+        self.pos_emb = nn.Linear(pos_size, word_size)
+        self.input_emb = nn.Linear(word_size, hidden_size)
+        self.gru = nn.GRU(hidden_size, hidden_size, batch_first=True, bidirectional=True)
+        # self.linear2 = nn.Linear(hidden_size, output_size)
+
+        self.input_emb.apply(init_weight)
+        self.pos_emb.apply(init_weight)
+        # self.linear2.apply(init_weight)
+        # self.batch_size = batch_size
+        self.hidden_size = hidden_size
+        self.hidden = nn.Parameter(torch.randn((2, 1, self.hidden_size), requires_grad=True))
+
+    # input(batch_size, seq_len, dim)
+    def forward(self, word_embs, pos_onehot, cap_lens):
+        num_samples = word_embs.shape[0]
+
+        pos_embs = self.pos_emb(pos_onehot)
+        inputs = word_embs + pos_embs
+        input_embs = self.input_emb(inputs)
+        hidden = self.hidden.repeat(1, num_samples, 1)
+
+        cap_lens = cap_lens.data.tolist()
+        emb = pack_padded_sequence(input_embs, cap_lens, batch_first=True)
+
+        gru_seq, gru_last = self.gru(emb, hidden)
+
+        gru_last = torch.cat([gru_last[0], gru_last[1]], dim=-1)
+        gru_seq = pad_packed_sequence(gru_seq, batch_first=True)[0]
+        forward_seq = gru_seq[..., :self.hidden_size]
+        backward_seq = gru_seq[..., self.hidden_size:].clone()
+
+        # Concate the forward and backward word embeddings
+        for i, length in enumerate(cap_lens):
+            backward_seq[i:i+1, :length] = torch.flip(backward_seq[i:i+1, :length].clone(), dims=[1])
+        gru_seq = torch.cat([forward_seq, backward_seq], dim=-1)
+
+        return gru_seq, gru_last
+
+
+class TextEncoderBiGRUCo(nn.Module):
+    def __init__(self, word_size, pos_size, hidden_size, output_size, device):
+        super(TextEncoderBiGRUCo, self).__init__()
+        self.device = device
+
+        self.pos_emb = nn.Linear(pos_size, word_size)
+        self.input_emb = nn.Linear(word_size, hidden_size)
+        self.gru = nn.GRU(hidden_size, hidden_size, batch_first=True, bidirectional=True)
+        self.output_net = nn.Sequential(
+            nn.Linear(hidden_size * 2, hidden_size),
+            nn.LayerNorm(hidden_size),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Linear(hidden_size, output_size)
+        )
+
+        self.input_emb.apply(init_weight)
+        self.pos_emb.apply(init_weight)
+        self.output_net.apply(init_weight)
+        # self.linear2.apply(init_weight)
+        # self.batch_size = batch_size
+        self.hidden_size = hidden_size
+        self.hidden = nn.Parameter(torch.randn((2, 1, self.hidden_size), requires_grad=True))
+
+    # input(batch_size, seq_len, dim)
+    def forward(self, word_embs, pos_onehot, cap_lens):
+        num_samples = word_embs.shape[0]
+
+        pos_embs = self.pos_emb(pos_onehot)
+        inputs = word_embs + pos_embs
+        input_embs = self.input_emb(inputs)
+        hidden = self.hidden.repeat(1, num_samples, 1)
+
+        cap_lens = cap_lens.data.tolist()
+        emb = pack_padded_sequence(input_embs, cap_lens, batch_first=True)
+
+        gru_seq, gru_last = self.gru(emb, hidden)
+
+        gru_last = torch.cat([gru_last[0], gru_last[1]], dim=-1)
+
+        return self.output_net(gru_last)
+
+
+class MotionEncoderBiGRUCo(nn.Module):
+    def __init__(self, input_size, hidden_size, output_size, device):
+        super(MotionEncoderBiGRUCo, self).__init__()
+        self.device = device
+
+        self.input_emb = nn.Linear(input_size, hidden_size)
+        self.gru = nn.GRU(hidden_size, hidden_size, batch_first=True, bidirectional=True)
+        self.output_net = nn.Sequential(
+            nn.Linear(hidden_size*2, hidden_size),
+            nn.LayerNorm(hidden_size),
+            nn.LeakyReLU(0.2, inplace=True),
+            nn.Linear(hidden_size, output_size)
+        )
+
+        self.input_emb.apply(init_weight)
+        self.output_net.apply(init_weight)
+        self.hidden_size = hidden_size
+        self.hidden = nn.Parameter(torch.randn((2, 1, self.hidden_size), requires_grad=True))
+
+    # input(batch_size, seq_len, dim)
+    def forward(self, inputs, m_lens):
+        num_samples = inputs.shape[0]
+
+        input_embs = self.input_emb(inputs)
+        hidden = self.hidden.repeat(1, num_samples, 1)
+
+        cap_lens = m_lens.data.tolist()
+        emb = pack_padded_sequence(input_embs, cap_lens, batch_first=True)
+
+        gru_seq, gru_last = self.gru(emb, hidden)
+
+        gru_last = torch.cat([gru_last[0], gru_last[1]], dim=-1)
+
+        return self.output_net(gru_last)
+
+
+class MotionLenEstimatorBiGRU(nn.Module):
+    def __init__(self, word_size, pos_size, hidden_size, output_size):
+        super(MotionLenEstimatorBiGRU, self).__init__()
+
+        self.pos_emb = nn.Linear(pos_size, word_size)
+        self.input_emb = nn.Linear(word_size, hidden_size)
+        self.gru = nn.GRU(hidden_size, hidden_size, batch_first=True, bidirectional=True)
+        nd = 512
+        self.output = nn.Sequential(
+            nn.Linear(hidden_size*2, nd),
+            nn.LayerNorm(nd),
+            nn.LeakyReLU(0.2, inplace=True),
+
+            nn.Linear(nd, nd // 2),
+            nn.LayerNorm(nd // 2),
+            nn.LeakyReLU(0.2, inplace=True),
+
+            nn.Linear(nd // 2, nd // 4),
+            nn.LayerNorm(nd // 4),
+            nn.LeakyReLU(0.2, inplace=True),
+
+            nn.Linear(nd // 4, output_size)
+        )
+        # self.linear2 = nn.Linear(hidden_size, output_size)
+
+        self.input_emb.apply(init_weight)
+        self.pos_emb.apply(init_weight)
+        self.output.apply(init_weight)
+        # self.linear2.apply(init_weight)
+        # self.batch_size = batch_size
+        self.hidden_size = hidden_size
+        self.hidden = nn.Parameter(torch.randn((2, 1, self.hidden_size), requires_grad=True))
+
+    # input(batch_size, seq_len, dim)
+    def forward(self, word_embs, pos_onehot, cap_lens):
+        num_samples = word_embs.shape[0]
+
+        pos_embs = self.pos_emb(pos_onehot)
+        inputs = word_embs + pos_embs
+        input_embs = self.input_emb(inputs)
+        hidden = self.hidden.repeat(1, num_samples, 1)
+
+        cap_lens = cap_lens.data.tolist()
+        emb = pack_padded_sequence(input_embs, cap_lens, batch_first=True)
+
+        gru_seq, gru_last = self.gru(emb, hidden)
+
+        gru_last = torch.cat([gru_last[0], gru_last[1]], dim=-1)
+
+        return self.output(gru_last)

models/__init__.py

@@ -0,0 +1,4 @@
+from .transformer import MotionTransformer
+from .gaussian_diffusion import GaussianDiffusion
+
+__all__ = ['MotionTransformer', 'GaussianDiffusion']

models/gaussian_diffusion.py

@@ -0,0 +1,1145 @@
+"""
+This code is borrowed from https://github.com/openai/guided-diffusion/blob/main/guided_diffusion/gaussian_diffusion.py
+"""
+
+import enum
+import math
+
+import numpy as np
+import torch as th
+
+
+from abc import ABC, abstractmethod
+import torch.distributed as dist
+
+
+def create_named_schedule_sampler(name, diffusion):
+    """
+    Create a ScheduleSampler from a library of pre-defined samplers.
+    :param name: the name of the sampler.
+    :param diffusion: the diffusion object to sample for.
+    """
+    if name == "uniform":
+        return UniformSampler(diffusion)
+    elif name == "loss-second-moment":
+        return LossSecondMomentResampler(diffusion)
+    else:
+        raise NotImplementedError(f"unknown schedule sampler: {name}")
+
+
+class ScheduleSampler(ABC):
+    """
+    A distribution over timesteps in the diffusion process, intended to reduce
+    variance of the objective.
+    By default, samplers perform unbiased importance sampling, in which the
+    objective's mean is unchanged.
+    However, subclasses may override sample() to change how the resampled
+    terms are reweighted, allowing for actual changes in the objective.
+    """
+
+    @abstractmethod
+    def weights(self):
+        """
+        Get a numpy array of weights, one per diffusion step.
+        The weights needn't be normalized, but must be positive.
+        """
+
+    def sample(self, batch_size, device):
+        """
+        Importance-sample timesteps for a batch.
+        :param batch_size: the number of timesteps.
+        :param device: the torch device to save to.
+        :return: a tuple (timesteps, weights):
+                 - timesteps: a tensor of timestep indices.
+                 - weights: a tensor of weights to scale the resulting losses.
+        """
+        w = self.weights()
+        p = w / np.sum(w)
+        indices_np = np.random.choice(len(p), size=(batch_size,), p=p)
+        indices = th.from_numpy(indices_np).long().to(device)
+        weights_np = 1 / (len(p) * p[indices_np])
+        weights = th.from_numpy(weights_np).float().to(device)
+        return indices, weights
+
+
+class UniformSampler(ScheduleSampler):
+    def __init__(self, diffusion):
+        self.diffusion = diffusion
+        self._weights = np.ones([diffusion.num_timesteps])
+
+    def weights(self):
+        return self._weights
+
+
+class LossAwareSampler(ScheduleSampler):
+    def update_with_local_losses(self, local_ts, local_losses):
+        """
+        Update the reweighting using losses from a model.
+        Call this method from each rank with a batch of timesteps and the
+        corresponding losses for each of those timesteps.
+        This method will perform synchronization to make sure all of the ranks
+        maintain the exact same reweighting.
+        :param local_ts: an integer Tensor of timesteps.
+        :param local_losses: a 1D Tensor of losses.
+        """
+        batch_sizes = [
+            th.tensor([0], dtype=th.int32, device=local_ts.device)
+            for _ in range(dist.get_world_size())
+        ]
+        dist.all_gather(
+            batch_sizes,
+            th.tensor([len(local_ts)], dtype=th.int32, device=local_ts.device),
+        )
+
+        # Pad all_gather batches to be the maximum batch size.
+        batch_sizes = [x.item() for x in batch_sizes]
+        max_bs = max(batch_sizes)
+
+        timestep_batches = [th.zeros(max_bs).to(local_ts) for bs in batch_sizes]
+        loss_batches = [th.zeros(max_bs).to(local_losses) for bs in batch_sizes]
+        dist.all_gather(timestep_batches, local_ts)
+        dist.all_gather(loss_batches, local_losses)
+        timesteps = [
+            x.item() for y, bs in zip(timestep_batches, batch_sizes) for x in y[:bs]
+        ]
+        losses = [x.item() for y, bs in zip(loss_batches, batch_sizes) for x in y[:bs]]
+        self.update_with_all_losses(timesteps, losses)
+
+    @abstractmethod
+    def update_with_all_losses(self, ts, losses):
+        """
+        Update the reweighting using losses from a model.
+        Sub-classes should override this method to update the reweighting
+        using losses from the model.
+        This method directly updates the reweighting without synchronizing
+        between workers. It is called by update_with_local_losses from all
+        ranks with identical arguments. Thus, it should have deterministic
+        behavior to maintain state across workers.
+        :param ts: a list of int timesteps.
+        :param losses: a list of float losses, one per timestep.
+        """
+
+
+class LossSecondMomentResampler(LossAwareSampler):
+    def __init__(self, diffusion, history_per_term=10, uniform_prob=0.001):
+        self.diffusion = diffusion
+        self.history_per_term = history_per_term
+        self.uniform_prob = uniform_prob
+        self._loss_history = np.zeros(
+            [diffusion.num_timesteps, history_per_term], dtype=np.float64
+        )
+        self._loss_counts = np.zeros([diffusion.num_timesteps], dtype=np.int)
+
+    def weights(self):
+        if not self._warmed_up():
+            return np.ones([self.diffusion.num_timesteps], dtype=np.float64)
+        weights = np.sqrt(np.mean(self._loss_history ** 2, axis=-1))
+        weights /= np.sum(weights)
+        weights *= 1 - self.uniform_prob
+        weights += self.uniform_prob / len(weights)
+        return weights
+
+    def update_with_all_losses(self, ts, losses):
+        for t, loss in zip(ts, losses):
+            if self._loss_counts[t] == self.history_per_term:
+                # Shift out the oldest loss term.
+                self._loss_history[t, :-1] = self._loss_history[t, 1:]
+                self._loss_history[t, -1] = loss
+            else:
+                self._loss_history[t, self._loss_counts[t]] = loss
+                self._loss_counts[t] += 1
+
+    def _warmed_up(self):
+        return (self._loss_counts == self.history_per_term).all()
+
+
+def mean_flat(tensor):
+    """
+    Take the mean over all non-batch dimensions.
+    """
+    return tensor.mean(dim=list(range(1, len(tensor.shape))))
+
+
+def normal_kl(mean1, logvar1, mean2, logvar2):
+    """
+    Compute the KL divergence between two gaussians.
+    Shapes are automatically broadcasted, so batches can be compared to
+    scalars, among other use cases.
+    """
+    tensor = None
+    for obj in (mean1, logvar1, mean2, logvar2):
+        if isinstance(obj, th.Tensor):
+            tensor = obj
+            break
+    assert tensor is not None, "at least one argument must be a Tensor"
+
+    # Force variances to be Tensors. Broadcasting helps convert scalars to
+    # Tensors, but it does not work for th.exp().
+    logvar1, logvar2 = [
+        x if isinstance(x, th.Tensor) else th.tensor(x).to(tensor)
+        for x in (logvar1, logvar2)
+    ]
+
+    return 0.5 * (
+        -1.0
+        + logvar2
+        - logvar1
+        + th.exp(logvar1 - logvar2)
+        + ((mean1 - mean2) ** 2) * th.exp(-logvar2)
+    )
+
+
+def approx_standard_normal_cdf(x):
+    """
+    A fast approximation of the cumulative distribution function of the
+    standard normal.
+    """
+    return 0.5 * (1.0 + th.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * th.pow(x, 3))))
+
+
+def discretized_gaussian_log_likelihood(x, *, means, log_scales):
+    """
+    Compute the log-likelihood of a Gaussian distribution discretizing to a
+    given image.
+    :param x: the target images. It is assumed that this was uint8 values,
+              rescaled to the range [-1, 1].
+    :param means: the Gaussian mean Tensor.
+    :param log_scales: the Gaussian log stddev Tensor.
+    :return: a tensor like x of log probabilities (in nats).
+    """
+    assert x.shape == means.shape == log_scales.shape
+    centered_x = x - means
+    inv_stdv = th.exp(-log_scales)
+    plus_in = inv_stdv * (centered_x + 1.0 / 255.0)
+    cdf_plus = approx_standard_normal_cdf(plus_in)
+    min_in = inv_stdv * (centered_x - 1.0 / 255.0)
+    cdf_min = approx_standard_normal_cdf(min_in)
+    log_cdf_plus = th.log(cdf_plus.clamp(min=1e-12))
+    log_one_minus_cdf_min = th.log((1.0 - cdf_min).clamp(min=1e-12))
+    cdf_delta = cdf_plus - cdf_min
+    log_probs = th.where(
+        x < -0.999,
+        log_cdf_plus,
+        th.where(x > 0.999, log_one_minus_cdf_min, th.log(cdf_delta.clamp(min=1e-12))),
+    )
+    assert log_probs.shape == x.shape
+    return log_probs
+
+
+def get_named_beta_schedule(schedule_name, num_diffusion_timesteps):
+    """
+    Get a pre-defined beta schedule for the given name.
+
+    The beta schedule library consists of beta schedules which remain similar
+    in the limit of num_diffusion_timesteps.
+    Beta schedules may be added, but should not be removed or changed once
+    they are committed to maintain backwards compatibility.
+    """
+    if schedule_name == "linear":
+        # Linear schedule from Ho et al, extended to work for any number of
+        # diffusion steps.
+        scale = 1000 / num_diffusion_timesteps
+        beta_start = scale * 0.0001
+        beta_end = scale * 0.02
+        return np.linspace(
+            beta_start, beta_end, num_diffusion_timesteps, dtype=np.float64
+        )
+    elif schedule_name == "cosine":
+        return betas_for_alpha_bar(
+            num_diffusion_timesteps,
+            lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2,
+        )
+    else:
+        raise NotImplementedError(f"unknown beta schedule: {schedule_name}")
+
+
+def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
+    """
+    Create a beta schedule that discretizes the given alpha_t_bar function,
+    which defines the cumulative product of (1-beta) over time from t = [0,1].
+
+    :param num_diffusion_timesteps: the number of betas to produce.
+    :param alpha_bar: a lambda that takes an argument t from 0 to 1 and
+                      produces the cumulative product of (1-beta) up to that
+                      part of the diffusion process.
+    :param max_beta: the maximum beta to use; use values lower than 1 to
+                     prevent singularities.
+    """
+    betas = []
+    for i in range(num_diffusion_timesteps):
+        t1 = i / num_diffusion_timesteps
+        t2 = (i + 1) / num_diffusion_timesteps
+        betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
+    return np.array(betas)
+
+
+class ModelMeanType(enum.Enum):
+    """
+    Which type of output the model predicts.
+    """
+
+    PREVIOUS_X = enum.auto()  # the model predicts x_{t-1}
+    START_X = enum.auto()  # the model predicts x_0
+    EPSILON = enum.auto()  # the model predicts epsilon
+
+
+class ModelVarType(enum.Enum):
+    """
+    What is used as the model's output variance.
+
+    The LEARNED_RANGE option has been added to allow the model to predict
+    values between FIXED_SMALL and FIXED_LARGE, making its job easier.
+    """
+
+    LEARNED = enum.auto()
+    FIXED_SMALL = enum.auto()
+    FIXED_LARGE = enum.auto()
+    LEARNED_RANGE = enum.auto()
+
+
+class LossType(enum.Enum):
+    MSE = enum.auto()  # use raw MSE loss (and KL when learning variances)
+    RESCALED_MSE = (
+        enum.auto()
+    )  # use raw MSE loss (with RESCALED_KL when learning variances)
+    KL = enum.auto()  # use the variational lower-bound
+    RESCALED_KL = enum.auto()  # like KL, but rescale to estimate the full VLB
+
+    def is_vb(self):
+        return self == LossType.KL or self == LossType.RESCALED_KL
+
+
+class GaussianDiffusion:
+    """
+    Utilities for training and sampling diffusion models.
+
+    Ported directly from here, and then adapted over time to further experimentation.
+    https://github.com/hojonathanho/diffusion/blob/1e0dceb3b3495bbe19116a5e1b3596cd0706c543/diffusion_tf/diffusion_utils_2.py#L42
+
+    :param betas: a 1-D numpy array of betas for each diffusion timestep,
+                  starting at T and going to 1.
+    :param model_mean_type: a ModelMeanType determining what the model outputs.
+    :param model_var_type: a ModelVarType determining how variance is output.
+    :param loss_type: a LossType determining the loss function to use.
+    :param rescale_timesteps: if True, pass floating point timesteps into the
+                              model so that they are always scaled like in the
+                              original paper (0 to 1000).
+    """
+
+    def __init__(
+        self,
+        *,
+        betas,
+        model_mean_type,
+        model_var_type,
+        loss_type,
+        rescale_timesteps=False,
+    ):
+        self.model_mean_type = model_mean_type
+        self.model_var_type = model_var_type
+        self.loss_type = loss_type
+        self.rescale_timesteps = rescale_timesteps
+
+        # Use float64 for accuracy.
+        betas = np.array(betas, dtype=np.float64)
+        self.betas = betas
+        assert len(betas.shape) == 1, "betas must be 1-D"
+        assert (betas > 0).all() and (betas <= 1).all()
+
+        self.num_timesteps = int(betas.shape[0])
+
+        alphas = 1.0 - betas
+        self.alphas_cumprod = np.cumprod(alphas, axis=0)
+        self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1])
+        self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0)
+        assert self.alphas_cumprod_prev.shape == (self.num_timesteps,)
+
+        # calculations for diffusion q(x_t | x_{t-1}) and others
+        self.sqrt_alphas_cumprod = np.sqrt(self.alphas_cumprod)
+        self.sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - self.alphas_cumprod)
+        self.log_one_minus_alphas_cumprod = np.log(1.0 - self.alphas_cumprod)
+        self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod)
+        self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - 1)
+
+        # calculations for posterior q(x_{t-1} | x_t, x_0)
+        self.posterior_variance = (
+            betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        )
+        # log calculation clipped because the posterior variance is 0 at the
+        # beginning of the diffusion chain.
+        self.posterior_log_variance_clipped = np.log(
+            np.append(self.posterior_variance[1], self.posterior_variance[1:])
+        )
+        self.posterior_mean_coef1 = (
+            betas * np.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
+        )
+        self.posterior_mean_coef2 = (
+            (1.0 - self.alphas_cumprod_prev)
+            * np.sqrt(alphas)
+            / (1.0 - self.alphas_cumprod)
+        )
+
+    def q_mean_variance(self, x_start, t):
+        """
+        Get the distribution q(x_t | x_0).
+
+        :param x_start: the [N x C x ...] tensor of noiseless inputs.
+        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+        :return: A tuple (mean, variance, log_variance), all of x_start's shape.
+        """
+        mean = (
+            _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+        )
+        variance = _extract_into_tensor(1.0 - self.alphas_cumprod, t, x_start.shape)
+        log_variance = _extract_into_tensor(
+            self.log_one_minus_alphas_cumprod, t, x_start.shape
+        )
+        return mean, variance, log_variance
+
+    def q_sample(self, x_start, t, noise=None):
+        """
+        Diffuse the data for a given number of diffusion steps.
+
+        In other words, sample from q(x_t | x_0).
+
+        :param x_start: the initial data batch.
+        :param t: the number of diffusion steps (minus 1). Here, 0 means one step.
+        :param noise: if specified, the split-out normal noise.
+        :return: A noisy version of x_start.
+        """
+        if noise is None:
+            noise = th.randn_like(x_start)
+        assert noise.shape == x_start.shape
+        return (
+            _extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
+            + _extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape)
+            * noise
+        )
+
+    def q_posterior_mean_variance(self, x_start, x_t, t):
+        """
+        Compute the mean and variance of the diffusion posterior:
+
+            q(x_{t-1} | x_t, x_0)
+
+        """
+        assert x_start.shape == x_t.shape
+        posterior_mean = (
+            _extract_into_tensor(self.posterior_mean_coef1, t, x_t.shape) * x_start
+            + _extract_into_tensor(self.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        posterior_variance = _extract_into_tensor(self.posterior_variance, t, x_t.shape)
+        posterior_log_variance_clipped = _extract_into_tensor(
+            self.posterior_log_variance_clipped, t, x_t.shape
+        )
+        assert (
+            posterior_mean.shape[0]
+            == posterior_variance.shape[0]
+            == posterior_log_variance_clipped.shape[0]
+            == x_start.shape[0]
+        )
+        return posterior_mean, posterior_variance, posterior_log_variance_clipped
+
+    def p_mean_variance(
+        self, model, x, t, clip_denoised=True, denoised_fn=None, model_kwargs=None
+    ):
+        """
+        Apply the model to get p(x_{t-1} | x_t), as well as a prediction of
+        the initial x, x_0.
+
+        :param model: the model, which takes a signal and a batch of timesteps
+                      as input.
+        :param x: the [N x C x ...] tensor at time t.
+        :param t: a 1-D Tensor of timesteps.
+        :param clip_denoised: if True, clip the denoised signal into [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample. Applies before
+            clip_denoised.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :return: a dict with the following keys:
+                 - 'mean': the model mean output.
+                 - 'variance': the model variance output.
+                 - 'log_variance': the log of 'variance'.
+                 - 'pred_xstart': the prediction for x_0.
+        """
+        if model_kwargs is None:
+            model_kwargs = {}
+
+        B, C = x.shape[:2]
+        assert t.shape == (B,)
+        model_output = model(x, self._scale_timesteps(t), **model_kwargs)
+
+        if self.model_var_type in [ModelVarType.LEARNED, ModelVarType.LEARNED_RANGE]:
+            assert model_output.shape == (B, 2 * C, *x.shape[2:])
+            model_output, model_var_values = th.split(model_output, C, dim=1)
+            if self.model_var_type == ModelVarType.LEARNED:
+                model_log_variance = model_var_values
+                model_variance = th.exp(model_log_variance)
+            else:
+                min_log = _extract_into_tensor(
+                    self.posterior_log_variance_clipped, t, x.shape
+                )
+                max_log = _extract_into_tensor(np.log(self.betas), t, x.shape)
+                # The model_var_values is [-1, 1] for [min_var, max_var].
+                frac = (model_var_values + 1) / 2
+                model_log_variance = frac * max_log + (1 - frac) * min_log
+                model_variance = th.exp(model_log_variance)
+        else:
+            model_variance, model_log_variance = {
+                # for fixedlarge, we set the initial (log-)variance like so
+                # to get a better decoder log likelihood.
+                ModelVarType.FIXED_LARGE: (
+                    np.append(self.posterior_variance[1], self.betas[1:]),
+                    np.log(np.append(self.posterior_variance[1], self.betas[1:])),
+                ),
+                ModelVarType.FIXED_SMALL: (
+                    self.posterior_variance,
+                    self.posterior_log_variance_clipped,
+                ),
+            }[self.model_var_type]
+            model_variance = _extract_into_tensor(model_variance, t, x.shape)
+            model_log_variance = _extract_into_tensor(model_log_variance, t, x.shape)
+
+        def process_xstart(x):
+            if denoised_fn is not None:
+                x = denoised_fn(x)
+            if clip_denoised:
+                return x.clamp(-1, 1)
+            return x
+
+        if self.model_mean_type == ModelMeanType.PREVIOUS_X:
+            pred_xstart = process_xstart(
+                self._predict_xstart_from_xprev(x_t=x, t=t, xprev=model_output)
+            )
+            model_mean = model_output
+        elif self.model_mean_type in [ModelMeanType.START_X, ModelMeanType.EPSILON]:
+            if self.model_mean_type == ModelMeanType.START_X:
+                pred_xstart = process_xstart(model_output)
+            else:
+                pred_xstart = process_xstart(
+                    self._predict_xstart_from_eps(x_t=x, t=t, eps=model_output)
+                )
+            model_mean, _, _ = self.q_posterior_mean_variance(
+                x_start=pred_xstart, x_t=x, t=t
+            )
+        else:
+            raise NotImplementedError(self.model_mean_type)
+
+        assert (
+            model_mean.shape == model_log_variance.shape == pred_xstart.shape == x.shape
+        )
+        return {
+            "mean": model_mean,
+            "variance": model_variance,
+            "log_variance": model_log_variance,
+            "pred_xstart": pred_xstart,
+        }
+
+    def _predict_xstart_from_eps(self, x_t, t, eps):
+        assert x_t.shape == eps.shape
+        return (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
+            - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape) * eps
+        )
+
+    def _predict_xstart_from_xprev(self, x_t, t, xprev):
+        assert x_t.shape == xprev.shape
+        return (  # (xprev - coef2*x_t) / coef1
+            _extract_into_tensor(1.0 / self.posterior_mean_coef1, t, x_t.shape) * xprev
+            - _extract_into_tensor(
+                self.posterior_mean_coef2 / self.posterior_mean_coef1, t, x_t.shape
+            )
+            * x_t
+        )
+
+    def _predict_eps_from_xstart(self, x_t, t, pred_xstart):
+        return (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t
+            - pred_xstart
+        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
+
+    def _scale_timesteps(self, t):
+        if self.rescale_timesteps:
+            return t.float() * (1000.0 / self.num_timesteps)
+        return t
+
+    def condition_mean(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
+        """
+        Compute the mean for the previous step, given a function cond_fn that
+        computes the gradient of a conditional log probability with respect to
+        x. In particular, cond_fn computes grad(log(p(y|x))), and we want to
+        condition on y.
+
+        This uses the conditioning strategy from Sohl-Dickstein et al. (2015).
+        """
+        gradient = cond_fn(x, self._scale_timesteps(t), **model_kwargs)
+        new_mean = (
+            p_mean_var["mean"].float() + p_mean_var["variance"] * gradient.float()
+        )
+        return new_mean
+
+    def condition_score(self, cond_fn, p_mean_var, x, t, model_kwargs=None):
+        """
+        Compute what the p_mean_variance output would have been, should the
+        model's score function be conditioned by cond_fn.
+
+        See condition_mean() for details on cond_fn.
+
+        Unlike condition_mean(), this instead uses the conditioning strategy
+        from Song et al (2020).
+        """
+        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+
+        eps = self._predict_eps_from_xstart(x, t, p_mean_var["pred_xstart"])
+        eps = eps - (1 - alpha_bar).sqrt() * cond_fn(
+            x, self._scale_timesteps(t), **model_kwargs
+        )
+
+        out = p_mean_var.copy()
+        out["pred_xstart"] = self._predict_xstart_from_eps(x, t, eps)
+        out["mean"], _, _ = self.q_posterior_mean_variance(
+            x_start=out["pred_xstart"], x_t=x, t=t
+        )
+        return out
+
+    def p_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        pre_seq=None,
+        transl_req=None,
+        model_kwargs=None,
+    ):
+        """
+        Sample x_{t-1} from the model at the given timestep.
+
+        :param model: the model to sample from.
+        :param x: the current tensor at x_{t-1}.
+        :param t: the value of t, starting at 0 for the first diffusion step.
+        :param clip_denoised: if True, clip the x_start prediction to [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample.
+        :param cond_fn: if not None, this is a gradient function that acts
+                        similarly to the model.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :return: a dict containing the following keys:
+                 - 'sample': a random sample from the model.
+                 - 'pred_xstart': a prediction of x_0.
+        """
+        # concat seq
+        if pre_seq is not None:
+            T = pre_seq.shape[2]
+            noise = th.randn_like(pre_seq)
+            x_t = self.q_sample(pre_seq, t, noise=noise)
+            x[:, :, :T] = x_t
+
+        if transl_req is not None:
+            for item in transl_req:
+                noise = th.randn(2).type_as(x)
+                transl = th.Tensor(item[1:]).type_as(x)
+                x_t = self.q_sample(transl, t, noise=noise)
+                x[:, :2, item[0]] = x_t
+
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        noise = th.randn_like(x)
+        nonzero_mask = (
+            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+        )  # no noise when t == 0
+        if cond_fn is not None:
+            out["mean"] = self.condition_mean(
+                cond_fn, out, x, t, model_kwargs=model_kwargs
+            )
+        sample = out["mean"] + nonzero_mask * th.exp(0.5 * out["log_variance"]) * noise
+        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
+
+    def p_sample_loop(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        pre_seq=None,
+        transl_req=None,
+        progress=False,
+    ):
+        """
+        Generate samples from the model.
+
+        :param model: the model module.
+        :param shape: the shape of the samples, (N, C, H, W).
+        :param noise: if specified, the noise from the encoder to sample.
+                      Should be of the same shape as `shape`.
+        :param clip_denoised: if True, clip x_start predictions to [-1, 1].
+        :param denoised_fn: if not None, a function which applies to the
+            x_start prediction before it is used to sample.
+        :param cond_fn: if not None, this is a gradient function that acts
+                        similarly to the model.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :param device: if specified, the device to create the samples on.
+                       If not specified, use a model parameter's device.
+        :param progress: if True, show a tqdm progress bar.
+        :return: a non-differentiable batch of samples.
+        """
+        final = None
+        for sample in self.p_sample_loop_progressive(
+            model,
+            shape,
+            noise=noise,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            cond_fn=cond_fn,
+            model_kwargs=model_kwargs,
+            device=device,
+            pre_seq=pre_seq,
+            transl_req=transl_req,
+            progress=progress,
+        ):
+            final = sample
+        return final["sample"]
+
+    def p_sample_loop_progressive(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        pre_seq=None,
+        transl_req=None,
+        progress=False,
+    ):
+        """
+        Generate samples from the model and yield intermediate samples from
+        each timestep of diffusion.
+
+        Arguments are the same as p_sample_loop().
+        Returns a generator over dicts, where each dict is the return value of
+        p_sample().
+        """
+        if device is None:
+            device = next(model.parameters()).device
+        assert isinstance(shape, (tuple, list))
+        if noise is not None:
+            img = noise
+        else:
+            img = th.randn(*shape, device=device)
+        indices = list(range(self.num_timesteps))[::-1]
+        if progress:
+            # Lazy import so that we don't depend on tqdm.
+            from tqdm.auto import tqdm
+
+            indices = tqdm(indices)
+
+        for i in indices:
+            t = th.tensor([i] * shape[0], device=device)
+            with th.no_grad():
+                out = self.p_sample(
+                    model,
+                    img,
+                    t,
+                    clip_denoised=clip_denoised,
+                    denoised_fn=denoised_fn,
+                    cond_fn=cond_fn,
+                    model_kwargs=model_kwargs,
+                    pre_seq=pre_seq,
+                    transl_req=transl_req
+                )
+                yield out
+                img = out["sample"]
+
+    def ddim_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        eta=0.0,
+    ):
+        """
+        Sample x_{t-1} from the model using DDIM.
+
+        Same usage as p_sample().
+        """
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        if cond_fn is not None:
+            out = self.condition_score(cond_fn, out, x, t, model_kwargs=model_kwargs)
+
+        # Usually our model outputs epsilon, but we re-derive it
+        # in case we used x_start or x_prev prediction.
+        eps = self._predict_eps_from_xstart(x, t, out["pred_xstart"])
+
+        alpha_bar = _extract_into_tensor(self.alphas_cumprod, t, x.shape)
+        alpha_bar_prev = _extract_into_tensor(self.alphas_cumprod_prev, t, x.shape)
+        sigma = (
+            eta
+            * th.sqrt((1 - alpha_bar_prev) / (1 - alpha_bar))
+            * th.sqrt(1 - alpha_bar / alpha_bar_prev)
+        )
+        # Equation 12.
+        noise = th.randn_like(x)
+        mean_pred = (
+            out["pred_xstart"] * th.sqrt(alpha_bar_prev)
+            + th.sqrt(1 - alpha_bar_prev - sigma ** 2) * eps
+        )
+        nonzero_mask = (
+            (t != 0).float().view(-1, *([1] * (len(x.shape) - 1)))
+        )  # no noise when t == 0
+        sample = mean_pred + nonzero_mask * sigma * noise
+        return {"sample": sample, "pred_xstart": out["pred_xstart"]}
+
+    def ddim_reverse_sample(
+        self,
+        model,
+        x,
+        t,
+        clip_denoised=True,
+        denoised_fn=None,
+        model_kwargs=None,
+        eta=0.0,
+    ):
+        """
+        Sample x_{t+1} from the model using DDIM reverse ODE.
+        """
+        assert eta == 0.0, "Reverse ODE only for deterministic path"
+        out = self.p_mean_variance(
+            model,
+            x,
+            t,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            model_kwargs=model_kwargs,
+        )
+        # Usually our model outputs epsilon, but we re-derive it
+        # in case we used x_start or x_prev prediction.
+        eps = (
+            _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, x.shape) * x
+            - out["pred_xstart"]
+        ) / _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, x.shape)
+        alpha_bar_next = _extract_into_tensor(self.alphas_cumprod_next, t, x.shape)
+
+        # Equation 12. reversed
+        mean_pred = (
+            out["pred_xstart"] * th.sqrt(alpha_bar_next)
+            + th.sqrt(1 - alpha_bar_next) * eps
+        )
+
+        return {"sample": mean_pred, "pred_xstart": out["pred_xstart"]}
+
+    def ddim_sample_loop(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        eta=0.0,
+    ):
+        """
+        Generate samples from the model using DDIM.
+
+        Same usage as p_sample_loop().
+        """
+        final = None
+        for sample in self.ddim_sample_loop_progressive(
+            model,
+            shape,
+            noise=noise,
+            clip_denoised=clip_denoised,
+            denoised_fn=denoised_fn,
+            cond_fn=cond_fn,
+            model_kwargs=model_kwargs,
+            device=device,
+            progress=progress,
+            eta=eta,
+        ):
+            final = sample
+        return final["sample"]
+
+    def ddim_sample_loop_progressive(
+        self,
+        model,
+        shape,
+        noise=None,
+        clip_denoised=True,
+        denoised_fn=None,
+        cond_fn=None,
+        model_kwargs=None,
+        device=None,
+        progress=False,
+        eta=0.0,
+    ):
+        """
+        Use DDIM to sample from the model and yield intermediate samples from
+        each timestep of DDIM.
+
+        Same usage as p_sample_loop_progressive().
+        """
+        if device is None:
+            device = next(model.parameters()).device
+        assert isinstance(shape, (tuple, list))
+        if noise is not None:
+            img = noise
+        else:
+            img = th.randn(*shape, device=device)
+        indices = list(range(self.num_timesteps))[::-1]
+
+        if progress:
+            # Lazy import so that we don't depend on tqdm.
+            from tqdm.auto import tqdm
+
+            indices = tqdm(indices)
+
+        for i in indices:
+            t = th.tensor([i] * shape[0], device=device)
+            with th.no_grad():
+                out = self.ddim_sample(
+                    model,
+                    img,
+                    t,
+                    clip_denoised=clip_denoised,
+                    denoised_fn=denoised_fn,
+                    cond_fn=cond_fn,
+                    model_kwargs=model_kwargs,
+                    eta=eta,
+                )
+                yield out
+                img = out["sample"]
+
+    def _vb_terms_bpd(
+        self, model, x_start, x_t, t, clip_denoised=True, model_kwargs=None
+    ):
+        """
+        Get a term for the variational lower-bound.
+
+        The resulting units are bits (rather than nats, as one might expect).
+        This allows for comparison to other papers.
+
+        :return: a dict with the following keys:
+                 - 'output': a shape [N] tensor of NLLs or KLs.
+                 - 'pred_xstart': the x_0 predictions.
+        """
+        true_mean, _, true_log_variance_clipped = self.q_posterior_mean_variance(
+            x_start=x_start, x_t=x_t, t=t
+        )
+        out = self.p_mean_variance(
+            model, x_t, t, clip_denoised=clip_denoised, model_kwargs=model_kwargs
+        )
+        kl = normal_kl(
+            true_mean, true_log_variance_clipped, out["mean"], out["log_variance"]
+        )
+        kl = mean_flat(kl) / np.log(2.0)
+
+        decoder_nll = -discretized_gaussian_log_likelihood(
+            x_start, means=out["mean"], log_scales=0.5 * out["log_variance"]
+        )
+        assert decoder_nll.shape == x_start.shape
+        decoder_nll = mean_flat(decoder_nll) / np.log(2.0)
+
+        # At the first timestep return the decoder NLL,
+        # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
+        output = th.where((t == 0), decoder_nll, kl)
+        return {"output": output, "pred_xstart": out["pred_xstart"]}
+
+    def training_losses(self, model, x_start, t, model_kwargs=None, noise=None):
+        """
+        Compute training losses for a single timestep.
+
+        :param model: the model to evaluate loss on.
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :param t: a batch of timestep indices.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+        :param noise: if specified, the specific Gaussian noise to try to remove.
+        :return: a dict with the key "loss" containing a tensor of shape [N].
+                 Some mean or variance settings may also have other keys.
+        """
+        if model_kwargs is None:
+            model_kwargs = {}
+        if noise is None:
+            noise = th.randn_like(x_start)
+        x_t = self.q_sample(x_start, t, noise=noise)
+
+        terms = {}
+
+        if self.loss_type == LossType.KL or self.loss_type == LossType.RESCALED_KL:
+            terms["loss"] = self._vb_terms_bpd(
+                model=model,
+                x_start=x_start,
+                x_t=x_t,
+                t=t,
+                clip_denoised=False,
+                model_kwargs=model_kwargs,
+            )["output"]
+            if self.loss_type == LossType.RESCALED_KL:
+                terms["loss"] *= self.num_timesteps
+        elif self.loss_type == LossType.MSE or self.loss_type == LossType.RESCALED_MSE:
+            model_output = model(x_t, self._scale_timesteps(t), **model_kwargs)
+
+            if self.model_var_type in [
+                ModelVarType.LEARNED,
+                ModelVarType.LEARNED_RANGE,
+            ]:
+                B, C = x_t.shape[:2]
+                assert model_output.shape == (B, C * 2, *x_t.shape[2:])
+                model_output, model_var_values = th.split(model_output, C, dim=1)
+                # Learn the variance using the variational bound, but don't let
+                # it affect our mean prediction.
+                frozen_out = th.cat([model_output.detach(), model_var_values], dim=1)
+                terms["vb"] = self._vb_terms_bpd(
+                    model=lambda *args, r=frozen_out: r,
+                    x_start=x_start,
+                    x_t=x_t,
+                    t=t,
+                    clip_denoised=False,
+                )["output"]
+                if self.loss_type == LossType.RESCALED_MSE:
+                    # Divide by 1000 for equivalence with initial implementation.
+                    # Without a factor of 1/1000, the VB term hurts the MSE term.
+                    terms["vb"] *= self.num_timesteps / 1000.0
+
+            target = {
+                ModelMeanType.PREVIOUS_X: self.q_posterior_mean_variance(
+                    x_start=x_start, x_t=x_t, t=t
+                )[0],
+                ModelMeanType.START_X: x_start,
+                ModelMeanType.EPSILON: noise,
+            }[self.model_mean_type]
+            assert model_output.shape == target.shape == x_start.shape
+            terms["mse"] = mean_flat((target - model_output) ** 2).view(-1, 1).mean(-1)
+            # if "vb" in terms:
+            #     terms["loss"] = terms["mse"] + terms["vb"]
+            # else:
+            #     terms["loss"] = terms["mse"]
+            terms["target"] = target
+            terms["pred"] = model_output
+        else:
+            raise NotImplementedError(self.loss_type)
+
+        return terms
+
+    def _prior_bpd(self, x_start):
+        """
+        Get the prior KL term for the variational lower-bound, measured in
+        bits-per-dim.
+
+        This term can't be optimized, as it only depends on the encoder.
+
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :return: a batch of [N] KL values (in bits), one per batch element.
+        """
+        batch_size = x_start.shape[0]
+        t = th.tensor([self.num_timesteps - 1] * batch_size, device=x_start.device)
+        qt_mean, _, qt_log_variance = self.q_mean_variance(x_start, t)
+        kl_prior = normal_kl(
+            mean1=qt_mean, logvar1=qt_log_variance, mean2=0.0, logvar2=0.0
+        )
+        return mean_flat(kl_prior) / np.log(2.0)
+
+    def calc_bpd_loop(self, model, x_start, clip_denoised=True, model_kwargs=None):
+        """
+        Compute the entire variational lower-bound, measured in bits-per-dim,
+        as well as other related quantities.
+
+        :param model: the model to evaluate loss on.
+        :param x_start: the [N x C x ...] tensor of inputs.
+        :param clip_denoised: if True, clip denoised samples.
+        :param model_kwargs: if not None, a dict of extra keyword arguments to
+            pass to the model. This can be used for conditioning.
+
+        :return: a dict containing the following keys:
+                 - total_bpd: the total variational lower-bound, per batch element.
+                 - prior_bpd: the prior term in the lower-bound.
+                 - vb: an [N x T] tensor of terms in the lower-bound.
+                 - xstart_mse: an [N x T] tensor of x_0 MSEs for each timestep.
+                 - mse: an [N x T] tensor of epsilon MSEs for each timestep.
+        """
+        device = x_start.device
+        batch_size = x_start.shape[0]
+
+        vb = []
+        xstart_mse = []
+        mse = []
+        for t in list(range(self.num_timesteps))[::-1]:
+            t_batch = th.tensor([t] * batch_size, device=device)
+            noise = th.randn_like(x_start)
+            x_t = self.q_sample(x_start=x_start, t=t_batch, noise=noise)
+            # Calculate VLB term at the current timestep
+            with th.no_grad():
+                out = self._vb_terms_bpd(
+                    model,
+                    x_start=x_start,
+                    x_t=x_t,
+                    t=t_batch,
+                    clip_denoised=clip_denoised,
+                    model_kwargs=model_kwargs,
+                )
+            vb.append(out["output"])
+            xstart_mse.append(mean_flat((out["pred_xstart"] - x_start) ** 2))
+            eps = self._predict_eps_from_xstart(x_t, t_batch, out["pred_xstart"])
+            mse.append(mean_flat((eps - noise) ** 2))
+
+        vb = th.stack(vb, dim=1)
+        xstart_mse = th.stack(xstart_mse, dim=1)
+        mse = th.stack(mse, dim=1)
+
+        prior_bpd = self._prior_bpd(x_start)
+        total_bpd = vb.sum(dim=1) + prior_bpd
+        return {
+            "total_bpd": total_bpd,
+            "prior_bpd": prior_bpd,
+            "vb": vb,
+            "xstart_mse": xstart_mse,
+            "mse": mse,
+        }
+
+
+def _extract_into_tensor(arr, timesteps, broadcast_shape):
+    """
+    Extract values from a 1-D numpy array for a batch of indices.
+
+    :param arr: the 1-D numpy array.
+    :param timesteps: a tensor of indices into the array to extract.
+    :param broadcast_shape: a larger shape of K dimensions with the batch
+                            dimension equal to the length of timesteps.
+    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
+    """
+    res = th.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
+    while len(res.shape) < len(broadcast_shape):
+        res = res[..., None]
+    return res.expand(broadcast_shape)
+

models/transformer.py

@@ -0,0 +1,426 @@
+"""
+Copyright 2021 S-Lab
+"""
+
+from cv2 import norm
+import torch
+import torch.nn.functional as F
+from torch import layer_norm, nn
+import numpy as np
+import clip
+
+import math
+
+
+def timestep_embedding(timesteps, dim, max_period=10000):
+    """
+    Create sinusoidal timestep embeddings.
+    :param timesteps: a 1-D Tensor of N indices, one per batch element.
+                      These may be fractional.
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an [N x dim] Tensor of positional embeddings.
+    """
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+    ).to(device=timesteps.device)
+    args = timesteps[:, None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+    return embedding
+
+
+def set_requires_grad(nets, requires_grad=False):
+    """Set requies_grad for all the networks.
+
+    Args:
+        nets (nn.Module | list[nn.Module]): A list of networks or a single
+            network.
+        requires_grad (bool): Whether the networks require gradients or not
+    """
+    if not isinstance(nets, list):
+        nets = [nets]
+    for net in nets:
+        if net is not None:
+            for param in net.parameters():
+                param.requires_grad = requires_grad
+
+
+def zero_module(module):
+    """
+    Zero out the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+class StylizationBlock(nn.Module):
+
+    def __init__(self, latent_dim, time_embed_dim, dropout):
+        super().__init__()
+        self.emb_layers = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(time_embed_dim, 2 * latent_dim),
+        )
+        self.norm = nn.LayerNorm(latent_dim)
+        self.out_layers = nn.Sequential(
+            nn.SiLU(),
+            nn.Dropout(p=dropout),
+            zero_module(nn.Linear(latent_dim, latent_dim)),
+        )
+
+    def forward(self, h, emb):
+        """
+        h: B, T, D
+        emb: B, D
+        """
+        # B, 1, 2D
+        emb_out = self.emb_layers(emb).unsqueeze(1)
+        # scale: B, 1, D / shift: B, 1, D
+        scale, shift = torch.chunk(emb_out, 2, dim=2)
+        h = self.norm(h) * (1 + scale) + shift
+        h = self.out_layers(h)
+        return h
+
+
+class LinearTemporalSelfAttention(nn.Module):
+
+    def __init__(self, seq_len, latent_dim, num_head, dropout, time_embed_dim):
+        super().__init__()
+        self.num_head = num_head
+        self.norm = nn.LayerNorm(latent_dim)
+        self.query = nn.Linear(latent_dim, latent_dim)
+        self.key = nn.Linear(latent_dim, latent_dim)
+        self.value = nn.Linear(latent_dim, latent_dim)
+        self.dropout = nn.Dropout(dropout)
+        self.proj_out = StylizationBlock(latent_dim, time_embed_dim, dropout)
+    
+    def forward(self, x, emb, src_mask):
+        """
+        x: B, T, D
+        """
+        B, T, D = x.shape
+        H = self.num_head
+        # B, T, D
+        query = self.query(self.norm(x))
+        # B, T, D
+        key = (self.key(self.norm(x)) + (1 - src_mask) * -1000000)
+        query = F.softmax(query.view(B, T, H, -1), dim=-1)
+        key = F.softmax(key.view(B, T, H, -1), dim=1)
+        # B, T, H, HD
+        value = (self.value(self.norm(x)) * src_mask).view(B, T, H, -1)
+        # B, H, HD, HD
+        attention = torch.einsum('bnhd,bnhl->bhdl', key, value)
+        y = torch.einsum('bnhd,bhdl->bnhl', query, attention).reshape(B, T, D)
+        y = x + self.proj_out(y, emb)
+        return y
+
+
+class LinearTemporalCrossAttention(nn.Module):
+
+    def __init__(self, seq_len, latent_dim, text_latent_dim, num_head, dropout, time_embed_dim):
+        super().__init__()
+        self.num_head = num_head
+        self.norm = nn.LayerNorm(latent_dim)
+        self.text_norm = nn.LayerNorm(text_latent_dim)
+        self.query = nn.Linear(latent_dim, latent_dim)
+        self.key = nn.Linear(text_latent_dim, latent_dim)
+        self.value = nn.Linear(text_latent_dim, latent_dim)
+        self.dropout = nn.Dropout(dropout)
+        self.proj_out = StylizationBlock(latent_dim, time_embed_dim, dropout)
+    
+    def forward(self, x, xf, emb):
+        """
+        x: B, T, D
+        xf: B, N, L
+        """
+        B, T, D = x.shape
+        N = xf.shape[1]
+        H = self.num_head
+        # B, T, D
+        query = self.query(self.norm(x))
+        # B, N, D
+        key = self.key(self.text_norm(xf))
+        query = F.softmax(query.view(B, T, H, -1), dim=-1)
+        key = F.softmax(key.view(B, N, H, -1), dim=1)
+        # B, N, H, HD
+        value = self.value(self.text_norm(xf)).view(B, N, H, -1)
+        # B, H, HD, HD
+        attention = torch.einsum('bnhd,bnhl->bhdl', key, value)
+        y = torch.einsum('bnhd,bhdl->bnhl', query, attention).reshape(B, T, D)
+        y = x + self.proj_out(y, emb)
+        return y
+
+class FFN(nn.Module):
+
+    def __init__(self, latent_dim, ffn_dim, dropout, time_embed_dim):
+        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)
+        self.proj_out = StylizationBlock(latent_dim, time_embed_dim, dropout)
+
+    def forward(self, x, emb):
+        y = self.linear2(self.dropout(self.activation(self.linear1(x))))
+        y = x + self.proj_out(y, emb)
+        return y
+    
+
+class LinearTemporalDiffusionTransformerDecoderLayer(nn.Module):
+
+    def __init__(self,
+                 seq_len=60,
+                 latent_dim=32,
+                 text_latent_dim=512,
+                 time_embed_dim=128,
+                 ffn_dim=256,
+                 num_head=4,
+                 dropout=0.1):
+        super().__init__()
+        self.sa_block = LinearTemporalSelfAttention(
+            seq_len, latent_dim, num_head, dropout, time_embed_dim)
+        self.ca_block = LinearTemporalCrossAttention(
+            seq_len, latent_dim, text_latent_dim, num_head, dropout, time_embed_dim)
+        self.ffn = FFN(latent_dim, ffn_dim, dropout, time_embed_dim)
+
+    def forward(self, x, xf, emb, src_mask):
+        x = self.sa_block(x, emb, src_mask)
+        x = self.ca_block(x, xf, emb)
+        x = self.ffn(x, emb)
+        return x
+
+class TemporalSelfAttention(nn.Module):
+
+    def __init__(self, seq_len, latent_dim, num_head, dropout, time_embed_dim):
+        super().__init__()
+        self.num_head = num_head
+        self.norm = nn.LayerNorm(latent_dim)
+        self.query = nn.Linear(latent_dim, latent_dim)
+        self.key = nn.Linear(latent_dim, latent_dim)
+        self.value = nn.Linear(latent_dim, latent_dim)
+        self.dropout = nn.Dropout(dropout)
+        self.proj_out = StylizationBlock(latent_dim, time_embed_dim, dropout)
+    
+    def forward(self, x, emb, src_mask):
+        """
+        x: B, T, D
+        """
+        B, T, D = x.shape
+        H = self.num_head
+        # B, T, 1, D
+        query = self.query(self.norm(x)).unsqueeze(2)
+        # B, 1, T, D
+        key = self.key(self.norm(x)).unsqueeze(1)
+        query = query.view(B, T, H, -1)
+        key = key.view(B, T, H, -1)
+        # B, T, T, H
+        attention = torch.einsum('bnhd,bmhd->bnmh', query, key) / math.sqrt(D // H)
+        attention = attention + (1 - src_mask.unsqueeze(-1)) * -100000
+        weight = self.dropout(F.softmax(attention, dim=2))
+        value = self.value(self.norm(x)).view(B, T, H, -1)
+        y = torch.einsum('bnmh,bmhd->bnhd', weight, value).reshape(B, T, D)
+        y = x + self.proj_out(y, emb)
+        return y
+
+class TemporalCrossAttention(nn.Module):
+
+    def __init__(self, seq_len, latent_dim, text_latent_dim, num_head, dropout, time_embed_dim):
+        super().__init__()
+        self.num_head = num_head
+        self.norm = nn.LayerNorm(latent_dim)
+        self.text_norm = nn.LayerNorm(text_latent_dim)
+        self.query = nn.Linear(latent_dim, latent_dim)
+        self.key = nn.Linear(text_latent_dim, latent_dim)
+        self.value = nn.Linear(text_latent_dim, latent_dim)
+        self.dropout = nn.Dropout(dropout)
+        self.proj_out = StylizationBlock(latent_dim, time_embed_dim, dropout)
+    
+    def forward(self, x, xf, emb):
+        """
+        x: B, T, D
+        xf: B, N, L
+        """
+        B, T, D = x.shape
+        N = xf.shape[1]
+        H = self.num_head
+        # B, T, 1, D
+        query = self.query(self.norm(x)).unsqueeze(2)
+        # B, 1, N, D
+        key = self.key(self.text_norm(xf)).unsqueeze(1)
+        query = query.view(B, T, H, -1)
+        key = key.view(B, N, H, -1)
+        # B, T, N, H
+        attention = torch.einsum('bnhd,bmhd->bnmh', query, key) / math.sqrt(D // H)
+        weight = self.dropout(F.softmax(attention, dim=2))
+        value = self.value(self.text_norm(xf)).view(B, N, H, -1)
+        y = torch.einsum('bnmh,bmhd->bnhd', weight, value).reshape(B, T, D)
+        y = x + self.proj_out(y, emb)
+        return y
+
+class TemporalDiffusionTransformerDecoderLayer(nn.Module):
+
+    def __init__(self,
+                 seq_len=60,
+                 latent_dim=32,
+                 text_latent_dim=512,
+                 time_embed_dim=128,
+                 ffn_dim=256,
+                 num_head=4,
+                 dropout=0.1):
+        super().__init__()
+        self.sa_block = TemporalSelfAttention(
+            seq_len, latent_dim, num_head, dropout, time_embed_dim)
+        self.ca_block = TemporalCrossAttention(
+            seq_len, latent_dim, text_latent_dim, num_head, dropout, time_embed_dim)
+        self.ffn = FFN(latent_dim, ffn_dim, dropout, time_embed_dim)
+
+    def forward(self, x, xf, emb, src_mask):
+        x = self.sa_block(x, emb, src_mask)
+        x = self.ca_block(x, xf, emb)
+        x = self.ffn(x, emb)
+        return x
+
+
+class MotionTransformer(nn.Module):
+    def __init__(self,
+                 input_feats,
+                 num_frames=240,
+                 latent_dim=512,
+                 ff_size=1024,
+                 num_layers=8,
+                 num_heads=8,
+                 dropout=0,
+                 activation="gelu", 
+                 num_text_layers=4,
+                 text_latent_dim=256,
+                 text_ff_size=2048,
+                 text_num_heads=4,
+                 no_clip=False,
+                 no_eff=False,
+                 **kargs):
+        super().__init__()
+        
+        self.num_frames = num_frames
+        self.latent_dim = latent_dim
+        self.ff_size = ff_size
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.activation = activation  
+        self.input_feats = input_feats
+        self.time_embed_dim = latent_dim * 4
+        self.sequence_embedding = nn.Parameter(torch.randn(num_frames, latent_dim))
+
+        # Text Transformer
+        self.clip, _ = clip.load('ViT-B/32', "cpu")
+        if no_clip:
+            self.clip.initialize_parameters()
+        else:
+            set_requires_grad(self.clip, False)
+        if text_latent_dim != 512:
+            self.text_pre_proj = nn.Linear(512, text_latent_dim)
+        else:
+            self.text_pre_proj = nn.Identity()
+        textTransEncoderLayer = nn.TransformerEncoderLayer(
+            d_model=text_latent_dim,
+            nhead=text_num_heads,
+            dim_feedforward=text_ff_size,
+            dropout=dropout,
+            activation=activation)
+        self.textTransEncoder = nn.TransformerEncoder(
+            textTransEncoderLayer,
+            num_layers=num_text_layers)
+        self.text_ln = nn.LayerNorm(text_latent_dim)
+        self.text_proj = nn.Sequential(
+            nn.Linear(text_latent_dim, self.time_embed_dim)
+        )
+
+        # Input Embedding
+        self.joint_embed = nn.Linear(self.input_feats, self.latent_dim)
+
+        self.time_embed = nn.Sequential(
+            nn.Linear(self.latent_dim, self.time_embed_dim),
+            nn.SiLU(),
+            nn.Linear(self.time_embed_dim, self.time_embed_dim),
+        )
+        self.temporal_decoder_blocks = nn.ModuleList()
+        for i in range(num_layers):
+            if no_eff:
+                self.temporal_decoder_blocks.append(
+                    TemporalDiffusionTransformerDecoderLayer(
+                        seq_len=num_frames,
+                        latent_dim=latent_dim,
+                        text_latent_dim=text_latent_dim,
+                        time_embed_dim=self.time_embed_dim,
+                        ffn_dim=ff_size,
+                        num_head=num_heads,
+                        dropout=dropout
+                    )
+                )
+            else:
+                self.temporal_decoder_blocks.append(
+                    LinearTemporalDiffusionTransformerDecoderLayer(
+                        seq_len=num_frames,
+                        latent_dim=latent_dim,
+                        text_latent_dim=text_latent_dim,
+                        time_embed_dim=self.time_embed_dim,
+                        ffn_dim=ff_size,
+                        num_head=num_heads,
+                        dropout=dropout
+                    )
+                )
+        
+        # Output Module
+        self.out = zero_module(nn.Linear(self.latent_dim, self.input_feats))
+        
+    def encode_text(self, text, device):
+        with torch.no_grad():
+            text = clip.tokenize(text, truncate=True).to(device)
+            x = self.clip.token_embedding(text).type(self.clip.dtype)  # [batch_size, n_ctx, d_model]
+
+            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)
+
+        # T, B, D
+        x = self.text_pre_proj(x)
+        xf_out = self.textTransEncoder(x)
+        xf_out = self.text_ln(xf_out)
+        xf_proj = self.text_proj(xf_out[text.argmax(dim=-1), torch.arange(xf_out.shape[1])])
+        # B, T, D
+        xf_out = xf_out.permute(1, 0, 2)
+        return xf_proj, xf_out
+
+    def generate_src_mask(self, T, length):
+        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, x, timesteps, length=None, text=None, xf_proj=None, xf_out=None):
+        """
+        x: B, T, D
+        """
+        B, T = x.shape[0], x.shape[1]
+        if xf_proj is None or xf_out is None:
+            xf_proj, xf_out = self.encode_text(text, x.device)
+
+        emb = self.time_embed(timestep_embedding(timesteps, self.latent_dim)) + xf_proj
+
+        # B, T, latent_dim
+        h = self.joint_embed(x)
+        h = h + self.sequence_embedding.unsqueeze(0)[:, :T, :]
+
+        src_mask = self.generate_src_mask(T, length).to(x.device).unsqueeze(-1)
+        for module in self.temporal_decoder_blocks:
+            h = module(h, xf_out, emb, src_mask)
+
+        output = self.out(h).view(B, T, -1).contiguous()
+        return output

options/base_options.py

@@ -0,0 +1,86 @@
+import argparse
+import os
+import torch
+from mmcv.runner import get_dist_info
+import torch.distributed as dist
+
+
+class BaseOptions():
+    def __init__(self):
+        self.parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+        self.initialized = False
+
+    def initialize(self):
+        self.parser.add_argument('--name', type=str, default="test", help='Name of this trial')
+        self.parser.add_argument('--decomp_name', type=str, default="Decomp_SP001_SM001_H512", help='Name of autoencoder model')
+
+        self.parser.add_argument("--gpu_id", type=int, default=-1, help='GPU id')
+        self.parser.add_argument("--distributed", action="store_true", help='Weather to use DDP training')
+
+        self.parser.add_argument('--dataset_name', type=str, default='t2m', help='Dataset Name')
+        self.parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')
+
+        self.parser.add_argument("--unit_length", type=int, default=4, help="Motions are cropped to the maximum times of unit_length")
+        self.parser.add_argument("--max_text_len", type=int, default=20, help="Maximum length of text description")
+
+        self.parser.add_argument('--text_enc_mod', type=str, default='bigru')
+        self.parser.add_argument('--estimator_mod', type=str, default='bigru')
+
+        self.parser.add_argument('--dim_text_hidden', type=int, default=512, help='Dimension of hidden unit in text encoder')
+        self.parser.add_argument('--dim_att_vec', type=int, default=512, help='Dimension of attention vector')
+        self.parser.add_argument('--dim_z', type=int, default=128, help='Dimension of latent Gaussian vector')
+
+        self.parser.add_argument('--n_layers_pri', type=int, default=1, help='Number of layers in prior network')
+        self.parser.add_argument('--n_layers_pos', type=int, default=1, help='Number of layers in posterior network')
+        self.parser.add_argument('--n_layers_dec', type=int, default=1, help='Number of layers in generator')
+
+        self.parser.add_argument('--dim_pri_hidden', type=int, default=1024, help='Dimension of hidden unit in prior network')
+        self.parser.add_argument('--dim_pos_hidden', type=int, default=1024, help='Dimension of hidden unit in posterior network')
+        self.parser.add_argument('--dim_dec_hidden', type=int, default=1024, help='Dimension of hidden unit in generator')
+
+        self.parser.add_argument('--dim_movement_enc_hidden', type=int, default=512,
+                                 help='Dimension of hidden in AutoEncoder(encoder)')
+        self.parser.add_argument('--dim_movement_dec_hidden', type=int, default=512,
+                                 help='Dimension of hidden in AutoEncoder(decoder)')
+        self.parser.add_argument('--dim_movement_latent', type=int, default=512, help='Dimension of motion snippet')
+
+        self.initialized = True
+
+
+
+    def parse(self):
+        if not self.initialized:
+            self.initialize()
+
+        self.opt = self.parser.parse_args()
+
+        self.opt.is_train = self.is_train
+
+        if self.opt.gpu_id != -1:
+            # self.opt.gpu_id = int(self.opt.gpu_id)
+            torch.cuda.set_device(self.opt.gpu_id)
+
+        args = vars(self.opt)
+
+        if args["distributed"]:
+            init_dist('slurm')
+        rank, world_size = get_dist_info()
+        if rank == 0:
+            print('------------ Options -------------')
+            for k, v in sorted(args.items()):
+                print('%s: %s' % (str(k), str(v)))
+            print('-------------- End ----------------')
+            if self.is_train:
+                # save to the disk
+                expr_dir = os.path.join(self.opt.checkpoints_dir, self.opt.dataset_name, self.opt.name)
+                if not os.path.exists(expr_dir):
+                    os.makedirs(expr_dir)
+                file_name = os.path.join(expr_dir, 'opt.txt')
+                with open(file_name, 'wt') as opt_file:
+                    opt_file.write('------------ Options -------------\n')
+                    for k, v in sorted(args.items()):
+                        opt_file.write('%s: %s\n' % (str(k), str(v)))
+                    opt_file.write('-------------- End ----------------\n')
+        if world_size > 1:
+            dist.barrier()
+        return self.opt

options/evaluate_options.py

@@ -0,0 +1,27 @@
+from options.base_options import BaseOptions
+
+
+class TestOptions(BaseOptions):
+    def initialize(self):
+        BaseOptions.initialize(self)
+        self.parser.add_argument('--batch_size', type=int, default=1, help='Batch size')
+        self.parser.add_argument('--start_mov_len', type=int, default=10)
+        self.parser.add_argument('--est_length', action="store_true", help="Whether to use sampled motion length")
+        self.parser.add_argument('--num_layers', type=int, default=8, help='num_layers of transformer')
+        self.parser.add_argument('--latent_dim', type=int, default=512, help='latent_dim of transformer')
+        self.parser.add_argument('--diffusion_steps', type=int, default=1000, help='diffusion_steps of transformer')
+        self.parser.add_argument('--no_clip', action='store_true', help='whether use clip pretrain')
+        self.parser.add_argument('--no_eff', action='store_true', help='whether use efficient attention')
+
+
+        self.parser.add_argument('--repeat_times', type=int, default=3, help="Number of generation rounds for each text description")
+        self.parser.add_argument('--split_file', type=str, default='test.txt')
+        self.parser.add_argument('--text', type=str, default="", help='Text description for motion generation')
+        self.parser.add_argument('--motion_length', type=int, default=0, help='Number of framese for motion generation')
+        self.parser.add_argument('--text_file', type=str, default="", help='Path of text description for motion generation')
+        self.parser.add_argument('--which_epoch', type=str, default="latest", help='Checkpoint that will be used')
+        self.parser.add_argument('--result_path', type=str, default="./eval_results/", help='Path to save generation results')
+        self.parser.add_argument('--num_results', type=int, default=40, help='Number of descriptions that will be used')
+        self.parser.add_argument('--ext', type=str, default='default', help='Save file path extension')
+
+        self.is_train = False

options/train_options.py

@@ -0,0 +1,26 @@
+from options.base_options import BaseOptions
+import argparse
+
+class TrainCompOptions(BaseOptions):
+    def initialize(self):
+        BaseOptions.initialize(self)
+        self.parser.add_argument('--num_layers', type=int, default=8, help='num_layers of transformer')
+        self.parser.add_argument('--latent_dim', type=int, default=512, help='latent_dim of transformer')
+        self.parser.add_argument('--diffusion_steps', type=int, default=1000, help='diffusion_steps of transformer')
+        self.parser.add_argument('--no_clip', action='store_true', help='whether use clip pretrain')
+        self.parser.add_argument('--no_eff', action='store_true', help='whether use efficient attention')
+
+        self.parser.add_argument('--num_epochs', type=int, default=50, help='Number of epochs')
+        self.parser.add_argument('--lr', type=float, default=2e-4, help='Learning rate')
+        self.parser.add_argument('--batch_size', type=int, default=32, help='Batch size per GPU')
+        self.parser.add_argument('--times', type=int, default=1, help='times of dataset')
+
+        self.parser.add_argument('--feat_bias', type=float, default=5, help='Scales for global motion features and foot contact')
+
+        self.parser.add_argument('--is_continue', action="store_true", help='Is this trail continued from previous trail?')
+
+        self.parser.add_argument('--log_every', type=int, default=50, help='Frequency of printing training progress (by iteration)')
+        self.parser.add_argument('--save_every_e', type=int, default=5, help='Frequency of saving models (by epoch)')
+        self.parser.add_argument('--eval_every_e', type=int, default=5, help='Frequency of animation results (by epoch)')
+        self.parser.add_argument('--save_latest', type=int, default=500, help='Frequency of saving models (by iteration)')
+        self.is_train = True

requirements.txt

@@ -0,0 +1,4 @@
+mmcv
+matplotlib==3.3.1
+torch==1.7.1
+git+https://github.com/openai/CLIP.git

tools/evaluation.py

@@ -0,0 +1,278 @@
+from datetime import datetime
+import numpy as np
+import torch
+from datasets import get_dataset_motion_loader, get_motion_loader
+from models import MotionTransformer
+from utils.get_opt import get_opt
+from utils.metrics import *
+from datasets import EvaluatorModelWrapper
+from collections import OrderedDict
+from utils.plot_script import *
+from utils import paramUtil
+from utils.utils import *
+from trainers import DDPMTrainer
+
+from os.path import join as pjoin
+import sys
+
+
+def build_models(opt, dim_pose):
+    encoder = MotionTransformer(
+        input_feats=dim_pose,
+        num_frames=opt.max_motion_length,
+        num_layers=opt.num_layers,
+        latent_dim=opt.latent_dim,
+        no_clip=opt.no_clip,
+        no_eff=opt.no_eff)
+    return encoder
+
+
+torch.multiprocessing.set_sharing_strategy('file_system')
+
+
+def evaluate_matching_score(motion_loaders, file):
+    match_score_dict = OrderedDict({})
+    R_precision_dict = OrderedDict({})
+    activation_dict = OrderedDict({})
+    # print(motion_loaders.keys())
+    print('========== Evaluating Matching Score ==========')
+    for motion_loader_name, motion_loader in motion_loaders.items():
+        all_motion_embeddings = []
+        score_list = []
+        all_size = 0
+        matching_score_sum = 0
+        top_k_count = 0
+        # print(motion_loader_name)
+        with torch.no_grad():
+            for idx, batch in enumerate(motion_loader):
+                word_embeddings, pos_one_hots, _, sent_lens, motions, m_lens, _ = batch
+                text_embeddings, motion_embeddings = eval_wrapper.get_co_embeddings(
+                    word_embs=word_embeddings,
+                    pos_ohot=pos_one_hots,
+                    cap_lens=sent_lens,
+                    motions=motions,
+                    m_lens=m_lens
+                )
+                dist_mat = euclidean_distance_matrix(text_embeddings.cpu().numpy(),
+                                                     motion_embeddings.cpu().numpy())
+                matching_score_sum += dist_mat.trace()
+
+                argsmax = np.argsort(dist_mat, axis=1)
+                top_k_mat = calculate_top_k(argsmax, top_k=3)
+                top_k_count += top_k_mat.sum(axis=0)
+
+                all_size += text_embeddings.shape[0]
+
+                all_motion_embeddings.append(motion_embeddings.cpu().numpy())
+
+            all_motion_embeddings = np.concatenate(all_motion_embeddings, axis=0)
+            matching_score = matching_score_sum / all_size
+            R_precision = top_k_count / all_size
+            match_score_dict[motion_loader_name] = matching_score
+            R_precision_dict[motion_loader_name] = R_precision
+            activation_dict[motion_loader_name] = all_motion_embeddings
+
+        print(f'---> [{motion_loader_name}] Matching Score: {matching_score:.4f}')
+        print(f'---> [{motion_loader_name}] Matching Score: {matching_score:.4f}', file=file, flush=True)
+
+        line = f'---> [{motion_loader_name}] R_precision: '
+        for i in range(len(R_precision)):
+            line += '(top %d): %.4f ' % (i+1, R_precision[i])
+        print(line)
+        print(line, file=file, flush=True)
+
+    return match_score_dict, R_precision_dict, activation_dict
+
+
+def evaluate_fid(groundtruth_loader, activation_dict, file):
+    eval_dict = OrderedDict({})
+    gt_motion_embeddings = []
+    print('========== Evaluating FID ==========')
+    with torch.no_grad():
+        for idx, batch in enumerate(groundtruth_loader):
+            _, _, _, sent_lens, motions, m_lens, _ = batch
+            motion_embeddings = eval_wrapper.get_motion_embeddings(
+                motions=motions,
+                m_lens=m_lens
+            )
+            gt_motion_embeddings.append(motion_embeddings.cpu().numpy())
+    gt_motion_embeddings = np.concatenate(gt_motion_embeddings, axis=0)
+    gt_mu, gt_cov = calculate_activation_statistics(gt_motion_embeddings)
+
+    # print(gt_mu)
+    for model_name, motion_embeddings in activation_dict.items():
+        mu, cov = calculate_activation_statistics(motion_embeddings)
+        # print(mu)
+        fid = calculate_frechet_distance(gt_mu, gt_cov, mu, cov)
+        print(f'---> [{model_name}] FID: {fid:.4f}')
+        print(f'---> [{model_name}] FID: {fid:.4f}', file=file, flush=True)
+        eval_dict[model_name] = fid
+    return eval_dict
+
+
+def evaluate_diversity(activation_dict, file):
+    eval_dict = OrderedDict({})
+    print('========== Evaluating Diversity ==========')
+    for model_name, motion_embeddings in activation_dict.items():
+        diversity = calculate_diversity(motion_embeddings, diversity_times)
+        eval_dict[model_name] = diversity
+        print(f'---> [{model_name}] Diversity: {diversity:.4f}')
+        print(f'---> [{model_name}] Diversity: {diversity:.4f}', file=file, flush=True)
+    return eval_dict
+
+
+def evaluate_multimodality(mm_motion_loaders, file):
+    eval_dict = OrderedDict({})
+    print('========== Evaluating MultiModality ==========')
+    for model_name, mm_motion_loader in mm_motion_loaders.items():
+        mm_motion_embeddings = []
+        with torch.no_grad():
+            for idx, batch in enumerate(mm_motion_loader):
+                # (1, mm_replications, dim_pos)
+                motions, m_lens = batch
+                motion_embedings = eval_wrapper.get_motion_embeddings(motions[0], m_lens[0])
+                mm_motion_embeddings.append(motion_embedings.unsqueeze(0))
+        if len(mm_motion_embeddings) == 0:
+            multimodality = 0
+        else:
+            mm_motion_embeddings = torch.cat(mm_motion_embeddings, dim=0).cpu().numpy()
+            multimodality = calculate_multimodality(mm_motion_embeddings, mm_num_times)
+        print(f'---> [{model_name}] Multimodality: {multimodality:.4f}')
+        print(f'---> [{model_name}] Multimodality: {multimodality:.4f}', file=file, flush=True)
+        eval_dict[model_name] = multimodality
+    return eval_dict
+
+
+def get_metric_statistics(values):
+    mean = np.mean(values, axis=0)
+    std = np.std(values, axis=0)
+    conf_interval = 1.96 * std / np.sqrt(replication_times)
+    return mean, conf_interval
+
+
+def evaluation(log_file):
+    with open(log_file, 'w') as f:
+        all_metrics = OrderedDict({'Matching Score': OrderedDict({}),
+                                   'R_precision': OrderedDict({}),
+                                   'FID': OrderedDict({}),
+                                   'Diversity': OrderedDict({}),
+                                   'MultiModality': OrderedDict({})})
+        for replication in range(replication_times):
+            motion_loaders = {}
+            mm_motion_loaders = {}
+            motion_loaders['ground truth'] = gt_loader
+            for motion_loader_name, motion_loader_getter in eval_motion_loaders.items():
+                motion_loader, mm_motion_loader = motion_loader_getter()
+                motion_loaders[motion_loader_name] = motion_loader
+                mm_motion_loaders[motion_loader_name] = mm_motion_loader
+
+            print(f'==================== Replication {replication} ====================')
+            print(f'==================== Replication {replication} ====================', file=f, flush=True)
+            print(f'Time: {datetime.now()}')
+            print(f'Time: {datetime.now()}', file=f, flush=True)
+            mat_score_dict, R_precision_dict, acti_dict = evaluate_matching_score(motion_loaders, f)
+
+            print(f'Time: {datetime.now()}')
+            print(f'Time: {datetime.now()}', file=f, flush=True)
+            fid_score_dict = evaluate_fid(gt_loader, acti_dict, f)
+
+            print(f'Time: {datetime.now()}')
+            print(f'Time: {datetime.now()}', file=f, flush=True)
+            div_score_dict = evaluate_diversity(acti_dict, f)
+
+            print(f'Time: {datetime.now()}')
+            print(f'Time: {datetime.now()}', file=f, flush=True)
+            mm_score_dict = evaluate_multimodality(mm_motion_loaders, f)
+
+            print(f'!!! DONE !!!')
+            print(f'!!! DONE !!!', file=f, flush=True)
+
+            for key, item in mat_score_dict.items():
+                if key not in all_metrics['Matching Score']:
+                    all_metrics['Matching Score'][key] = [item]
+                else:
+                    all_metrics['Matching Score'][key] += [item]
+
+            for key, item in R_precision_dict.items():
+                if key not in all_metrics['R_precision']:
+                    all_metrics['R_precision'][key] = [item]
+                else:
+                    all_metrics['R_precision'][key] += [item]
+
+            for key, item in fid_score_dict.items():
+                if key not in all_metrics['FID']:
+                    all_metrics['FID'][key] = [item]
+                else:
+                    all_metrics['FID'][key] += [item]
+
+            for key, item in div_score_dict.items():
+                if key not in all_metrics['Diversity']:
+                    all_metrics['Diversity'][key] = [item]
+                else:
+                    all_metrics['Diversity'][key] += [item]
+
+            for key, item in mm_score_dict.items():
+                if key not in all_metrics['MultiModality']:
+                    all_metrics['MultiModality'][key] = [item]
+                else:
+                    all_metrics['MultiModality'][key] += [item]
+
+
+        # print(all_metrics['Diversity'])
+        for metric_name, metric_dict in all_metrics.items():
+            print('========== %s Summary ==========' % metric_name)
+            print('========== %s Summary ==========' % metric_name, file=f, flush=True)
+
+            for model_name, values in metric_dict.items():
+                # print(metric_name, model_name)
+                mean, conf_interval = get_metric_statistics(np.array(values))
+                # print(mean, mean.dtype)
+                if isinstance(mean, np.float64) or isinstance(mean, np.float32):
+                    print(f'---> [{model_name}] Mean: {mean:.4f} CInterval: {conf_interval:.4f}')
+                    print(f'---> [{model_name}] Mean: {mean:.4f} CInterval: {conf_interval:.4f}', file=f, flush=True)
+                elif isinstance(mean, np.ndarray):
+                    line = f'---> [{model_name}]'
+                    for i in range(len(mean)):
+                        line += '(top %d) Mean: %.4f CInt: %.4f;' % (i+1, mean[i], conf_interval[i])
+                    print(line)
+                    print(line, file=f, flush=True)
+
+
+if __name__ == '__main__':
+    mm_num_samples = 100
+    mm_num_repeats = 30
+    mm_num_times = 10
+
+    diversity_times = 300
+    replication_times = 1
+    batch_size = 32
+    opt_path = sys.argv[1]
+    dataset_opt_path = opt_path
+
+    try:
+        device_id = int(sys.argv[2])
+    except:
+        device_id = 0
+    device = torch.device('cuda:%d' % device_id if torch.cuda.is_available() else 'cpu')
+    torch.cuda.set_device(device_id)
+
+    gt_loader, gt_dataset = get_dataset_motion_loader(dataset_opt_path, batch_size, device)
+    wrapper_opt = get_opt(dataset_opt_path, device)
+    eval_wrapper = EvaluatorModelWrapper(wrapper_opt)
+
+    opt = get_opt(opt_path, device)
+    encoder = build_models(opt, opt.dim_pose)
+    trainer = DDPMTrainer(opt, encoder)
+    eval_motion_loaders = {
+        'text2motion': lambda: get_motion_loader(
+            opt,
+            batch_size,
+            trainer,
+            gt_dataset,
+            mm_num_samples,
+            mm_num_repeats
+        )
+    }
+
+    log_file = './t2m_evaluation.log'
+    evaluation(log_file)

tools/train.py

@@ -0,0 +1,89 @@
+import os
+from os.path import join as pjoin
+
+import utils.paramUtil as paramUtil
+from options.train_options import TrainCompOptions
+from utils.plot_script import *
+
+from models import MotionTransformer
+from trainers import DDPMTrainer
+from datasets import Text2MotionDataset
+
+from mmcv.runner import get_dist_info, init_dist
+from mmcv.parallel import MMDistributedDataParallel
+import torch
+import torch.distributed as dist
+
+
+def build_models(opt, dim_pose):
+    encoder = MotionTransformer(
+        input_feats=dim_pose,
+        num_frames=opt.max_motion_length,
+        num_layers=opt.num_layers,
+        latent_dim=opt.latent_dim,
+        no_clip=opt.no_clip,
+        no_eff=opt.no_eff)
+    return encoder
+
+
+if __name__ == '__main__':
+    parser = TrainCompOptions()
+    opt = parser.parse()
+    rank, world_size = get_dist_info()
+
+    opt.device = torch.device("cuda")
+    torch.autograd.set_detect_anomaly(True)
+
+    opt.save_root = pjoin(opt.checkpoints_dir, opt.dataset_name, opt.name)
+    opt.model_dir = pjoin(opt.save_root, 'model')
+    opt.meta_dir = pjoin(opt.save_root, 'meta')
+
+    if rank == 0:
+        os.makedirs(opt.model_dir, exist_ok=True)
+        os.makedirs(opt.meta_dir, exist_ok=True)
+    if world_size > 1:
+        dist.barrier()
+
+    if opt.dataset_name == 't2m':
+        opt.data_root = './data/HumanML3D'
+        opt.motion_dir = pjoin(opt.data_root, 'new_joint_vecs')
+        opt.text_dir = pjoin(opt.data_root, 'texts')
+        opt.joints_num = 22
+        radius = 4
+        fps = 20
+        opt.max_motion_length = 196
+        dim_pose = 263
+        kinematic_chain = paramUtil.t2m_kinematic_chain
+    elif opt.dataset_name == 'kit':
+        opt.data_root = './data/KIT-ML'
+        opt.motion_dir = pjoin(opt.data_root, 'new_joint_vecs')
+        opt.text_dir = pjoin(opt.data_root, 'texts')
+        opt.joints_num = 21
+        radius = 240 * 8
+        fps = 12.5
+        dim_pose = 251
+        opt.max_motion_length = 196
+        kinematic_chain = paramUtil.kit_kinematic_chain
+
+    else:
+        raise KeyError('Dataset Does Not Exist')
+
+    dim_word = 300
+    mean = np.load(pjoin(opt.data_root, 'Mean.npy'))
+    std = np.load(pjoin(opt.data_root, 'Std.npy'))
+
+    train_split_file = pjoin(opt.data_root, 'train.txt')
+
+    encoder = build_models(opt, dim_pose)
+    if world_size > 1:
+        encoder = MMDistributedDataParallel(
+            encoder.cuda(),
+            device_ids=[torch.cuda.current_device()],
+            broadcast_buffers=False,
+            find_unused_parameters=True)
+    else:
+        encoder = encoder.cuda()
+
+    trainer = DDPMTrainer(opt, encoder)
+    train_dataset = Text2MotionDataset(opt, mean, std, train_split_file, opt.times)
+    trainer.train(train_dataset)

tools/visualization.py

@@ -0,0 +1,30 @@
+import os
+import torch
+import argparse
+
+import utils.paramUtil as paramUtil
+from torch.utils.data import DataLoader
+from utils.plot_script import *
+
+from utils.utils import *
+from utils.motion_process import recover_from_ric
+
+
+def plot_t2m(opt, data, result_path, caption):
+    joint = recover_from_ric(torch.from_numpy(data).float(), opt.joints_num).numpy()
+    # joint = motion_temporal_filter(joint, sigma=1)
+    plot_3d_motion(result_path, paramUtil.t2m_kinematic_chain, joint, title=caption, fps=20)
+
+
+def process(trainer, opt, device, mean, std, text, motion_length, result_path):
+
+    result_dict = {}
+    with torch.no_grad():
+        if motion_length != -1:
+            caption = [text]
+            m_lens = torch.LongTensor([motion_length]).to(device)
+            pred_motions = trainer.generate(caption, m_lens, opt.dim_pose)
+            motion = pred_motions[0].cpu().numpy()
+            motion = motion * std + mean
+            title = text + " #%d" % motion.shape[0]
+            plot_t2m(opt, motion, result_path, title)

trainers/__init__.py

@@ -0,0 +1,4 @@
+from .ddpm_trainer import DDPMTrainer
+
+
+__all__ = ['DDPMTrainer']

trainers/ddpm_trainer.py

@@ -0,0 +1,222 @@
+import torch
+import torch.nn.functional as F
+import random
+import time
+from models.transformer import MotionTransformer
+from torch.utils.data import DataLoader
+import torch.optim as optim
+from torch.nn.utils import clip_grad_norm_
+from collections import OrderedDict
+from utils.utils import print_current_loss
+from os.path import join as pjoin
+import codecs as cs
+import torch.distributed as dist
+
+
+from mmcv.runner import get_dist_info
+from models.gaussian_diffusion import (
+    GaussianDiffusion,
+    get_named_beta_schedule,
+    create_named_schedule_sampler,
+    ModelMeanType,
+    ModelVarType,
+    LossType
+)
+
+from datasets import build_dataloader
+
+
+class DDPMTrainer(object):
+
+    def __init__(self, args, encoder):
+        self.opt = args
+        self.device = args.device
+        self.encoder = encoder
+        self.diffusion_steps = args.diffusion_steps
+        sampler = 'uniform'
+        beta_scheduler = 'linear'
+        betas = get_named_beta_schedule(beta_scheduler, self.diffusion_steps)
+        self.diffusion = GaussianDiffusion(
+            betas=betas,
+            model_mean_type=ModelMeanType.EPSILON,
+            model_var_type=ModelVarType.FIXED_SMALL,
+            loss_type=LossType.MSE
+        )
+        self.sampler = create_named_schedule_sampler(sampler, self.diffusion)
+        self.sampler_name = sampler
+
+        if args.is_train:
+            self.mse_criterion = torch.nn.MSELoss(reduction='none')
+        self.to(self.device)
+
+    @staticmethod
+    def zero_grad(opt_list):
+        for opt in opt_list:
+            opt.zero_grad()
+
+    @staticmethod
+    def clip_norm(network_list):
+        for network in network_list:
+            clip_grad_norm_(network.parameters(), 0.5)
+
+    @staticmethod
+    def step(opt_list):
+        for opt in opt_list:
+            opt.step()
+
+    def forward(self, batch_data, eval_mode=False):
+        caption, motions, m_lens = batch_data
+        motions = motions.detach().to(self.device).float()
+
+        self.caption = caption
+        self.motions = motions
+        x_start = motions
+        B, T = x_start.shape[:2]
+        cur_len = torch.LongTensor([min(T, m_len) for m_len in  m_lens]).to(self.device)
+        t, _ = self.sampler.sample(B, x_start.device)
+        output = self.diffusion.training_losses(
+            model=self.encoder,
+            x_start=x_start,
+            t=t,
+            model_kwargs={"text": caption, "length": cur_len}
+        )
+
+        self.real_noise = output['target']
+        self.fake_noise = output['pred']
+        try:
+            self.src_mask = self.encoder.module.generate_src_mask(T, cur_len).to(x_start.device)
+        except:
+            self.src_mask = self.encoder.generate_src_mask(T, cur_len).to(x_start.device)
+
+    def generate_batch(self, caption, m_lens, dim_pose):
+        xf_proj, xf_out = self.encoder.encode_text(caption, self.device)
+        
+        B = len(caption)
+        T = min(m_lens.max(), self.encoder.num_frames)
+        output = self.diffusion.p_sample_loop(
+            self.encoder,
+            (B, T, dim_pose),
+            clip_denoised=False,
+            progress=True,
+            model_kwargs={
+                'xf_proj': xf_proj,
+                'xf_out': xf_out,
+                'length': m_lens
+            })
+        return output
+
+    def generate(self, caption, m_lens, dim_pose, batch_size=1024):
+        N = len(caption)
+        cur_idx = 0
+        self.encoder.eval()
+        all_output = []
+        while cur_idx < N:
+            if cur_idx + batch_size >= N:
+                batch_caption = caption[cur_idx:]
+                batch_m_lens = m_lens[cur_idx:]
+            else:
+                batch_caption = caption[cur_idx: cur_idx + batch_size]
+                batch_m_lens = m_lens[cur_idx: cur_idx + batch_size]
+            output = self.generate_batch(batch_caption, batch_m_lens, dim_pose)
+            B = output.shape[0]
+
+            for i in range(B):
+                all_output.append(output[i])
+            cur_idx += batch_size
+        return all_output
+
+    def backward_G(self):
+        loss_mot_rec = self.mse_criterion(self.fake_noise, self.real_noise).mean(dim=-1)
+        loss_mot_rec = (loss_mot_rec * self.src_mask).sum() / self.src_mask.sum()
+        self.loss_mot_rec = loss_mot_rec
+        loss_logs = OrderedDict({})
+        loss_logs['loss_mot_rec'] = self.loss_mot_rec.item()
+        return loss_logs
+
+    def update(self):
+        self.zero_grad([self.opt_encoder])
+        loss_logs = self.backward_G()
+        self.loss_mot_rec.backward()
+        self.clip_norm([self.encoder])
+        self.step([self.opt_encoder])
+
+        return loss_logs
+
+    def to(self, device):
+        if self.opt.is_train:
+            self.mse_criterion.to(device)
+        self.encoder = self.encoder.to(device)
+
+    def train_mode(self):
+        self.encoder.train()
+
+    def eval_mode(self):
+        self.encoder.eval()
+
+    def save(self, file_name, ep, total_it):
+        state = {
+            'opt_encoder': self.opt_encoder.state_dict(),
+            'ep': ep,
+            'total_it': total_it
+        }
+        try:
+            state['encoder'] = self.encoder.module.state_dict()
+        except:
+            state['encoder'] = self.encoder.state_dict()
+        torch.save(state, file_name)
+        return
+
+    def load(self, model_dir):
+        checkpoint = torch.load(model_dir, map_location=self.device)
+        if self.opt.is_train:
+            self.opt_encoder.load_state_dict(checkpoint['opt_encoder'])
+        self.encoder.load_state_dict(checkpoint['encoder'], strict=True)
+        return checkpoint['ep'], checkpoint.get('total_it', 0)
+
+    def train(self, train_dataset):
+        rank, world_size = get_dist_info()
+        self.to(self.device)
+        self.opt_encoder = optim.Adam(self.encoder.parameters(), lr=self.opt.lr)
+        it = 0
+        cur_epoch = 0
+        if self.opt.is_continue:
+            model_dir = pjoin(self.opt.model_dir, 'latest.tar')
+            cur_epoch, it = self.load(model_dir)
+
+        start_time = time.time()
+
+        train_loader = build_dataloader(
+            train_dataset,
+            samples_per_gpu=self.opt.batch_size,
+            drop_last=True,
+            workers_per_gpu=4,
+            shuffle=True)
+
+        logs = OrderedDict()
+        for epoch in range(cur_epoch, self.opt.num_epochs):
+            self.train_mode()
+            for i, batch_data in enumerate(train_loader):
+                self.forward(batch_data)
+                log_dict = self.update()
+                for k, v in log_dict.items():
+                    if k not in logs:
+                        logs[k] = v
+                    else:
+                        logs[k] += v
+                it += 1
+                if it % self.opt.log_every == 0 and rank == 0:
+                    mean_loss = OrderedDict({})
+                    for tag, value in logs.items():
+                        mean_loss[tag] = value / self.opt.log_every
+                    logs = OrderedDict()
+                    print_current_loss(start_time, it, mean_loss, epoch, inner_iter=i)
+
+                if it % self.opt.save_latest == 0 and rank == 0:
+                    self.save(pjoin(self.opt.model_dir, 'latest.tar'), epoch, it)
+
+            if rank == 0:
+                self.save(pjoin(self.opt.model_dir, 'latest.tar'), epoch, it)
+
+            if epoch % self.opt.save_every_e == 0 and rank == 0:
+                self.save(pjoin(self.opt.model_dir, 'ckpt_e%03d.tar'%(epoch)),
+                            epoch, total_it=it)

utils/get_opt.py

@@ -0,0 +1,91 @@
+import os
+from argparse import Namespace
+import re
+from os.path import join as pjoin
+from utils.word_vectorizer import POS_enumerator
+
+
+def is_float(numStr):
+    flag = False
+    numStr = str(numStr).strip().lstrip('-').lstrip('+')
+    try:
+        reg = re.compile(r'^[-+]?[0-9]+\.[0-9]+$')
+        res = reg.match(str(numStr))
+        if res:
+            flag = True
+    except Exception as ex:
+        print("is_float() - error: " + str(ex))
+    return flag
+
+
+def is_number(numStr):
+    flag = False
+    numStr = str(numStr).strip().lstrip('-').lstrip('+')
+    if str(numStr).isdigit():
+        flag = True
+    return flag
+
+
+def get_opt(opt_path, device):
+    opt = Namespace()
+    opt_dict = vars(opt)
+
+    skip = ('-------------- End ----------------',
+            '------------ Options -------------',
+            '\n')
+    print('Reading', opt_path)
+    with open(opt_path) as f:
+        for line in f:
+            if line.strip() not in skip:
+                # print(line.strip())
+                key, value = line.strip().split(': ')
+                if value in ('True', 'False'):
+                    opt_dict[key] = True if value == 'True' else False
+                elif is_float(value):
+                    opt_dict[key] = float(value)
+                elif is_number(value):
+                    opt_dict[key] = int(value)
+                else:
+                    opt_dict[key] = str(value)
+
+    opt_dict['which_epoch'] = 'latest'
+    if 'num_layers' not in opt_dict:
+        opt_dict['num_layers'] = 8
+    if 'latent_dim' not in opt_dict:
+        opt_dict['latent_dim'] = 512
+    if 'diffusion_steps' not in opt_dict:
+        opt_dict['diffusion_steps'] = 1000
+    if 'no_clip' not in opt_dict:
+        opt_dict['no_clip'] = False
+    if 'no_eff' not in opt_dict:
+        opt_dict['no_eff'] = False
+
+    opt.save_root = pjoin(opt.checkpoints_dir, opt.dataset_name, opt.name)
+    opt.model_dir = pjoin(opt.save_root, 'model')
+    opt.meta_dir = pjoin(opt.save_root, 'meta')
+
+    if opt.dataset_name == 't2m':
+        opt.data_root = './data/HumanML3D'
+        opt.motion_dir = pjoin(opt.data_root, 'new_joint_vecs')
+        opt.text_dir = pjoin(opt.data_root, 'texts')
+        opt.joints_num = 22
+        opt.dim_pose = 263
+        opt.max_motion_length = 196
+    elif opt.dataset_name == 'kit':
+        opt.data_root = './data/KIT-ML'
+        opt.motion_dir = pjoin(opt.data_root, 'new_joint_vecs')
+        opt.text_dir = pjoin(opt.data_root, 'texts')
+        opt.joints_num = 21
+        opt.dim_pose = 251
+        opt.max_motion_length = 196
+    else:
+        raise KeyError('Dataset not recognized')
+
+    opt.dim_word = 300
+    opt.num_classes = 200 // opt.unit_length
+    opt.dim_pos_ohot = len(POS_enumerator)
+    opt.is_train = False
+    opt.is_continue = False
+    opt.device = device
+
+    return opt

utils/metrics.py

@@ -0,0 +1,146 @@
+import numpy as np
+from scipy import linalg
+
+
+# (X - X_train)*(X - X_train) = -2X*X_train + X*X + X_train*X_train
+def euclidean_distance_matrix(matrix1, matrix2):
+    """
+        Params:
+        -- matrix1: N1 x D
+        -- matrix2: N2 x D
+        Returns:
+        -- dist: N1 x N2
+        dist[i, j] == distance(matrix1[i], matrix2[j])
+    """
+    assert matrix1.shape[1] == matrix2.shape[1]
+    d1 = -2 * np.dot(matrix1, matrix2.T)    # shape (num_test, num_train)
+    d2 = np.sum(np.square(matrix1), axis=1, keepdims=True)    # shape (num_test, 1)
+    d3 = np.sum(np.square(matrix2), axis=1)     # shape (num_train, )
+    dists = np.sqrt(d1 + d2 + d3)  # broadcasting
+    return dists
+
+def calculate_top_k(mat, top_k):
+    size = mat.shape[0]
+    gt_mat = np.expand_dims(np.arange(size), 1).repeat(size, 1)
+    bool_mat = (mat == gt_mat)
+    correct_vec = False
+    top_k_list = []
+    for i in range(top_k):
+#         print(correct_vec, bool_mat[:, i])
+        correct_vec = (correct_vec | bool_mat[:, i])
+        # print(correct_vec)
+        top_k_list.append(correct_vec[:, None])
+    top_k_mat = np.concatenate(top_k_list, axis=1)
+    return top_k_mat
+
+
+def calculate_R_precision(embedding1, embedding2, top_k, sum_all=False):
+    dist_mat = euclidean_distance_matrix(embedding1, embedding2)
+    argmax = np.argsort(dist_mat, axis=1)
+    top_k_mat = calculate_top_k(argmax, top_k)
+    if sum_all:
+        return top_k_mat.sum(axis=0)
+    else:
+        return top_k_mat
+
+
+def calculate_matching_score(embedding1, embedding2, sum_all=False):
+    assert len(embedding1.shape) == 2
+    assert embedding1.shape[0] == embedding2.shape[0]
+    assert embedding1.shape[1] == embedding2.shape[1]
+
+    dist = linalg.norm(embedding1 - embedding2, axis=1)
+    if sum_all:
+        return dist.sum(axis=0)
+    else:
+        return dist
+
+
+
+def calculate_activation_statistics(activations):
+    """
+    Params:
+    -- activation: num_samples x dim_feat
+    Returns:
+    -- mu: dim_feat
+    -- sigma: dim_feat x dim_feat
+    """
+    mu = np.mean(activations, axis=0)
+    cov = np.cov(activations, rowvar=False)
+    return mu, cov
+
+
+def calculate_diversity(activation, diversity_times):
+    assert len(activation.shape) == 2
+    assert activation.shape[0] > diversity_times
+    num_samples = activation.shape[0]
+
+    first_indices = np.random.choice(num_samples, diversity_times, replace=False)
+    second_indices = np.random.choice(num_samples, diversity_times, replace=False)
+    dist = linalg.norm(activation[first_indices] - activation[second_indices], axis=1)
+    return dist.mean()
+
+
+def calculate_multimodality(activation, multimodality_times):
+    assert len(activation.shape) == 3
+    assert activation.shape[1] > multimodality_times
+    num_per_sent = activation.shape[1]
+
+    first_dices = np.random.choice(num_per_sent, multimodality_times, replace=False)
+    second_dices = np.random.choice(num_per_sent, multimodality_times, replace=False)
+    dist = linalg.norm(activation[:, first_dices] - activation[:, second_dices], axis=2)
+    return dist.mean()
+
+
+def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
+    """Numpy implementation of the Frechet Distance.
+    The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
+    and X_2 ~ N(mu_2, C_2) is
+            d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
+    Stable version by Dougal J. Sutherland.
+    Params:
+    -- mu1   : Numpy array containing the activations of a layer of the
+               inception net (like returned by the function 'get_predictions')
+               for generated samples.
+    -- mu2   : The sample mean over activations, precalculated on an
+               representative data set.
+    -- sigma1: The covariance matrix over activations for generated samples.
+    -- sigma2: The covariance matrix over activations, precalculated on an
+               representative data set.
+    Returns:
+    --   : The Frechet Distance.
+    """
+
+    mu1 = np.atleast_1d(mu1)
+    mu2 = np.atleast_1d(mu2)
+
+    sigma1 = np.atleast_2d(sigma1)
+    sigma2 = np.atleast_2d(sigma2)
+
+    assert mu1.shape == mu2.shape, \
+        'Training and test mean vectors have different lengths'
+    assert sigma1.shape == sigma2.shape, \
+        'Training and test covariances have different dimensions'
+
+    diff = mu1 - mu2
+
+    # Product might be almost singular
+    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
+    if not np.isfinite(covmean).all():
+        msg = ('fid calculation produces singular product; '
+               'adding %s to diagonal of cov estimates') % eps
+        print(msg)
+        offset = np.eye(sigma1.shape[0]) * eps
+        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
+
+    # Numerical error might give slight imaginary component
+    if np.iscomplexobj(covmean):
+        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
+            m = np.max(np.abs(covmean.imag))
+            raise ValueError('Imaginary component {}'.format(m))
+        covmean = covmean.real
+
+    tr_covmean = np.trace(covmean)
+
+    return (diff.dot(diff) + np.trace(sigma1) +
+            np.trace(sigma2) - 2 * tr_covmean)

utils/motion_process.py

@@ -0,0 +1,515 @@
+from os.path import join as pjoin
+
+import numpy as np
+import os
+from utils.quaternion import *
+from utils.skeleton import Skeleton
+from utils.paramUtil import *
+
+import torch
+from tqdm import tqdm
+
+# positions (batch, joint_num, 3)
+def uniform_skeleton(positions, target_offset):
+    src_skel = Skeleton(n_raw_offsets, kinematic_chain, 'cpu')
+    src_offset = src_skel.get_offsets_joints(torch.from_numpy(positions[0]))
+    src_offset = src_offset.numpy()
+    tgt_offset = target_offset.numpy()
+    # print(src_offset)
+    # print(tgt_offset)
+    '''Calculate Scale Ratio as the ratio of legs'''
+    src_leg_len = np.abs(src_offset[l_idx1]).max() + np.abs(src_offset[l_idx2]).max()
+    tgt_leg_len = np.abs(tgt_offset[l_idx1]).max() + np.abs(tgt_offset[l_idx2]).max()
+
+    scale_rt = tgt_leg_len / src_leg_len
+    # print(scale_rt)
+    src_root_pos = positions[:, 0]
+    tgt_root_pos = src_root_pos * scale_rt
+
+    '''Inverse Kinematics'''
+    quat_params = src_skel.inverse_kinematics_np(positions, face_joint_indx)
+    # print(quat_params.shape)
+
+    '''Forward Kinematics'''
+    src_skel.set_offset(target_offset)
+    new_joints = src_skel.forward_kinematics_np(quat_params, tgt_root_pos)
+    return new_joints
+
+
+def extract_features(positions, feet_thre, n_raw_offsets, kinematic_chain, face_joint_indx, fid_r, fid_l):
+    global_positions = positions.copy()
+    """ Get Foot Contacts """
+
+    def foot_detect(positions, thres):
+        velfactor, heightfactor = np.array([thres, thres]), np.array([3.0, 2.0])
+
+        feet_l_x = (positions[1:, fid_l, 0] - positions[:-1, fid_l, 0]) ** 2
+        feet_l_y = (positions[1:, fid_l, 1] - positions[:-1, fid_l, 1]) ** 2
+        feet_l_z = (positions[1:, fid_l, 2] - positions[:-1, fid_l, 2]) ** 2
+        #     feet_l_h = positions[:-1,fid_l,1]
+        #     feet_l = (((feet_l_x + feet_l_y + feet_l_z) < velfactor) & (feet_l_h < heightfactor)).astype(np.float)
+        feet_l = ((feet_l_x + feet_l_y + feet_l_z) < velfactor).astype(np.float)
+
+        feet_r_x = (positions[1:, fid_r, 0] - positions[:-1, fid_r, 0]) ** 2
+        feet_r_y = (positions[1:, fid_r, 1] - positions[:-1, fid_r, 1]) ** 2
+        feet_r_z = (positions[1:, fid_r, 2] - positions[:-1, fid_r, 2]) ** 2
+        #     feet_r_h = positions[:-1,fid_r,1]
+        #     feet_r = (((feet_r_x + feet_r_y + feet_r_z) < velfactor) & (feet_r_h < heightfactor)).astype(np.float)
+        feet_r = (((feet_r_x + feet_r_y + feet_r_z) < velfactor)).astype(np.float)
+        return feet_l, feet_r
+
+    #
+    feet_l, feet_r = foot_detect(positions, feet_thre)
+    # feet_l, feet_r = foot_detect(positions, 0.002)
+
+    '''Quaternion and Cartesian representation'''
+    r_rot = None
+
+    def get_rifke(positions):
+        '''Local pose'''
+        positions[..., 0] -= positions[:, 0:1, 0]
+        positions[..., 2] -= positions[:, 0:1, 2]
+        '''All pose face Z+'''
+        positions = qrot_np(np.repeat(r_rot[:, None], positions.shape[1], axis=1), positions)
+        return positions
+
+    def get_quaternion(positions):
+        skel = Skeleton(n_raw_offsets, kinematic_chain, "cpu")
+        # (seq_len, joints_num, 4)
+        quat_params = skel.inverse_kinematics_np(positions, face_joint_indx, smooth_forward=False)
+
+        '''Fix Quaternion Discontinuity'''
+        quat_params = qfix(quat_params)
+        # (seq_len, 4)
+        r_rot = quat_params[:, 0].copy()
+        #     print(r_rot[0])
+        '''Root Linear Velocity'''
+        # (seq_len - 1, 3)
+        velocity = (positions[1:, 0] - positions[:-1, 0]).copy()
+        #     print(r_rot.shape, velocity.shape)
+        velocity = qrot_np(r_rot[1:], velocity)
+        '''Root Angular Velocity'''
+        # (seq_len - 1, 4)
+        r_velocity = qmul_np(r_rot[1:], qinv_np(r_rot[:-1]))
+        quat_params[1:, 0] = r_velocity
+        # (seq_len, joints_num, 4)
+        return quat_params, r_velocity, velocity, r_rot
+
+    def get_cont6d_params(positions):
+        skel = Skeleton(n_raw_offsets, kinematic_chain, "cpu")
+        # (seq_len, joints_num, 4)
+        quat_params = skel.inverse_kinematics_np(positions, face_joint_indx, smooth_forward=True)
+
+        '''Quaternion to continuous 6D'''
+        cont_6d_params = quaternion_to_cont6d_np(quat_params)
+        # (seq_len, 4)
+        r_rot = quat_params[:, 0].copy()
+        #     print(r_rot[0])
+        '''Root Linear Velocity'''
+        # (seq_len - 1, 3)
+        velocity = (positions[1:, 0] - positions[:-1, 0]).copy()
+        #     print(r_rot.shape, velocity.shape)
+        velocity = qrot_np(r_rot[1:], velocity)
+        '''Root Angular Velocity'''
+        # (seq_len - 1, 4)
+        r_velocity = qmul_np(r_rot[1:], qinv_np(r_rot[:-1]))
+        # (seq_len, joints_num, 4)
+        return cont_6d_params, r_velocity, velocity, r_rot
+
+    cont_6d_params, r_velocity, velocity, r_rot = get_cont6d_params(positions)
+    positions = get_rifke(positions)
+
+    #     trejec = np.cumsum(np.concatenate([np.array([[0, 0, 0]]), velocity], axis=0), axis=0)
+    #     r_rotations, r_pos = recover_ric_glo_np(r_velocity, velocity[:, [0, 2]])
+
+    # plt.plot(positions_b[:, 0, 0], positions_b[:, 0, 2], marker='*')
+    # plt.plot(ground_positions[:, 0, 0], ground_positions[:, 0, 2], marker='o', color='r')
+    # plt.plot(trejec[:, 0], trejec[:, 2], marker='^', color='g')
+    # plt.plot(r_pos[:, 0], r_pos[:, 2], marker='s', color='y')
+    # plt.xlabel('x')
+    # plt.ylabel('z')
+    # plt.axis('equal')
+    # plt.show()
+
+    '''Root height'''
+    root_y = positions[:, 0, 1:2]
+
+    '''Root rotation and linear velocity'''
+    # (seq_len-1, 1) rotation velocity along y-axis
+    # (seq_len-1, 2) linear velovity on xz plane
+    r_velocity = np.arcsin(r_velocity[:, 2:3])
+    l_velocity = velocity[:, [0, 2]]
+    #     print(r_velocity.shape, l_velocity.shape, root_y.shape)
+    root_data = np.concatenate([r_velocity, l_velocity, root_y[:-1]], axis=-1)
+
+    '''Get Joint Rotation Representation'''
+    # (seq_len, (joints_num-1) *6) quaternion for skeleton joints
+    rot_data = cont_6d_params[:, 1:].reshape(len(cont_6d_params), -1)
+
+    '''Get Joint Rotation Invariant Position Represention'''
+    # (seq_len, (joints_num-1)*3) local joint position
+    ric_data = positions[:, 1:].reshape(len(positions), -1)
+
+    '''Get Joint Velocity Representation'''
+    # (seq_len-1, joints_num*3)
+    local_vel = qrot_np(np.repeat(r_rot[:-1, None], global_positions.shape[1], axis=1),
+                        global_positions[1:] - global_positions[:-1])
+    local_vel = local_vel.reshape(len(local_vel), -1)
+
+    data = root_data
+    data = np.concatenate([data, ric_data[:-1]], axis=-1)
+    data = np.concatenate([data, rot_data[:-1]], axis=-1)
+    #     print(data.shape, local_vel.shape)
+    data = np.concatenate([data, local_vel], axis=-1)
+    data = np.concatenate([data, feet_l, feet_r], axis=-1)
+
+    return data
+
+
+def process_file(positions, feet_thre):
+    # (seq_len, joints_num, 3)
+    #     '''Down Sample'''
+    #     positions = positions[::ds_num]
+
+    '''Uniform Skeleton'''
+    positions = uniform_skeleton(positions, tgt_offsets)
+
+    '''Put on Floor'''
+    floor_height = positions.min(axis=0).min(axis=0)[1]
+    positions[:, :, 1] -= floor_height
+    #     print(floor_height)
+
+    #     plot_3d_motion("./positions_1.mp4", kinematic_chain, positions, 'title', fps=20)
+
+    '''XZ at origin'''
+    root_pos_init = positions[0]
+    root_pose_init_xz = root_pos_init[0] * np.array([1, 0, 1])
+    positions = positions - root_pose_init_xz
+
+    # '''Move the first pose to origin '''
+    # root_pos_init = positions[0]
+    # positions = positions - root_pos_init[0]
+
+    '''All initially face Z+'''
+    r_hip, l_hip, sdr_r, sdr_l = face_joint_indx
+    across1 = root_pos_init[r_hip] - root_pos_init[l_hip]
+    across2 = root_pos_init[sdr_r] - root_pos_init[sdr_l]
+    across = across1 + across2
+    across = across / np.sqrt((across ** 2).sum(axis=-1))[..., np.newaxis]
+
+    # forward (3,), rotate around y-axis
+    forward_init = np.cross(np.array([[0, 1, 0]]), across, axis=-1)
+    # forward (3,)
+    forward_init = forward_init / np.sqrt((forward_init ** 2).sum(axis=-1))[..., np.newaxis]
+
+    #     print(forward_init)
+
+    target = np.array([[0, 0, 1]])
+    root_quat_init = qbetween_np(forward_init, target)
+    root_quat_init = np.ones(positions.shape[:-1] + (4,)) * root_quat_init
+
+    positions_b = positions.copy()
+
+    positions = qrot_np(root_quat_init, positions)
+
+    #     plot_3d_motion("./positions_2.mp4", kinematic_chain, positions, 'title', fps=20)
+
+    '''New ground truth positions'''
+    global_positions = positions.copy()
+
+    # plt.plot(positions_b[:, 0, 0], positions_b[:, 0, 2], marker='*')
+    # plt.plot(positions[:, 0, 0], positions[:, 0, 2], marker='o', color='r')
+    # plt.xlabel('x')
+    # plt.ylabel('z')
+    # plt.axis('equal')
+    # plt.show()
+
+    """ Get Foot Contacts """
+
+    def foot_detect(positions, thres):
+        velfactor, heightfactor = np.array([thres, thres]), np.array([3.0, 2.0])
+
+        feet_l_x = (positions[1:, fid_l, 0] - positions[:-1, fid_l, 0]) ** 2
+        feet_l_y = (positions[1:, fid_l, 1] - positions[:-1, fid_l, 1]) ** 2
+        feet_l_z = (positions[1:, fid_l, 2] - positions[:-1, fid_l, 2]) ** 2
+        #     feet_l_h = positions[:-1,fid_l,1]
+        #     feet_l = (((feet_l_x + feet_l_y + feet_l_z) < velfactor) & (feet_l_h < heightfactor)).astype(np.float)
+        feet_l = ((feet_l_x + feet_l_y + feet_l_z) < velfactor).astype(np.float)
+
+        feet_r_x = (positions[1:, fid_r, 0] - positions[:-1, fid_r, 0]) ** 2
+        feet_r_y = (positions[1:, fid_r, 1] - positions[:-1, fid_r, 1]) ** 2
+        feet_r_z = (positions[1:, fid_r, 2] - positions[:-1, fid_r, 2]) ** 2
+        #     feet_r_h = positions[:-1,fid_r,1]
+        #     feet_r = (((feet_r_x + feet_r_y + feet_r_z) < velfactor) & (feet_r_h < heightfactor)).astype(np.float)
+        feet_r = (((feet_r_x + feet_r_y + feet_r_z) < velfactor)).astype(np.float)
+        return feet_l, feet_r
+    #
+    feet_l, feet_r = foot_detect(positions, feet_thre)
+    # feet_l, feet_r = foot_detect(positions, 0.002)
+
+    '''Quaternion and Cartesian representation'''
+    r_rot = None
+
+    def get_rifke(positions):
+        '''Local pose'''
+        positions[..., 0] -= positions[:, 0:1, 0]
+        positions[..., 2] -= positions[:, 0:1, 2]
+        '''All pose face Z+'''
+        positions = qrot_np(np.repeat(r_rot[:, None], positions.shape[1], axis=1), positions)
+        return positions
+
+    def get_quaternion(positions):
+        skel = Skeleton(n_raw_offsets, kinematic_chain, "cpu")
+        # (seq_len, joints_num, 4)
+        quat_params = skel.inverse_kinematics_np(positions, face_joint_indx, smooth_forward=False)
+
+        '''Fix Quaternion Discontinuity'''
+        quat_params = qfix(quat_params)
+        # (seq_len, 4)
+        r_rot = quat_params[:, 0].copy()
+        #     print(r_rot[0])
+        '''Root Linear Velocity'''
+        # (seq_len - 1, 3)
+        velocity = (positions[1:, 0] - positions[:-1, 0]).copy()
+        #     print(r_rot.shape, velocity.shape)
+        velocity = qrot_np(r_rot[1:], velocity)
+        '''Root Angular Velocity'''
+        # (seq_len - 1, 4)
+        r_velocity = qmul_np(r_rot[1:], qinv_np(r_rot[:-1]))
+        quat_params[1:, 0] = r_velocity
+        # (seq_len, joints_num, 4)
+        return quat_params, r_velocity, velocity, r_rot
+
+    def get_cont6d_params(positions):
+        skel = Skeleton(n_raw_offsets, kinematic_chain, "cpu")
+        # (seq_len, joints_num, 4)
+        quat_params = skel.inverse_kinematics_np(positions, face_joint_indx, smooth_forward=True)
+
+        '''Quaternion to continuous 6D'''
+        cont_6d_params = quaternion_to_cont6d_np(quat_params)
+        # (seq_len, 4)
+        r_rot = quat_params[:, 0].copy()
+        #     print(r_rot[0])
+        '''Root Linear Velocity'''
+        # (seq_len - 1, 3)
+        velocity = (positions[1:, 0] - positions[:-1, 0]).copy()
+        #     print(r_rot.shape, velocity.shape)
+        velocity = qrot_np(r_rot[1:], velocity)
+        '''Root Angular Velocity'''
+        # (seq_len - 1, 4)
+        r_velocity = qmul_np(r_rot[1:], qinv_np(r_rot[:-1]))
+        # (seq_len, joints_num, 4)
+        return cont_6d_params, r_velocity, velocity, r_rot
+
+    cont_6d_params, r_velocity, velocity, r_rot = get_cont6d_params(positions)
+    positions = get_rifke(positions)
+
+    #     trejec = np.cumsum(np.concatenate([np.array([[0, 0, 0]]), velocity], axis=0), axis=0)
+    #     r_rotations, r_pos = recover_ric_glo_np(r_velocity, velocity[:, [0, 2]])
+
+    # plt.plot(positions_b[:, 0, 0], positions_b[:, 0, 2], marker='*')
+    # plt.plot(ground_positions[:, 0, 0], ground_positions[:, 0, 2], marker='o', color='r')
+    # plt.plot(trejec[:, 0], trejec[:, 2], marker='^', color='g')
+    # plt.plot(r_pos[:, 0], r_pos[:, 2], marker='s', color='y')
+    # plt.xlabel('x')
+    # plt.ylabel('z')
+    # plt.axis('equal')
+    # plt.show()
+
+    '''Root height'''
+    root_y = positions[:, 0, 1:2]
+
+    '''Root rotation and linear velocity'''
+    # (seq_len-1, 1) rotation velocity along y-axis
+    # (seq_len-1, 2) linear velovity on xz plane
+    r_velocity = np.arcsin(r_velocity[:, 2:3])
+    l_velocity = velocity[:, [0, 2]]
+    #     print(r_velocity.shape, l_velocity.shape, root_y.shape)
+    root_data = np.concatenate([r_velocity, l_velocity, root_y[:-1]], axis=-1)
+
+    '''Get Joint Rotation Representation'''
+    # (seq_len, (joints_num-1) *6) quaternion for skeleton joints
+    rot_data = cont_6d_params[:, 1:].reshape(len(cont_6d_params), -1)
+
+    '''Get Joint Rotation Invariant Position Represention'''
+    # (seq_len, (joints_num-1)*3) local joint position
+    ric_data = positions[:, 1:].reshape(len(positions), -1)
+
+    '''Get Joint Velocity Representation'''
+    # (seq_len-1, joints_num*3)
+    local_vel = qrot_np(np.repeat(r_rot[:-1, None], global_positions.shape[1], axis=1),
+                        global_positions[1:] - global_positions[:-1])
+    local_vel = local_vel.reshape(len(local_vel), -1)
+
+    data = root_data
+    data = np.concatenate([data, ric_data[:-1]], axis=-1)
+    data = np.concatenate([data, rot_data[:-1]], axis=-1)
+    #     print(data.shape, local_vel.shape)
+    data = np.concatenate([data, local_vel], axis=-1)
+    data = np.concatenate([data, feet_l, feet_r], axis=-1)
+
+    return data, global_positions, positions, l_velocity
+
+
+# Recover global angle and positions for rotation data
+# root_rot_velocity (B, seq_len, 1)
+# root_linear_velocity (B, seq_len, 2)
+# root_y (B, seq_len, 1)
+# ric_data (B, seq_len, (joint_num - 1)*3)
+# rot_data (B, seq_len, (joint_num - 1)*6)
+# local_velocity (B, seq_len, joint_num*3)
+# foot contact (B, seq_len, 4)
+def recover_root_rot_pos(data):
+    rot_vel = data[..., 0]
+    r_rot_ang = torch.zeros_like(rot_vel).to(data.device)
+    '''Get Y-axis rotation from rotation velocity'''
+    r_rot_ang[..., 1:] = rot_vel[..., :-1]
+    r_rot_ang = torch.cumsum(r_rot_ang, dim=-1)
+
+    r_rot_quat = torch.zeros(data.shape[:-1] + (4,)).to(data.device)
+    r_rot_quat[..., 0] = torch.cos(r_rot_ang)
+    r_rot_quat[..., 2] = torch.sin(r_rot_ang)
+
+    r_pos = torch.zeros(data.shape[:-1] + (3,)).to(data.device)
+    r_pos[..., 1:, [0, 2]] = data[..., :-1, 1:3]
+    '''Add Y-axis rotation to root position'''
+    r_pos = qrot(qinv(r_rot_quat), r_pos)
+
+    r_pos = torch.cumsum(r_pos, dim=-2)
+
+    r_pos[..., 1] = data[..., 3]
+    return r_rot_quat, r_pos
+
+
+def recover_from_rot(data, joints_num, skeleton):
+    r_rot_quat, r_pos = recover_root_rot_pos(data)
+
+    r_rot_cont6d = quaternion_to_cont6d(r_rot_quat)
+
+    start_indx = 1 + 2 + 1 + (joints_num - 1) * 3
+    end_indx = start_indx + (joints_num - 1) * 6
+    cont6d_params = data[..., start_indx:end_indx]
+    #     print(r_rot_cont6d.shape, cont6d_params.shape, r_pos.shape)
+    cont6d_params = torch.cat([r_rot_cont6d, cont6d_params], dim=-1)
+    cont6d_params = cont6d_params.view(-1, joints_num, 6)
+
+    positions = skeleton.forward_kinematics_cont6d(cont6d_params, r_pos)
+
+    return positions
+
+
+def recover_from_ric(data, joints_num):
+    r_rot_quat, r_pos = recover_root_rot_pos(data)
+    positions = data[..., 4:(joints_num - 1) * 3 + 4]
+    positions = positions.view(positions.shape[:-1] + (-1, 3))
+
+    '''Add Y-axis rotation to local joints'''
+    positions = qrot(qinv(r_rot_quat[..., None, :]).expand(positions.shape[:-1] + (4,)), positions)
+
+    '''Add root XZ to joints'''
+    positions[..., 0] += r_pos[..., 0:1]
+    positions[..., 2] += r_pos[..., 2:3]
+
+    '''Concate root and joints'''
+    positions = torch.cat([r_pos.unsqueeze(-2), positions], dim=-2)
+
+    return positions
+'''
+For Text2Motion Dataset
+'''
+'''
+if __name__ == "__main__":
+    example_id = "000021"
+    # Lower legs
+    l_idx1, l_idx2 = 5, 8
+    # Right/Left foot
+    fid_r, fid_l = [8, 11], [7, 10]
+    # Face direction, r_hip, l_hip, sdr_r, sdr_l
+    face_joint_indx = [2, 1, 17, 16]
+    # l_hip, r_hip
+    r_hip, l_hip = 2, 1
+    joints_num = 22
+    # ds_num = 8
+    data_dir = '../dataset/pose_data_raw/joints/'
+    save_dir1 = '../dataset/pose_data_raw/new_joints/'
+    save_dir2 = '../dataset/pose_data_raw/new_joint_vecs/'
+
+    n_raw_offsets = torch.from_numpy(t2m_raw_offsets)
+    kinematic_chain = t2m_kinematic_chain
+
+    # Get offsets of target skeleton
+    example_data = np.load(os.path.join(data_dir, example_id + '.npy'))
+    example_data = example_data.reshape(len(example_data), -1, 3)
+    example_data = torch.from_numpy(example_data)
+    tgt_skel = Skeleton(n_raw_offsets, kinematic_chain, 'cpu')
+    # (joints_num, 3)
+    tgt_offsets = tgt_skel.get_offsets_joints(example_data[0])
+    # print(tgt_offsets)
+
+    source_list = os.listdir(data_dir)
+    frame_num = 0
+    for source_file in tqdm(source_list):
+        source_data = np.load(os.path.join(data_dir, source_file))[:, :joints_num]
+        try:
+            data, ground_positions, positions, l_velocity = process_file(source_data, 0.002)
+            rec_ric_data = recover_from_ric(torch.from_numpy(data).unsqueeze(0).float(), joints_num)
+            np.save(pjoin(save_dir1, source_file), rec_ric_data.squeeze().numpy())
+            np.save(pjoin(save_dir2, source_file), data)
+            frame_num += data.shape[0]
+        except Exception as e:
+            print(source_file)
+            print(e)
+
+    print('Total clips: %d, Frames: %d, Duration: %fm' %
+          (len(source_list), frame_num, frame_num / 20 / 60))
+'''
+
+if __name__ == "__main__":
+    example_id = "03950_gt"
+    # Lower legs
+    l_idx1, l_idx2 = 17, 18
+    # Right/Left foot
+    fid_r, fid_l = [14, 15], [19, 20]
+    # Face direction, r_hip, l_hip, sdr_r, sdr_l
+    face_joint_indx = [11, 16, 5, 8]
+    # l_hip, r_hip
+    r_hip, l_hip = 11, 16
+    joints_num = 21
+    # ds_num = 8
+    data_dir = '../dataset/kit_mocap_dataset/joints/'
+    save_dir1 = '../dataset/kit_mocap_dataset/new_joints/'
+    save_dir2 = '../dataset/kit_mocap_dataset/new_joint_vecs/'
+
+    n_raw_offsets = torch.from_numpy(kit_raw_offsets)
+    kinematic_chain = kit_kinematic_chain
+
+    '''Get offsets of target skeleton'''
+    example_data = np.load(os.path.join(data_dir, example_id + '.npy'))
+    example_data = example_data.reshape(len(example_data), -1, 3)
+    example_data = torch.from_numpy(example_data)
+    tgt_skel = Skeleton(n_raw_offsets, kinematic_chain, 'cpu')
+    # (joints_num, 3)
+    tgt_offsets = tgt_skel.get_offsets_joints(example_data[0])
+    # print(tgt_offsets)
+
+    source_list = os.listdir(data_dir)
+    frame_num = 0
+    '''Read source data'''
+    for source_file in tqdm(source_list):
+        source_data = np.load(os.path.join(data_dir, source_file))[:, :joints_num]
+        try:
+            name = ''.join(source_file[:-7].split('_')) + '.npy'
+            data, ground_positions, positions, l_velocity = process_file(source_data, 0.05)
+            rec_ric_data = recover_from_ric(torch.from_numpy(data).unsqueeze(0).float(), joints_num)
+            if np.isnan(rec_ric_data.numpy()).any():
+                print(source_file)
+                continue
+            np.save(pjoin(save_dir1, name), rec_ric_data.squeeze().numpy())
+            np.save(pjoin(save_dir2, name), data)
+            frame_num += data.shape[0]
+        except Exception as e:
+            print(source_file)
+            print(e)
+
+    print('Total clips: %d, Frames: %d, Duration: %fm' %
+          (len(source_list), frame_num, frame_num / 12.5 / 60))

utils/paramUtil.py

@@ -0,0 +1,63 @@
+import numpy as np
+
+# Define a kinematic tree for the skeletal struture
+kit_kinematic_chain = [[0, 11, 12, 13, 14, 15], [0, 16, 17, 18, 19, 20], [0, 1, 2, 3, 4], [3, 5, 6, 7], [3, 8, 9, 10]]
+
+kit_raw_offsets = np.array(
+    [
+        [0, 0, 0],
+        [0, 1, 0],
+        [0, 1, 0],
+        [0, 1, 0],
+        [0, 1, 0],
+        [1, 0, 0],
+        [0, -1, 0],
+        [0, -1, 0],
+        [-1, 0, 0],
+        [0, -1, 0],
+        [0, -1, 0],
+        [1, 0, 0],
+        [0, -1, 0],
+        [0, -1, 0],
+        [0, 0, 1],
+        [0, 0, 1],
+        [-1, 0, 0],
+        [0, -1, 0],
+        [0, -1, 0],
+        [0, 0, 1],
+        [0, 0, 1]
+    ]
+)
+
+t2m_raw_offsets = np.array([[0,0,0],
+                           [1,0,0],
+                           [-1,0,0],
+                           [0,1,0],
+                           [0,-1,0],
+                           [0,-1,0],
+                           [0,1,0],
+                           [0,-1,0],
+                           [0,-1,0],
+                           [0,1,0],
+                           [0,0,1],
+                           [0,0,1],
+                           [0,1,0],
+                           [1,0,0],
+                           [-1,0,0],
+                           [0,0,1],
+                           [0,-1,0],
+                           [0,-1,0],
+                           [0,-1,0],
+                           [0,-1,0],
+                           [0,-1,0],
+                           [0,-1,0]])
+
+t2m_kinematic_chain = [[0, 2, 5, 8, 11], [0, 1, 4, 7, 10], [0, 3, 6, 9, 12, 15], [9, 14, 17, 19, 21], [9, 13, 16, 18, 20]]
+t2m_left_hand_chain = [[20, 22, 23, 24], [20, 34, 35, 36], [20, 25, 26, 27], [20, 31, 32, 33], [20, 28, 29, 30]]
+t2m_right_hand_chain = [[21, 43, 44, 45], [21, 46, 47, 48], [21, 40, 41, 42], [21, 37, 38, 39], [21, 49, 50, 51]]
+
+
+kit_tgt_skel_id = '03950'
+
+t2m_tgt_skel_id = '000021'
+

utils/plot_script.py

@@ -0,0 +1,115 @@
+import math
+import numpy as np
+import matplotlib
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
+from matplotlib.animation import FuncAnimation, FFMpegFileWriter
+from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+import mpl_toolkits.mplot3d.axes3d as p3
+# import cv2
+
+
+def list_cut_average(ll, intervals):
+    if intervals == 1:
+        return ll
+
+    bins = math.ceil(len(ll) * 1.0 / intervals)
+    ll_new = []
+    for i in range(bins):
+        l_low = intervals * i
+        l_high = l_low + intervals
+        l_high = l_high if l_high < len(ll) else len(ll)
+        ll_new.append(np.mean(ll[l_low:l_high]))
+    return ll_new
+
+
+def plot_3d_motion(save_path, kinematic_tree, joints, title, figsize=(10, 10), fps=120, radius=4):
+    matplotlib.use('Agg')
+
+    title_sp = title.split(' ')
+    if len(title_sp) > 20:
+        title = '\n'.join([' '.join(title_sp[:10]), ' '.join(title_sp[10:20]), ' '.join(title_sp[20:])])
+    elif len(title_sp) > 10:
+        title = '\n'.join([' '.join(title_sp[:10]), ' '.join(title_sp[10:])])
+
+    def init():
+        ax.set_xlim3d([-radius / 4, radius / 4])
+        ax.set_ylim3d([0, radius / 2])
+        ax.set_zlim3d([0, radius / 2])
+        # print(title)
+        fig.suptitle(title, fontsize=20)
+        ax.grid(b=False)
+
+    def plot_xzPlane(minx, maxx, miny, minz, maxz):
+        ## Plot a plane XZ
+        verts = [
+            [minx, miny, minz],
+            [minx, miny, maxz],
+            [maxx, miny, maxz],
+            [maxx, miny, minz]
+        ]
+        xz_plane = Poly3DCollection([verts])
+        xz_plane.set_facecolor((0.5, 0.5, 0.5, 0.5))
+        ax.add_collection3d(xz_plane)
+
+    #         return ax
+
+    # (seq_len, joints_num, 3)
+    data = joints.copy().reshape(len(joints), -1, 3)
+    fig = plt.figure(figsize=figsize)
+    ax = p3.Axes3D(fig)
+    init()
+    MINS = data.min(axis=0).min(axis=0)
+    MAXS = data.max(axis=0).max(axis=0)
+    colors = ['red', 'blue', 'black', 'red', 'blue',
+              'darkblue', 'darkblue', 'darkblue', 'darkblue', 'darkblue',
+              'darkred', 'darkred', 'darkred', 'darkred', 'darkred']
+    frame_number = data.shape[0]
+    #     print(data.shape)
+
+    height_offset = MINS[1]
+    data[:, :, 1] -= height_offset
+    trajec = data[:, 0, [0, 2]]
+
+    data[..., 0] -= data[:, 0:1, 0]
+    data[..., 2] -= data[:, 0:1, 2]
+
+    #     print(trajec.shape)
+
+    def update(index):
+        #         print(index)
+        ax.lines = []
+        ax.collections = []
+        ax.view_init(elev=120, azim=-90)
+        ax.dist = 7.5
+        #         ax =
+        plot_xzPlane(MINS[0] - trajec[index, 0], MAXS[0] - trajec[index, 0], 0, MINS[2] - trajec[index, 1],
+                     MAXS[2] - trajec[index, 1])
+        #         ax.scatter(data[index, :22, 0], data[index, :22, 1], data[index, :22, 2], color='black', s=3)
+
+        if index > 1:
+            ax.plot3D(trajec[:index, 0] - trajec[index, 0], np.zeros_like(trajec[:index, 0]),
+                      trajec[:index, 1] - trajec[index, 1], linewidth=1.0,
+                      color='blue')
+        #             ax = plot_xzPlane(ax, MINS[0], MAXS[0], 0, MINS[2], MAXS[2])
+
+        for i, (chain, color) in enumerate(zip(kinematic_tree, colors)):
+            #             print(color)
+            if i < 5:
+                linewidth = 4.0
+            else:
+                linewidth = 2.0
+            ax.plot3D(data[index, chain, 0], data[index, chain, 1], data[index, chain, 2], linewidth=linewidth,
+                      color=color)
+        #         print(trajec[:index, 0].shape)
+
+        plt.axis('off')
+        ax.set_xticklabels([])
+        ax.set_yticklabels([])
+        ax.set_zticklabels([])
+
+    ani = FuncAnimation(fig, update, frames=frame_number, interval=1000 / fps, repeat=False)
+
+    # writer = FFMpegFileWriter(fps=fps)
+    ani.save(save_path, fps=fps)
+    plt.close()

utils/quaternion.py

@@ -0,0 +1,423 @@
+# Copyright (c) 2018-present, Facebook, Inc.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+
+import torch
+import numpy as np
+
+_EPS4 = np.finfo(float).eps * 4.0
+
+_FLOAT_EPS = np.finfo(np.float).eps
+
+# PyTorch-backed implementations
+def qinv(q):
+    assert q.shape[-1] == 4, 'q must be a tensor of shape (*, 4)'
+    mask = torch.ones_like(q)
+    mask[..., 1:] = -mask[..., 1:]
+    return q * mask
+
+
+def qinv_np(q):
+    assert q.shape[-1] == 4, 'q must be a tensor of shape (*, 4)'
+    return qinv(torch.from_numpy(q).float()).numpy()
+
+
+def qnormalize(q):
+    assert q.shape[-1] == 4, 'q must be a tensor of shape (*, 4)'
+    return q / torch.norm(q, dim=-1, keepdim=True)
+
+
+def qmul(q, r):
+    """
+    Multiply quaternion(s) q with quaternion(s) r.
+    Expects two equally-sized tensors of shape (*, 4), where * denotes any number of dimensions.
+    Returns q*r as a tensor of shape (*, 4).
+    """
+    assert q.shape[-1] == 4
+    assert r.shape[-1] == 4
+
+    original_shape = q.shape
+
+    # Compute outer product
+    terms = torch.bmm(r.view(-1, 4, 1), q.view(-1, 1, 4))
+
+    w = terms[:, 0, 0] - terms[:, 1, 1] - terms[:, 2, 2] - terms[:, 3, 3]
+    x = terms[:, 0, 1] + terms[:, 1, 0] - terms[:, 2, 3] + terms[:, 3, 2]
+    y = terms[:, 0, 2] + terms[:, 1, 3] + terms[:, 2, 0] - terms[:, 3, 1]
+    z = terms[:, 0, 3] - terms[:, 1, 2] + terms[:, 2, 1] + terms[:, 3, 0]
+    return torch.stack((w, x, y, z), dim=1).view(original_shape)
+
+
+def qrot(q, v):
+    """
+    Rotate vector(s) v about the rotation described by quaternion(s) q.
+    Expects a tensor of shape (*, 4) for q and a tensor of shape (*, 3) for v,
+    where * denotes any number of dimensions.
+    Returns a tensor of shape (*, 3).
+    """
+    assert q.shape[-1] == 4
+    assert v.shape[-1] == 3
+    assert q.shape[:-1] == v.shape[:-1]
+
+    original_shape = list(v.shape)
+    # print(q.shape)
+    q = q.contiguous().view(-1, 4)
+    v = v.contiguous().view(-1, 3)
+
+    qvec = q[:, 1:]
+    uv = torch.cross(qvec, v, dim=1)
+    uuv = torch.cross(qvec, uv, dim=1)
+    return (v + 2 * (q[:, :1] * uv + uuv)).view(original_shape)
+
+
+def qeuler(q, order, epsilon=0, deg=True):
+    """
+    Convert quaternion(s) q to Euler angles.
+    Expects a tensor of shape (*, 4), where * denotes any number of dimensions.
+    Returns a tensor of shape (*, 3).
+    """
+    assert q.shape[-1] == 4
+
+    original_shape = list(q.shape)
+    original_shape[-1] = 3
+    q = q.view(-1, 4)
+
+    q0 = q[:, 0]
+    q1 = q[:, 1]
+    q2 = q[:, 2]
+    q3 = q[:, 3]
+
+    if order == 'xyz':
+        x = torch.atan2(2 * (q0 * q1 - q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2))
+        y = torch.asin(torch.clamp(2 * (q1 * q3 + q0 * q2), -1 + epsilon, 1 - epsilon))
+        z = torch.atan2(2 * (q0 * q3 - q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3))
+    elif order == 'yzx':
+        x = torch.atan2(2 * (q0 * q1 - q2 * q3), 1 - 2 * (q1 * q1 + q3 * q3))
+        y = torch.atan2(2 * (q0 * q2 - q1 * q3), 1 - 2 * (q2 * q2 + q3 * q3))
+        z = torch.asin(torch.clamp(2 * (q1 * q2 + q0 * q3), -1 + epsilon, 1 - epsilon))
+    elif order == 'zxy':
+        x = torch.asin(torch.clamp(2 * (q0 * q1 + q2 * q3), -1 + epsilon, 1 - epsilon))
+        y = torch.atan2(2 * (q0 * q2 - q1 * q3), 1 - 2 * (q1 * q1 + q2 * q2))
+        z = torch.atan2(2 * (q0 * q3 - q1 * q2), 1 - 2 * (q1 * q1 + q3 * q3))
+    elif order == 'xzy':
+        x = torch.atan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q3 * q3))
+        y = torch.atan2(2 * (q0 * q2 + q1 * q3), 1 - 2 * (q2 * q2 + q3 * q3))
+        z = torch.asin(torch.clamp(2 * (q0 * q3 - q1 * q2), -1 + epsilon, 1 - epsilon))
+    elif order == 'yxz':
+        x = torch.asin(torch.clamp(2 * (q0 * q1 - q2 * q3), -1 + epsilon, 1 - epsilon))
+        y = torch.atan2(2 * (q1 * q3 + q0 * q2), 1 - 2 * (q1 * q1 + q2 * q2))
+        z = torch.atan2(2 * (q1 * q2 + q0 * q3), 1 - 2 * (q1 * q1 + q3 * q3))
+    elif order == 'zyx':
+        x = torch.atan2(2 * (q0 * q1 + q2 * q3), 1 - 2 * (q1 * q1 + q2 * q2))
+        y = torch.asin(torch.clamp(2 * (q0 * q2 - q1 * q3), -1 + epsilon, 1 - epsilon))
+        z = torch.atan2(2 * (q0 * q3 + q1 * q2), 1 - 2 * (q2 * q2 + q3 * q3))
+    else:
+        raise
+
+    if deg:
+        return torch.stack((x, y, z), dim=1).view(original_shape) * 180 / np.pi
+    else:
+        return torch.stack((x, y, z), dim=1).view(original_shape)
+
+
+# Numpy-backed implementations
+
+def qmul_np(q, r):
+    q = torch.from_numpy(q).contiguous().float()
+    r = torch.from_numpy(r).contiguous().float()
+    return qmul(q, r).numpy()
+
+
+def qrot_np(q, v):
+    q = torch.from_numpy(q).contiguous().float()
+    v = torch.from_numpy(v).contiguous().float()
+    return qrot(q, v).numpy()
+
+
+def qeuler_np(q, order, epsilon=0, use_gpu=False):
+    if use_gpu:
+        q = torch.from_numpy(q).cuda().float()
+        return qeuler(q, order, epsilon).cpu().numpy()
+    else:
+        q = torch.from_numpy(q).contiguous().float()
+        return qeuler(q, order, epsilon).numpy()
+
+
+def qfix(q):
+    """
+    Enforce quaternion continuity across the time dimension by selecting
+    the representation (q or -q) with minimal distance (or, equivalently, maximal dot product)
+    between two consecutive frames.
+
+    Expects a tensor of shape (L, J, 4), where L is the sequence length and J is the number of joints.
+    Returns a tensor of the same shape.
+    """
+    assert len(q.shape) == 3
+    assert q.shape[-1] == 4
+
+    result = q.copy()
+    dot_products = np.sum(q[1:] * q[:-1], axis=2)
+    mask = dot_products < 0
+    mask = (np.cumsum(mask, axis=0) % 2).astype(bool)
+    result[1:][mask] *= -1
+    return result
+
+
+def euler2quat(e, order, deg=True):
+    """
+    Convert Euler angles to quaternions.
+    """
+    assert e.shape[-1] == 3
+
+    original_shape = list(e.shape)
+    original_shape[-1] = 4
+
+    e = e.view(-1, 3)
+
+    ## if euler angles in degrees
+    if deg:
+        e = e * np.pi / 180.
+
+    x = e[:, 0]
+    y = e[:, 1]
+    z = e[:, 2]
+
+    rx = torch.stack((torch.cos(x / 2), torch.sin(x / 2), torch.zeros_like(x), torch.zeros_like(x)), dim=1)
+    ry = torch.stack((torch.cos(y / 2), torch.zeros_like(y), torch.sin(y / 2), torch.zeros_like(y)), dim=1)
+    rz = torch.stack((torch.cos(z / 2), torch.zeros_like(z), torch.zeros_like(z), torch.sin(z / 2)), dim=1)
+
+    result = None
+    for coord in order:
+        if coord == 'x':
+            r = rx
+        elif coord == 'y':
+            r = ry
+        elif coord == 'z':
+            r = rz
+        else:
+            raise
+        if result is None:
+            result = r
+        else:
+            result = qmul(result, r)
+
+    # Reverse antipodal representation to have a non-negative "w"
+    if order in ['xyz', 'yzx', 'zxy']:
+        result *= -1
+
+    return result.view(original_shape)
+
+
+def expmap_to_quaternion(e):
+    """
+    Convert axis-angle rotations (aka exponential maps) to quaternions.
+    Stable formula from "Practical Parameterization of Rotations Using the Exponential Map".
+    Expects a tensor of shape (*, 3), where * denotes any number of dimensions.
+    Returns a tensor of shape (*, 4).
+    """
+    assert e.shape[-1] == 3
+
+    original_shape = list(e.shape)
+    original_shape[-1] = 4
+    e = e.reshape(-1, 3)
+
+    theta = np.linalg.norm(e, axis=1).reshape(-1, 1)
+    w = np.cos(0.5 * theta).reshape(-1, 1)
+    xyz = 0.5 * np.sinc(0.5 * theta / np.pi) * e
+    return np.concatenate((w, xyz), axis=1).reshape(original_shape)
+
+
+def euler_to_quaternion(e, order):
+    """
+    Convert Euler angles to quaternions.
+    """
+    assert e.shape[-1] == 3
+
+    original_shape = list(e.shape)
+    original_shape[-1] = 4
+
+    e = e.reshape(-1, 3)
+
+    x = e[:, 0]
+    y = e[:, 1]
+    z = e[:, 2]
+
+    rx = np.stack((np.cos(x / 2), np.sin(x / 2), np.zeros_like(x), np.zeros_like(x)), axis=1)
+    ry = np.stack((np.cos(y / 2), np.zeros_like(y), np.sin(y / 2), np.zeros_like(y)), axis=1)
+    rz = np.stack((np.cos(z / 2), np.zeros_like(z), np.zeros_like(z), np.sin(z / 2)), axis=1)
+
+    result = None
+    for coord in order:
+        if coord == 'x':
+            r = rx
+        elif coord == 'y':
+            r = ry
+        elif coord == 'z':
+            r = rz
+        else:
+            raise
+        if result is None:
+            result = r
+        else:
+            result = qmul_np(result, r)
+
+    # Reverse antipodal representation to have a non-negative "w"
+    if order in ['xyz', 'yzx', 'zxy']:
+        result *= -1
+
+    return result.reshape(original_shape)
+
+
+def quaternion_to_matrix(quaternions):
+    """
+    Convert rotations given as quaternions to rotation matrices.
+    Args:
+        quaternions: quaternions with real part first,
+            as tensor of shape (..., 4).
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    r, i, j, k = torch.unbind(quaternions, -1)
+    two_s = 2.0 / (quaternions * quaternions).sum(-1)
+
+    o = torch.stack(
+        (
+            1 - two_s * (j * j + k * k),
+            two_s * (i * j - k * r),
+            two_s * (i * k + j * r),
+            two_s * (i * j + k * r),
+            1 - two_s * (i * i + k * k),
+            two_s * (j * k - i * r),
+            two_s * (i * k - j * r),
+            two_s * (j * k + i * r),
+            1 - two_s * (i * i + j * j),
+        ),
+        -1,
+    )
+    return o.reshape(quaternions.shape[:-1] + (3, 3))
+
+
+def quaternion_to_matrix_np(quaternions):
+    q = torch.from_numpy(quaternions).contiguous().float()
+    return quaternion_to_matrix(q).numpy()
+
+
+def quaternion_to_cont6d_np(quaternions):
+    rotation_mat = quaternion_to_matrix_np(quaternions)
+    cont_6d = np.concatenate([rotation_mat[..., 0], rotation_mat[..., 1]], axis=-1)
+    return cont_6d
+
+
+def quaternion_to_cont6d(quaternions):
+    rotation_mat = quaternion_to_matrix(quaternions)
+    cont_6d = torch.cat([rotation_mat[..., 0], rotation_mat[..., 1]], dim=-1)
+    return cont_6d
+
+
+def cont6d_to_matrix(cont6d):
+    assert cont6d.shape[-1] == 6, "The last dimension must be 6"
+    x_raw = cont6d[..., 0:3]
+    y_raw = cont6d[..., 3:6]
+
+    x = x_raw / torch.norm(x_raw, dim=-1, keepdim=True)
+    z = torch.cross(x, y_raw, dim=-1)
+    z = z / torch.norm(z, dim=-1, keepdim=True)
+
+    y = torch.cross(z, x, dim=-1)
+
+    x = x[..., None]
+    y = y[..., None]
+    z = z[..., None]
+
+    mat = torch.cat([x, y, z], dim=-1)
+    return mat
+
+
+def cont6d_to_matrix_np(cont6d):
+    q = torch.from_numpy(cont6d).contiguous().float()
+    return cont6d_to_matrix(q).numpy()
+
+
+def qpow(q0, t, dtype=torch.float):
+    ''' q0 : tensor of quaternions
+    t: tensor of powers
+    '''
+    q0 = qnormalize(q0)
+    theta0 = torch.acos(q0[..., 0])
+
+    ## if theta0 is close to zero, add epsilon to avoid NaNs
+    mask = (theta0 <= 10e-10) * (theta0 >= -10e-10)
+    theta0 = (1 - mask) * theta0 + mask * 10e-10
+    v0 = q0[..., 1:] / torch.sin(theta0).view(-1, 1)
+
+    if isinstance(t, torch.Tensor):
+        q = torch.zeros(t.shape + q0.shape)
+        theta = t.view(-1, 1) * theta0.view(1, -1)
+    else:  ## if t is a number
+        q = torch.zeros(q0.shape)
+        theta = t * theta0
+
+    q[..., 0] = torch.cos(theta)
+    q[..., 1:] = v0 * torch.sin(theta).unsqueeze(-1)
+
+    return q.to(dtype)
+
+
+def qslerp(q0, q1, t):
+    '''
+    q0: starting quaternion
+    q1: ending quaternion
+    t: array of points along the way
+
+    Returns:
+    Tensor of Slerps: t.shape + q0.shape
+    '''
+
+    q0 = qnormalize(q0)
+    q1 = qnormalize(q1)
+    q_ = qpow(qmul(q1, qinv(q0)), t)
+
+    return qmul(q_,
+                q0.contiguous().view(torch.Size([1] * len(t.shape)) + q0.shape).expand(t.shape + q0.shape).contiguous())
+
+
+def qbetween(v0, v1):
+    '''
+    find the quaternion used to rotate v0 to v1
+    '''
+    assert v0.shape[-1] == 3, 'v0 must be of the shape (*, 3)'
+    assert v1.shape[-1] == 3, 'v1 must be of the shape (*, 3)'
+
+    v = torch.cross(v0, v1)
+    w = torch.sqrt((v0 ** 2).sum(dim=-1, keepdim=True) * (v1 ** 2).sum(dim=-1, keepdim=True)) + (v0 * v1).sum(dim=-1,
+                                                                                                              keepdim=True)
+    return qnormalize(torch.cat([w, v], dim=-1))
+
+
+def qbetween_np(v0, v1):
+    '''
+    find the quaternion used to rotate v0 to v1
+    '''
+    assert v0.shape[-1] == 3, 'v0 must be of the shape (*, 3)'
+    assert v1.shape[-1] == 3, 'v1 must be of the shape (*, 3)'
+
+    v0 = torch.from_numpy(v0).float()
+    v1 = torch.from_numpy(v1).float()
+    return qbetween(v0, v1).numpy()
+
+
+def lerp(p0, p1, t):
+    if not isinstance(t, torch.Tensor):
+        t = torch.Tensor([t])
+
+    new_shape = t.shape + p0.shape
+    new_view_t = t.shape + torch.Size([1] * len(p0.shape))
+    new_view_p = torch.Size([1] * len(t.shape)) + p0.shape
+    p0 = p0.view(new_view_p).expand(new_shape)
+    p1 = p1.view(new_view_p).expand(new_shape)
+    t = t.view(new_view_t).expand(new_shape)
+
+    return p0 + t * (p1 - p0)

utils/skeleton.py

@@ -0,0 +1,199 @@
+from utils.quaternion import *
+import scipy.ndimage.filters as filters
+
+class Skeleton(object):
+    def __init__(self, offset, kinematic_tree, device):
+        self.device = device
+        self._raw_offset_np = offset.numpy()
+        self._raw_offset = offset.clone().detach().to(device).float()
+        self._kinematic_tree = kinematic_tree
+        self._offset = None
+        self._parents = [0] * len(self._raw_offset)
+        self._parents[0] = -1
+        for chain in self._kinematic_tree:
+            for j in range(1, len(chain)):
+                self._parents[chain[j]] = chain[j-1]
+
+    def njoints(self):
+        return len(self._raw_offset)
+
+    def offset(self):
+        return self._offset
+
+    def set_offset(self, offsets):
+        self._offset = offsets.clone().detach().to(self.device).float()
+
+    def kinematic_tree(self):
+        return self._kinematic_tree
+
+    def parents(self):
+        return self._parents
+
+    # joints (batch_size, joints_num, 3)
+    def get_offsets_joints_batch(self, joints):
+        assert len(joints.shape) == 3
+        _offsets = self._raw_offset.expand(joints.shape[0], -1, -1).clone()
+        for i in range(1, self._raw_offset.shape[0]):
+            _offsets[:, i] = torch.norm(joints[:, i] - joints[:, self._parents[i]], p=2, dim=1)[:, None] * _offsets[:, i]
+
+        self._offset = _offsets.detach()
+        return _offsets
+
+    # joints (joints_num, 3)
+    def get_offsets_joints(self, joints):
+        assert len(joints.shape) == 2
+        _offsets = self._raw_offset.clone()
+        for i in range(1, self._raw_offset.shape[0]):
+            # print(joints.shape)
+            _offsets[i] = torch.norm(joints[i] - joints[self._parents[i]], p=2, dim=0) * _offsets[i]
+
+        self._offset = _offsets.detach()
+        return _offsets
+
+    # face_joint_idx should follow the order of right hip, left hip, right shoulder, left shoulder
+    # joints (batch_size, joints_num, 3)
+    def inverse_kinematics_np(self, joints, face_joint_idx, smooth_forward=False):
+        assert len(face_joint_idx) == 4
+        '''Get Forward Direction'''
+        l_hip, r_hip, sdr_r, sdr_l = face_joint_idx
+        across1 = joints[:, r_hip] - joints[:, l_hip]
+        across2 = joints[:, sdr_r] - joints[:, sdr_l]
+        across = across1 + across2
+        across = across / np.sqrt((across**2).sum(axis=-1))[:, np.newaxis]
+        # print(across1.shape, across2.shape)
+
+        # forward (batch_size, 3)
+        forward = np.cross(np.array([[0, 1, 0]]), across, axis=-1)
+        if smooth_forward:
+            forward = filters.gaussian_filter1d(forward, 20, axis=0, mode='nearest')
+            # forward (batch_size, 3)
+        forward = forward / np.sqrt((forward**2).sum(axis=-1))[..., np.newaxis]
+
+        '''Get Root Rotation'''
+        target = np.array([[0,0,1]]).repeat(len(forward), axis=0)
+        root_quat = qbetween_np(forward, target)
+
+        '''Inverse Kinematics'''
+        # quat_params (batch_size, joints_num, 4)
+        # print(joints.shape[:-1])
+        quat_params = np.zeros(joints.shape[:-1] + (4,))
+        # print(quat_params.shape)
+        root_quat[0] = np.array([[1.0, 0.0, 0.0, 0.0]])
+        quat_params[:, 0] = root_quat
+        # quat_params[0, 0] = np.array([[1.0, 0.0, 0.0, 0.0]])
+        for chain in self._kinematic_tree:
+            R = root_quat
+            for j in range(len(chain) - 1):
+                # (batch, 3)
+                u = self._raw_offset_np[chain[j+1]][np.newaxis,...].repeat(len(joints), axis=0)
+                # print(u.shape)
+                # (batch, 3)
+                v = joints[:, chain[j+1]] - joints[:, chain[j]]
+                v = v / np.sqrt((v**2).sum(axis=-1))[:, np.newaxis]
+                # print(u.shape, v.shape)
+                rot_u_v = qbetween_np(u, v)
+
+                R_loc = qmul_np(qinv_np(R), rot_u_v)
+
+                quat_params[:,chain[j + 1], :] = R_loc
+                R = qmul_np(R, R_loc)
+
+        return quat_params
+
+    # Be sure root joint is at the beginning of kinematic chains
+    def forward_kinematics(self, quat_params, root_pos, skel_joints=None, do_root_R=True):
+        # quat_params (batch_size, joints_num, 4)
+        # joints (batch_size, joints_num, 3)
+        # root_pos (batch_size, 3)
+        if skel_joints is not None:
+            offsets = self.get_offsets_joints_batch(skel_joints)
+        if len(self._offset.shape) == 2:
+            offsets = self._offset.expand(quat_params.shape[0], -1, -1)
+        joints = torch.zeros(quat_params.shape[:-1] + (3,)).to(self.device)
+        joints[:, 0] = root_pos
+        for chain in self._kinematic_tree:
+            if do_root_R:
+                R = quat_params[:, 0]
+            else:
+                R = torch.tensor([[1.0, 0.0, 0.0, 0.0]]).expand(len(quat_params), -1).detach().to(self.device)
+            for i in range(1, len(chain)):
+                R = qmul(R, quat_params[:, chain[i]])
+                offset_vec = offsets[:, chain[i]]
+                joints[:, chain[i]] = qrot(R, offset_vec) + joints[:, chain[i-1]]
+        return joints
+
+    # Be sure root joint is at the beginning of kinematic chains
+    def forward_kinematics_np(self, quat_params, root_pos, skel_joints=None, do_root_R=True):
+        # quat_params (batch_size, joints_num, 4)
+        # joints (batch_size, joints_num, 3)
+        # root_pos (batch_size, 3)
+        if skel_joints is not None:
+            skel_joints = torch.from_numpy(skel_joints)
+            offsets = self.get_offsets_joints_batch(skel_joints)
+        if len(self._offset.shape) == 2:
+            offsets = self._offset.expand(quat_params.shape[0], -1, -1)
+        offsets = offsets.numpy()
+        joints = np.zeros(quat_params.shape[:-1] + (3,))
+        joints[:, 0] = root_pos
+        for chain in self._kinematic_tree:
+            if do_root_R:
+                R = quat_params[:, 0]
+            else:
+                R = np.array([[1.0, 0.0, 0.0, 0.0]]).repeat(len(quat_params), axis=0)
+            for i in range(1, len(chain)):
+                R = qmul_np(R, quat_params[:, chain[i]])
+                offset_vec = offsets[:, chain[i]]
+                joints[:, chain[i]] = qrot_np(R, offset_vec) + joints[:, chain[i - 1]]
+        return joints
+
+    def forward_kinematics_cont6d_np(self, cont6d_params, root_pos, skel_joints=None, do_root_R=True):
+        # cont6d_params (batch_size, joints_num, 6)
+        # joints (batch_size, joints_num, 3)
+        # root_pos (batch_size, 3)
+        if skel_joints is not None:
+            skel_joints = torch.from_numpy(skel_joints)
+            offsets = self.get_offsets_joints_batch(skel_joints)
+        if len(self._offset.shape) == 2:
+            offsets = self._offset.expand(cont6d_params.shape[0], -1, -1)
+        offsets = offsets.numpy()
+        joints = np.zeros(cont6d_params.shape[:-1] + (3,))
+        joints[:, 0] = root_pos
+        for chain in self._kinematic_tree:
+            if do_root_R:
+                matR = cont6d_to_matrix_np(cont6d_params[:, 0])
+            else:
+                matR = np.eye(3)[np.newaxis, :].repeat(len(cont6d_params), axis=0)
+            for i in range(1, len(chain)):
+                matR = np.matmul(matR, cont6d_to_matrix_np(cont6d_params[:, chain[i]]))
+                offset_vec = offsets[:, chain[i]][..., np.newaxis]
+                # print(matR.shape, offset_vec.shape)
+                joints[:, chain[i]] = np.matmul(matR, offset_vec).squeeze(-1) + joints[:, chain[i-1]]
+        return joints
+
+    def forward_kinematics_cont6d(self, cont6d_params, root_pos, skel_joints=None, do_root_R=True):
+        # cont6d_params (batch_size, joints_num, 6)
+        # joints (batch_size, joints_num, 3)
+        # root_pos (batch_size, 3)
+        if skel_joints is not None:
+            # skel_joints = torch.from_numpy(skel_joints)
+            offsets = self.get_offsets_joints_batch(skel_joints)
+        if len(self._offset.shape) == 2:
+            offsets = self._offset.expand(cont6d_params.shape[0], -1, -1)
+        joints = torch.zeros(cont6d_params.shape[:-1] + (3,)).to(cont6d_params.device)
+        joints[..., 0, :] = root_pos
+        for chain in self._kinematic_tree:
+            if do_root_R:
+                matR = cont6d_to_matrix(cont6d_params[:, 0])
+            else:
+                matR = torch.eye(3).expand((len(cont6d_params), -1, -1)).detach().to(cont6d_params.device)
+            for i in range(1, len(chain)):
+                matR = torch.matmul(matR, cont6d_to_matrix(cont6d_params[:, chain[i]]))
+                offset_vec = offsets[:, chain[i]].unsqueeze(-1)
+                # print(matR.shape, offset_vec.shape)
+                joints[:, chain[i]] = torch.matmul(matR, offset_vec).squeeze(-1) + joints[:, chain[i-1]]
+        return joints
+
+
+
+
+

utils/utils.py

@@ -0,0 +1,131 @@
+import os
+import numpy as np
+# import cv2
+from PIL import Image
+from utils import paramUtil
+import math
+import time
+import matplotlib.pyplot as plt
+from scipy.ndimage import gaussian_filter
+
+
+def mkdir(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+COLORS = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0],
+          [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255],
+          [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]
+
+MISSING_VALUE = -1
+
+def save_image(image_numpy, image_path):
+    img_pil = Image.fromarray(image_numpy)
+    img_pil.save(image_path)
+
+
+def save_logfile(log_loss, save_path):
+    with open(save_path, 'wt') as f:
+        for k, v in log_loss.items():
+            w_line = k
+            for digit in v:
+                w_line += ' %.3f' % digit
+            f.write(w_line + '\n')
+
+
+def print_current_loss(start_time, niter_state, losses, epoch=None, inner_iter=None):
+
+    def as_minutes(s):
+        m = math.floor(s / 60)
+        s -= m * 60
+        return '%dm %ds' % (m, s)
+
+    def time_since(since, percent):
+        now = time.time()
+        s = now - since
+        es = s / percent
+        rs = es - s
+        return '%s (- %s)' % (as_minutes(s), as_minutes(rs))
+
+    if epoch is not None:
+        print('epoch: %3d niter: %6d  inner_iter: %4d' % (epoch, niter_state, inner_iter), end=" ")
+
+    now = time.time()
+    message = '%s'%(as_minutes(now - start_time))
+
+    for k, v in losses.items():
+        message += ' %s: %.4f ' % (k, v)
+    print(message)
+
+
+def compose_gif_img_list(img_list, fp_out, duration):
+    img, *imgs = [Image.fromarray(np.array(image)) for image in img_list]
+    img.save(fp=fp_out, format='GIF', append_images=imgs, optimize=False,
+             save_all=True, loop=0, duration=duration)
+
+
+def save_images(visuals, image_path):
+    if not os.path.exists(image_path):
+        os.makedirs(image_path)
+
+    for i, (label, img_numpy) in enumerate(visuals.items()):
+        img_name = '%d_%s.jpg' % (i, label)
+        save_path = os.path.join(image_path, img_name)
+        save_image(img_numpy, save_path)
+
+
+def save_images_test(visuals, image_path, from_name, to_name):
+    if not os.path.exists(image_path):
+        os.makedirs(image_path)
+
+    for i, (label, img_numpy) in enumerate(visuals.items()):
+        img_name = "%s_%s_%s" % (from_name, to_name, label)
+        save_path = os.path.join(image_path, img_name)
+        save_image(img_numpy, save_path)
+
+
+def compose_and_save_img(img_list, save_dir, img_name, col=4, row=1, img_size=(256, 200)):
+    # print(col, row)
+    compose_img = compose_image(img_list, col, row, img_size)
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+    img_path = os.path.join(save_dir, img_name)
+    # print(img_path)
+    compose_img.save(img_path)
+
+
+def compose_image(img_list, col, row, img_size):
+    to_image = Image.new('RGB', (col * img_size[0], row * img_size[1]))
+    for y in range(0, row):
+        for x in range(0, col):
+            from_img = Image.fromarray(img_list[y * col + x])
+            # print((x * img_size[0], y*img_size[1],
+            #                           (x + 1) * img_size[0], (y + 1) * img_size[1]))
+            paste_area = (x * img_size[0], y*img_size[1],
+                                      (x + 1) * img_size[0], (y + 1) * img_size[1])
+            to_image.paste(from_img, paste_area)
+            # to_image[y*img_size[1]:(y + 1) * img_size[1], x * img_size[0] :(x + 1) * img_size[0]] = from_img
+    return to_image
+
+
+def list_cut_average(ll, intervals):
+    if intervals == 1:
+        return ll
+
+    bins = math.ceil(len(ll) * 1.0 / intervals)
+    ll_new = []
+    for i in range(bins):
+        l_low = intervals * i
+        l_high = l_low + intervals
+        l_high = l_high if l_high < len(ll) else len(ll)
+        ll_new.append(np.mean(ll[l_low:l_high]))
+    return ll_new
+
+
+def motion_temporal_filter(motion, sigma=1):
+    motion = motion.reshape(motion.shape[0], -1)
+    # print(motion.shape)
+    for i in range(motion.shape[1]):
+        motion[:, i] = gaussian_filter(motion[:, i], sigma=sigma, mode="nearest")
+    return motion.reshape(motion.shape[0], -1, 3)
+

utils/word_vectorizer.py

@@ -0,0 +1,80 @@
+import numpy as np
+import pickle
+from os.path import join as pjoin
+
+POS_enumerator = {
+    'VERB': 0,
+    'NOUN': 1,
+    'DET': 2,
+    'ADP': 3,
+    'NUM': 4,
+    'AUX': 5,
+    'PRON': 6,
+    'ADJ': 7,
+    'ADV': 8,
+    'Loc_VIP': 9,
+    'Body_VIP': 10,
+    'Obj_VIP': 11,
+    'Act_VIP': 12,
+    'Desc_VIP': 13,
+    'OTHER': 14,
+}
+
+Loc_list = ('left', 'right', 'clockwise', 'counterclockwise', 'anticlockwise', 'forward', 'back', 'backward',
+            'up', 'down', 'straight', 'curve')
+
+Body_list = ('arm', 'chin', 'foot', 'feet', 'face', 'hand', 'mouth', 'leg', 'waist', 'eye', 'knee', 'shoulder', 'thigh')
+
+Obj_List = ('stair', 'dumbbell', 'chair', 'window', 'floor', 'car', 'ball', 'handrail', 'baseball', 'basketball')
+
+Act_list = ('walk', 'run', 'swing', 'pick', 'bring', 'kick', 'put', 'squat', 'throw', 'hop', 'dance', 'jump', 'turn',
+            'stumble', 'dance', 'stop', 'sit', 'lift', 'lower', 'raise', 'wash', 'stand', 'kneel', 'stroll',
+            'rub', 'bend', 'balance', 'flap', 'jog', 'shuffle', 'lean', 'rotate', 'spin', 'spread', 'climb')
+
+Desc_list = ('slowly', 'carefully', 'fast', 'careful', 'slow', 'quickly', 'happy', 'angry', 'sad', 'happily',
+             'angrily', 'sadly')
+
+VIP_dict = {
+    'Loc_VIP': Loc_list,
+    'Body_VIP': Body_list,
+    'Obj_VIP': Obj_List,
+    'Act_VIP': Act_list,
+    'Desc_VIP': Desc_list,
+}
+
+
+class WordVectorizer(object):
+    def __init__(self, meta_root, prefix):
+        vectors = np.load(pjoin(meta_root, '%s_data.npy'%prefix))
+        words = pickle.load(open(pjoin(meta_root, '%s_words.pkl'%prefix), 'rb'))
+        word2idx = pickle.load(open(pjoin(meta_root, '%s_idx.pkl'%prefix), 'rb'))
+        self.word2vec = {w: vectors[word2idx[w]] for w in words}
+
+    def _get_pos_ohot(self, pos):
+        pos_vec = np.zeros(len(POS_enumerator))
+        if pos in POS_enumerator:
+            pos_vec[POS_enumerator[pos]] = 1
+        else:
+            pos_vec[POS_enumerator['OTHER']] = 1
+        return pos_vec
+
+    def __len__(self):
+        return len(self.word2vec)
+
+    def __getitem__(self, item):
+        word, pos = item.split('/')
+        if word in self.word2vec:
+            word_vec = self.word2vec[word]
+            vip_pos = None
+            for key, values in VIP_dict.items():
+                if word in values:
+                    vip_pos = key
+                    break
+            if vip_pos is not None:
+                pos_vec = self._get_pos_ohot(vip_pos)
+            else:
+                pos_vec = self._get_pos_ohot(pos)
+        else:
+            word_vec = self.word2vec['unk']
+            pos_vec = self._get_pos_ohot('OTHER')
+        return word_vec, pos_vec