Upload 473 files

.gitattributes

@@ -32,3 +32,4 @@ 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
+generate_human_motion/VQTrans/output/results.gif filter=lfs diff=lfs merge=lfs -text

generate_human_motion/VQTrans/GPT_eval_multi.py

@@ -0,0 +1,121 @@
+import os 
+import torch
+import numpy as np
+from torch.utils.tensorboard import SummaryWriter
+import json
+import clip
+
+import options.option_transformer as option_trans
+import models.vqvae as vqvae
+import utils.utils_model as utils_model
+import utils.eval_trans as eval_trans
+from dataset import dataset_TM_eval
+import models.t2m_trans as trans
+from options.get_eval_option import get_opt
+from models.evaluator_wrapper import EvaluatorModelWrapper
+import warnings
+warnings.filterwarnings('ignore')
+
+##### ---- Exp dirs ---- #####
+args = option_trans.get_args_parser()
+torch.manual_seed(args.seed)
+
+args.out_dir = os.path.join(args.out_dir, f'{args.exp_name}')
+os.makedirs(args.out_dir, exist_ok = True)
+
+##### ---- Logger ---- #####
+logger = utils_model.get_logger(args.out_dir)
+writer = SummaryWriter(args.out_dir)
+logger.info(json.dumps(vars(args), indent=4, sort_keys=True))
+
+from utils.word_vectorizer import WordVectorizer
+w_vectorizer = WordVectorizer('./glove', 'our_vab')
+val_loader = dataset_TM_eval.DATALoader(args.dataname, True, 32, w_vectorizer)
+
+dataset_opt_path = 'checkpoints/kit/Comp_v6_KLD005/opt.txt' if args.dataname == 'kit' else 'checkpoints/t2m/Comp_v6_KLD005/opt.txt'
+
+wrapper_opt = get_opt(dataset_opt_path, torch.device('cuda'))
+eval_wrapper = EvaluatorModelWrapper(wrapper_opt)
+
+##### ---- Network ---- #####
+
+## load clip model and datasets
+clip_model, clip_preprocess = clip.load("ViT-B/32", device=torch.device('cuda'), jit=False, download_root='/apdcephfs_cq2/share_1290939/maelyszhang/.cache/clip')  # Must set jit=False for training
+clip.model.convert_weights(clip_model)  # Actually this line is unnecessary since clip by default already on float16
+clip_model.eval()
+for p in clip_model.parameters():
+    p.requires_grad = False
+
+net = vqvae.HumanVQVAE(args, ## use args to define different parameters in different quantizers
+                       args.nb_code,
+                       args.code_dim,
+                       args.output_emb_width,
+                       args.down_t,
+                       args.stride_t,
+                       args.width,
+                       args.depth,
+                       args.dilation_growth_rate)
+
+
+trans_encoder = trans.Text2Motion_Transformer(num_vq=args.nb_code, 
+                                embed_dim=args.embed_dim_gpt, 
+                                clip_dim=args.clip_dim, 
+                                block_size=args.block_size, 
+                                num_layers=args.num_layers, 
+                                n_head=args.n_head_gpt, 
+                                drop_out_rate=args.drop_out_rate, 
+                                fc_rate=args.ff_rate)
+
+
+print ('loading checkpoint from {}'.format(args.resume_pth))
+ckpt = torch.load(args.resume_pth, map_location='cpu')
+net.load_state_dict(ckpt['net'], strict=True)
+net.eval()
+net.cuda()
+
+if args.resume_trans is not None:
+    print ('loading transformer checkpoint from {}'.format(args.resume_trans))
+    ckpt = torch.load(args.resume_trans, map_location='cpu')
+    trans_encoder.load_state_dict(ckpt['trans'], strict=True)
+trans_encoder.train()
+trans_encoder.cuda()
+
+
+fid = []
+div = []
+top1 = []
+top2 = []
+top3 = []
+matching = []
+multi = []
+repeat_time = 20
+
+        
+for i in range(repeat_time):
+    best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching, best_multi, writer, logger = eval_trans.evaluation_transformer_test(args.out_dir, val_loader, net, trans_encoder, logger, writer, 0, best_fid=1000, best_iter=0, best_div=100, best_top1=0, best_top2=0, best_top3=0, best_matching=100, best_multi=0, clip_model=clip_model, eval_wrapper=eval_wrapper, draw=False, savegif=False, save=False, savenpy=(i==0))
+    fid.append(best_fid)
+    div.append(best_div)
+    top1.append(best_top1)
+    top2.append(best_top2)
+    top3.append(best_top3)
+    matching.append(best_matching)
+    multi.append(best_multi)
+
+print('final result:')
+print('fid: ', sum(fid)/repeat_time)
+print('div: ', sum(div)/repeat_time)
+print('top1: ', sum(top1)/repeat_time)
+print('top2: ', sum(top2)/repeat_time)
+print('top3: ', sum(top3)/repeat_time)
+print('matching: ', sum(matching)/repeat_time)
+print('multi: ', sum(multi)/repeat_time)
+
+fid = np.array(fid)
+div = np.array(div)
+top1 = np.array(top1)
+top2 = np.array(top2)
+top3 = np.array(top3)
+matching = np.array(matching)
+multi = np.array(multi)
+msg_final = f"FID. {np.mean(fid):.3f}, conf. {np.std(fid)*1.96/np.sqrt(repeat_time):.3f}, Diversity. {np.mean(div):.3f}, conf. {np.std(div)*1.96/np.sqrt(repeat_time):.3f}, TOP1. {np.mean(top1):.3f}, conf. {np.std(top1)*1.96/np.sqrt(repeat_time):.3f}, TOP2. {np.mean(top2):.3f}, conf. {np.std(top2)*1.96/np.sqrt(repeat_time):.3f}, TOP3. {np.mean(top3):.3f}, conf. {np.std(top3)*1.96/np.sqrt(repeat_time):.3f}, Matching. {np.mean(matching):.3f}, conf. {np.std(matching)*1.96/np.sqrt(repeat_time):.3f}, Multi. {np.mean(multi):.3f}, conf. {np.std(multi)*1.96/np.sqrt(repeat_time):.3f}"
+logger.info(msg_final)

generate_human_motion/VQTrans/VQ_eval.py

@@ -0,0 +1,95 @@
+import os
+import json
+
+import torch
+from torch.utils.tensorboard import SummaryWriter
+import numpy as np
+import models.vqvae as vqvae
+import options.option_vq as option_vq
+import utils.utils_model as utils_model
+from dataset import dataset_TM_eval
+import utils.eval_trans as eval_trans
+from options.get_eval_option import get_opt
+from models.evaluator_wrapper import EvaluatorModelWrapper
+import warnings
+warnings.filterwarnings('ignore')
+import numpy as np
+##### ---- Exp dirs ---- #####
+args = option_vq.get_args_parser()
+torch.manual_seed(args.seed)
+
+args.out_dir = os.path.join(args.out_dir, f'{args.exp_name}')
+os.makedirs(args.out_dir, exist_ok = True)
+
+##### ---- Logger ---- #####
+logger = utils_model.get_logger(args.out_dir)
+writer = SummaryWriter(args.out_dir)
+logger.info(json.dumps(vars(args), indent=4, sort_keys=True))
+
+
+from utils.word_vectorizer import WordVectorizer
+w_vectorizer = WordVectorizer('./glove', 'our_vab')
+
+
+dataset_opt_path = 'checkpoints/kit/Comp_v6_KLD005/opt.txt' if args.dataname == 'kit' else 'checkpoints/t2m/Comp_v6_KLD005/opt.txt'
+
+wrapper_opt = get_opt(dataset_opt_path, torch.device('cuda'))
+eval_wrapper = EvaluatorModelWrapper(wrapper_opt)
+
+
+##### ---- Dataloader ---- #####
+args.nb_joints = 21 if args.dataname == 'kit' else 22
+
+val_loader = dataset_TM_eval.DATALoader(args.dataname, True, 32, w_vectorizer, unit_length=2**args.down_t)
+
+##### ---- Network ---- #####
+net = vqvae.HumanVQVAE(args, ## use args to define different parameters in different quantizers
+                       args.nb_code,
+                       args.code_dim,
+                       args.output_emb_width,
+                       args.down_t,
+                       args.stride_t,
+                       args.width,
+                       args.depth,
+                       args.dilation_growth_rate,
+                       args.vq_act,
+                       args.vq_norm)
+
+if args.resume_pth : 
+    logger.info('loading checkpoint from {}'.format(args.resume_pth))
+    ckpt = torch.load(args.resume_pth, map_location='cpu')
+    net.load_state_dict(ckpt['net'], strict=True)
+net.train()
+net.cuda()
+
+fid = []
+div = []
+top1 = []
+top2 = []
+top3 = []
+matching = []
+repeat_time = 20
+for i in range(repeat_time):
+    best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching, writer, logger = eval_trans.evaluation_vqvae(args.out_dir, val_loader, net, logger, writer, 0, best_fid=1000, best_iter=0, best_div=100, best_top1=0, best_top2=0, best_top3=0, best_matching=100, eval_wrapper=eval_wrapper, draw=False, save=False, savenpy=(i==0))
+    fid.append(best_fid)
+    div.append(best_div)
+    top1.append(best_top1)
+    top2.append(best_top2)
+    top3.append(best_top3)
+    matching.append(best_matching)
+print('final result:')
+print('fid: ', sum(fid)/repeat_time)
+print('div: ', sum(div)/repeat_time)
+print('top1: ', sum(top1)/repeat_time)
+print('top2: ', sum(top2)/repeat_time)
+print('top3: ', sum(top3)/repeat_time)
+print('matching: ', sum(matching)/repeat_time)
+
+fid = np.array(fid)
+div = np.array(div)
+top1 = np.array(top1)
+top2 = np.array(top2)
+top3 = np.array(top3)
+matching = np.array(matching)
+msg_final = f"FID. {np.mean(fid):.3f}, conf. {np.std(fid)*1.96/np.sqrt(repeat_time):.3f}, Diversity. {np.mean(div):.3f}, conf. {np.std(div)*1.96/np.sqrt(repeat_time):.3f}, TOP1. {np.mean(top1):.3f}, conf. {np.std(top1)*1.96/np.sqrt(repeat_time):.3f}, TOP2. {np.mean(top2):.3f}, conf. {np.std(top2)*1.96/np.sqrt(repeat_time):.3f}, TOP3. {np.mean(top3):.3f}, conf. {np.std(top3)*1.96/np.sqrt(repeat_time):.3f}, Matching. {np.mean(matching):.3f}, conf. {np.std(matching)*1.96/np.sqrt(repeat_time):.3f}"
+logger.info(msg_final)

generate_human_motion/VQTrans/ViT-B-32.pt

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

generate_human_motion/VQTrans/body_models/smpl/J_regressor_extra.npy

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

generate_human_motion/VQTrans/body_models/smpl/SMPL_NEUTRAL.pkl

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

generate_human_motion/VQTrans/body_models/smpl/kintree_table.pkl

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

generate_human_motion/VQTrans/body_models/smpl/smplfaces.npy

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

generate_human_motion/VQTrans/checkpoints/kit.zip

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

generate_human_motion/VQTrans/checkpoints/t2m.zip

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

generate_human_motion/VQTrans/checkpoints/train_vq.py

