chencws/Stable-text-to-motion
0
1import sys2import os3import OpenGL.GL as gl4os.environ["PYOPENGL_PLATFORM"] = "egl"5os.environ["MESA_GL_VERSION_OVERRIDE"] = "4.1"6os.system('pip install /home/user/app/pyrender')7 8sys.argv = ['VQ-Trans/GPT_eval_multi.py']9os.chdir('VQ-Trans')10 11sys.path.append('/home/user/app/VQ-Trans')12sys.path.append('/home/user/app/pyrender')13 14import options.option_transformer as option_trans15from huggingface_hub import snapshot_download16model_path = snapshot_download(repo_id="vumichien/T2M-GPT")17 18args = option_trans.get_args_parser()19 20args.dataname = 't2m'21args.resume_pth = f'{model_path}/VQVAE/net_last.pth'22args.resume_trans = f'{model_path}/VQTransformer_corruption05/net_best_fid.pth'23args.down_t = 224args.depth = 325args.block_size = 5126 27import clip28import torch29import numpy as np30import models.vqvae as vqvae31import models.t2m_trans as trans32from utils.motion_process import recover_from_ric33import visualization.plot_3d_global as plot_3d34from models.rotation2xyz import Rotation2xyz35import numpy as np36from trimesh import Trimesh37import gc38 39import torch40from visualize.simplify_loc2rot import joints2smpl41import pyrender42# import matplotlib.pyplot as plt43 44import io45import imageio46from shapely import geometry47import trimesh48from pyrender.constants import RenderFlags49import math50# import ffmpeg51# from PIL import Image52import hashlib53import gradio as gr54import moviepy.editor as mp55 56## load clip model and datasets57is_cuda = torch.cuda.is_available()58device = torch.device("cuda" if is_cuda else "cpu")59print(device)60clip_model, clip_preprocess = clip.load("ViT-B/32", device=device, jit=False, download_root='./') # Must set jit=False for training61 62if is_cuda:63 clip.model.convert_weights(clip_model)64 65clip_model.eval()66for p in clip_model.parameters():67 p.requires_grad = False68 69net = vqvae.HumanVQVAE(args, ## use args to define different parameters in different quantizers70 args.nb_code,71 args.code_dim,72 args.output_emb_width,73 args.down_t,74 args.stride_t,75 args.width,76 args.depth,77 args.dilation_growth_rate)78 79 80trans_encoder = trans.Text2Motion_Transformer(num_vq=args.nb_code, 81 embed_dim=1024, 82 clip_dim=args.clip_dim, 83 block_size=args.block_size, 84 num_layers=9, 85 n_head=16, 86 drop_out_rate=args.drop_out_rate, 87 fc_rate=args.ff_rate)88 89 90print('loading checkpoint from {}'.format(args.resume_pth))91ckpt = torch.load(args.resume_pth, map_location='cpu')92net.load_state_dict(ckpt['net'], strict=True)93net.eval()94 95print('loading transformer checkpoint from {}'.format(args.resume_trans))96ckpt = torch.load(args.resume_trans, map_location='cpu')97trans_encoder.load_state_dict(ckpt['trans'], strict=True)98trans_encoder.eval()99 100mean = torch.from_numpy(np.load(f'{model_path}/meta/mean.npy'))101std = torch.from_numpy(np.load(f'{model_path}/meta/std.npy'))102 103if is_cuda:104 net.cuda()105 trans_encoder.cuda()106 mean = mean.cuda()107 std = std.cuda()108 109def render(motions, device_id=0, name='test_vis'):110 frames, njoints, nfeats = motions.shape111 MINS = motions.min(axis=0).min(axis=0)112 MAXS = motions.max(axis=0).max(axis=0)113 114 height_offset = MINS[1]115 motions[:, :, 1] -= height_offset116 trajec = motions[:, 0, [0, 2]]117 is_cuda = torch.cuda.is_available()118 # device = torch.device("cuda" if is_cuda else "cpu")119 j2s = joints2smpl(num_frames=frames, device_id=0, cuda=is_cuda)120 rot2xyz = Rotation2xyz(device=device)121 faces = rot2xyz.smpl_model.faces122 123 if not os.path.exists(f'output/{name}_pred.pt'): 124 print(f'Running SMPLify, it may take a few minutes.')125 motion_tensor, opt_dict = j2s.joint2smpl(motions) # [nframes, njoints, 3]126 127 vertices = rot2xyz(torch.tensor(motion_tensor).clone(), mask=None,128 pose_rep='rot6d', translation=True, glob=True,129 jointstype='vertices',130 vertstrans=True)131 vertices = vertices.detach().cpu()132 torch.save(vertices, f'output/{name}_pred.pt')133 else:134 vertices = torch.load(f'output/{name}_pred.pt')135 frames = vertices.shape[3] # shape: 1, nb_frames, 3, nb_joints136 print(vertices.shape)137 MINS = torch.min(torch.min(vertices[0], axis=0)[0], axis=1)[0]138 MAXS = torch.max(torch.max(vertices[0], axis=0)[0], axis=1)[0]139 140 out_list = []141 142 minx = MINS[0] - 0.5143 maxx = MAXS[0] + 0.5144 minz = MINS[2] - 0.5 145 maxz = MAXS[2] + 0.5146 polygon = geometry.Polygon([[minx, minz], [minx, maxz], [maxx, maxz], [maxx, minz]])147 polygon_mesh = trimesh.creation.extrude_polygon(polygon, 1e-5)148 149 vid = []150 for i in range(frames):151 if i % 10 == 0:152 print(i)153 154 mesh = Trimesh(vertices=vertices[0, :, :, i].squeeze().tolist(), faces=faces)155 156 base_color = (0.11, 0.53, 0.8, 0.5)157 ## OPAQUE rendering without alpha158 ## BLEND rendering consider alpha 159 material = pyrender.MetallicRoughnessMaterial(160 metallicFactor=0.7,161 alphaMode='OPAQUE',162 baseColorFactor=base_color163 )164 165 166 mesh = pyrender.Mesh.from_trimesh(mesh, material=material)167 168 polygon_mesh.visual.face_colors = [0, 0, 0, 0.21]169 polygon_render = pyrender.Mesh.from_trimesh(polygon_mesh, smooth=False)170 171 bg_color = [1, 1, 1, 0.8]172 scene = pyrender.Scene(bg_color=bg_color, ambient_light=(0.4, 0.4, 0.4))173 174 sx, sy, tx, ty = [0.75, 0.75, 0, 0.10]175 176 camera = pyrender.PerspectiveCamera(yfov=(np.pi / 3.0))177 178 light = pyrender.DirectionalLight(color=[1,1,1], intensity=300)179 180 scene.add(mesh)181 182 c = np.pi / 2183 184 scene.add(polygon_render, pose=np.array([[ 1, 0, 0, 0],185 186 [ 0, np.cos(c), -np.sin(c), MINS[1].cpu().numpy()],187 188 [ 0, np.sin(c), np.cos(c), 0],189 190 [ 0, 0, 0, 1]]))191 192 light_pose = np.eye(4)193 light_pose[:3, 3] = [0, -1, 1]194 scene.add(light, pose=light_pose.copy())195 196 light_pose[:3, 3] = [0, 1, 1]197 scene.add(light, pose=light_pose.copy())198 199 light_pose[:3, 3] = [1, 1, 2]200 scene.add(light, pose=light_pose.copy())201 202 203 c = -np.pi / 6204 205 scene.add(camera, pose=[[ 1, 0, 0, (minx+maxx).cpu().numpy()/2],206 207 [ 0, np.cos(c), -np.sin(c), 1.5],208 209 [ 0, np.sin(c), np.cos(c), max(4, minz.cpu().numpy()+(1.5-MINS[1].cpu().numpy())*2, (maxx-minx).cpu().numpy())],210 211 [ 0, 0, 0, 1]212 ])213 214 # render scene215 r = pyrender.OffscreenRenderer(960, 960)216 217 color, _ = r.render(scene, flags=RenderFlags.RGBA)218 # Image.fromarray(color).save(outdir+'/'+name+'_'+str(i)+'.png')219 220 vid.append(color)221 222 r.delete()223 224 out = np.stack(vid, axis=0)225 imageio.mimwrite(f'output/results.gif', out, duration=50)226 out_video = mp.VideoFileClip(f'output/results.gif')227 out_video.write_videofile("output/results.mp4")228 del out, vertices229 return f'output/results.mp4'230 231def predict(clip_text, method='fast'):232 gc.collect()233 print('prompt text instruction: {}'.format(clip_text))234 if torch.cuda.is_available():235 text = clip.tokenize([clip_text], truncate=True).cuda()236 else:237 text = clip.tokenize([clip_text], truncate=True)238 feat_clip_text = clip_model.encode_text(text).float()239 index_motion = trans_encoder.sample(feat_clip_text[0:1], False)240 pred_pose = net.forward_decoder(index_motion)241 pred_xyz = recover_from_ric((pred_pose*std+mean).float(), 22)242 output_name = hashlib.md5(clip_text.encode()).hexdigest()243 if method == 'fast':244 xyz = pred_xyz.reshape(1, -1, 22, 3)245 pose_vis = plot_3d.draw_to_batch(xyz.detach().cpu().numpy(), title_batch=None, outname=[f'output/results.gif'])246 out_video = mp.VideoFileClip("output/results.gif")247 out_video.write_videofile("output/results.mp4")248 return f'output/results.mp4'249 elif method == 'slow':250 output_path = render(pred_xyz.detach().cpu().numpy().squeeze(axis=0), device_id=0, name=output_name)251 return output_path252 253 254# ---- Gradio Layout -----255video_out = gr.Video(label="Motion", mirror_webcam=False, interactive=False) 256demo = gr.Blocks()257demo.encrypt = False258 259with demo:260 gr.Markdown('''261 <div>262 <h1 style='text-align: center'>Generating Human Motion from Textual Descriptions (T2M-GPT)</h1>263 This space uses <a href='https://mael-zys.github.io/T2M-GPT/' target='_blank'><b>T2M-GPT models</b></a> based on Vector Quantised-Variational AutoEncoder (VQ-VAE) and Generative Pre-trained Transformer (GPT) for human motion generation from textural descriptions🤗264 </div>265 ''')266 with gr.Row():267 with gr.Column():268 gr.Markdown('''269 <figure>270 <img src="https://huggingface.co/vumichien/T2M-GPT/resolve/main/demo_slow1.gif" alt="Demo Slow", width="425", height=480/>271 <figcaption> a man starts off in an up right position with botg arms extended out by his sides, he then brings his arms down to his body and claps his hands together. after this he wals down amd the the left where he proceeds to sit on a seat272 </figcaption>273 </figure>274 ''')275 with gr.Column():276 gr.Markdown('''277 <figure>278 <img src="https://huggingface.co/vumichien/T2M-GPT/resolve/main/demo_slow2.gif" alt="Demo Slow 2", width="425", height=480/>279 <figcaption> a person puts their hands together, leans forwards slightly then swings the arms from right to left280 </figcaption>281 </figure>282 ''')283 with gr.Column():284 gr.Markdown('''285 <figure>286 <img src="https://huggingface.co/vumichien/T2M-GPT/resolve/main/demo_slow3.gif" alt="Demo Slow 3", width="425", height=480/>287 <figcaption> a man is practicing the waltz with a partner288 </figcaption>289 </figure>290 ''')291 with gr.Row():292 with gr.Column():293 gr.Markdown('''294 ### Generate human motion by **T2M-GPT**295 ##### Step 1. Give prompt text describing human motion296 ##### Step 2. Choice method to render output (Fast: Sketch skeleton; Slow: SMPL mesh, only work with GPU and running time around 2 mins)297 ##### Step 3. Generate output and enjoy298 ''')299 with gr.Column():300 with gr.Row():301 text_prompt = gr.Textbox(label="Text prompt", lines=1, interactive=True)302 method = gr.Dropdown(["slow", "fast"], label="Method", value="slow")303 with gr.Row():304 generate_btn = gr.Button("Generate")305 generate_btn.click(predict, [text_prompt, method], [video_out], api_name="generate")306 with gr.Row(): 307 video_out.render()308 with gr.Row(): 309 gr.Markdown('''310 ### You can test by following examples:311 ''')312 examples = gr.Examples(313 examples=[314 ["a person jogs in place, slowly at first, then increases speed. they then back up and squat down.", "slow"],315 ["a man steps forward and does a handstand", "slow"],316 ["a man rises from the ground, walks in a circle and sits back down on the ground", "slow"],317 ["a man starts off in an up right position with botg arms extended out by his sides, he then brings his arms down to his body and claps his hands together. after this he wals down amd the the left where he proceeds to sit on a seat", "slow"],318 ["a person puts their hands together, leans forwards slightly then swings the arms from right to left","slow"],319 ["a man is practicing the waltz with a partner","slow"],320 ],321 label="Examples", 322 inputs=[text_prompt, method],323 outputs=[video_out],324 fn=predict,325 cache_examples=True,326 )327 328 329demo.launch(debug=True)330