momask-codes 部署踩坑笔记

目录

依赖项

t2m_nlayer8_nhead6_ld384_ff1024_cdp0.1_rvq6ns

推理代码完善：

---

代码地址：

https://github.com/EricGuo5513/momask-codes

依赖项

pip install numpy==1.23

matplotlib 必须指定版本：pip install matplotlib==3.3.4

t2m_nlayer8_nhead6_ld384_ff1024_cdp0.1_rvq6ns

下载模型：

cd t2m 
echo -e "Downloading pretrained models for HumanML3D dataset"
gdown --fuzzy https://drive.google.com/file/d/1vXS7SHJBgWPt59wupQ5UUzhFObrnGkQ0/view?usp=sharing

echo -e "Unzipping humanml3d_models.zip"
unzip humanml3d_models.zip

推理代码完善：

text_motion/momask-codes/gen_t2m.py

# coding=utf-8
import sys
import os
os.chdir(os.path.dirname(os.path.abspath(__file__)))
current_dir = os.path.dirname(os.path.abspath(__file__))paths = [os.path.abspath(__file__).split('scripts')[0]]
print('current_dir',current_dir)
paths.append(os.path.abspath(os.path.join(current_dir, 'src')))for path in paths:sys.path.insert(0, path)os.environ['PYTHONPATH'] = (os.environ.get('PYTHONPATH', '') + ':' + path).strip(':')import os
from os.path import join as pjoinimport torch
import torch.nn.functional as Ffrom models.mask_transformer.transformer import MaskTransformer, ResidualTransformer
from models.vq.model import RVQVAE, LengthEstimator# from options.eval_option import EvalT2MOptionsfrom options.base_option import BaseOptionsclass EvalT2MOptions(BaseOptions):def initialize(self):BaseOptions.initialize(self)self.parser.add_argument('--which_epoch', type=str, default="latest", help='Checkpoint you want to use, {latest, net_best_fid, etc}')self.parser.add_argument('--batch_size', type=int, default=32, help='Batch size')self.parser.add_argument('--ext', type=str, default='text2motion', help='Extension of the result file or folder')self.parser.add_argument("--num_batch", default=2, type=int,help="Number of batch for generation")self.parser.add_argument("--repeat_times", default=1, type=int,help="Number of repetitions, per sample text prompt")self.parser.add_argument("--cond_scale", default=4, type=float,help="For classifier-free sampling - specifies the s parameter, as defined in the paper.")self.parser.add_argument("--temperature", default=1., type=float,help="Sampling Temperature.")self.parser.add_argument("--topkr", default=0.9, type=float,help="Filter out percentil low prop entries.")self.parser.add_argument("--time_steps", default=18, type=int,help="Mask Generate steps.")self.parser.add_argument("--seed", default=10107, type=int)self.parser.add_argument('--gumbel_sample', action="store_true", help='True: gumbel sampling, False: categorical sampling.')self.parser.add_argument('--use_res_model', action="store_true", help='Whether to use residual transformer.')# self.parser.add_argument('--est_length', action="store_true", help='Training iterations')self.parser.add_argument('--res_name', type=str, default='tres_nlayer8_ld384_ff1024_rvq6ns_cdp0.2_sw', help='Model name of residual transformer')self.parser.add_argument('--text_path', type=str, default="", help='Text prompt file')self.parser.add_argument('-msec', '--mask_edit_section', nargs='*', type=str, help='Indicate sections for editing, use comma to separate the start and end of a section''type int will specify the token frame, type float will specify the ratio of seq_len')self.parser.add_argument('--text_prompt', default='A person is running on a treadmill.', type=str, help="A text prompt to be generated. If empty, will take text prompts from dataset.")self.parser.add_argument('--source_motion', default='example_data/000612.npy', type=str, help="Source motion path for editing. (new_joint_vecs format .npy file)")self.parser.add_argument("--motion_length", default=0, type=int,help="Motion length for generation, only applicable with single text prompt.")self.is_train = Falsefrom utils.get_opt import get_optfrom utils.fixseed import fixseed
from visualization.joints2bvh import Joint2BVHConvertor
from torch.distributions.categorical import Categoricalfrom utils.motion_process import recover_from_ric
from utils.plot_script import plot_3d_motionfrom utils.paramUtil import t2m_kinematic_chainimport numpy as np
clip_version = 'ViT-B/32'def load_vq_model(vq_opt):# opt_path = pjoin(opt.checkpoints_dir, opt.dataset_name, opt.vq_name, 'opt.txt')vq_model = RVQVAE(vq_opt,vq_opt.dim_pose,vq_opt.nb_code,vq_opt.code_dim,vq_opt.output_emb_width,vq_opt.down_t,vq_opt.stride_t,vq_opt.width,vq_opt.depth,vq_opt.dilation_growth_rate,vq_opt.vq_act,vq_opt.vq_norm)ckpt = torch.load(pjoin(vq_opt.checkpoints_dir, vq_opt.dataset_name, vq_opt.name, 'model', 'net_best_fid.tar'),map_location='cpu')model_key = 'vq_model' if 'vq_model' in ckpt else 'net'vq_model.load_state_dict(ckpt[model_key])print(f'Loading VQ Model {vq_opt.name} Completed!')return vq_model, vq_optdef load_trans_model(model_opt, opt, which_model):t2m_transformer = MaskTransformer(code_dim=model_opt.code_dim,cond_mode='text',latent_dim=model_opt.latent_dim,ff_size=model_opt.ff_size,num_layers=model_opt.n_layers,num_heads=model_opt.n_heads,dropout=model_opt.dropout,clip_dim=512,cond_drop_prob=model_opt.cond_drop_prob,clip_version=clip_version,opt=model_opt)ckpt = torch.load(pjoin(model_opt.checkpoints_dir, model_opt.dataset_name, model_opt.name, 'model', which_model),map_location='cpu')model_key = 't2m_transformer' if 't2m_transformer' in ckpt else 'trans'# print(ckpt.keys())missing_keys, unexpected_keys = t2m_transformer.load_state_dict(ckpt[model_key], strict=False)assert len(unexpected_keys) == 0assert all([k.startswith('clip_model.') for k in missing_keys])print(f'Loading Transformer {opt.name} from epoch {ckpt["ep"]}!')return t2m_transformerdef load_res_model(res_opt, vq_opt, opt):res_opt.num_quantizers = vq_opt.num_quantizersres_opt.num_tokens = vq_opt.nb_coderes_transformer = ResidualTransformer(code_dim=vq_opt.code_dim,cond_mode='text',latent_dim=res_opt.latent_dim,ff_size=res_opt.ff_size,num_layers=res_opt.n_layers,num_heads=res_opt.n_heads,dropout=res_opt.dropout,clip_dim=512,shared_codebook=vq_opt.shared_codebook,cond_drop_prob=res_opt.cond_drop_prob,# codebook=vq_model.quantizer.codebooks[0] if opt.fix_token_emb else None,share_weight=res_opt.share_weight,clip_version=clip_version,opt=res_opt)ckpt = torch.load(pjoin(res_opt.checkpoints_dir, res_opt.dataset_name, res_opt.name, 'model', 'net_best_fid.tar'),map_location=opt.device)missing_keys, unexpected_keys = res_transformer.load_state_dict(ckpt['res_transformer'], strict=False)assert len(unexpected_keys) == 0assert all([k.startswith('clip_model.') for k in missing_keys])print(f'Loading Residual Transformer {res_opt.name} from epoch {ckpt["ep"]}!')return res_transformerdef load_len_estimator(opt):model = LengthEstimator(512, 50)ckpt = torch.load(pjoin(opt.checkpoints_dir, opt.dataset_name, 'length_estimator', 'model', 'finest.tar'),map_location=opt.device)model.load_state_dict(ckpt['estimator'])print(f'Loading Length Estimator from epoch {ckpt["epoch"]}!')return modelif __name__ == '__main__':parser = EvalT2MOptions()opt = parser.parse()fixseed(opt.seed)opt.device = torch.device("cuda:1")torch.autograd.set_detect_anomaly(True)dim_pose = 251 if opt.dataset_name == 'kit' else 263# out_dir = pjoin(opt.check)root_dir = pjoin(opt.checkpoints_dir, opt.dataset_name, opt.name)model_dir = pjoin(root_dir, 'model')result_dir = pjoin('./generation', opt.ext)joints_dir = pjoin(result_dir, 'joints')animation_dir = pjoin(result_dir, 'animations')os.makedirs(joints_dir, exist_ok=True)os.makedirs(animation_dir,exist_ok=True)model_opt_path = pjoin(root_dir, 'opt.txt')model_opt = get_opt(model_opt_path, device=opt.device)#############################Loading RVQ#############################vq_opt_path = pjoin(opt.checkpoints_dir, opt.dataset_name, model_opt.vq_name, 'opt.txt')vq_opt = get_opt(vq_opt_path, device=opt.device)vq_opt.dim_pose = dim_posevq_model, vq_opt = load_vq_model(vq_opt)model_opt.num_tokens = vq_opt.nb_codemodel_opt.num_quantizers = vq_opt.num_quantizersmodel_opt.code_dim = vq_opt.code_dim#######################################Loading R-Transformer#######################################res_opt_path = pjoin(opt.checkpoints_dir, opt.dataset_name, opt.res_name, 'opt.txt')res_opt = get_opt(res_opt_path, device=opt.device)res_model = load_res_model(res_opt, vq_opt, opt)assert res_opt.vq_name == model_opt.vq_name#######################################Loading M-Transformer#######################################t2m_transformer = load_trans_model(model_opt, opt, 'latest.tar')#######################################Loading Length Predictor#######################################length_estimator = load_len_estimator(model_opt)t2m_transformer.eval()vq_model.eval()res_model.eval()length_estimator.eval()res_model.to(opt.device)t2m_transformer.to(opt.device)vq_model.to(opt.device)length_estimator.to(opt.device)##### ---- Dataloader ---- #####opt.nb_joints = 21 if opt.dataset_name == 'kit' else 22mean = np.load(pjoin(opt.checkpoints_dir, opt.dataset_name, model_opt.vq_name, 'meta', 'mean.npy'))std = np.load(pjoin(opt.checkpoints_dir, opt.dataset_name, model_opt.vq_name, 'meta', 'std.npy'))def inv_transform(data):return data * std + meanprompt_list = []length_list = []est_length = Falseif opt.text_prompt != "":prompt_list.append(opt.text_prompt)if opt.motion_length == 0:est_length = Trueelse:length_list.append(opt.motion_length)elif opt.text_path != "":with open(opt.text_path, 'r') as f:lines = f.readlines()for line in lines:infos = line.split('#')prompt_list.append(infos[0])if len(infos) == 1 or (not infos[1].isdigit()):est_length = Truelength_list = []else:length_list.append(int(infos[-1]))else:raise "A text prompt, or a file a text prompts are required!!!"# print('loading checkpoint {}'.format(file))if est_length:print("Since no motion length are specified, we will use estimated motion lengthes!!")text_embedding = t2m_transformer.encode_text(prompt_list)pred_dis = length_estimator(text_embedding)probs = F.softmax(pred_dis, dim=-1)  # (b, ntoken)token_lens = Categorical(probs).sample()  # (b, seqlen)# lengths = torch.multinomial()else:token_lens = torch.LongTensor(length_list) // 4token_lens = token_lens.to(opt.device).long()m_length = token_lens * 4captions = prompt_listsample = 0kinematic_chain = t2m_kinematic_chainconverter = Joint2BVHConvertor()for r in range(opt.repeat_times):print("-->Repeat %d"%r)with torch.no_grad():mids = t2m_transformer.generate(captions, token_lens,timesteps=opt.time_steps,cond_scale=opt.cond_scale,temperature=opt.temperature,topk_filter_thres=opt.topkr,gsample=opt.gumbel_sample)# print(mids)# print(mids.shape)mids = res_model.generate(mids, captions, token_lens, temperature=1, cond_scale=5)pred_motions = vq_model.forward_decoder(mids)pred_motions = pred_motions.detach().cpu().numpy()data = inv_transform(pred_motions)for k, (caption, joint_data)  in enumerate(zip(captions, data)):print("---->Sample %d: %s %d"%(k, caption, m_length[k]))animation_path = pjoin(animation_dir, str(k))joint_path = pjoin(joints_dir, str(k))print('save_dir',animation_path)os.makedirs(animation_path, exist_ok=True)os.makedirs(joint_path, exist_ok=True)joint_data = joint_data[:m_length[k]]joint = recover_from_ric(torch.from_numpy(joint_data).float(), 22).numpy()bvh_path = pjoin(animation_path, "sample%d_repeat%d_len%d_ik.bvh"%(k, r, m_length[k]))_, ik_joint = converter.convert(joint, filename=bvh_path, iterations=100)bvh_path = pjoin(animation_path, "sample%d_repeat%d_len%d.bvh" % (k, r, m_length[k]))_, joint = converter.convert(joint, filename=bvh_path, iterations=100, foot_ik=False)save_path = pjoin(animation_path, "sample%d_repeat%d_len%d.mp4"%(k, r, m_length[k]))ik_save_path = pjoin(animation_path, "sample%d_repeat%d_len%d_ik.mp4"%(k, r, m_length[k]))plot_3d_motion(ik_save_path, kinematic_chain, ik_joint, title=caption, fps=20)plot_3d_motion(save_path, kinematic_chain, joint, title=caption, fps=20)np.save(pjoin(joint_path, "sample%d_repeat%d_len%d.npy"%(k, r, m_length[k])), joint)np.save(pjoin(joint_path, "sample%d_repeat%d_len%d_ik.npy"%(k, r, m_length[k])), ik_joint)