Update app.py

app.py

@@ -1,22 +1,43 @@
 import gradio as gr
+import clip
 import torch
 import numpy as np
+import models.vqvae as vqvae
+import models.t2m_trans as trans
 import matplotlib.pyplot as plt
 from matplotlib.animation import FuncAnimation
 from mpl_toolkits.mplot3d import Axes3D
+import sys
+import os
+from matplotlib.animation import FuncAnimation
+from mpl_toolkits.mplot3d import Axes3D
 from mpl_toolkits.mplot3d.art3d import Poly3DCollection
 
+import warnings
+warnings.filterwarnings('ignore')
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-vqvae_model = VQVAE().to(device)
-transformer_model = Transformer().to(device)
+vqvae_model = vqvae.VQVAE().to(device)
+transformer_model = trans.Text2Motion_Transformer().to(device)
 vqvae_model.load_state_dict(torch.load("output/VQVAE_imp_resnet_100k_hml3d/net_last.pth", map_location=device))
 transformer_model.load_state_dict(torch.load("output/net_best_fid.pth", map_location=device))
 vqvae_model.eval()
 transformer_model.eval()
 
+mean = torch.from_numpy(np.load('output/Mean.npy', allow_pickle=True)).to(device)
+std = torch.from_numpy(np.load('output/Std.npy', allow_pickle=True)).to(device)
+
+def generate_motion(text, vqvae_model, transformer_model):
+    clip_text = [text]
+    text_encoded = clip.tokenize(clip_text, truncate=True).to(device)
+    with torch.no_grad():
+        motion_indices = transformer_model.sample(text_encoded, False)
+        pred_pose = vqvae_model.forward_decoder(motion_indices)
+        pred_xyz = recover_from_ric((pred_pose * std + mean).float(), 22)
+    return pred_xyz.cpu().numpy().reshape(-1, 22, 3)
 
-def create_animation(joints, title="3D Motion", save_path="animation.gif"):
+def create_animation(joints, title="3D Motion", save_path="static/animation.gif"):
     fig = plt.figure(figsize=(10, 10))
     ax = fig.add_subplot(111, projection='3d')
     data = np.array(joints).T
@@ -28,31 +49,31 @@ def create_animation(joints, title="3D Motion", save_path="animation.gif"):
         return line,
 
     ani = FuncAnimation(fig, update, frames=len(joints), fargs=(data, line), interval=50, blit=True)
-    ani.save(save_path, writer='pillow')
+    ani.save(save_path, writer=PillowWriter(fps=20))
     plt.close(fig)
     return save_path
-
-def infer(text):
-    motion_data = generate_motion(text, vqvae_model, transformer_model)
-    gif_path = create_animation(motion_data, kinematic_tree)
-    return gif_path
-
+    
 examples = [
     "Person doing yoga",
     "A person is dancing ballet",
 ]
 
+def infer(text):
+    motion_data = generate_motion(text, vqvae_model, transformer_model)
+    gif_path = create_animation(motion_data)
+    return gif_path
+
 with gr.Blocks(css=".container { max-width: 800px; margin: auto; }") as demo:
     with gr.Column():
-        gr.Markdown("## 3D Motion Generation")
-        text_input = gr.Textbox(label="Describe the action", placeholder="Enter text description of the action here...")
+        gr.Markdown("## 3D Motion Generation on " + ("GPU" if device == "cuda" else "CPU"))
+        text_input = gr.Textbox(label="Describe the action", placeholder="Enter text description for the action here...")
         output_image = gr.Image(label="Generated Motion Animation")
         submit_button = gr.Button("Generate Motion")
         
         submit_button.click(
             fn=infer,
-            inputs=text_input,
-            outputs=output_image
+            inputs=[text_input],
+            outputs=[output_image]
         )
 
 if __name__ == "__main__":