mingyuan/MotionDiffuse

1import os2import sys3import gradio as gr4 5os.makedirs("outputs", exist_ok=True) 6sys.path.insert(0, '.')7 8 9from utils.get_opt import get_opt10from os.path import join as pjoin11import numpy as np12from trainers import DDPMTrainer13from models import MotionTransformer14 15device = 'cpu'16opt = get_opt("checkpoints/t2m/t2m_motiondiffuse/opt.txt", device)17opt.do_denoise = True18 19assert opt.dataset_name == "t2m"20opt.data_root = './dataset/HumanML3D'21opt.motion_dir = pjoin(opt.data_root, 'new_joint_vecs')22opt.text_dir = pjoin(opt.data_root, 'texts')23opt.joints_num = 2224opt.dim_pose = 26325 26mean = np.load(pjoin(opt.meta_dir, 'mean.npy'))27std = np.load(pjoin(opt.meta_dir, 'std.npy'))28 29 30def build_models(opt):31    encoder = MotionTransformer(32        input_feats=opt.dim_pose,33        num_frames=opt.max_motion_length,34        num_layers=opt.num_layers,35        latent_dim=opt.latent_dim,36        no_clip=opt.no_clip,37        no_eff=opt.no_eff)38    return encoder39 40 41encoder = build_models(opt).to(device)42trainer = DDPMTrainer(opt, encoder)43trainer.load(pjoin(opt.model_dir, opt.which_epoch + '.tar'))44 45trainer.eval_mode()46trainer.to(opt.device)47 48def generate(prompt, length):49    from tools.visualization import process50    result_path = "outputs/" + str(hash(prompt)) + ".mp4"51    process(trainer, opt, device, mean, std, prompt, int(length), result_path)52    return result_path53 54demo = gr.Interface(55    fn=generate,56    inputs=["text", gr.Slider(20, 196, value=60)],57    examples=[58        ["the man throws a punch with each hand.", 58],59        ["a person spins quickly and takes off running.", 29],60        ["a person quickly waves with their right hand", 46],61        ["a person performing a slight bow", 89],62    ],63    outputs="video",64    title="MotionDiffuse: Text-Driven Human Motion Generation with Diffusion Model",65    description="This is an interactive demo for MotionDiffuse. For more information, feel free to visit our project page(https://mingyuan-zhang.github.io/projects/MotionDiffuse.html).")66 67demo.launch(enable_queue=True)