@@ -0,0 +1,171 @@
+import os
+import json
+
+import torch
+import torch.optim as optim
+from torch.utils.tensorboard import SummaryWriter
+
+import models.vqvae as vqvae
+import utils.losses as losses 
+import options.option_vq as option_vq
+import utils.utils_model as utils_model
+from dataset import dataset_VQ, dataset_TM_eval
+import utils.eval_trans as eval_trans
+from options.get_eval_option import get_opt
+from models.evaluator_wrapper import EvaluatorModelWrapper
+import warnings
+warnings.filterwarnings('ignore')
+from utils.word_vectorizer import WordVectorizer
+
+def update_lr_warm_up(optimizer, nb_iter, warm_up_iter, lr):
+
+    current_lr = lr * (nb_iter + 1) / (warm_up_iter + 1)
+    for param_group in optimizer.param_groups:
+        param_group["lr"] = current_lr
+
+    return optimizer, current_lr
+
+##### ---- Exp dirs ---- #####
+args = option_vq.get_args_parser()
+torch.manual_seed(args.seed)
+
+args.out_dir = os.path.join(args.out_dir, f'{args.exp_name}')
+os.makedirs(args.out_dir, exist_ok = True)
+
+##### ---- Logger ---- #####
+logger = utils_model.get_logger(args.out_dir)
+writer = SummaryWriter(args.out_dir)
+logger.info(json.dumps(vars(args), indent=4, sort_keys=True))
+
+
+
+w_vectorizer = WordVectorizer('./glove', 'our_vab')
+
+if args.dataname == 'kit' : 
+    dataset_opt_path = 'checkpoints/kit/Comp_v6_KLD005/opt.txt'  
+    args.nb_joints = 21
+    
+else :
+    dataset_opt_path = 'checkpoints/t2m/Comp_v6_KLD005/opt.txt'
+    args.nb_joints = 22
+
+logger.info(f'Training on {args.dataname}, motions are with {args.nb_joints} joints')
+
+wrapper_opt = get_opt(dataset_opt_path, torch.device('cuda'))
+eval_wrapper = EvaluatorModelWrapper(wrapper_opt)
+
+
+##### ---- Dataloader ---- #####
+train_loader = dataset_VQ.DATALoader(args.dataname,
+                                        args.batch_size,
+                                        window_size=args.window_size,
+                                        unit_length=2**args.down_t)
+
+train_loader_iter = dataset_VQ.cycle(train_loader)
+
+val_loader = dataset_TM_eval.DATALoader(args.dataname, False,
+                                        32,
+                                        w_vectorizer,
+                                        unit_length=2**args.down_t)
+
+##### ---- Network ---- #####
+net = vqvae.HumanVQVAE(args, ## use args to define different parameters in different quantizers
+                       args.nb_code,
+                       args.code_dim,
+                       args.output_emb_width,
+                       args.down_t,
+                       args.stride_t,
+                       args.width,
+                       args.depth,
+                       args.dilation_growth_rate,
+                       args.vq_act,
+                       args.vq_norm)
+
+
+if args.resume_pth : 
+    logger.info('loading checkpoint from {}'.format(args.resume_pth))
+    ckpt = torch.load(args.resume_pth, map_location='cpu')
+    net.load_state_dict(ckpt['net'], strict=True)
+net.train()
+net.cuda()
+
+##### ---- Optimizer & Scheduler ---- #####
+optimizer = optim.AdamW(net.parameters(), lr=args.lr, betas=(0.9, 0.99), weight_decay=args.weight_decay)
+scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.lr_scheduler, gamma=args.gamma)
+  
+
+Loss = losses.ReConsLoss(args.recons_loss, args.nb_joints)
+
+##### ------ warm-up ------- #####
+avg_recons, avg_perplexity, avg_commit = 0., 0., 0.
+
+for nb_iter in range(1, args.warm_up_iter):
+    
+    optimizer, current_lr = update_lr_warm_up(optimizer, nb_iter, args.warm_up_iter, args.lr)
+    
+    gt_motion = next(train_loader_iter)
+    gt_motion = gt_motion.cuda().float() # (bs, 64, dim)
+
+    pred_motion, loss_commit, perplexity = net(gt_motion)
+    loss_motion = Loss(pred_motion, gt_motion)
+    loss_vel = Loss.forward_vel(pred_motion, gt_motion)
+    
+    loss = loss_motion + args.commit * loss_commit + args.loss_vel * loss_vel
+    
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+    avg_recons += loss_motion.item()
+    avg_perplexity += perplexity.item()
+    avg_commit += loss_commit.item()
+    
+    if nb_iter % args.print_iter ==  0 :
+        avg_recons /= args.print_iter
+        avg_perplexity /= args.print_iter
+        avg_commit /= args.print_iter
+        
+        logger.info(f"Warmup. Iter {nb_iter} :  lr {current_lr:.5f} \t Commit. {avg_commit:.5f} \t PPL. {avg_perplexity:.2f} \t Recons.  {avg_recons:.5f}")
+        
+        avg_recons, avg_perplexity, avg_commit = 0., 0., 0.
+
+##### ---- Training ---- #####
+avg_recons, avg_perplexity, avg_commit = 0., 0., 0.
+best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching, writer, logger = eval_trans.evaluation_vqvae(args.out_dir, val_loader, net, logger, writer, 0, best_fid=1000, best_iter=0, best_div=100, best_top1=0, best_top2=0, best_top3=0, best_matching=100, eval_wrapper=eval_wrapper)
+
+for nb_iter in range(1, args.total_iter + 1):
+    
+    gt_motion = next(train_loader_iter)
+    gt_motion = gt_motion.cuda().float() # bs, nb_joints, joints_dim, seq_len
+    
+    pred_motion, loss_commit, perplexity = net(gt_motion)
+    loss_motion = Loss(pred_motion, gt_motion)
+    loss_vel = Loss.forward_vel(pred_motion, gt_motion)
+    
+    loss = loss_motion + args.commit * loss_commit + args.loss_vel * loss_vel
+    
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+    scheduler.step()
+    
+    avg_recons += loss_motion.item()
+    avg_perplexity += perplexity.item()
+    avg_commit += loss_commit.item()
+    
+    if nb_iter % args.print_iter ==  0 :
+        avg_recons /= args.print_iter
+        avg_perplexity /= args.print_iter
+        avg_commit /= args.print_iter
+        
+        writer.add_scalar('./Train/L1', avg_recons, nb_iter)
+        writer.add_scalar('./Train/PPL', avg_perplexity, nb_iter)
+        writer.add_scalar('./Train/Commit', avg_commit, nb_iter)
+        
+        logger.info(f"Train. Iter {nb_iter} : \t Commit. {avg_commit:.5f} \t PPL. {avg_perplexity:.2f} \t Recons.  {avg_recons:.5f}")
+        
+        avg_recons, avg_perplexity, avg_commit = 0., 0., 0.,
+
+    if nb_iter % args.eval_iter==0 :
+        best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching, writer, logger = eval_trans.evaluation_vqvae(args.out_dir, val_loader, net, logger, writer, nb_iter, best_fid, best_iter, best_div, best_top1, best_top2, best_top3, best_matching, eval_wrapper=eval_wrapper)
+        

generate_human_motion/VQTrans/dataset/dataset_TM_eval.py

