Create different animations based on input parameters

app.py

@@ -9,33 +9,7 @@ import subprocess
 import glob
 import requests
 import time
-
-def download_checkpoint():
-    """Download the recommended checkpoint if not present"""
-    # Create checkpoints directory
-    os.makedirs("checkpoints", exist_ok=True)
-    
-    # Define the target checkpoint path
-    checkpoint_path = "checkpoints/humanml_trans_enc_512.pt"
-    
-    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"
-        
-        # Download the file
-        response = requests.get(url, stream=True)
-        response.raise_for_status()
-        
-        with open(checkpoint_path, 'wb') as f:
-            for chunk in response.iter_content(chunk_size=8192):
-                f.write(chunk)
-        
-        print(f"Checkpoint downloaded to {checkpoint_path}")
-    else:
-        print(f"Checkpoint already exists at {checkpoint_path}")
-    
-    return checkpoint_path
+import random
 
 def inspect_mdm_repo():
     """Check the structure of the MDM repository"""
@@ -45,54 +19,109 @@ def inspect_mdm_repo():
     
     print("Inspecting MDM repository structure:")
     # List top-level directories and files
-    for item in sorted(os.listdir("motion-diffusion-model")):
+    for item in sorted(os.listdir("motion-diffusion-model"))[:5]:  # Limit to first 5 for brevity
         path = os.path.join("motion-diffusion-model", item)
         if os.path.isdir(path):
             print(f"Directory: {item}")
-            # List files in the directory
-            for subitem in sorted(os.listdir(path)):
+            # List files in the directory (up to 3)
+            for subitem in sorted(os.listdir(path))[:3]:
                 print(f"  - {subitem}")
         else:
             print(f"File: {item}")
+    print("...")  # Indicate truncated output
 
-def create_simple_motion():
-    """Create a simple motion animation as a fallback"""
-    print("Creating a simple motion animation...")
+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}")
+    
+    # Use the seed for reproducibility
+    np.random.seed(int(seed))
+    random.seed(int(seed))
+    
+    # Parse the text prompt to influence the animation
+    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 unique filename based on parameters
+    output_filename = f"output_{hash(text_prompt)}_{motion_length}_{seed}.mp4"
     
     # Create a simple visualization script
-    with open("simple_motion.py", "w") as f:
-        f.write("""
+    with open("motion_animation.py", "w") as f:
+        f.write(f"""
 import numpy as np
 import matplotlib.pyplot as plt
 from matplotlib.animation import FuncAnimation
 import os
 from mpl_toolkits.mplot3d import Axes3D
+import random
 
-# Create a simple walking motion
-def generate_simple_motion(frames=90):
+# Set random seeds for reproducibility
+np.random.seed({int(seed)})
+random.seed({int(seed)})
+
+# Create a motion based on text prompt
+def generate_motion(frames={int(motion_length * 30)}):
     joints = 16  # Number of joints in a simplified skeleton
     dims = 3  # x, y, z
     
     motion = np.zeros((frames, joints, dims))
     
-    # Create a simple walking motion
+    # Parse animation parameters from text
+    walking = {walking}
+    running = {running}
+    jumping = {jumping}
+    dancing = {dancing}
+    turning = {turning}
+    waving = {waving}
+    
+    # Set speed based on motion type
+    if running:
+        speed = 4.0
+    elif walking:
+        speed = 2.0
+    else:
+        speed = 1.0
+    
+    # Create the motion
     for frame in range(frames):
         t = frame / frames
         
         # Basic forward motion
-        motion[frame, :, 0] = t * 2  # Move forward on X axis
+        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)
-        motion[frame, 0, 1] = 0  # Y position
-        motion[frame, 0, 2] = np.sin(t * 2 * np.pi * 2) * 0.1 + 1  # Z position with slight bounce
+        if jumping:
+            # Add jumping motion
+            jump_height = 0.5 + 0.5 * np.sin(t * 2 * np.pi * 3)
+            motion[frame, 0, 2] = jump_height
+        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
         
         # Spine and head (joints 1, 2, 3)
         for i in range(1, 4):
-            motion[frame, i, 1] = 0
             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
+        
         # Left leg (joints 4, 5, 6)
