Create app.py

app.py

@@ -0,0 +1,288 @@
+import sys
+import os
+import OpenGL.GL as gl
+os.environ["PYOPENGL_PLATFORM"] = "egl"
+sys.argv = ['VQ-Trans/GPT_eval_multi.py']
+os.makedirs('output', exist_ok=True)
+os.chdir('VQ-Trans')
+os.makedirs('checkpoints', exist_ok=True)
+os.system('gdown --fuzzy https://drive.google.com/file/d/1o7RTDQcToJjTm9_mNWTyzvZvjTWpZfug/view -O checkpoints')
+os.system('gdown --fuzzy https://drive.google.com/file/d/1tX79xk0fflp07EZ660Xz1RAFE33iEyJR/view -O checkpoints')
+os.system('unzip checkpoints/t2m.zip')
+os.system('unzip checkpoints/kit.zip')
+os.system('rm checkpoints/t2m.zip')
+os.system('rm checkpoints/kit.zip')
+sys.path.append('/home/user/app/VQ-Trans')
+
+import options.option_transformer as option_trans
+from huggingface_hub import snapshot_download
+model_path = snapshot_download(repo_id="vumichien/T2M-GPT")
+
+args = option_trans.get_args_parser()
+
+args.dataname = 't2m'
+args.resume_pth = f'{model_path}/VQVAE/net_last.pth'
+args.resume_trans = f'{model_path}/VQTransformer_corruption05/net_best_fid.pth'
+args.down_t = 2
+args.depth = 3
+args.block_size = 51
+
+import clip
+import torch
+import numpy as np
+import models.vqvae as vqvae
+import models.t2m_trans as trans
+from utils.motion_process import recover_from_ric
+import visualization.plot_3d_global as plot_3d
+from models.rotation2xyz import Rotation2xyz
+import numpy as np
+from trimesh import Trimesh
+import gc
+
+import torch
+from visualize.simplify_loc2rot import joints2smpl
+import pyrender
+# import matplotlib.pyplot as plt
+
+import io
+import imageio
+from shapely import geometry
+import trimesh
+from pyrender.constants import RenderFlags
+import math
+# import ffmpeg
+# from PIL import Image
+import hashlib
+import gradio as gr
+
+## load clip model and datasets
+is_cuda = torch.cuda.is_available()
+device = torch.device("cuda" if is_cuda else "cpu")
+print(device)
+clip_model, clip_preprocess = clip.load("ViT-B/32", device=device, jit=False, download_root='./')  # 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=1024, 
+                                clip_dim=args.clip_dim, 
+                                block_size=args.block_size, 
+                                num_layers=9, 
+                                n_head=16, 
+                                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()
+    
+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.eval()
+mean = torch.from_numpy(np.load('./checkpoints/t2m/VQVAEV3_CB1024_CMT_H1024_NRES3/meta/mean.npy'))
+std = torch.from_numpy(np.load('./checkpoints/t2m/VQVAEV3_CB1024_CMT_H1024_NRES3/meta/std.npy'))
+if is_cuda:
+    net.cuda()
+    trans_encoder.cuda()
+    mean = mean.cuda()
+    std = std.cuda()
+
+
+def render(motions, device_id=0, name='test_vis'):
+    frames, njoints, nfeats = motions.shape
+    MINS = motions.min(axis=0).min(axis=0)
+    MAXS = motions.max(axis=0).max(axis=0)
+
+    height_offset = MINS[1]
+    motions[:, :, 1] -= height_offset
+    trajec = motions[:, 0, [0, 2]]
+    is_cuda = torch.cuda.is_available()
+    # device = torch.device("cuda" if is_cuda else "cpu")
+    j2s = joints2smpl(num_frames=frames, device_id=0, cuda=is_cuda)
+    rot2xyz = Rotation2xyz(device=device)
+    faces = rot2xyz.smpl_model.faces
+
+    if not os.path.exists(f'output/{name}_pred.pt'): 
+        print(f'Running SMPLify, it may take a few minutes.')
+        motion_tensor, opt_dict = j2s.joint2smpl(motions)  # [nframes, njoints, 3]
+
+        vertices = rot2xyz(torch.tensor(motion_tensor).clone(), mask=None,
+                                        pose_rep='rot6d', translation=True, glob=True,
+                                        jointstype='vertices',
+                                        vertstrans=True)
+        vertices = vertices.detach().cpu()
+        torch.save(vertices, f'output/{name}_pred.pt')
+    else:
+        vertices = torch.load(f'output/{name}_pred.pt')
+    frames = vertices.shape[3] # shape: 1, nb_frames, 3, nb_joints
+    print(vertices.shape)
+    MINS = torch.min(torch.min(vertices[0], axis=0)[0], axis=1)[0]
+    MAXS = torch.max(torch.max(vertices[0], axis=0)[0], axis=1)[0]
+
+    out_list = []
+    
+    minx = MINS[0] - 0.5
+    maxx = MAXS[0] + 0.5
+    minz = MINS[2] - 0.5 
+    maxz = MAXS[2] + 0.5
+    polygon = geometry.Polygon([[minx, minz], [minx, maxz], [maxx, maxz], [maxx, minz]])
+    polygon_mesh = trimesh.creation.extrude_polygon(polygon, 1e-5)
+
+    vid = []
+    for i in range(frames):
+        if i % 10 == 0:
+            print(i)
+
+        mesh = Trimesh(vertices=vertices[0, :, :, i].squeeze().tolist(), faces=faces)
+
+        base_color = (0.11, 0.53, 0.8, 0.5)
+        ## OPAQUE rendering without alpha
+        ## BLEND rendering consider alpha 
+        material = pyrender.MetallicRoughnessMaterial(
+            metallicFactor=0.7,
+            alphaMode='OPAQUE',
+            baseColorFactor=base_color
+        )
+
+
+        mesh = pyrender.Mesh.from_trimesh(mesh, material=material)
+
+        polygon_mesh.visual.face_colors = [0, 0, 0, 0.21]
+        polygon_render = pyrender.Mesh.from_trimesh(polygon_mesh, smooth=False)
+
+        bg_color = [1, 1, 1, 0.8]
+        scene = pyrender.Scene(bg_color=bg_color, ambient_light=(0.4, 0.4, 0.4))
+        
+        sx, sy, tx, ty = [0.75, 0.75, 0, 0.10]
+
+        camera = pyrender.PerspectiveCamera(yfov=(np.pi / 3.0))
+
+        light = pyrender.DirectionalLight(color=[1,1,1], intensity=300)
+
+        scene.add(mesh)
+
+        c = np.pi / 2
+
+        scene.add(polygon_render, pose=np.array([[ 1, 0, 0, 0],
+
+        [ 0, np.cos(c), -np.sin(c), MINS[1].cpu().numpy()],
+
+        [ 0, np.sin(c), np.cos(c), 0],
+
+        [ 0, 0, 0, 1]]))
+
+        light_pose = np.eye(4)
+        light_pose[:3, 3] = [0, -1, 1]
+        scene.add(light, pose=light_pose.copy())
+
+        light_pose[:3, 3] = [0, 1, 1]
+        scene.add(light, pose=light_pose.copy())
+
+        light_pose[:3, 3] = [1, 1, 2]
+        scene.add(light, pose=light_pose.copy())
+
+
+        c = -np.pi / 6
+
+        scene.add(camera, pose=[[ 1, 0, 0, (minx+maxx).cpu().numpy()/2],
+
+                                [ 0, np.cos(c), -np.sin(c), 1.5],
+
+                                [ 0, np.sin(c), np.cos(c), max(4, minz.cpu().numpy()+(1.5-MINS[1].cpu().numpy())*2, (maxx-minx).cpu().numpy())],
+
+                                [ 0, 0, 0, 1]
+                                ])
+        
+        # render scene
+        r = pyrender.OffscreenRenderer(960, 960)
+
+        color, _ = r.render(scene, flags=RenderFlags.RGBA)
+        # Image.fromarray(color).save(outdir+'/'+name+'_'+str(i)+'.png')
+
+        vid.append(color)
+
+        r.delete()
+
+    out = np.stack(vid, axis=0)
+    imageio.mimwrite(f'output/results.gif', out, fps=20)
+    del out, vertices
+    return f'output/results.gif'
+
+def predict(clip_text, method='fast'):
+    gc.collect()
+    if torch.cuda.is_available():
+        text = clip.tokenize([clip_text], truncate=True).cuda()
+    else:
+        text = clip.tokenize([clip_text], truncate=True)
+    feat_clip_text = clip_model.encode_text(text).float()
+    index_motion = trans_encoder.sample(feat_clip_text[0:1], False)
+    pred_pose = net.forward_decoder(index_motion)
+    pred_xyz = recover_from_ric((pred_pose*std+mean).float(), 22)
+    output_name = hashlib.md5(clip_text.encode()).hexdigest()
+    if method == 'fast':
+        xyz = pred_xyz.reshape(1, -1, 22, 3)
+        pose_vis = plot_3d.draw_to_batch(xyz.detach().cpu().numpy(), title_batch=None, outname=[f'output/results.gif'])
+        return f'output/results.gif'
+    elif method == 'slow':
+        output_path = render(pred_xyz.detach().cpu().numpy().squeeze(axis=0), device_id=0, name=output_name)
+        return output_path
+
+        
+# ---- Gradio Layout -----
+text_prompt = gr.Textbox(label="Text prompt", lines=1, interactive=True)
+video_out = gr.Video(label="Motion", mirror_webcam=False, interactive=False) 
+demo = gr.Blocks()
+demo.encrypt = False
+
+with demo:
+    gr.Markdown('''
+            <div>
+            <h1 style='text-align: center'>Generating Human Motion from Textual Descriptions with Discrete Representations (T2M-GPT)</h1>
+            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🤗
+            </div>
+        ''')
+    with gr.Row():
+            gr.Markdown('''
+            ### Generate human motion by **T2M-GPT**
+            ##### Step 1. Give prompt text describing human motion
+            ##### Step 2. Choice method to generate output (Fast: Sketch skeleton; Slow: SMPL mesh)
+            ##### Step 3. Generate output and enjoy
+            ''')
+    with gr.Row():         
+        gr.Markdown('''
+            ### You can test by following examples:
+            ''')
+    examples = gr.Examples(examples=
+            [ "a person jogs in place, slowly at first, then increases speed. they then back up and squat down.", 
+              "a man steps forward and does a handstand", 
+              "a man rises from the ground, walks in a circle and sits back down on the ground"],
+          label="Examples", inputs=[text_prompt])
+    
+    with gr.Column():
+        with gr.Row():
+            text_prompt.render()
+            method = gr.Dropdown(["slow", "fast"], label="Method", value="fast")
+        with gr.Row():
+            generate_btn = gr.Button("Generate")
+            generate_btn.click(predict, [text_prompt, method], [video_out])
+            print(video_out)
+        with gr.Row():  
+            video_out.render()
+        
+demo.launch(debug=True)