Update app.py

app.py

@@ -20,8 +20,33 @@ warnings.filterwarnings('ignore')
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-vqvae_model = vqvae.VQVAE().to(device)
-transformer_model = trans.Text2Motion_Transformer().to(device)
+# Create an args object for initializing the HumanVQVAE
+class Args:
+    def __init__(self):
+        self.dataname = 't2m'  # example property, adjust as needed
+        self.quantizer = 'ema'  # Set the type of quantizer used in VQVAE_251
+        # Add other properties required by HumanVQVAE initialization
+
+args = Args()
+
+
+vqvae_model = vqvae.HumanVQVAE(args,
+                       args.nb_code,
+                       args.code_dim,
+                       args.output_emb_width,
+                       args.down_t,
+                       args.stride_t,
+                       args.width,
+                       args.depth,
+                       args.dilation_growth_rate).to(device)
+transformer_model = 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).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()