Update api/mdm_loader.py

api/mdm_loader.py

@@ -60,16 +60,17 @@ def generate_motion(prompt: str, num_frames: int, style: str = "default"):
     # Setup device
     dist_util.setup_dist(args.device)
 
-    # Load model and diffusion
-    model, diffusion = create_model_and_diffusion(args, data=None)
-    load_saved_model(model, args.model_path, use_avg=getattr(args, 'use_ema', False))
-    model.to(dist_util.dev())
-    model.eval()
-
     # Prepare dataset and model_kwargs
     max_frames = 196 if args.dataset in ['kit', 'humanml'] else 60
     n_frames = min(max_frames, int(args.motion_length * 20))
     data = get_dataset_loader(name=args.dataset, batch_size=1, num_frames=max_frames, split='test', hml_mode='text_only', fixed_len=0, pred_len=0, device=dist_util.dev())
+
+    # Load model and diffusion with dataset
+    model, diffusion = create_model_and_diffusion(args, data)
+    load_saved_model(model, args.model_path, use_avg=getattr(args, 'use_ema', False))
+    model.to(dist_util.dev())
+    model.eval()
+
     iterator = iter(data)
     _, model_kwargs = next(iterator)
     model_kwargs['y']['text'] = [prompt]