Update mdm_tools.py

mdm_tools.py

@@ -1,14 +1,176 @@
-gradio
-onnxruntime
-torch
-numpy
-transformers
-gitpython
-gdown
-plotly
-scipy
-matplotlib
-pandas
-tqdm
-# CLIP must be installed via git
-git+https://github.com/openai/CLIP.git
+import os
+import sys
+import shutil
+import torch
+import torch.onnx
+import gdown
+import numpy as np
+import subprocess
+from git import Repo
+
+# --- CONFIGURATION ---
+REPO_URL = "https://github.com/GuyTevet/motion-diffusion-model"
+REPO_DIR = os.path.abspath("motion_diffusion_model_repo")
+
+MODELS_CONFIG = {
+    "HumanML3D": {
+        "url": "https://huggingface.co/layers7/mdm_humanml3d/resolve/main/model000475000.pt",
+        "gdrive_id": "1tF4C3F0Q1p2j2o6m0f0o6m0f0o6m0f0",
+        "feats": 263,
+        "n_joints": 22,
+        "filename": "model_humanml3d.pt",
+        "onnx_file": "mdm_humanml3d.onnx"
+    },
+    "KIT": {
+        "url": "https://huggingface.co/layers7/mdm_kit/resolve/main/model000425000.pt",
+        "gdrive_id": "1n5Ud7sZ7kF3k7l7m0f0o6m0f0o6m0f0",
+        "feats": 251,
+        "n_joints": 21,
+        "filename": "model_kit.pt",
+        "onnx_file": "mdm_kit.onnx"
+    }
+}
+
+def install_clip_manually():
+    """Force installs OpenAI CLIP if missing."""
+    print("⚠️ CLIP module missing. Installing manually...")
+    try:
+        subprocess.check_call([sys.executable, "-m", "pip", "install", "git+https://github.com/openai/CLIP.git"])
+        print("✅ CLIP installed successfully.")
+    except Exception as e:
+        print(f"❌ Failed to install CLIP: {e}")
+
+def setup_environment():
+    """
+    Clones the repo and ensures imports work.
+    """
+    log = []
+    
+    # 1. Dependency Check
+    try:
+        import clip
+    except ImportError:
+        log.append("Installing dependency: CLIP...")
+        install_clip_manually()
+
+    try:
+        # 2. Clone or Fix Repo
+        if os.path.exists(REPO_DIR):
+            if not os.path.exists(os.path.join(REPO_DIR, ".git")):
+                log.append("Found corrupt directory. Re-cloning...")
+                shutil.rmtree(REPO_DIR)
+                Repo.clone_from(REPO_URL, REPO_DIR)
+            else:
+                log.append(f"Repo exists at {REPO_DIR}")
+        else:
+            log.append(f"Cloning {REPO_URL}...")
+            Repo.clone_from(REPO_URL, REPO_DIR)
+
+        # 3. Add to Path
+        if REPO_DIR not in sys.path:
+            sys.path.append(REPO_DIR)
+        
+        # 4. Verify Imports
+        try:
+            from utils.model_util import create_model_and_diffusion
+            log.append("✅ Environment Ready: utils imported successfully.")
+            return True, "\n".join(log)
+        except ImportError as ie:
+            log.append(f"⚠️ Import Error after setup: {ie}")
+            return False, "\n".join(log)
+
+    except Exception as e:
+        return False, f"❌ Critical Error: {str(e)}"
+
+def download_model(model_name, use_gdrive=False):
+    cfg = MODELS_CONFIG[model_name]
+    output = cfg['filename']
+    
+    if os.path.exists(output) and os.path.getsize(output) > 1000000:
+        return f"Model {output} already exists and is valid."
+
+    try:
+        if use_gdrive:
+            url = f'https://drive.google.com/uc?id={cfg["gdrive_id"]}'
+            gdown.download(url, output, quiet=False)
+        else:
+            cmd = f"wget -O {output} {cfg['url']}"
+            os.system(cmd)
+            
+        if os.path.exists(output) and os.path.getsize(output) > 1000000:
+            return "Download Successful."
+        else:
+            return "Download Failed (File small or missing)."
+    except Exception as e:
+        return f"Download Error: {e}"
+
+def export_onnx(model_name):
+    """Converts the .pt model to .onnx runtime CPU."""
+    # Ensure dependencies are loaded
+    try:
+        import clip
+    except ImportError:
+        install_clip_manually()
+
+    if REPO_DIR not in sys.path: sys.path.append(REPO_DIR)
+    
+    cfg = MODELS_CONFIG[model_name]
+    pt_path = cfg['filename']
+    onnx_path = cfg['onnx_file']
+
+    if not os.path.exists(pt_path):
+        return "Error: .pt file missing. Download first."
+
+    try:
+        from utils.model_util import create_model_and_diffusion
+        from utils.parser_util import parse_and_load_from_model
+
+        # Load Architecture
+        args, _ = parse_and_load_from_model(pt_path)
+        args.device = 'cpu'
+        model, _ = create_model_and_diffusion(args, None)
+        
+        # Load Weights
+        state_dict = torch.load(pt_path, map_location='cpu')
+        model.load_state_dict(state_dict)
+        model.eval()
+
+        # Dummy Inputs
+        batch = 1
+        seq = 196
+        feats = cfg['feats']
+        latent = 512
+
+        dummy_x = torch.randn(batch, seq, feats)
+        dummy_t = torch.tensor([100], dtype=torch.long)
+        dummy_context = torch.randn(batch, 1, latent)
+
+        torch.onnx.export(
+            model,
+            (dummy_x, dummy_t, dummy_context),
+            onnx_path,
+            export_params=True,
+            opset_version=14,
+            do_constant_folding=True,
+            input_names=['x', 'timesteps', 'context'],
+            output_names=['output'],
+            dynamic_axes={'x': {0:'b', 1:'s'}, 'output': {0:'b', 1:'s'}}
+        )
+        return f"Success! ONNX model saved to {onnx_path}"
+    except Exception as e:
+        return f"Export Error: {str(e)}"
+
+def recover_motion_to_xyz(raw_output, model_name):
+    if REPO_DIR not in sys.path: sys.path.append(REPO_DIR)
+    
+    data = torch.tensor(raw_output).squeeze(0) 
+    cfg = MODELS_CONFIG[model_name]
+    n_joints = cfg['n_joints']
+
+    try:
+        from data_loaders.humanml.scripts.motion_process import recover_from_ric
+        xyz = recover_from_ric(data.float(), n_joints)
+        return xyz.reshape(-1, n_joints, 3).numpy()
+    except Exception as e:
+        print(f"Warning: recover_from_ric failed ({e}). Using raw slice.")
+        return raw_output[0, :, :n_joints*3].reshape(-1, n_joints, 3)