Update api/mdm_loader.py

api/mdm_loader.py

@@ -1,23 +1,21 @@
 import torch
-from utils.model_util import create_model_and_diffusion, load_saved_model
-from utils.parser_util import generate_args
+from mdm_repo.utils.model_util import create_model_and_diffusion, load_saved_model
+from mdm_repo.utils.parser_util import generate_args
 import os
 
-# Load once on startup
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Load args.json from models directory
-args_path = os.path.join(os.path.dirname(__file__), "..", "models", "args.json")
-args = generate_args()
-if os.path.exists(args_path):
-    import json
-    with open(args_path, "r") as f:
-        args.__dict__.update(json.load(f))
-
-model, diffusion = create_model_and_diffusion(args, data=None)
-model = load_saved_model(model, "models/model000475000.pt")
-model.to(device)
-model.eval()
+def load_model_and_args():
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    args_path = os.path.join(os.path.dirname(__file__), "..", "models", "args.json")
+    args = generate_args()
+    if os.path.exists(args_path):
+        import json
+        with open(args_path, "r") as f:
+            args.__dict__.update(json.load(f))
+    model, diffusion = create_model_and_diffusion(args, data=None)
+    model = load_saved_model(model, "models/model000475000.pt")
+    model.to(device)
+    model.eval()
+    return model, diffusion, args, device
 
 def generate_motion(prompt: str, num_frames: int, style: str = "default"):
     # TODO: Implement motion generation using MDM repo's sampling logic