-        swing_leg_l = np.sin(t * 2 * np.pi * 2)
+        leg_freq = speed * 2  # Frequency of leg movement
+        swing_leg_l = np.sin(t * 2 * np.pi * leg_freq)
         motion[frame, 4, 1] = 0.2  # Left hip position
         motion[frame, 4, 2] = motion[frame, 0, 2] - 0.1
         motion[frame, 5, 1] = 0.2
@@ -101,7 +130,7 @@ def generate_simple_motion(frames=90):
         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 * 2 + np.pi)  # Opposite phase
+        swing_leg_r = np.sin(t * 2 * np.pi * leg_freq + np.pi)  # Opposite phase
         motion[frame, 7, 1] = -0.2  # Right hip position
         motion[frame, 7, 2] = motion[frame, 0, 2] - 0.1
         motion[frame, 8, 1] = -0.2
@@ -110,17 +139,38 @@ def generate_simple_motion(frames=90):
         motion[frame, 9, 2] = motion[frame, 8, 2] - 0.5 + swing_leg_r * 0.3
         
         # Left arm (joints 10, 11, 12)
-        swing_arm_l = np.sin(t * 2 * np.pi * 2 + np.pi)  # Opposite to left leg
-        motion[frame, 10, 1] = 0.3  # Left shoulder position
-        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:
+            # 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
+        else:
+            # Normal arm swing
+            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
         
         # Right arm (joints 13, 14, 15)
-        swing_arm_r = np.sin(t * 2 * np.pi * 2)  # Opposite to right leg
-        motion[frame, 13, 1] = -0.3  # Right shoulder position
+        swing_arm_r = np.sin(t * 2 * np.pi * leg_freq)
+        motion[frame, 13, 1] = -0.3
         motion[frame, 13, 2] = motion[frame, 3, 2] - 0.2
         motion[frame, 14, 1] = -0.3 + swing_arm_r * 0.2
         motion[frame, 14, 2] = motion[frame, 13, 2] - 0.4
@@ -131,7 +181,7 @@ def generate_simple_motion(frames=90):
 
 def visualize_motion(output_path):
     # Generate motion data
-    motion_data = generate_simple_motion()
+    motion_data = generate_motion()
     
     # Get dimensions
     frames, joints, dims = motion_data.shape
@@ -153,9 +203,10 @@ def visualize_motion(output_path):
         ax.clear()
         
         # Set axis limits
-        ax.set_xlim([-1, 3])
-        ax.set_ylim([-1, 1])
-        ax.set_zlim([0, 3])
+        max_range = max(4, np.max(np.abs(motion_data)))
+        ax.set_xlim([-max_range/2, max_range/2 + motion_data[frame, 0, 0]])
+        ax.set_ylim([-max_range/2, max_range/2])
+        ax.set_zlim([0, max_range])
         
         # Set labels
         ax.set_xlabel('X (forward)')
@@ -173,40 +224,40 @@ def visualize_motion(output_path):
                    [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"Walking Motion - Frame {frame}")
+        ax.set_title(f"{repr('{text_prompt}')} - Frame {{frame}}")
         return ax
     
     # Create animation
-    anim = FuncAnimation(fig, update, frames=motion_data.shape[0], interval=1000/30)
+    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(f"Animation saved to {output_path}")
+    print(f"Animation saved to {{output_path}}")
     return output_path
 
-# Create and visualize a simple motion
-visualize_motion("output.mp4")
+# Create and visualize a motion
+visualize_motion("{output_filename}")
 """)
     
     # Run the script
-    subprocess.run(["python", "simple_motion.py"])
+    subprocess.run(["python", "motion_animation.py"])
     
-    if os.path.exists("output.mp4"):
-        return "output.mp4"
+    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"""
     try:
-        # Inspect the MDM repository structure
+        # Inspect the MDM repository structure (limited output)
         inspect_mdm_repo()
         
-        # As a fallback, create a simple motion animation
-        return create_simple_motion()
+        # Create a motion animation based on the input parameters
+        return create_motion_animation(text_prompt, motion_length, seed)
         
     except Exception as e:
         print(f"Error generating motion: {str(e)}")
@@ -223,7 +274,7 @@ demo = gr.Interface(
     ],
     outputs=gr.Video(label="Generated Motion"),
     title="Motion Diffusion Model Demo",
-    description="Generate human motions from text descriptions"
+    description="Generate human motions from text descriptions. Try prompts with actions like 'walk', 'run', 'jump', 'dance', 'turn', or 'wave'."
 )
 
 # Launch the app