Update app.py

app.py

@@ -49,12 +49,12 @@ transformer_model = trans.Text2Motion_Transformer(num_vq=args.nb_code,
                                                   fc_rate=args.ff_rate).to(device)
 
 vqvae_checkpoint = torch.load("output/VQVAE_imp_resnet_100k_hml3d/net_last.pth", map_location=device)
-transformed_state_dict = {}
+transformed_vqvae_state_dict = {}
 for k, v in vqvae_checkpoint['net'].items():
-    new_key = k.replace("vqvae.", "")
-    transformed_state_dict[new_key] = v
+    new_key = k.replace("vqvae.", "") 
+    transformed_vqvae_state_dict[new_key] = v
 
-vqvae_model.load_state_dict(transformed_state_dict, strict=False)
+vqvae_model.load_state_dict(transformed_vqvae_state_dict, strict=False)
 
 transformer_checkpoint = torch.load("output/VQVAE_imp_resnet_100k_hml3d/net_best_fid.pth", map_location=device)
 transformed_transformer_state_dict = {}
@@ -74,7 +74,7 @@ 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)
+        motion_indices = transformer_model.sample(text_encoded.float(), 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)
@@ -94,11 +94,6 @@ def create_animation(joints, title="3D Motion", save_path="static/animation.gif"
     ani.save(save_path, writer=PillowWriter(fps=20))
     plt.close(fig)
     return save_path
-    
-examples = [
-    "A person doing a kick",
-    "A person is dancing ballet",
-]
 
 def infer(text):
     motion_data = generate_motion(text, vqvae_model, transformer_model)