Update viz_utils.py

viz_utils.py

@@ -1,117 +1,176 @@
+import os
+import sys
+import shutil
+import torch
+import torch.onnx
+import gdown
 import numpy as np
-import plotly.graph_objects as go
+import subprocess
+from git import Repo
 
-# HumanML3D Skeleton Topology (Parent Indices)
-# -1 is Root
-CHAIN = [-1, 0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 9, 12, 13, 14, 16, 17, 18, 19]
+# --- CONFIGURATION ---
+REPO_URL = "https://github.com/GuyTevet/motion-diffusion-model"
+REPO_DIR = os.path.abspath("motion_diffusion_model_repo")
 
-# Joint Names matching HumanML3D order
-JOINT_NAMES = [
-    "Hips", "RightUpLeg", "RightLeg", "RightFoot", "LeftUpLeg", "LeftLeg", "LeftFoot",
-    "Spine", "Spine1", "Spine2", "Neck", "Head", "RightShoulder", "RightArm", "RightForeArm",
-    "RightHand", "LeftShoulder", "LeftArm", "LeftForeArm", "LeftHand", "RightToeBase", "LeftToeBase"
-]
+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 save_bvh_for_blender(xyz_data, filename):
+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():
     """
-    Exports a BVH with Position Channels on all nodes.
-    This ensures Blender imports the exact shape generated by the AI 
-    without requiring Inverse Kinematics.
-    
-    xyz_data: (Seq, Joints, 3)
+    Clones the repo and ensures imports work.
     """
-    seq_len, n_joints, _ = xyz_data.shape
+    log = []
     
-    with open(filename, 'w') as f:
-        f.write("HIERARCHY\n")
-        
-        # Recursive function to write hierarchy
-        def write_joint(idx, indent_level):
-            indent = "\t" * indent_level
-            name = JOINT_NAMES[idx] if idx < len(JOINT_NAMES) else f"Joint_{idx}"
-            
-            # Type
-            if idx == 0:
-                f.write(f"{indent}ROOT {name}\n")
-            else:
-                f.write(f"{indent}JOINT {name}\n")
-            
-            f.write(f"{indent}{{\n")
-            f.write(f"{indent}\tOFFSET 0.00 0.00 0.00\n") # We drive pos via channels, so offset is 0
-            
-            # CRITICAL: We enable Position channels for EVERY bone.
-            # Standard BVH only has rotation for children, but Blender supports this.
-            f.write(f"{indent}\tCHANNELS 3 Xposition Yposition Zposition\n")
-            
-            # Find children
-            children = [i for i, p in enumerate(CHAIN) if p == idx]
-            
-            if len(children) > 0:
-                for child in children:
-                    write_joint(child, indent_level + 1)
+    # 1. Dependency Check (Fix for 'No module named clip')
+    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:
-                # End Site
-                f.write(f"{indent}\tEnd Site\n")
-                f.write(f"{indent}\t{{\n")
-                f.write(f"{indent}\t\tOFFSET 0.00 0.00 0.00\n")
-                f.write(f"{indent}\t}}\n")
-                
-            f.write(f"{indent}}}\n")
-
-        # Start recursion
-        write_joint(0, 0)
-        
-        f.write("MOTION\n")
-        f.write(f"Frames: {seq_len}\n")
-        f.write("Frame Time: 0.050000\n") # 20 FPS standard for HumanML3D
+                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)
         
-        # Write Data
-        for i in range(seq_len):
-            # Flatten frame: Joint1_X J1_Y J1_Z Joint2_X ...
-            row = xyz_data[i].flatten()
-            line = " ".join([f"{val:.6f}" for val in row])
-            f.write(line + "\n")
-            
-    return filename
+        # 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)
 
-def plot_skeleton(xyz_data):
-    """Creates a Plotly 3D animation."""
-    # Downsample for web smoothness
-    xyz = xyz_data[::4] 
-    
-    # Center the motion on the first frame's root
-    root_offset = xyz[0:1, 0:1, :]
-    xyz = xyz - root_offset
+    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']
     
-    frames = []
-    for k in range(len(xyz)):
-        x_p, y_p, z_p = [], [], []
-        for child, parent in enumerate(CHAIN):
-            if parent == -1 or child >= xyz.shape[1]: continue
+    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)
             
-            # 3D Line segment
-            x_p.extend([xyz[k, parent, 0], xyz[k, child, 0], None])
-            # Swap Y-Z for visualization (Plotly Z is up, MDM Y is up)
-            y_p.extend([xyz[k, parent, 2], xyz[k, child, 2], None]) 
-            z_p.extend([xyz[k, parent, 1], xyz[k, child, 1], None])
+        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)
         
-        frames.append(go.Frame(data=[go.Scatter3d(
-            x=x_p, y=y_p, z=z_p, mode='lines+markers',
-            marker=dict(size=4, color='red'), line=dict(color='blue', width=4)
-        )]))
+        # 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)
     
-    fig = go.Figure(
-        data=frames[0].data,
-        layout=go.Layout(
-            title="Generated Motion Preview",
-            scene=dict(
-                xaxis=dict(range=[-2, 2], autorange=False),
-                yaxis=dict(range=[-2, 2], autorange=False),
-                zaxis=dict(range=[-0, 2], autorange=False),
-                aspectmode='cube'
-            ),
-            updatemenus=[dict(type="buttons", buttons=[dict(label="Play", method="animate", args=[None])])]
-        ),
-        frames=frames
-    )
-    return fig
+    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)