Update app to use proper MDM checkpoint and setup

app.py

@@ -9,157 +9,77 @@ import subprocess
 import glob
 import requests
 
-def setup_mdm():
-    """Set up the MDM repository and files"""
-    # Create a simple inference script
-    inference_script_content = """
-import torch
-import numpy as np
-import argparse
-import os
-import imageio
-import matplotlib.pyplot as plt
-from mpl_toolkits.mplot3d import Axes3D
-from matplotlib.animation import FuncAnimation
-
-# Parse arguments
-parser = argparse.ArgumentParser()
-parser.add_argument("--model_path", type=str, required=True)
-parser.add_argument("--text_prompt", type=str, required=True)
-parser.add_argument("--motion_length", type=float, default=3.0)
-parser.add_argument("--seed", type=int, default=0)
-args = parser.parse_args()
-
-# Mock function to generate simple motion data for testing
-def generate_mock_motion(text_prompt, motion_length, seed):
-    np.random.seed(seed)
-    print(f"Generating motion for: {text_prompt}")
-    # Create a simple walking motion
-    frames = int(motion_length * 30)  # 30 fps
-    joints = 24  # Number of joints in a typical skeleton
-    dimensions = 3  # x, y, z
+def download_checkpoint():
+    """Download the recommended checkpoint if not present"""
+    # Create checkpoints directory
+    os.makedirs("checkpoints", exist_ok=True)
     
-    motion = np.zeros((frames, joints, dimensions))
+    # Define the target checkpoint path
+    checkpoint_path = "checkpoints/humanml_trans_enc_512.pt"
     
-    # Create a simple walking motion - pendulum motion for legs and arms
-    for frame in range(frames):
-        t = frame / frames
-        
-        # Basic forward motion
-        motion[frame, :, 0] = t * 2  # Move forward on X axis
+    if not Path(checkpoint_path).exists():
+        print(f"Downloading checkpoint to {checkpoint_path}...")
+        # URL for the checkpoint from HuggingFace or direct link
+        url = "https://huggingface.co/spaces/mohaed/testMDM/resolve/main/checkpoints/mld_humanml.pt"
         
-        # Leg and arm swing
-        swing = np.sin(t * 2 * np.pi * 2)  # Two cycles over the motion length
+        # Download the file
+        response = requests.get(url, stream=True)
+        response.raise_for_status()
         
-        # Left leg, right leg, left arm, right arm
-        joint_indices = [4, 7, 16, 20]
-        for ji, joint_idx in enumerate(joint_indices):
-            # Alternate phase for left/right sides
-            phase = 0 if ji % 2 == 0 else np.pi
-            motion[frame, joint_idx, 2] = np.sin(t * 2 * np.pi * 2 + phase) * 0.5
+        with open(checkpoint_path, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
         
-    return motion
-
-# Visualize the motion
-def visualize_motion(motion_data, output_path):
-    frames, joints, dims = motion_data.shape
+        print(f"Checkpoint downloaded to {checkpoint_path}")
+    else:
+        print(f"Checkpoint already exists at {checkpoint_path}")
     
-    # Create a simple stick figure animation
-    fig = plt.figure(figsize=(10, 10))
-    ax = fig.add_subplot(111, projection='3d')
-    
-    # Define connections between joints (simplified)
-    connections = [
-        (0, 1), (1, 2), (2, 3),  # Spine
-        (0, 4), (4, 5), (5, 6),  # Left leg
-        (0, 7), (7, 8), (8, 9),  # Right leg
-        (3, 10), (10, 11), (11, 12),  # Left arm
-        (3, 13), (13, 14), (14, 15)   # Right arm
-    ]
-    
-    def update(frame):
-        ax.clear()
-        
-        # Set the axis limits
-        ax.set_xlim([-2, 4])
-        ax.set_ylim([-2, 2])
-        ax.set_zlim([-2, 2])
-        
-        # Plot joints
-        ax.scatter(motion_data[frame, :, 0], 
-                  motion_data[frame, :, 1], 
-                  motion_data[frame, :, 2], c='b', marker='o')
+    return checkpoint_path
+
+def clone_mdm_repo():
+    """Clone the MDM repository if not present"""
+    if not Path("motion-diffusion-model").exists():
+        print("Cloning Motion Diffusion Model repository...")
+        subprocess.run(["git", "clone", "https://github.com/GuyTevet/motion-diffusion-model.git"])
+        print("Installing Spacy language model...")
+        subprocess.run(["python", "-m", "spacy", "download", "en_core_web_sm"])
         
-        # Plot connections
-        for start, end in connections:
-            ax.plot([motion_data[frame, start, 0], motion_data[frame, end, 0]],
-                   [motion_data[frame, start, 1], motion_data[frame, end, 1]],
-                   [motion_data[frame, start, 2], motion_data[frame, end, 2]], 'r-')
+        # Download additional required files
+        print("Downloading additional required files...")
+        os.chdir("motion-diffusion-model")
+        subprocess.run(["bash", "prepare/download_smpl_files.sh"])
+        subprocess.run(["bash", "prepare/download_glove.sh"])
+        subprocess.run(["bash", "prepare/download_t2m_evaluators.sh"])
+        os.chdir("..")
         
-        ax.set_title(f"Frame {frame}")
-        return ax
-    
-    # Create animation
-    anim = FuncAnimation(fig, update, frames=frames, interval=1000/30)
-    
-    # Save animation to mp4
-    os.makedirs(os.path.dirname(output_path) or '.', exist_ok=True)
-    anim.save(output_path, writer='ffmpeg', fps=30)
-    plt.close()
-    
-    return output_path
-
-# Main function
-def main():
-    print(f"Processing text prompt: {args.text_prompt}")
-    print(f"Using model: {args.model_path}")
-    print(f"Motion length: {args.motion_length}")
-    print(f"Seed: {args.seed}")
-    
-    # Generate motion
-    motion_data = generate_mock_motion(
-        args.text_prompt, 
-        args.motion_length, 
-        args.seed
-    )
-    
-    # Visualize and save
-    output_path = "output.mp4"
-    visualize_motion(motion_data, output_path)
-    print(f"Saved animation to {output_path}")
-    
-    return output_path
-
-if __name__ == "__main__":
-    main()
-"""
-
-    # Write the inference script to file
-    with open("mdm_inference.py", "w") as f:
-        f.write(inference_script_content)
-    
-    print("Created MDM inference script")
+        print("MDM repository setup complete")
+    else:
+        print("MDM repository already exists")
 
 def text_to_motion(text_prompt, motion_length=3.0, seed=0):
     """Generate motion from text prompt using MDM"""
     try:
