Update app to use proper MDM checkpoint and setup

app.py

@@ -1,29 +1,4 @@
 import gradio as gr
-import torch
-import os
-import sys
-import numpy as np
-from pathlib import Path
-import traceback
-import subprocess
-import glob
-import requests
-import time
-import random
-
-def create_motion_animation(text_prompt, motion_length=3.0, seed=0):
-    """Create a motion animation based on input parameters"""
-    print(f"Creating animation for: '{text_prompt}', length: {motion_length}s, seed: {seed}")
-    
-    # Create a unique filename based on parameters
-    output_filename = f"output_{abs(hash(text_prompt) % 10000)}_{int(motion_length)}_{seed}.mp4"
-    
-    # Sanitize the text prompt for Python string
-    safe_prompt = text_prompt.replace('"', '')
-    
-    # Create a simple visualization script
-    with open("motion_animation.py", "w") as f:
-        f.write("""
 import numpy as np
 import matplotlib.pyplot as plt
 from matplotlib.animation import FuncAnimation
@@ -31,71 +6,60 @@ import os
 from mpl_toolkits.mplot3d import Axes3D
 import random
 
-# Set random seeds for reproducibility
-np.random.seed({})
-random.seed({})
-
-# Animation parameters
-text_prompt = "{}"
-walking = "walk" in text_prompt.lower()
-running = "run" in text_prompt.lower()
-jumping = "jump" in text_prompt.lower()
-dancing = "danc" in text_prompt.lower()
-turning = "turn" in text_prompt.lower() or "spin" in text_prompt.lower()
-waving = "wave" in text_prompt.lower()
-
-# Create a motion based on text prompt
-def generate_motion(frames={}):
-    joints = 16  # Number of joints in a simplified skeleton
-    dims = 3  # x, y, z
-    
+def generate_motion(text_prompt, motion_length, seed):
+    """Generate a motion animation based on the text prompt"""
+    # Use the seed for reproducibility
+    np.random.seed(seed)
+    random.seed(seed)
+    
+    # Parse the text prompt to detect actions
+    text_lower = text_prompt.lower()
+    walking = "walk" in text_lower
+    running = "run" in text_lower
+    jumping = "jump" in text_lower
+    dancing = "danc" in text_lower
+    turning = "turn" in text_lower or "spin" in text_lower
+    waving = "wave" in text_lower
+    
+    # Set speed and other parameters based on the action
+    speed = 4.0 if running else 2.0 if walking else 1.0
+    frames = int(motion_length * 30)  # 30 fps
+    
+    # Create motion data - 16 joints with 3D coordinates
+    joints = 16
+    dims = 3
     motion = np.zeros((frames, joints, dims))
     
-    # Set speed based on motion type
-    if running:
-        speed = 4.0
-    elif walking:
-        speed = 2.0
-    else:
-        speed = 1.0
-    
-    # Create the motion
+    # Generate the motion
     for frame in range(frames):
         t = frame / frames
         
-        # Basic forward motion
+        # Basic forward motion or turning
         if turning:
-            # Move in a circle
             angle = t * 2 * np.pi * 2
             motion[frame, :, 0] = np.cos(angle) * 2
             motion[frame, :, 1] = np.sin(angle) * 2
         else:
-            # Move forward
             motion[frame, :, 0] = t * speed * 4
         
-        # Root joint (pelvis)
+        # Root joint (pelvis) with jumping or bouncing
         if jumping:
-            # Add jumping motion
-            jump_height = 0.5 + 0.5 * np.sin(t * 2 * np.pi * 3)
-            motion[frame, 0, 2] = jump_height
+            motion[frame, 0, 2] = 0.5 + 0.5 * np.sin(t * 2 * np.pi * 3)
         else:
-            # Regular walking bounce
-            bounce = 0.1 * np.sin(t * 2 * np.pi * speed * 2) if walking or running else 0.05
-            motion[frame, 0, 2] = bounce + 1
+            motion[frame, 0, 2] = 0.1 * np.sin(t * 2 * np.pi * speed * 2) + 1 if walking or running else 0.05 + 1
         
         # Spine and head (joints 1, 2, 3)
         for i in range(1, 4):
-            motion[frame, i, 2] = motion[frame, 0, 2] + i * 0.2  # Stack joints vertically
-        
-        # Add dancing motion if needed
-        if dancing:
-            wiggle = 0.2 * np.sin(t * 2 * np.pi * 4 + np.pi * i/4)
-            motion[frame, 1:4, 1] = wiggle  # Side-to-side motion for upper body
+            motion[frame, i, 2] = motion[frame, 0, 2] + i * 0.2
+            
+            # Add dancing motion for upper body
+            if dancing:
+                motion[frame, i, 1] = 0.2 * np.sin(t * 2 * np.pi * 4 + np.pi * i/4)
         
         # Left leg (joints 4, 5, 6)
-        leg_freq = speed * 2  # Frequency of leg movement
+        leg_freq = speed * 2
         swing_leg_l = np.sin(t * 2 * np.pi * leg_freq)
-        motion[frame, 4, 1] = 0.2  # Left hip position
+        motion[frame, 4, 1] = 0.2
         motion[frame, 4, 2] = motion[frame, 0, 2] - 0.1
         motion[frame, 5, 1] = 0.2
         motion[frame, 5, 2] = motion[frame, 4, 2] - 0.5 + swing_leg_l * 0.3
@@ -103,8 +67,8 @@ def generate_motion(frames={}):
         motion[frame, 6, 2] = motion[frame, 5, 2] - 0.5 + swing_leg_l * 0.3
         
         # Right leg (joints 7, 8, 9)
-        swing_leg_r = np.sin(t * 2 * np.pi * leg_freq + np.pi)  # Opposite phase
-        motion[frame, 7, 1] = -0.2  # Right hip position
+        swing_leg_r = np.sin(t * 2 * np.pi * leg_freq + np.pi)
+        motion[frame, 7, 1] = -0.2
         motion[frame, 7, 2] = motion[frame, 0, 2] - 0.1
         motion[frame, 8, 1] = -0.2
         motion[frame, 8, 2] = motion[frame, 7, 2] - 0.5 + swing_leg_r * 0.3
@@ -112,25 +76,15 @@ def generate_motion(frames={}):
         motion[frame, 9, 2] = motion[frame, 8, 2] - 0.5 + swing_leg_r * 0.3
         
         # Left arm (joints 10, 11, 12)
-        if waving:
-            # Waving motion for left arm
-            if t > 0.3 and t < 0.7:
-                wave = 0.5 * np.sin(t * 2 * np.pi * 8)
-                motion[frame, 10, 1] = 0.3
-                motion[frame, 10, 2] = motion[frame, 3, 2] - 0.2
-                motion[frame, 11, 1] = 0.5
-                motion[frame, 11, 2] = motion[frame, 10, 2]
-                motion[frame, 12, 1] = 0.7
-                motion[frame, 12, 2] = motion[frame, 11, 2] + wave
-            else:
-                # Normal arm swing during non-waving
-                swing_arm_l = np.sin(t * 2 * np.pi * leg_freq + np.pi)
-                motion[frame, 10, 1] = 0.3
-                motion[frame, 10, 2] = motion[frame, 3, 2] - 0.2
-                motion[frame, 11, 1] = 0.3 + swing_arm_l * 0.2
-                motion[frame, 11, 2] = motion[frame, 10, 2] - 0.4
-                motion[frame, 12, 1] = 0.3 + swing_arm_l * 0.4
-                motion[frame, 12, 2] = motion[frame, 11, 2] - 0.4
+        if waving and t > 0.3 and t < 0.7:
+            # Waving motion
+            wave = 0.5 * np.sin(t * 2 * np.pi * 8)
+            motion[frame, 10, 1] = 0.3
+            motion[frame, 10, 2] = motion[frame, 3, 2] - 0.2
+            motion[frame, 11, 1] = 0.5
+            motion[frame, 11, 2] = motion[frame, 10, 2]
+            motion[frame, 12, 1] = 0.7
+            motion[frame, 12, 2] = motion[frame, 11, 2] + wave
         else:
             # Normal arm swing
             swing_arm_l = np.sin(t * 2 * np.pi * leg_freq + np.pi)
@@ -152,18 +106,13 @@ def generate_motion(frames={}):
     
     return motion
 
-def visualize_motion(output_path):
-    # Generate motion data
-    motion_data = generate_motion()
-    
-    # Get dimensions
-    frames, joints, dims = motion_data.shape
-    
+def visualize_motion(motion_data, output_path="output.mp4"):
+    """Visualize the motion data as a 3D animation"""
     # Create figure
     fig = plt.figure(figsize=(10, 6))
     ax = fig.add_subplot(111, projection='3d')
     
-    # Define connections between joints (simplified skeleton)
+    # Define connections between joints
     connections = [
         (0, 1), (1, 2), (2, 3),  # Spine and head
         (0, 4), (4, 5), (5, 6),  # Left leg
@@ -172,6 +121,7 @@ def visualize_motion(output_path):
         (3, 13), (13, 14), (14, 15)   # Right arm
     ]
     
+    # Animation update function
     def update(frame):
         ax.clear()
         
@@ -196,65 +146,34 @@ def visualize_motion(output_path):
             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-')
-        
-        # Set title
-        action_type = ""
-        if running:
-            action_type = "Running"
-        elif walking:
-            action_type = "Walking"
-        elif jumping:
-            action_type = "Jumping"
-        elif dancing:
-            action_type = "Dancing"
-        elif turning:
-            action_type = "Turning"
-        elif waving:
-            action_type = "Waving"
-        else:
-            action_type = "Moving"
             
-        ax.set_title(action_type + " Motion - Frame " + str(frame))
         return ax
     
     # Create animation
     anim = FuncAnimation(fig, update, frames=min(180, motion_data.shape[0]), interval=1000/30)
     
     # Save animation
-    os.makedirs(os.path.dirname(output_path) or '.', exist_ok=True)
     anim.save(output_path, writer='ffmpeg', fps=30)
     plt.close()
     
-    print("Animation saved to " + output_path)
     return output_path
 
-# Create and visualize a motion
-visualize_motion("{}")
-""".format(
-            int(seed),  # First {}
-            int(seed),  # Second {}
-            safe_prompt,  # Third {}
-            int(motion_length * 30),  # Fourth {}
-            output_filename  # Fifth {}
-        ))
-    
-    # Run the script
-    subprocess.run(["python", "motion_animation.py"])
-    
-    if os.path.exists(output_filename):
-        return output_filename
-    else:
-        return None
-
 def text_to_motion(text_prompt, motion_length=3.0, seed=0):
-    """Generate motion from text prompt using MDM"""
+    """Generate motion from text prompt"""
     try:
-        # Create a motion animation based on the input parameters
-        return create_motion_animation(text_prompt, motion_length, seed)
+        print(f"Generating motion for: '{text_prompt}', length: {motion_length}s, seed: {seed}")
+        
+        # Create a unique filename
+        output_path = f"output_{abs(hash(text_prompt) % 10000)}_{int(motion_length)}_{seed}.mp4"
+        
+        # Generate and visualize the motion
+        motion_data = generate_motion(text_prompt, motion_length, seed)
+        return visualize_motion(motion_data, output_path)
         
     except Exception as e:
         print(f"Error generating motion: {str(e)}")
-        print(traceback.format_exc())  # Print the full traceback
+        import traceback
+        print(traceback.format_exc())
         return None
 
 # Create the Gradio interface