@@ -0,0 +1,217 @@
+import torch
+from torch.utils import data
+import numpy as np
+from os.path import join as pjoin
+import random
+import codecs as cs
+from tqdm import tqdm
+
+import utils.paramUtil as paramUtil
+from torch.utils.data._utils.collate import default_collate
+
+
+def collate_fn(batch):
+    batch.sort(key=lambda x: x[3], reverse=True)
+    return default_collate(batch)
+
+
+'''For use of training text-2-motion generative model'''
+class Text2MotionDataset(data.Dataset):
+    def __init__(self, dataset_name, is_test, w_vectorizer, feat_bias = 5, max_text_len = 20, unit_length = 4):
+        
+        self.max_length = 20
+        self.pointer = 0
+        self.dataset_name = dataset_name
+        self.is_test = is_test
+        self.max_text_len = max_text_len
+        self.unit_length = unit_length
+        self.w_vectorizer = w_vectorizer
+        if dataset_name == 't2m':
+            self.data_root = './dataset/HumanML3D'
+            self.motion_dir = pjoin(self.data_root, 'new_joint_vecs')
+            self.text_dir = pjoin(self.data_root, 'texts')
+            self.joints_num = 22
+            radius = 4
+            fps = 20
+            self.max_motion_length = 196
+            dim_pose = 263
+            kinematic_chain = paramUtil.t2m_kinematic_chain
+            self.meta_dir = 'checkpoints/t2m/VQVAEV3_CB1024_CMT_H1024_NRES3/meta'
+        elif dataset_name == 'kit':
+            self.data_root = './dataset/KIT-ML'
+            self.motion_dir = pjoin(self.data_root, 'new_joint_vecs')
+            self.text_dir = pjoin(self.data_root, 'texts')
+            self.joints_num = 21
+            radius = 240 * 8
+            fps = 12.5
+            dim_pose = 251
+            self.max_motion_length = 196
+            kinematic_chain = paramUtil.kit_kinematic_chain
+            self.meta_dir = 'checkpoints/kit/VQVAEV3_CB1024_CMT_H1024_NRES3/meta'
+
+        mean = np.load(pjoin(self.meta_dir, 'mean.npy'))
+        std = np.load(pjoin(self.meta_dir, 'std.npy'))
+        
+        if is_test:
+            split_file = pjoin(self.data_root, 'test.txt')
+        else:
+            split_file = pjoin(self.data_root, 'val.txt')
+
+        min_motion_len = 40 if self.dataset_name =='t2m' else 24
+        # min_motion_len = 64
+
+        joints_num = self.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(self.motion_dir, name + '.npy'))
+                if (len(motion)) < min_motion_len or (len(motion) >= 200):
+                    continue
+                text_data = []
+                flag = False
+                with cs.open(pjoin(self.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*fps) : int(to_tag*fps)]
+                                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 Exception as e:
+                # print(e)
+                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 forward_transform(self, data):
+        return (data - self.mean) / self.std
+
+    def __len__(self):
+        return len(self.data_dict) - self.pointer
+
+    def __getitem__(self, item):
+        idx = self.pointer + item
+        name = self.name_list[idx]
+        data = self.data_dict[name]
+        # 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.max_text_len:
+            # pad with "unk"
+            tokens = ['sos/OTHER'] + tokens + ['eos/OTHER']
+            sent_len = len(tokens)
+            tokens = tokens + ['unk/OTHER'] * (self.max_text_len + 2 - sent_len)
+        else:
+            # crop
+            tokens = tokens[:self.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)
+
+        if self.unit_length < 10:
+            coin2 = np.random.choice(['single', 'single', 'double'])
+        else:
+            coin2 = 'single'
+
+        if coin2 == 'double':
+            m_length = (m_length // self.unit_length - 1) * self.unit_length
+        elif coin2 == 'single':
+            m_length = (m_length // self.unit_length) * self.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), name
+
+
+
+
+def DATALoader(dataset_name, is_test,
+                batch_size, w_vectorizer,
+                num_workers = 8, unit_length = 4) : 
+    
+    val_loader = torch.utils.data.DataLoader(Text2MotionDataset(dataset_name, is_test, w_vectorizer, unit_length=unit_length),
+                                              batch_size,
+                                              shuffle = True,
+                                              num_workers=num_workers,
+                                              collate_fn=collate_fn,
+                                              drop_last = True)
+    return val_loader
+
+
+def cycle(iterable):
+    while True:
+        for x in iterable:
+            yield x

generate_human_motion/VQTrans/dataset/dataset_TM_train.py

@@ -0,0 +1,161 @@
+import torch
+from torch.utils import data
+import numpy as np
+from os.path import join as pjoin
+import random
+import codecs as cs
+from tqdm import tqdm
+import utils.paramUtil as paramUtil
+from torch.utils.data._utils.collate import default_collate
+
+
+def collate_fn(batch):
+    batch.sort(key=lambda x: x[3], reverse=True)
+    return default_collate(batch)
+
+
+'''For use of training text-2-motion generative model'''
+class Text2MotionDataset(data.Dataset):
+    def __init__(self, dataset_name, feat_bias = 5, unit_length = 4, codebook_size = 1024, tokenizer_name=None):
+        
+        self.max_length = 64
+        self.pointer = 0
+        self.dataset_name = dataset_name
+
+        self.unit_length = unit_length
+        # self.mot_start_idx = codebook_size
+        self.mot_end_idx = codebook_size
+        self.mot_pad_idx = codebook_size + 1
+        if dataset_name == 't2m':
+            self.data_root = './dataset/HumanML3D'
+            self.motion_dir = pjoin(self.data_root, 'new_joint_vecs')
+            self.text_dir = pjoin(self.data_root, 'texts')
+            self.joints_num = 22
+            radius = 4
+            fps = 20
+            self.max_motion_length = 26 if unit_length == 8 else 51
+            dim_pose = 263
+            kinematic_chain = paramUtil.t2m_kinematic_chain
+        elif dataset_name == 'kit':
+            self.data_root = './dataset/KIT-ML'
+            self.motion_dir = pjoin(self.data_root, 'new_joint_vecs')
+            self.text_dir = pjoin(self.data_root, 'texts')
+            self.joints_num = 21
+            radius = 240 * 8
+            fps = 12.5
+            dim_pose = 251
+            self.max_motion_length = 26 if unit_length == 8 else 51
+            kinematic_chain = paramUtil.kit_kinematic_chain
+
+        split_file = pjoin(self.data_root, 'train.txt')
+
+
+        id_list = []
+        with cs.open(split_file, 'r') as f:
+            for line in f.readlines():
+                id_list.append(line.strip())
+
+        new_name_list = []
+        data_dict = {}
+        for name in tqdm(id_list):
+            try:
+                m_token_list = np.load(pjoin(self.data_root, tokenizer_name, '%s.npy'%name))
+
+                # Read text
+                with cs.open(pjoin(self.text_dir, name + '.txt')) as f:
+                    text_data = []
+                    flag = False
+                    lines = f.readlines()
+
+                    for line in lines:
+                        try:
+                            text_dict = {}
+                            line_split = line.strip().split('#')
+                            caption = line_split[0]
+                            t_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'] = t_tokens
+                            if f_tag == 0.0 and to_tag == 0.0:
+                                flag = True
+                                text_data.append(text_dict)
+                            else:
+                                m_token_list_new = [tokens[int(f_tag*fps/unit_length) : int(to_tag*fps/unit_length)] for tokens in m_token_list if int(f_tag*fps/unit_length) < int(to_tag*fps/unit_length)]
+
+                                if len(m_token_list_new) == 0:
+                                    continue
+                                new_name = '%s_%f_%f'%(name, f_tag, to_tag)
+
+                                data_dict[new_name] = {'m_token_list': m_token_list_new,
+                                                       'text':[text_dict]}
+                                new_name_list.append(new_name)
+                        except:
+                            pass
+
+                if flag:
+                    data_dict[name] = {'m_token_list': m_token_list,
+                                       'text':text_data}
+                    new_name_list.append(name)
+            except:
+                pass
+        self.data_dict = data_dict
+        self.name_list = new_name_list
+
+    def __len__(self):
+        return len(self.data_dict)
+
+    def __getitem__(self, item):
+        data = self.data_dict[self.name_list[item]]
+        m_token_list, text_list = data['m_token_list'], data['text']
+        m_tokens = random.choice(m_token_list)
+
+        text_data = random.choice(text_list)
+        caption= text_data['caption']
+
+        
+        coin = np.random.choice([False, False, True])
+        # print(len(m_tokens))
+        if coin:
+            # drop one token at the head or tail
+            coin2 = np.random.choice([True, False])
+            if coin2:
+                m_tokens = m_tokens[:-1]
+            else:
+                m_tokens = m_tokens[1:]
+        m_tokens_len = m_tokens.shape[0]
+
+        if m_tokens_len+1 < self.max_motion_length:
+            m_tokens = np.concatenate([m_tokens, np.ones((1), dtype=int) * self.mot_end_idx, np.ones((self.max_motion_length-1-m_tokens_len), dtype=int) * self.mot_pad_idx], axis=0)
+        else:
+            m_tokens = np.concatenate([m_tokens, np.ones((1), dtype=int) * self.mot_end_idx], axis=0)
+
+        return caption, m_tokens.reshape(-1), m_tokens_len
+
+
+
+
+def DATALoader(dataset_name,
+                batch_size, codebook_size, tokenizer_name, unit_length=4,
+                num_workers = 8) : 
+
+    train_loader = torch.utils.data.DataLoader(Text2MotionDataset(dataset_name, codebook_size = codebook_size, tokenizer_name = tokenizer_name, unit_length=unit_length),
+                                              batch_size,
+                                              shuffle=True,
+                                              num_workers=num_workers,
+                                              #collate_fn=collate_fn,
+                                              drop_last = True)
+    
+
+    return train_loader
+
+
+def cycle(iterable):
+    while True:
+        for x in iterable:
+            yield x
+
+

generate_human_motion/VQTrans/dataset/dataset_VQ.py

@@ -0,0 +1,109 @@
+import torch
+from torch.utils import data
+import numpy as np
+from os.path import join as pjoin
+import random
+import codecs as cs
+from tqdm import tqdm
+
+
+
+class VQMotionDataset(data.Dataset):
+    def __init__(self, dataset_name, window_size = 64, unit_length = 4):
+        self.window_size = window_size
+        self.unit_length = unit_length
+        self.dataset_name = dataset_name
+
+        if dataset_name == 't2m':
+            self.data_root = './dataset/HumanML3D'
+            self.motion_dir = pjoin(self.data_root, 'new_joint_vecs')
+            self.text_dir = pjoin(self.data_root, 'texts')
+            self.joints_num = 22
+            self.max_motion_length = 196
+            self.meta_dir = 'checkpoints/t2m/VQVAEV3_CB1024_CMT_H1024_NRES3/meta'
+
+        elif dataset_name == 'kit':
+            self.data_root = './dataset/KIT-ML'
+            self.motion_dir = pjoin(self.data_root, 'new_joint_vecs')
+            self.text_dir = pjoin(self.data_root, 'texts')
+            self.joints_num = 21
+
+            self.max_motion_length = 196
+            self.meta_dir = 'checkpoints/kit/VQVAEV3_CB1024_CMT_H1024_NRES3/meta'
+        
+        joints_num = self.joints_num
+
+        mean = np.load(pjoin(self.meta_dir, 'mean.npy'))
+        std = np.load(pjoin(self.meta_dir, 'std.npy'))
+
+        split_file = pjoin(self.data_root, 'train.txt')
+
+        self.data = []
+        self.lengths = []
+        id_list = []
+        with cs.open(split_file, 'r') as f:
+            for line in f.readlines():
+                id_list.append(line.strip())
+
+        for name in tqdm(id_list):
+            try:
+                motion = np.load(pjoin(self.motion_dir, name + '.npy'))
+                if motion.shape[0] < self.window_size:
+                    continue
+                self.lengths.append(motion.shape[0] - self.window_size)
+                self.data.append(motion)
+            except:
+                # Some motion may not exist in KIT dataset
+                pass
+
+            
+        self.mean = mean
+        self.std = std
+        print("Total number of motions {}".format(len(self.data)))
+
+    def inv_transform(self, data):
+        return data * self.std + self.mean
+    
+    def compute_sampling_prob(self) : 
+        
+        prob = np.array(self.lengths, dtype=np.float32)
+        prob /= np.sum(prob)
+        return prob
+    
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, item):
+        motion = self.data[item]
+        
+        idx = random.randint(0, len(motion) - self.window_size)
+
+        motion = motion[idx:idx+self.window_size]
+        "Z Normalization"
+        motion = (motion - self.mean) / self.std
+
+        return motion
+
+def DATALoader(dataset_name,
+               batch_size,
+               num_workers = 8,
+               window_size = 64,
+               unit_length = 4):
+    
+    trainSet = VQMotionDataset(dataset_name, window_size=window_size, unit_length=unit_length)
+    prob = trainSet.compute_sampling_prob()
+    sampler = torch.utils.data.WeightedRandomSampler(prob, num_samples = len(trainSet) * 1000, replacement=True)
+    train_loader = torch.utils.data.DataLoader(trainSet,
+                                              batch_size,
+                                              shuffle=True,
+                                              #sampler=sampler,
+                                              num_workers=num_workers,
+                                              #collate_fn=collate_fn,
+                                              drop_last = True)
+    
+    return train_loader
+
+def cycle(iterable):
+    while True:
+        for x in iterable:
+            yield x

generate_human_motion/VQTrans/dataset/dataset_tokenize.py

@@ -0,0 +1,117 @@
+import torch
+from torch.utils import data
+import numpy as np
+from os.path import join as pjoin
+import random
+import codecs as cs
+from tqdm import tqdm
+
+
+
+class VQMotionDataset(data.Dataset):
+    def __init__(self, dataset_name, feat_bias = 5, window_size = 64, unit_length = 8):
+        self.window_size = window_size
+        self.unit_length = unit_length
+        self.feat_bias = feat_bias
+
+        self.dataset_name = dataset_name
+        min_motion_len = 40 if dataset_name =='t2m' else 24
+        
+        if dataset_name == 't2m':
+            self.data_root = './dataset/HumanML3D'
+            self.motion_dir = pjoin(self.data_root, 'new_joint_vecs')
+            self.text_dir = pjoin(self.data_root, 'texts')
+            self.joints_num = 22
+            radius = 4
+            fps = 20
+            self.max_motion_length = 196
+            dim_pose = 263
+            self.meta_dir = 'checkpoints/t2m/VQVAEV3_CB1024_CMT_H1024_NRES3/meta'
+            #kinematic_chain = paramUtil.t2m_kinematic_chain
+        elif dataset_name == 'kit':
+            self.data_root = './dataset/KIT-ML'
+            self.motion_dir = pjoin(self.data_root, 'new_joint_vecs')
+            self.text_dir = pjoin(self.data_root, 'texts')
+            self.joints_num = 21
+            radius = 240 * 8
+            fps = 12.5
+            dim_pose = 251
+            self.max_motion_length = 196
+            self.meta_dir = 'checkpoints/kit/VQVAEV3_CB1024_CMT_H1024_NRES3/meta'
+            #kinematic_chain = paramUtil.kit_kinematic_chain
+        
+        joints_num = self.joints_num
+
+        mean = np.load(pjoin(self.meta_dir, 'mean.npy'))
+        std = np.load(pjoin(self.meta_dir, 'std.npy'))
+        
+        split_file = pjoin(self.data_root, 'train.txt')
+        
+        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(self.motion_dir, name + '.npy'))
+                if (len(motion)) < min_motion_len or (len(motion) >= 200):
+                    continue
+
+                data_dict[name] = {'motion': motion,
+                                   'length': len(motion),
+                                   'name': name}
+                new_name_list.append(name)
+                length_list.append(len(motion))
+            except:
+                # Some motion may not exist in KIT dataset
+                pass
+
+
+        self.mean = mean
+        self.std = std
+        self.length_arr = np.array(length_list)
+        self.data_dict = data_dict
+        self.name_list = new_name_list
+
+    def inv_transform(self, data):
+        return data * self.std + self.mean
+
+    def __len__(self):
+        return len(self.data_dict)
+
+    def __getitem__(self, item):
+        name = self.name_list[item]
+        data = self.data_dict[name]
+        motion, m_length = data['motion'], data['length']
+
+        m_length = (m_length // self.unit_length) * self.unit_length
+
+        idx = random.randint(0, len(motion) - m_length)
+        motion = motion[idx:idx+m_length]
+
+        "Z Normalization"
+        motion = (motion - self.mean) / self.std
+
+        return motion, name
+
+def DATALoader(dataset_name,
+                batch_size = 1,
+                num_workers = 8, unit_length = 4) : 
+    
+    train_loader = torch.utils.data.DataLoader(VQMotionDataset(dataset_name, unit_length=unit_length),
+                                              batch_size,
+                                              shuffle=True,
+                                              num_workers=num_workers,
+                                              #collate_fn=collate_fn,
+                                              drop_last = True)
+    
+    return train_loader
+
+def cycle(iterable):
+    while True:
+        for x in iterable:
+            yield x

generate_human_motion/VQTrans/dataset/prepare/download_extractor.sh

@@ -0,0 +1,15 @@
+rm -rf checkpoints
+mkdir checkpoints
+cd checkpoints
+echo -e "Downloading extractors"
+gdown --fuzzy https://drive.google.com/file/d/1o7RTDQcToJjTm9_mNWTyzvZvjTWpZfug/view
+gdown --fuzzy https://drive.google.com/file/d/1tX79xk0fflp07EZ660Xz1RAFE33iEyJR/view
+
+
+unzip t2m.zip
+unzip kit.zip
+
+echo -e "Cleaning\n"
+rm t2m.zip
+rm kit.zip
+echo -e "Downloading done!"

generate_human_motion/VQTrans/dataset/prepare/download_glove.sh

@@ -0,0 +1,9 @@
+echo -e "Downloading glove (in use by the evaluators)"
+gdown --fuzzy https://drive.google.com/file/d/1bCeS6Sh_mLVTebxIgiUHgdPrroW06mb6/view?usp=sharing
+rm -rf glove
+
+unzip glove.zip
+echo -e "Cleaning\n"
+rm glove.zip
+
+echo -e "Downloading done!"

generate_human_motion/VQTrans/dataset/prepare/download_model.sh

@@ -0,0 +1,12 @@
+
+mkdir -p pretrained
+cd pretrained/
+
+echo -e "The pretrained model files will be stored in the 'pretrained' folder\n"
+gdown 1LaOvwypF-jM2Axnq5dc-Iuvv3w_G-WDE
+
+unzip VQTrans_pretrained.zip
+echo -e "Cleaning\n"
+rm VQTrans_pretrained.zip
+
+echo -e "Downloading done!"

generate_human_motion/VQTrans/dataset/prepare/download_smpl.sh

@@ -0,0 +1,13 @@
+
+mkdir -p body_models
+cd body_models/
+
+echo -e "The smpl files will be stored in the 'body_models/smpl/' folder\n"
+gdown 1INYlGA76ak_cKGzvpOV2Pe6RkYTlXTW2
+rm -rf smpl
+
+unzip smpl.zip
+echo -e "Cleaning\n"
+rm smpl.zip
+
+echo -e "Downloading done!"

generate_human_motion/VQTrans/environment.yml

@@ -0,0 +1,121 @@
+name: VQTrans
+channels:
+  - pytorch
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=main
+  - _openmp_mutex=4.5=1_gnu
+  - blas=1.0=mkl
+  - bzip2=1.0.8=h7b6447c_0
+  - ca-certificates=2021.7.5=h06a4308_1
+  - certifi=2021.5.30=py38h06a4308_0
+  - cudatoolkit=10.1.243=h6bb024c_0
+  - ffmpeg=4.3=hf484d3e_0
+  - freetype=2.10.4=h5ab3b9f_0
+  - gmp=6.2.1=h2531618_2
+  - gnutls=3.6.15=he1e5248_0
+  - intel-openmp=2021.3.0=h06a4308_3350
+  - jpeg=9b=h024ee3a_2
+  - lame=3.100=h7b6447c_0
+  - lcms2=2.12=h3be6417_0
+  - ld_impl_linux-64=2.35.1=h7274673_9
+  - libffi=3.3=he6710b0_2
+  - libgcc-ng=9.3.0=h5101ec6_17
+  - libgomp=9.3.0=h5101ec6_17
+  - libiconv=1.15=h63c8f33_5
+  - libidn2=2.3.2=h7f8727e_0
+  - libpng=1.6.37=hbc83047_0
+  - libstdcxx-ng=9.3.0=hd4cf53a_17
+  - libtasn1=4.16.0=h27cfd23_0
+  - libtiff=4.2.0=h85742a9_0
+  - libunistring=0.9.10=h27cfd23_0
+  - libuv=1.40.0=h7b6447c_0
+  - libwebp-base=1.2.0=h27cfd23_0
+  - lz4-c=1.9.3=h295c915_1
+  - mkl=2021.3.0=h06a4308_520
+  - mkl-service=2.4.0=py38h7f8727e_0
+  - mkl_fft=1.3.0=py38h42c9631_2
+  - mkl_random=1.2.2=py38h51133e4_0
+  - ncurses=6.2=he6710b0_1
+  - nettle=3.7.3=hbbd107a_1
+  - ninja=1.10.2=hff7bd54_1
+  - numpy=1.20.3=py38hf144106_0
+  - numpy-base=1.20.3=py38h74d4b33_0
+  - olefile=0.46=py_0
+  - openh264=2.1.0=hd408876_0
+  - openjpeg=2.3.0=h05c96fa_1
+  - openssl=1.1.1k=h27cfd23_0
+  - pillow=8.3.1=py38h2c7a002_0
+  - pip=21.0.1=py38h06a4308_0
+  - python=3.8.11=h12debd9_0_cpython
+  - pytorch=1.8.1=py3.8_cuda10.1_cudnn7.6.3_0
+  - readline=8.1=h27cfd23_0
+  - setuptools=52.0.0=py38h06a4308_0
+  - six=1.16.0=pyhd3eb1b0_0
+  - sqlite=3.36.0=hc218d9a_0
+  - tk=8.6.10=hbc83047_0
+  - torchaudio=0.8.1=py38
+  - torchvision=0.9.1=py38_cu101
+  - typing_extensions=3.10.0.0=pyh06a4308_0
+  - wheel=0.37.0=pyhd3eb1b0_0
+  - xz=5.2.5=h7b6447c_0
+  - zlib=1.2.11=h7b6447c_3
+  - zstd=1.4.9=haebb681_0
+  - pip:
+    - absl-py==0.13.0
+    - backcall==0.2.0
+    - cachetools==4.2.2
+    - charset-normalizer==2.0.4
+    - chumpy==0.70
+    - cycler==0.10.0
+    - decorator==5.0.9
+    - google-auth==1.35.0
+    - google-auth-oauthlib==0.4.5
+    - grpcio==1.39.0
+    - idna==3.2
+    - imageio==2.9.0
+    - ipdb==0.13.9
+    - ipython==7.26.0
+    - ipython-genutils==0.2.0
+    - jedi==0.18.0
+    - joblib==1.0.1
+    - kiwisolver==1.3.1
+    - markdown==3.3.4
+    - matplotlib==3.4.3
+    - matplotlib-inline==0.1.2
+    - oauthlib==3.1.1
+    - pandas==1.3.2
+    - parso==0.8.2
+    - pexpect==4.8.0
+    - pickleshare==0.7.5
+    - prompt-toolkit==3.0.20
+    - protobuf==3.17.3
+    - ptyprocess==0.7.0
+    - pyasn1==0.4.8
+    - pyasn1-modules==0.2.8
+    - pygments==2.10.0
+    - pyparsing==2.4.7
+    - python-dateutil==2.8.2
+    - pytz==2021.1
+    - pyyaml==5.4.1
+    - requests==2.26.0
+    - requests-oauthlib==1.3.0
+    - rsa==4.7.2
+    - scikit-learn==0.24.2
+    - scipy==1.7.1
+    - sklearn==0.0
+    - smplx==0.1.28
+    - tensorboard==2.6.0
+    - tensorboard-data-server==0.6.1
+    - tensorboard-plugin-wit==1.8.0
+    - threadpoolctl==2.2.0
+    - toml==0.10.2
+    - tqdm==4.62.2
+    - traitlets==5.0.5
+    - urllib3==1.26.6
+    - wcwidth==0.2.5
+    - werkzeug==2.0.1
+    - git+https://github.com/openai/CLIP.git
+    - git+https://github.com/nghorbani/human_body_prior
+    - gdown
+    - moviepy

generate_human_motion/VQTrans/models/encdec.py

@@ -0,0 +1,67 @@
+import torch.nn as nn
+from resnet import Resnet1D
+
+class Encoder(nn.Module):
+    def __init__(self,
+                 input_emb_width = 3,
+                 output_emb_width = 512,
+                 down_t = 3,
+                 stride_t = 2,
+                 width = 512,
+                 depth = 3,
+                 dilation_growth_rate = 3,
+                 activation='relu',
+                 norm=None):
+        super().__init__()
+        
+        blocks = []
+        filter_t, pad_t = stride_t * 2, stride_t // 2
+        blocks.append(nn.Conv1d(input_emb_width, width, 3, 1, 1))
+        blocks.append(nn.ReLU())
+        
+        for i in range(down_t):
+            input_dim = width
+            block = nn.Sequential(
+                nn.Conv1d(input_dim, width, filter_t, stride_t, pad_t),
+                Resnet1D(width, depth, dilation_growth_rate, activation=activation, norm=norm),
+            )
+            blocks.append(block)
+        blocks.append(nn.Conv1d(width, output_emb_width, 3, 1, 1))
+        self.model = nn.Sequential(*blocks)
+
+    def forward(self, x):
+        return self.model(x)
+
+class Decoder(nn.Module):
+    def __init__(self,
+                 input_emb_width = 3,
+                 output_emb_width = 512,
+                 down_t = 3,
+                 stride_t = 2,
+                 width = 512,
+                 depth = 3,
+                 dilation_growth_rate = 3, 
+                 activation='relu',
+                 norm=None):
+        super().__init__()
+        blocks = []
+        
+        filter_t, pad_t = stride_t * 2, stride_t // 2
+        blocks.append(nn.Conv1d(output_emb_width, width, 3, 1, 1))
+        blocks.append(nn.ReLU())
+        for i in range(down_t):
+            out_dim = width
+            block = nn.Sequential(
+                Resnet1D(width, depth, dilation_growth_rate, reverse_dilation=True, activation=activation, norm=norm),
+                nn.Upsample(scale_factor=2, mode='nearest'),
+                nn.Conv1d(width, out_dim, 3, 1, 1)
+            )
+            blocks.append(block)
+        blocks.append(nn.Conv1d(width, width, 3, 1, 1))
+        blocks.append(nn.ReLU())
+        blocks.append(nn.Conv1d(width, input_emb_width, 3, 1, 1))
+        self.model = nn.Sequential(*blocks)
+
+    def forward(self, x):
+        return self.model(x)
+    

generate_human_motion/VQTrans/models/evaluator_wrapper.py

@@ -0,0 +1,92 @@
+
+import torch
+from os.path import join as pjoin
+import numpy as np
+from models.modules import MovementConvEncoder, TextEncoderBiGRUCo, MotionEncoderBiGRUCo
+from utils.word_vectorizer import POS_enumerator
+
+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(opt.checkpoints_dir, 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
+
+        # print(opt)
+
+        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()
+
+            '''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)
+        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

generate_human_motion/VQTrans/models/modules.py

@@ -0,0 +1,109 @@
+import torch
+import torch.nn as nn
+from torch.nn.utils.rnn import pack_padded_sequence
+
+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)
+
+            
+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 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.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, enforce_sorted=False)
+
+        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)

generate_human_motion/VQTrans/models/pos_encoding.py

@@ -0,0 +1,43 @@
+"""
+Various positional encodings for the transformer.
+"""
+import math
+import torch
+from torch import nn
+
+def PE1d_sincos(seq_length, dim):
+    """
+    :param d_model: dimension of the model
+    :param length: length of positions
+    :return: length*d_model position matrix
+    """
+    if dim % 2 != 0:
+        raise ValueError("Cannot use sin/cos positional encoding with "
+                         "odd dim (got dim={:d})".format(dim))
+    pe = torch.zeros(seq_length, dim)
+    position = torch.arange(0, seq_length).unsqueeze(1)
+    div_term = torch.exp((torch.arange(0, dim, 2, dtype=torch.float) *
+                         -(math.log(10000.0) / dim)))
+    pe[:, 0::2] = torch.sin(position.float() * div_term)
+    pe[:, 1::2] = torch.cos(position.float() * div_term)
+
+    return pe.unsqueeze(1)
+
+
+class PositionEmbedding(nn.Module):
+    """
+    Absolute pos embedding (standard), learned.
+    """
+    def __init__(self, seq_length, dim, dropout, grad=False):
+        super().__init__()
+        self.embed = nn.Parameter(data=PE1d_sincos(seq_length, dim), requires_grad=grad)
+        self.dropout = nn.Dropout(p=dropout)
+        
+    def forward(self, x):
+        # x.shape: bs, seq_len, feat_dim
+        l = x.shape[1]
+        x = x.permute(1, 0, 2) + self.embed[:l].expand(x.permute(1, 0, 2).shape)
+        x = self.dropout(x.permute(1, 0, 2))
+        return x
+
+ 

generate_human_motion/VQTrans/models/quantize_cnn.py

@@ -0,0 +1,415 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class QuantizeEMAReset(nn.Module):
+    def __init__(self, nb_code, code_dim, args):
+        super().__init__()
+        self.nb_code = nb_code
+        self.code_dim = code_dim
+        self.mu = args.mu
+        self.reset_codebook()
+        
+    def reset_codebook(self):
+        self.init = False
+        self.code_sum = None
+        self.code_count = None
+        if torch.cuda.is_available():
+            self.register_buffer('codebook', torch.zeros(self.nb_code, self.code_dim).cuda())
+        else:
+            self.register_buffer('codebook', torch.zeros(self.nb_code, self.code_dim))
+
+    def _tile(self, x):
+        nb_code_x, code_dim = x.shape
+        if nb_code_x < self.nb_code:
+            n_repeats = (self.nb_code + nb_code_x - 1) // nb_code_x
+            std = 0.01 / np.sqrt(code_dim)
+            out = x.repeat(n_repeats, 1)
+            out = out + torch.randn_like(out) * std
+        else :
+            out = x
+        return out
+
+    def init_codebook(self, x):
+        out = self._tile(x)
+        self.codebook = out[:self.nb_code]
+        self.code_sum = self.codebook.clone()
+        self.code_count = torch.ones(self.nb_code, device=self.codebook.device)
+        self.init = True
+        
+    @torch.no_grad()
+    def compute_perplexity(self, code_idx) : 
+        # Calculate new centres
+        code_onehot = torch.zeros(self.nb_code, code_idx.shape[0], device=code_idx.device)  # nb_code, N * L
+        code_onehot.scatter_(0, code_idx.view(1, code_idx.shape[0]), 1)
+
+        code_count = code_onehot.sum(dim=-1)  # nb_code
+        prob = code_count / torch.sum(code_count)  
+        perplexity = torch.exp(-torch.sum(prob * torch.log(prob + 1e-7)))
+        return perplexity
+    
+    @torch.no_grad()
+    def update_codebook(self, x, code_idx):
+        
+        code_onehot = torch.zeros(self.nb_code, x.shape[0], device=x.device)  # nb_code, N * L
+        code_onehot.scatter_(0, code_idx.view(1, x.shape[0]), 1)
+
+        code_sum = torch.matmul(code_onehot, x)  # nb_code, w
+        code_count = code_onehot.sum(dim=-1)  # nb_code
+
+        out = self._tile(x)
+        code_rand = out[:self.nb_code]
+
+        # Update centres
+        self.code_sum = self.mu * self.code_sum + (1. - self.mu) * code_sum  # w, nb_code
+        self.code_count = self.mu * self.code_count + (1. - self.mu) * code_count  # nb_code
+
+        usage = (self.code_count.view(self.nb_code, 1) >= 1.0).float()
+        code_update = self.code_sum.view(self.nb_code, self.code_dim) / self.code_count.view(self.nb_code, 1)
+
+        self.codebook = usage * code_update + (1 - usage) * code_rand
+        prob = code_count / torch.sum(code_count)  
+        perplexity = torch.exp(-torch.sum(prob * torch.log(prob + 1e-7)))
+
+            
+        return perplexity
+
+    def preprocess(self, x):
+        # NCT -> NTC -> [NT, C]
+        x = x.permute(0, 2, 1).contiguous()
+        x = x.view(-1, x.shape[-1])  
+        return x
+
+    def quantize(self, x):
+        # Calculate latent code x_l
+        k_w = self.codebook.t()
+        distance = torch.sum(x ** 2, dim=-1, keepdim=True) - 2 * torch.matmul(x, k_w) + torch.sum(k_w ** 2, dim=0,
+                                                                                            keepdim=True)  # (N * L, b)
+        _, code_idx = torch.min(distance, dim=-1)
+        return code_idx
+
+    def dequantize(self, code_idx):
+        x = F.embedding(code_idx, self.codebook)
+        return x
+
+    
+    def forward(self, x):
+        N, width, T = x.shape
+
+        # Preprocess
+        x = self.preprocess(x)
+
+        # Init codebook if not inited
+        if self.training and not self.init:
+            self.init_codebook(x)
+
+        # quantize and dequantize through bottleneck
+        code_idx = self.quantize(x)
+        x_d = self.dequantize(code_idx)
+
+        # Update embeddings
+        if self.training:
+            perplexity = self.update_codebook(x, code_idx)
+        else : 
+            perplexity = self.compute_perplexity(code_idx)
+        
+        # Loss
+        commit_loss = F.mse_loss(x, x_d.detach())
+
+        # Passthrough
+        x_d = x + (x_d - x).detach()
+
+        # Postprocess
+        x_d = x_d.view(N, T, -1).permute(0, 2, 1).contiguous()   #(N, DIM, T)
+        
+        return x_d, commit_loss, perplexity
+
+
+
+class Quantizer(nn.Module):
+    def __init__(self, n_e, e_dim, beta):
+        super(Quantizer, self).__init__()
+
+        self.e_dim = e_dim
+        self.n_e = n_e
+        self.beta = beta
+
+        self.embedding = nn.Embedding(self.n_e, self.e_dim)
+        self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e)
+
+    def forward(self, z):
+        
+        N, width, T = z.shape
+        z = self.preprocess(z)
+        assert z.shape[-1] == self.e_dim
+        z_flattened = z.contiguous().view(-1, self.e_dim)
+
+        # B x V
+        d = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \
+            torch.sum(self.embedding.weight**2, dim=1) - 2 * \
+            torch.matmul(z_flattened, self.embedding.weight.t())
+        # B x 1
+        min_encoding_indices = torch.argmin(d, dim=1)
+        z_q = self.embedding(min_encoding_indices).view(z.shape)
+
+        # compute loss for embedding
+        loss = torch.mean((z_q - z.detach())**2) + self.beta * \
+               torch.mean((z_q.detach() - z)**2)
+
+        # preserve gradients
+        z_q = z + (z_q - z).detach()
+        z_q = z_q.view(N, T, -1).permute(0, 2, 1).contiguous()   #(N, DIM, T)
+
+        min_encodings = F.one_hot(min_encoding_indices, self.n_e).type(z.dtype)
+        e_mean = torch.mean(min_encodings, dim=0)
+        perplexity = torch.exp(-torch.sum(e_mean*torch.log(e_mean + 1e-10)))
+        return z_q, loss, perplexity
+
+    def quantize(self, z):
+
+        assert z.shape[-1] == self.e_dim
+
+        # B x V
+        d = torch.sum(z ** 2, dim=1, keepdim=True) + \
+            torch.sum(self.embedding.weight ** 2, dim=1) - 2 * \
+            torch.matmul(z, self.embedding.weight.t())
+        # B x 1
+        min_encoding_indices = torch.argmin(d, dim=1)
+        return min_encoding_indices
+
+    def dequantize(self, indices):
+
+        index_flattened = indices.view(-1)
+        z_q = self.embedding(index_flattened)
+        z_q = z_q.view(indices.shape + (self.e_dim, )).contiguous()
+        return z_q
+
+    def preprocess(self, x):
+        # NCT -> NTC -> [NT, C]
+        x = x.permute(0, 2, 1).contiguous()
+        x = x.view(-1, x.shape[-1])  
+        return x
+
+
+
+class QuantizeReset(nn.Module):
+    def __init__(self, nb_code, code_dim, args):
+        super().__init__()
+        self.nb_code = nb_code
+        self.code_dim = code_dim
+        self.reset_codebook()
+        self.codebook = nn.Parameter(torch.randn(nb_code, code_dim))
+        
+    def reset_codebook(self):
+        self.init = False
+        self.code_count = None
+
+    def _tile(self, x):
+        nb_code_x, code_dim = x.shape
+        if nb_code_x < self.nb_code:
+            n_repeats = (self.nb_code + nb_code_x - 1) // nb_code_x
+            std = 0.01 / np.sqrt(code_dim)
+            out = x.repeat(n_repeats, 1)
+            out = out + torch.randn_like(out) * std
+        else :
+            out = x
+        return out
+
+    def init_codebook(self, x):
+        out = self._tile(x)
+        self.codebook = nn.Parameter(out[:self.nb_code])
+        self.code_count = torch.ones(self.nb_code, device=self.codebook.device)
+        self.init = True
+        
+    @torch.no_grad()
+    def compute_perplexity(self, code_idx) : 
+        # Calculate new centres
+        code_onehot = torch.zeros(self.nb_code, code_idx.shape[0], device=code_idx.device)  # nb_code, N * L
+        code_onehot.scatter_(0, code_idx.view(1, code_idx.shape[0]), 1)
+
+        code_count = code_onehot.sum(dim=-1)  # nb_code
+        prob = code_count / torch.sum(code_count)  
+        perplexity = torch.exp(-torch.sum(prob * torch.log(prob + 1e-7)))
+        return perplexity
+    
+    def update_codebook(self, x, code_idx):
+        
+        code_onehot = torch.zeros(self.nb_code, x.shape[0], device=x.device)  # nb_code, N * L
+        code_onehot.scatter_(0, code_idx.view(1, x.shape[0]), 1)
+
+        code_count = code_onehot.sum(dim=-1)  # nb_code
+
+        out = self._tile(x)
+        code_rand = out[:self.nb_code]
+
+        # Update centres
+        self.code_count = code_count  # nb_code
+        usage = (self.code_count.view(self.nb_code, 1) >= 1.0).float()
+
+        self.codebook.data = usage * self.codebook.data + (1 - usage) * code_rand
+        prob = code_count / torch.sum(code_count)  
+        perplexity = torch.exp(-torch.sum(prob * torch.log(prob + 1e-7)))
+
+            
+        return perplexity
+
+    def preprocess(self, x):
+        # NCT -> NTC -> [NT, C]
+        x = x.permute(0, 2, 1).contiguous()
+        x = x.view(-1, x.shape[-1])  
+        return x
+
+    def quantize(self, x):
+        # Calculate latent code x_l
+        k_w = self.codebook.t()
+        distance = torch.sum(x ** 2, dim=-1, keepdim=True) - 2 * torch.matmul(x, k_w) + torch.sum(k_w ** 2, dim=0,
+                                                                                            keepdim=True)  # (N * L, b)
+        _, code_idx = torch.min(distance, dim=-1)
+        return code_idx
+
+    def dequantize(self, code_idx):
+        x = F.embedding(code_idx, self.codebook)
+        return x
+
+    
+    def forward(self, x):
+        N, width, T = x.shape
+        # Preprocess
+        x = self.preprocess(x)
+        # Init codebook if not inited
+        if self.training and not self.init:
+            self.init_codebook(x)
+        # quantize and dequantize through bottleneck
+        code_idx = self.quantize(x)
+        x_d = self.dequantize(code_idx)
+        # Update embeddings
+        if self.training:
+            perplexity = self.update_codebook(x, code_idx)
+        else : 
+            perplexity = self.compute_perplexity(code_idx)
+        
+        # Loss
+        commit_loss = F.mse_loss(x, x_d.detach())
+
+        # Passthrough
+        x_d = x + (x_d - x).detach()
+
+        # Postprocess
+        x_d = x_d.view(N, T, -1).permute(0, 2, 1).contiguous()   #(N, DIM, T)
+        
+        return x_d, commit_loss, perplexity
+
+class QuantizeEMA(nn.Module):
+    def __init__(self, nb_code, code_dim, args):
+        super().__init__()
+        self.nb_code = nb_code
+        self.code_dim = code_dim
+        self.mu = 0.99
+        self.reset_codebook()
+        
+    def reset_codebook(self):
+        self.init = False
+        self.code_sum = None
+        self.code_count = None
+        self.register_buffer('codebook', torch.zeros(self.nb_code, self.code_dim).cuda())
+
+    def _tile(self, x):
+        nb_code_x, code_dim = x.shape
+        if nb_code_x < self.nb_code:
+            n_repeats = (self.nb_code + nb_code_x - 1) // nb_code_x
+            std = 0.01 / np.sqrt(code_dim)
+            out = x.repeat(n_repeats, 1)
+            out = out + torch.randn_like(out) * std
+        else :
+            out = x
+        return out
+
+    def init_codebook(self, x):
+        out = self._tile(x)
+        self.codebook = out[:self.nb_code]
+        self.code_sum = self.codebook.clone()
+        self.code_count = torch.ones(self.nb_code, device=self.codebook.device)
+        self.init = True
+        
+    @torch.no_grad()
+    def compute_perplexity(self, code_idx) : 
+        # Calculate new centres
+        code_onehot = torch.zeros(self.nb_code, code_idx.shape[0], device=code_idx.device)  # nb_code, N * L
+        code_onehot.scatter_(0, code_idx.view(1, code_idx.shape[0]), 1)
+
+        code_count = code_onehot.sum(dim=-1)  # nb_code
+        prob = code_count / torch.sum(code_count)  
+        perplexity = torch.exp(-torch.sum(prob * torch.log(prob + 1e-7)))
+        return perplexity
+    
+    @torch.no_grad()
+    def update_codebook(self, x, code_idx):
+        
+        code_onehot = torch.zeros(self.nb_code, x.shape[0], device=x.device)  # nb_code, N * L
+        code_onehot.scatter_(0, code_idx.view(1, x.shape[0]), 1)
+
+        code_sum = torch.matmul(code_onehot, x)  # nb_code, w
+        code_count = code_onehot.sum(dim=-1)  # nb_code
+
+        # Update centres
+        self.code_sum = self.mu * self.code_sum + (1. - self.mu) * code_sum  # w, nb_code
+        self.code_count = self.mu * self.code_count + (1. - self.mu) * code_count  # nb_code
+
+        code_update = self.code_sum.view(self.nb_code, self.code_dim) / self.code_count.view(self.nb_code, 1)
+
+        self.codebook = code_update
+        prob = code_count / torch.sum(code_count)  
+        perplexity = torch.exp(-torch.sum(prob * torch.log(prob + 1e-7)))
+            
+        return perplexity
+
+    def preprocess(self, x):
+        # NCT -> NTC -> [NT, C]
+        x = x.permute(0, 2, 1).contiguous()
+        x = x.view(-1, x.shape[-1])  
+        return x
+
+    def quantize(self, x):
+        # Calculate latent code x_l
+        k_w = self.codebook.t()
+        distance = torch.sum(x ** 2, dim=-1, keepdim=True) - 2 * torch.matmul(x, k_w) + torch.sum(k_w ** 2, dim=0,
+                                                                                            keepdim=True)  # (N * L, b)
+        _, code_idx = torch.min(distance, dim=-1)
+        return code_idx
+
+    def dequantize(self, code_idx):
+        x = F.embedding(code_idx, self.codebook)
+        return x
+
+    
+    def forward(self, x):
+        N, width, T = x.shape
+
+        # Preprocess
+        x = self.preprocess(x)
+
+        # Init codebook if not inited
+        if self.training and not self.init:
+            self.init_codebook(x)
+
+        # quantize and dequantize through bottleneck
+        code_idx = self.quantize(x)
+        x_d = self.dequantize(code_idx)
+
+        # Update embeddings
+        if self.training:
+            perplexity = self.update_codebook(x, code_idx)
+        else : 
+            perplexity = self.compute_perplexity(code_idx)
+        
+        # Loss
+        commit_loss = F.mse_loss(x, x_d.detach())
+
+        # Passthrough
+        x_d = x + (x_d - x).detach()
+
+        # Postprocess
+        x_d = x_d.view(N, T, -1).permute(0, 2, 1).contiguous()   #(N, DIM, T)
+        
+        return x_d, commit_loss, perplexity

generate_human_motion/VQTrans/models/resnet.py

@@ -0,0 +1,82 @@
+import torch.nn as nn
+import torch
+
+class nonlinearity(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        # swish
+        return x * torch.sigmoid(x)
+
+class ResConv1DBlock(nn.Module):
+    def __init__(self, n_in, n_state, dilation=1, activation='silu', norm=None, dropout=None):
+        super().__init__()
+        padding = dilation
+        self.norm = norm
+        if norm == "LN":
+            self.norm1 = nn.LayerNorm(n_in)
+            self.norm2 = nn.LayerNorm(n_in)
+        elif norm == "GN":
+            self.norm1 = nn.GroupNorm(num_groups=32, num_channels=n_in, eps=1e-6, affine=True)
+            self.norm2 = nn.GroupNorm(num_groups=32, num_channels=n_in, eps=1e-6, affine=True)
+        elif norm == "BN":
+            self.norm1 = nn.BatchNorm1d(num_features=n_in, eps=1e-6, affine=True)
+            self.norm2 = nn.BatchNorm1d(num_features=n_in, eps=1e-6, affine=True)
+        
+        else:
+            self.norm1 = nn.Identity()
+            self.norm2 = nn.Identity()
+
+        if activation == "relu":
+            self.activation1 = nn.ReLU()
+            self.activation2 = nn.ReLU()
+            
+        elif activation == "silu":
+            self.activation1 = nonlinearity()
+            self.activation2 = nonlinearity()
+            
+        elif activation == "gelu":
+            self.activation1 = nn.GELU()
+            self.activation2 = nn.GELU()
+            
+        
+
+        self.conv1 = nn.Conv1d(n_in, n_state, 3, 1, padding, dilation)
+        self.conv2 = nn.Conv1d(n_state, n_in, 1, 1, 0,)     
+
+
+    def forward(self, x):
+        x_orig = x
+        if self.norm == "LN":
+            x = self.norm1(x.transpose(-2, -1))
+            x = self.activation1(x.transpose(-2, -1))
+        else:
+            x = self.norm1(x)
+            x = self.activation1(x)
+            
+        x = self.conv1(x)
+
+        if self.norm == "LN":
+            x = self.norm2(x.transpose(-2, -1))
+            x = self.activation2(x.transpose(-2, -1))
+        else:
+            x = self.norm2(x)
+            x = self.activation2(x)
+
+        x = self.conv2(x)
+        x = x + x_orig
+        return x
+
+class Resnet1D(nn.Module):
+    def __init__(self, n_in, n_depth, dilation_growth_rate=1, reverse_dilation=True, activation='relu', norm=None):
+        super().__init__()
+        
+        blocks = [ResConv1DBlock(n_in, n_in, dilation=dilation_growth_rate ** depth, activation=activation, norm=norm) for depth in range(n_depth)]
+        if reverse_dilation:
+            blocks = blocks[::-1]
+        
+        self.model = nn.Sequential(*blocks)
+
+    def forward(self, x):        
+        return self.model(x)

generate_human_motion/VQTrans/models/rotation2xyz.py

@@ -0,0 +1,92 @@
+# This code is based on https://github.com/Mathux/ACTOR.git
+import torch
+import VQTrans.utils.rotation_conversions as geometry
+
+
+from smpl import SMPL, JOINTSTYPE_ROOT
+# from .get_model import JOINTSTYPES
+JOINTSTYPES = ["a2m", "a2mpl", "smpl", "vibe", "vertices"]
+
+
+class Rotation2xyz:
+    def __init__(self, device, dataset='amass'):
+        self.device = device
+        self.dataset = dataset
+        self.smpl_model = SMPL().eval().to(device)
+
+    def __call__(self, x, mask, pose_rep, translation, glob,
+                 jointstype, vertstrans, betas=None, beta=0,
+                 glob_rot=None, get_rotations_back=False, **kwargs):
+        if pose_rep == "xyz":
+            return x
+
+        if mask is None:
+            mask = torch.ones((x.shape[0], x.shape[-1]), dtype=bool, device=x.device)
+
+        if not glob and glob_rot is None:
+            raise TypeError("You must specify global rotation if glob is False")
+
+        if jointstype not in JOINTSTYPES:
+            raise NotImplementedError("This jointstype is not implemented.")
+
+        if translation:
+            x_translations = x[:, -1, :3]
+            x_rotations = x[:, :-1]
+        else:
+            x_rotations = x
+
+        x_rotations = x_rotations.permute(0, 3, 1, 2)
+        nsamples, time, njoints, feats = x_rotations.shape
+
+        # Compute rotations (convert only masked sequences output)
+        if pose_rep == "rotvec":
+            rotations = geometry.axis_angle_to_matrix(x_rotations[mask])
+        elif pose_rep == "rotmat":
+            rotations = x_rotations[mask].view(-1, njoints, 3, 3)
+        elif pose_rep == "rotquat":
+            rotations = geometry.quaternion_to_matrix(x_rotations[mask])
+        elif pose_rep == "rot6d":
+            rotations = geometry.rotation_6d_to_matrix(x_rotations[mask])
+        else:
+            raise NotImplementedError("No geometry for this one.")
+
+        if not glob:
+            global_orient = torch.tensor(glob_rot, device=x.device)
+            global_orient = geometry.axis_angle_to_matrix(global_orient).view(1, 1, 3, 3)
+            global_orient = global_orient.repeat(len(rotations), 1, 1, 1)
+        else:
+            global_orient = rotations[:, 0]
+            rotations = rotations[:, 1:]
+
+        if betas is None:
+            betas = torch.zeros([rotations.shape[0], self.smpl_model.num_betas],
+                                dtype=rotations.dtype, device=rotations.device)
+            betas[:, 1] = beta
+            # import ipdb; ipdb.set_trace()
+        out = self.smpl_model(body_pose=rotations, global_orient=global_orient, betas=betas)
+
+        # get the desirable joints
+        joints = out[jointstype]
+
+        x_xyz = torch.empty(nsamples, time, joints.shape[1], 3, device=x.device, dtype=x.dtype)
+        x_xyz[~mask] = 0
+        x_xyz[mask] = joints
+
+        x_xyz = x_xyz.permute(0, 2, 3, 1).contiguous()
+
+        # the first translation root at the origin on the prediction
+        if jointstype != "vertices":
+            rootindex = JOINTSTYPE_ROOT[jointstype]
+            x_xyz = x_xyz - x_xyz[:, [rootindex], :, :]
+
+        if translation and vertstrans:
+            # the first translation root at the origin
+            x_translations = x_translations - x_translations[:, :, [0]]
+
+            # add the translation to all the joints
+            x_xyz = x_xyz + x_translations[:, None, :, :]
+
+        if get_rotations_back:
+            return x_xyz, rotations, global_orient
+        else:
+            return x_xyz

generate_human_motion/VQTrans/models/smpl.py

@@ -0,0 +1,97 @@
+# This code is based on https://github.com/Mathux/ACTOR.git
+import numpy as np
+import torch
+
+import contextlib
+
+from smplx import SMPLLayer as _SMPLLayer
+from smplx.lbs import vertices2joints
+
+
+# action2motion_joints = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 21, 24, 38]
+# change 0 and 8
+action2motion_joints = [8, 1, 2, 3, 4, 5, 6, 7, 0, 9, 10, 11, 12, 13, 14, 21, 24, 38]
+
+from VQTrans.utils.config import SMPL_MODEL_PATH, JOINT_REGRESSOR_TRAIN_EXTRA
+
+JOINTSTYPE_ROOT = {"a2m": 0, # action2motion
+                   "smpl": 0,
+                   "a2mpl": 0, # set(smpl, a2m)
+                   "vibe": 8}  # 0 is the 8 position: OP MidHip below
+
+JOINT_MAP = {
+    'OP Nose': 24, 'OP Neck': 12, 'OP RShoulder': 17,
+    'OP RElbow': 19, 'OP RWrist': 21, 'OP LShoulder': 16,
+    'OP LElbow': 18, 'OP LWrist': 20, 'OP MidHip': 0,
+    'OP RHip': 2, 'OP RKnee': 5, 'OP RAnkle': 8,
+    'OP LHip': 1, 'OP LKnee': 4, 'OP LAnkle': 7,
+    'OP REye': 25, 'OP LEye': 26, 'OP REar': 27,
+    'OP LEar': 28, 'OP LBigToe': 29, 'OP LSmallToe': 30,
+    'OP LHeel': 31, 'OP RBigToe': 32, 'OP RSmallToe': 33, 'OP RHeel': 34,
+    'Right Ankle': 8, 'Right Knee': 5, 'Right Hip': 45,
+    'Left Hip': 46, 'Left Knee': 4, 'Left Ankle': 7,
+    'Right Wrist': 21, 'Right Elbow': 19, 'Right Shoulder': 17,
+    'Left Shoulder': 16, 'Left Elbow': 18, 'Left Wrist': 20,
+    'Neck (LSP)': 47, 'Top of Head (LSP)': 48,
+    'Pelvis (MPII)': 49, 'Thorax (MPII)': 50,
+    'Spine (H36M)': 51, 'Jaw (H36M)': 52,
+    'Head (H36M)': 53, 'Nose': 24, 'Left Eye': 26,
+    'Right Eye': 25, 'Left Ear': 28, 'Right Ear': 27
+}
+
+JOINT_NAMES = [
+    'OP Nose', 'OP Neck', 'OP RShoulder',
+    'OP RElbow', 'OP RWrist', 'OP LShoulder',
+    'OP LElbow', 'OP LWrist', 'OP MidHip',
+    'OP RHip', 'OP RKnee', 'OP RAnkle',
+    'OP LHip', 'OP LKnee', 'OP LAnkle',
+    'OP REye', 'OP LEye', 'OP REar',
+    'OP LEar', 'OP LBigToe', 'OP LSmallToe',
+    'OP LHeel', 'OP RBigToe', 'OP RSmallToe', 'OP RHeel',
+    'Right Ankle', 'Right Knee', 'Right Hip',
+    'Left Hip', 'Left Knee', 'Left Ankle',
+    'Right Wrist', 'Right Elbow', 'Right Shoulder',
+    'Left Shoulder', 'Left Elbow', 'Left Wrist',
+    'Neck (LSP)', 'Top of Head (LSP)',
+    'Pelvis (MPII)', 'Thorax (MPII)',
+    'Spine (H36M)', 'Jaw (H36M)',
+    'Head (H36M)', 'Nose', 'Left Eye',
+    'Right Eye', 'Left Ear', 'Right Ear'
+]
+
+
+# adapted from VIBE/SPIN to output smpl_joints, vibe joints and action2motion joints
+class SMPL(_SMPLLayer):
+    """ Extension of the official SMPL implementation to support more joints """
+
+    def __init__(self, model_path=SMPL_MODEL_PATH, **kwargs):
+        kwargs["model_path"] = model_path
+
+        # remove the verbosity for the 10-shapes beta parameters
+        with contextlib.redirect_stdout(None):
+            super(SMPL, self).__init__(**kwargs)
+            
+        J_regressor_extra = np.load(JOINT_REGRESSOR_TRAIN_EXTRA)
+        self.register_buffer('J_regressor_extra', torch.tensor(J_regressor_extra, dtype=torch.float32))
+        vibe_indexes = np.array([JOINT_MAP[i] for i in JOINT_NAMES])
+        a2m_indexes = vibe_indexes[action2motion_joints]
+        smpl_indexes = np.arange(24)
+        a2mpl_indexes = np.unique(np.r_[smpl_indexes, a2m_indexes])
+
+        self.maps = {"vibe": vibe_indexes,
+                     "a2m": a2m_indexes,
+                     "smpl": smpl_indexes,
+                     "a2mpl": a2mpl_indexes}
+        
+    def forward(self, *args, **kwargs):
+        smpl_output = super(SMPL, self).forward(*args, **kwargs)
+        
+        extra_joints = vertices2joints(self.J_regressor_extra, smpl_output.vertices)
+        all_joints = torch.cat([smpl_output.joints, extra_joints], dim=1)
+
+        output = {"vertices": smpl_output.vertices}
+
+        for joinstype, indexes in self.maps.items():
+            output[joinstype] = all_joints[:, indexes]
+            
+        return output

generate_human_motion/VQTrans/models/t2m_trans.py

@@ -0,0 +1,211 @@
+import math
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+from torch.distributions import Categorical
+import pos_encoding as pos_encoding
+
+class Text2Motion_Transformer(nn.Module):
+
+    def __init__(self, 
+                num_vq=1024, 
+                embed_dim=512, 
+                clip_dim=512, 
+                block_size=16, 
+                num_layers=2, 
+                n_head=8, 
+                drop_out_rate=0.1, 
+                fc_rate=4):
+        super().__init__()
+        self.trans_base = CrossCondTransBase(num_vq, embed_dim, clip_dim, block_size, num_layers, n_head, drop_out_rate, fc_rate)
+        self.trans_head = CrossCondTransHead(num_vq, embed_dim, block_size, num_layers, n_head, drop_out_rate, fc_rate)
+        self.block_size = block_size
+        self.num_vq = num_vq
+
+    def get_block_size(self):
+        return self.block_size
+
+    def forward(self, idxs, clip_feature):
+        feat = self.trans_base(idxs, clip_feature)
+        logits = self.trans_head(feat)
+        return logits
+
+    def sample(self, clip_feature, if_categorial=False):
+        for k in range(self.block_size):
+            if k == 0:
+                x = []
+            else:
+                x = xs
+            logits = self.forward(x, clip_feature)
+            logits = logits[:, -1, :]
+            probs = F.softmax(logits, dim=-1)
+            if if_categorial:
+                dist = Categorical(probs)
+                idx = dist.sample()
+                if idx == self.num_vq:
+                    break
+                idx = idx.unsqueeze(-1)
+            else:
+                _, idx = torch.topk(probs, k=1, dim=-1)
+                if idx[0] == self.num_vq:
+                    break
+            # append to the sequence and continue
+            if k == 0:
+                xs = idx
+            else:
+                xs = torch.cat((xs, idx), dim=1)
+            
+            if k == self.block_size - 1:
+                return xs[:, :-1]
+        return xs
+
+class CausalCrossConditionalSelfAttention(nn.Module):
+
+    def __init__(self, embed_dim=512, block_size=16, n_head=8, drop_out_rate=0.1):
+        super().__init__()
+        assert embed_dim % 8 == 0
+        # key, query, value projections for all heads
+        self.key = nn.Linear(embed_dim, embed_dim)
+        self.query = nn.Linear(embed_dim, embed_dim)
+        self.value = nn.Linear(embed_dim, embed_dim)
+
+        self.attn_drop = nn.Dropout(drop_out_rate)
+        self.resid_drop = nn.Dropout(drop_out_rate)
+
+        self.proj = nn.Linear(embed_dim, embed_dim)
+        # causal mask to ensure that attention is only applied to the left in the input sequence
+        self.register_buffer("mask", torch.tril(torch.ones(block_size, block_size)).view(1, 1, block_size, block_size))
+        self.n_head = n_head
+
+    def forward(self, x):
+        B, T, C = x.size() 
+
+        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+        k = self.key(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        v = self.value(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
+        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+        att = att.masked_fill(self.mask[:,:,:T,:T] == 0, float('-inf'))
+        att = F.softmax(att, dim=-1)
+        att = self.attn_drop(att)
+        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
+
+        # output projection
+        y = self.resid_drop(self.proj(y))
+        return y
+
+class Block(nn.Module):
+
+    def __init__(self, embed_dim=512, block_size=16, n_head=8, drop_out_rate=0.1, fc_rate=4):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(embed_dim)
+        self.ln2 = nn.LayerNorm(embed_dim)
+        self.attn = CausalCrossConditionalSelfAttention(embed_dim, block_size, n_head, drop_out_rate)
+        self.mlp = nn.Sequential(
+            nn.Linear(embed_dim, fc_rate * embed_dim),
+            nn.GELU(),
+            nn.Linear(fc_rate * embed_dim, embed_dim),
+            nn.Dropout(drop_out_rate),
+        )
+
+    def forward(self, x):
+        x = x + self.attn(self.ln1(x))
+        x = x + self.mlp(self.ln2(x))
+        return x
+
+class CrossCondTransBase(nn.Module):
+
+    def __init__(self, 
+                num_vq=1024, 
+                embed_dim=512, 
+                clip_dim=512, 
+                block_size=16, 
+                num_layers=2, 
+                n_head=8, 
+                drop_out_rate=0.1, 
+                fc_rate=4):
+        super().__init__()
+        self.tok_emb = nn.Embedding(num_vq + 2, embed_dim)
+        self.cond_emb = nn.Linear(clip_dim, embed_dim)
+        self.pos_embedding = nn.Embedding(block_size, embed_dim)
+        self.drop = nn.Dropout(drop_out_rate)
+        # transformer block
+        self.blocks = nn.Sequential(*[Block(embed_dim, block_size, n_head, drop_out_rate, fc_rate) for _ in range(num_layers)])
+        self.pos_embed = pos_encoding.PositionEmbedding(block_size, embed_dim, 0.0, False)
+
+        self.block_size = block_size
+
+        self.apply(self._init_weights)
+
+    def get_block_size(self):
+        return self.block_size
+
+    def _init_weights(self, module):
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+            if isinstance(module, nn.Linear) and module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+    
+    def forward(self, idx, clip_feature):
+        if len(idx) == 0:
+            token_embeddings = self.cond_emb(clip_feature).unsqueeze(1)
+        else:
+            b, t = idx.size()
+            assert t <= self.block_size, "Cannot forward, model block size is exhausted."
+            # forward the Trans model
+            token_embeddings = self.tok_emb(idx)
+            token_embeddings = torch.cat([self.cond_emb(clip_feature).unsqueeze(1), token_embeddings], dim=1)
+            
+        x = self.pos_embed(token_embeddings)
+        x = self.blocks(x)
+
+        return x
+
+
+class CrossCondTransHead(nn.Module):
+
+    def __init__(self, 
+                num_vq=1024, 
+                embed_dim=512, 
+                block_size=16, 
+                num_layers=2, 
+                n_head=8, 
+                drop_out_rate=0.1, 
+                fc_rate=4):
+        super().__init__()
+
+        self.blocks = nn.Sequential(*[Block(embed_dim, block_size, n_head, drop_out_rate, fc_rate) for _ in range(num_layers)])
+        self.ln_f = nn.LayerNorm(embed_dim)
+        self.head = nn.Linear(embed_dim, num_vq + 1, bias=False)
+        self.block_size = block_size
+
+        self.apply(self._init_weights)
+
+    def get_block_size(self):
+        return self.block_size
+
+    def _init_weights(self, module):
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+            if isinstance(module, nn.Linear) and module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+    def forward(self, x):
+        x = self.blocks(x)
+        x = self.ln_f(x)
+        logits = self.head(x)
+        return logits
+
+    
+
+
+        
+

generate_human_motion/VQTrans/models/vqvae.py

@@ -0,0 +1,118 @@
+import torch.nn as nn
+from encdec import Encoder, Decoder
+from quantize_cnn import QuantizeEMAReset, Quantizer, QuantizeEMA, QuantizeReset
+
+
+class VQVAE_251(nn.Module):
+    def __init__(self,
+                 args,
+                 nb_code=1024,
+                 code_dim=512,
+                 output_emb_width=512,
+                 down_t=3,
+                 stride_t=2,
+                 width=512,
+                 depth=3,
+                 dilation_growth_rate=3,
+                 activation='relu',
+                 norm=None):
+        
+        super().__init__()
+        self.code_dim = code_dim
+        self.num_code = nb_code
+        self.quant = args.quantizer
+        self.encoder = Encoder(251 if args.dataname == 'kit' else 263, output_emb_width, down_t, stride_t, width, depth, dilation_growth_rate, activation=activation, norm=norm)
+        self.decoder = Decoder(251 if args.dataname == 'kit' else 263, output_emb_width, down_t, stride_t, width, depth, dilation_growth_rate, activation=activation, norm=norm)
+        if args.quantizer == "ema_reset":
+            self.quantizer = QuantizeEMAReset(nb_code, code_dim, args)
+        elif args.quantizer == "orig":
+            self.quantizer = Quantizer(nb_code, code_dim, 1.0)
+        elif args.quantizer == "ema":
+            self.quantizer = QuantizeEMA(nb_code, code_dim, args)
+        elif args.quantizer == "reset":
+            self.quantizer = QuantizeReset(nb_code, code_dim, args)
+
+
+    def preprocess(self, x):
+        # (bs, T, Jx3) -> (bs, Jx3, T)
+        x = x.permute(0,2,1).float()
+        return x
+
+
+    def postprocess(self, x):
+        # (bs, Jx3, T) ->  (bs, T, Jx3)
+        x = x.permute(0,2,1)
+        return x
+
+
+    def encode(self, x):
+        N, T, _ = x.shape
+        x_in = self.preprocess(x)
+        x_encoder = self.encoder(x_in)
+        x_encoder = self.postprocess(x_encoder)
+        x_encoder = x_encoder.contiguous().view(-1, x_encoder.shape[-1])  # (NT, C)
+        code_idx = self.quantizer.quantize(x_encoder)
+        code_idx = code_idx.view(N, -1)
+        return code_idx
+
+
+    def forward(self, x):
+        
+        x_in = self.preprocess(x)
+        # Encode
+        x_encoder = self.encoder(x_in)
+        
+        ## quantization
+        x_quantized, loss, perplexity  = self.quantizer(x_encoder)
+
+        ## decoder
+        x_decoder = self.decoder(x_quantized)
+        x_out = self.postprocess(x_decoder)
+        return x_out, loss, perplexity
+
+
+    def forward_decoder(self, x):
+        x_d = self.quantizer.dequantize(x)
+        x_d = x_d.view(1, -1, self.code_dim).permute(0, 2, 1).contiguous()
+        
+        # decoder
+        x_decoder = self.decoder(x_d)
+        x_out = self.postprocess(x_decoder)
+        return x_out
+
+
+
+class HumanVQVAE(nn.Module):
+    def __init__(self,
+                 args,
+                 nb_code=512,
+                 code_dim=512,
+                 output_emb_width=512,
+                 down_t=3,
+                 stride_t=2,
+                 width=512,
+                 depth=3,
+                 dilation_growth_rate=3,
+                 activation='relu',
+                 norm=None):
+        
+        super().__init__()
+        
+        self.nb_joints = 21 if args.dataname == 'kit' else 22
+        self.vqvae = VQVAE_251(args, nb_code, code_dim, output_emb_width, down_t, stride_t, width, depth, dilation_growth_rate, activation=activation, norm=norm)
+
+    def encode(self, x):
+        b, t, c = x.size()
+        quants = self.vqvae.encode(x) # (N, T)
+        return quants
+
+    def forward(self, x):
+
+        x_out, loss, perplexity = self.vqvae(x)
+        
+        return x_out, loss, perplexity
+
+    def forward_decoder(self, x):
+        x_out = self.vqvae.forward_decoder(x)
+        return x_out
+        

generate_human_motion/VQTrans/options/get_eval_option.py

@@ -0,0 +1,83 @@
+from argparse import Namespace
+import re
+from os.path import join as pjoin
+
+
+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] = (value == 'True')
+                #     print(key, value)
+                elif is_float(value):
+                    opt_dict[key] = float(value)
+                elif is_number(value):
+                    opt_dict[key] = int(value)
+                else:
+                    opt_dict[key] = str(value)
+
+    # print(opt)
+    opt_dict['which_epoch'] = 'finest'
+    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 = './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
+        opt.max_motion_length = 196
+        opt.max_motion_frame = 196
+        opt.max_motion_token = 55
+    elif opt.dataset_name == 'kit':
+        opt.data_root = './dataset/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
+        opt.max_motion_frame = 196
+        opt.max_motion_token = 55
+    else:
+        raise KeyError('Dataset not recognized')
+
+    opt.dim_word = 300
+    opt.num_classes = 200 // opt.unit_length
+    opt.is_train = False
+    opt.is_continue = False
+    opt.device = device
+
+    return opt

generate_human_motion/VQTrans/options/option_transformer.py

@@ -0,0 +1,68 @@
+import argparse
+
+def get_args_parser():
+    parser = argparse.ArgumentParser(description='Optimal Transport AutoEncoder training for Amass',
+                                     add_help=True,
+                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+    
+    ## dataloader
+    
+    parser.add_argument('--dataname', type=str, default='kit', help='dataset directory')
+    parser.add_argument('--batch-size', default=128, type=int, help='batch size')
+    parser.add_argument('--fps', default=[20], nargs="+", type=int, help='frames per second')
+    parser.add_argument('--seq-len', type=int, default=64, help='training motion length')
+    
+    ## optimization
+    parser.add_argument('--total-iter', default=100000, type=int, help='number of total iterations to run')
+    parser.add_argument('--warm-up-iter', default=1000, type=int, help='number of total iterations for warmup')
+    parser.add_argument('--lr', default=2e-4, type=float, help='max learning rate')
+    parser.add_argument('--lr-scheduler', default=[60000], nargs="+", type=int, help="learning rate schedule (iterations)")
+    parser.add_argument('--gamma', default=0.05, type=float, help="learning rate decay")
+    
+    parser.add_argument('--weight-decay', default=1e-6, type=float, help='weight decay') 
+    parser.add_argument('--decay-option',default='all', type=str, choices=['all', 'noVQ'], help='disable weight decay on codebook')
+    parser.add_argument('--optimizer',default='adamw', type=str, choices=['adam', 'adamw'], help='disable weight decay on codebook')
+    
+    ## vqvae arch
+    parser.add_argument("--code-dim", type=int, default=512, help="embedding dimension")
+    parser.add_argument("--nb-code", type=int, default=512, help="nb of embedding")
+    parser.add_argument("--mu", type=float, default=0.99, help="exponential moving average to update the codebook")
+    parser.add_argument("--down-t", type=int, default=3, help="downsampling rate")
+    parser.add_argument("--stride-t", type=int, default=2, help="stride size")
+    parser.add_argument("--width", type=int, default=512, help="width of the network")
+    parser.add_argument("--depth", type=int, default=3, help="depth of the network")
+    parser.add_argument("--dilation-growth-rate", type=int, default=3, help="dilation growth rate")
+    parser.add_argument("--output-emb-width", type=int, default=512, help="output embedding width")
+    parser.add_argument('--vq-act', type=str, default='relu', choices = ['relu', 'silu', 'gelu'], help='dataset directory')
+
+    ## gpt arch
+    parser.add_argument("--block-size", type=int, default=25, help="seq len")
+    parser.add_argument("--embed-dim-gpt", type=int, default=512, help="embedding dimension")
+    parser.add_argument("--clip-dim", type=int, default=512, help="latent dimension in the clip feature")
+    parser.add_argument("--num-layers", type=int, default=2, help="nb of transformer layers")
+    parser.add_argument("--n-head-gpt", type=int, default=8, help="nb of heads")
+    parser.add_argument("--ff-rate", type=int, default=4, help="feedforward size")
+    parser.add_argument("--drop-out-rate", type=float, default=0.1, help="dropout ratio in the pos encoding")
+    
+    ## quantizer
+    parser.add_argument("--quantizer", type=str, default='ema_reset', choices = ['ema', 'orig', 'ema_reset', 'reset'], help="eps for optimal transport")
+    parser.add_argument('--quantbeta', type=float, default=1.0, help='dataset directory')
+
+    ## resume
+    parser.add_argument("--resume-pth", type=str, default=None, help='resume vq pth')
+    parser.add_argument("--resume-trans", type=str, default=None, help='resume gpt pth')
+    
+    
+    ## output directory 
+    parser.add_argument('--out-dir', type=str, default='output_GPT_Final/', help='output directory')
+    parser.add_argument('--exp-name', type=str, default='exp_debug', help='name of the experiment, will create a file inside out-dir')
+    parser.add_argument('--vq-name', type=str, default='exp_debug', help='name of the generated dataset .npy, will create a file inside out-dir')
+    ## other
+    parser.add_argument('--print-iter', default=200, type=int, help='print frequency')
+    parser.add_argument('--eval-iter', default=5000, type=int, help='evaluation frequency')
+    parser.add_argument('--seed', default=123, type=int, help='seed for initializing training. ')
+    parser.add_argument("--if-maxtest", action='store_true', help="test in max")
+    parser.add_argument('--pkeep', type=float, default=1.0, help='keep rate for gpt training')
+    
+    
+    return parser.parse_args()

generate_human_motion/VQTrans/options/option_vq.py

@@ -0,0 +1,61 @@
+import argparse
+
+def get_args_parser():
+    parser = argparse.ArgumentParser(description='Optimal Transport AutoEncoder training for AIST',
+                                     add_help=True,
+                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+
+    ## dataloader  
+    parser.add_argument('--dataname', type=str, default='kit', help='dataset directory')
+    parser.add_argument('--batch-size', default=128, type=int, help='batch size')
+    parser.add_argument('--window-size', type=int, default=64, help='training motion length')
+
+    ## optimization
+    parser.add_argument('--total-iter', default=200000, type=int, help='number of total iterations to run')
+    parser.add_argument('--warm-up-iter', default=1000, type=int, help='number of total iterations for warmup')
+    parser.add_argument('--lr', default=2e-4, type=float, help='max learning rate')
+    parser.add_argument('--lr-scheduler', default=[50000, 400000], nargs="+", type=int, help="learning rate schedule (iterations)")
+    parser.add_argument('--gamma', default=0.05, type=float, help="learning rate decay")
+
+    parser.add_argument('--weight-decay', default=0.0, type=float, help='weight decay')
+    parser.add_argument("--commit", type=float, default=0.02, help="hyper-parameter for the commitment loss")
+    parser.add_argument('--loss-vel', type=float, default=0.1, help='hyper-parameter for the velocity loss')
+    parser.add_argument('--recons-loss', type=str, default='l2', help='reconstruction loss')
+    
+    ## vqvae arch
+    parser.add_argument("--code-dim", type=int, default=512, help="embedding dimension")
+    parser.add_argument("--nb-code", type=int, default=512, help="nb of embedding")
+    parser.add_argument("--mu", type=float, default=0.99, help="exponential moving average to update the codebook")
+    parser.add_argument("--down-t", type=int, default=2, help="downsampling rate")
+    parser.add_argument("--stride-t", type=int, default=2, help="stride size")
+    parser.add_argument("--width", type=int, default=512, help="width of the network")
+    parser.add_argument("--depth", type=int, default=3, help="depth of the network")
+    parser.add_argument("--dilation-growth-rate", type=int, default=3, help="dilation growth rate")
+    parser.add_argument("--output-emb-width", type=int, default=512, help="output embedding width")
+    parser.add_argument('--vq-act', type=str, default='relu', choices = ['relu', 'silu', 'gelu'], help='dataset directory')
+    parser.add_argument('--vq-norm', type=str, default=None, help='dataset directory')
+    
+    ## quantizer
+    parser.add_argument("--quantizer", type=str, default='ema_reset', choices = ['ema', 'orig', 'ema_reset', 'reset'], help="eps for optimal transport")
+    parser.add_argument('--beta', type=float, default=1.0, help='commitment loss in standard VQ')
+
+    ## resume
+    parser.add_argument("--resume-pth", type=str, default=None, help='resume pth for VQ')
+    parser.add_argument("--resume-gpt", type=str, default=None, help='resume pth for GPT')
+    
+    
+    ## output directory 
+    parser.add_argument('--out-dir', type=str, default='output_vqfinal/', help='output directory')
+    parser.add_argument('--results-dir', type=str, default='visual_results/', help='output directory')
+    parser.add_argument('--visual-name', type=str, default='baseline', help='output directory')
+    parser.add_argument('--exp-name', type=str, default='exp_debug', help='name of the experiment, will create a file inside out-dir')
+    ## other
+    parser.add_argument('--print-iter', default=200, type=int, help='print frequency')
+    parser.add_argument('--eval-iter', default=1000, type=int, help='evaluation frequency')
+    parser.add_argument('--seed', default=123, type=int, help='seed for initializing training.')
+    
+    parser.add_argument('--vis-gt', action='store_true', help='whether visualize GT motions')
+    parser.add_argument('--nb-vis', default=20, type=int, help='nb of visualizations')
+    
+    
+    return parser.parse_args()

generate_human_motion/VQTrans/output/02ab4ad275eda92f352e2ed8d942eeef_pred.pt

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

generate_human_motion/VQTrans/output/06c27c738e874b23067c006f52e18ebc_pred.pt

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

generate_human_motion/VQTrans/output/0edd5f692aeec051d748dee0844f94e1_pred.pt

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

generate_human_motion/VQTrans/output/23cb7d0e26bb1646b3d386331971449c_pred.pt

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