-        # Ensure MDM scripts are set up
-        setup_mdm()
+        # Clone the MDM repository
+        clone_mdm_repo()
+        
+        # Download the recommended checkpoint
+        checkpoint_path = download_checkpoint()
+        
+        # Write a simple run script that properly sets up the environment
+        run_script = """
+#!/bin/bash
+cd motion-diffusion-model
+export PYTHONPATH=$PYTHONPATH:$(pwd)
+python -m sample.generate --model_path ../checkpoints/humanml_trans_enc_512.pt --text_prompt "$1" --motion_length $2 --seed $3 --num_samples 1
+"""
         
-        # Get absolute path to the checkpoint
-        checkpoint_path = os.path.abspath("checkpoints/mld_humanml.pt")
-        print(f"Checkpoint path: {checkpoint_path}")
+        with open("run_mdm.sh", "w") as f:
+            f.write(run_script)
         
-        # Run the inference script
-        cmd = [
-            "python", 
-            "mdm_inference.py",
-            "--model_path", checkpoint_path,
-            "--text_prompt", text_prompt,
-            "--motion_length", str(motion_length),
-            "--seed", str(int(seed))
-        ]
+        subprocess.run(["chmod", "+x", "run_mdm.sh"])
         
+        # Run the script
+        cmd = ["./run_mdm.sh", text_prompt, str(motion_length), str(int(seed))]
         print(f"Running command: {' '.join(cmd)}")
         result = subprocess.run(cmd, capture_output=True, text=True)
         
@@ -168,12 +88,16 @@ def text_to_motion(text_prompt, motion_length=3.0, seed=0):
         if result.stderr:
             print("Command error:", result.stderr)
         
-        # Check if the output file exists
-        if Path("output.mp4").exists():
-            print("Found output.mp4 file")
-            return "output.mp4"
+        # Check for output files
+        print("Checking for output files:")
+        for root, dirs, files in os.walk("."):
+            for file in files:
+                if file.endswith(".mp4"):
+                    path = os.path.join(root, file)
+                    print(f"Found video file: {path}")
+                    return path
         
-        print("No output.mp4 file found.")
+        print("No MP4 output files found.")
         return None
     except Exception as e:
         print(f"Error generating motion: {str(e)}")