Create standalone inference script for MDM demo

app.py

@@ -7,47 +7,156 @@ from pathlib import Path
 import traceback
 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
+    
+    motion = np.zeros((frames, joints, dimensions))
+    
+    # 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
+        
+        # Leg and arm swing
+        swing = np.sin(t * 2 * np.pi * 2)  # Two cycles over the motion length
+        
+        # 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
+        
+    return motion
+
+# Visualize the motion
+def visualize_motion(motion_data, output_path):
+    frames, joints, dims = motion_data.shape
+    
+    # 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')
+        
+        # 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-')
+        
+        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")
 
 def text_to_motion(text_prompt, motion_length=3.0, seed=0):
     """Generate motion from text prompt using MDM"""
     try:
-        # First, check if we need to clone the repository
-        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"])
-        
-        # Let's examine the repository structure
-        print("Repository contents:")
-        for path in glob.glob("motion-diffusion-model/*"):
-            print(f"- {path}")
-        
-        if Path("motion-diffusion-model/sample").exists():
-            print("Sample directory contents:")
-            for path in glob.glob("motion-diffusion-model/sample/*"):
-                print(f"- {path}")
+        # Ensure MDM scripts are set up
+        setup_mdm()
         
         # Get absolute path to the checkpoint
         checkpoint_path = os.path.abspath("checkpoints/mld_humanml.pt")
         print(f"Checkpoint path: {checkpoint_path}")
         
-        # Set the working directory to the repository
-        os.chdir("motion-diffusion-model")
-        
-        # Let's see what files are in the current directory
-        print("Current directory contents:")
-        for path in glob.glob("*"):
-            print(f"- {path}")
-        
-        # Try to use the python module directly
+        # Run the inference script
         cmd = [
             "python", 
-            "-m", "sample.sample_text", 
+            "mdm_inference.py",
             "--model_path", checkpoint_path,
             "--text_prompt", text_prompt,
             "--motion_length", str(motion_length),
-            "--num_samples", "1",
             "--seed", str(int(seed))
         ]
         
@@ -59,19 +168,12 @@ def text_to_motion(text_prompt, motion_length=3.0, seed=0):
         if result.stderr:
             print("Command error:", result.stderr)
         
-        # Return to the original directory
-        os.chdir("..")
-        
-        # 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
+        # Check if the output file exists
+        if Path("output.mp4").exists():
+            print("Found output.mp4 file")
+            return "output.mp4"
         
-        print("No MP4 output files found.")
+        print("No output.mp4 file found.")
         return None
     except Exception as e:
         print(f"Error generating motion: {str(e)}")