Improve MDM integration for better animation quality

app.py

@@ -1,206 +1,149 @@
+# app.py
+"""
+Motion Diffusion Demo on Hugging Face Spaces
+-------------------------------------------
+Generates human motion from a text prompt using the Motion‑Diffusion‑Model (MDM)
+checkpoint that is already stored in this Space.
+"""
+
+from __future__ import annotations
+
+import sys
+import subprocess
+import traceback
+from pathlib import Path
+from typing import Optional
+
 import gradio as gr
-import numpy as np
-import matplotlib.pyplot as plt
-from matplotlib.animation import FuncAnimation
-import os
-from mpl_toolkits.mplot3d import Axes3D
-
-def create_motion(text_prompt, motion_length, seed):
-    """Create a motion animation based on the text prompt"""
-    print(f"Creating motion for: '{text_prompt}', length: {motion_length}s, seed: {seed}")
-    
-    # Create output directory
-    os.makedirs("output", exist_ok=True)
-    output_path = f"output/motion_{abs(hash(text_prompt) % 10000)}_{int(motion_length)}_{seed}.mp4"
-    
-    # Use the seed for reproducibility
-    np.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))
-    
-    # Generate the motion
-    for frame in range(frames):
-        t = frame / frames
-        
-        # Basic forward motion or turning
-        if turning:
-            angle = t * 2 * np.pi * 2
-            motion[frame, :, 0] = np.cos(angle) * 2
-            motion[frame, :, 1] = np.sin(angle) * 2
-        else:
-            motion[frame, :, 0] = t * speed * 4
-        
-        # Root joint (pelvis) with jumping or bouncing
-        if jumping:
-            motion[frame, 0, 2] = 0.5 + 0.5 * np.sin(t * 2 * np.pi * 3)
-        else:
-            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
-            
-            # 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
-        swing_leg_l = np.sin(t * 2 * np.pi * leg_freq)
-        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
-        motion[frame, 6, 1] = 0.2
-        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)
-        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
-        motion[frame, 9, 1] = -0.2
-        motion[frame, 9, 2] = motion[frame, 8, 2] - 0.5 + swing_leg_r * 0.3
-        
-        # Left arm (joints 10, 11, 12)
-        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)
-            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 * 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
-        motion[frame, 15, 1] = -0.3 + swing_arm_r * 0.4
-        motion[frame, 15, 2] = motion[frame, 14, 2] - 0.4
-    
-    # Create figure
-    fig = plt.figure(figsize=(10, 6))
-    ax = fig.add_subplot(111, projection='3d')
-    
-    # Define connections between joints
-    connections = [
-        (0, 1), (1, 2), (2, 3),  # Spine and head
-        (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
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+REPO_DIR = "motion-diffusion-model"               # cloned repo folder
+CHECKPOINT_PATH = "checkpoints/opt000750000.pt"   # unchanged (user‑supplied)
+OUTPUT_DIR = "output"                             # where final MP4s are moved
+MAX_LEN_SEC = 9.8                                  # model’s hard limit
+
+# ---------------------------------------------------------------------------
+# Setup helpers
+# ---------------------------------------------------------------------------
+
+def ensure_repo_ready() -> None:
+    """Clone repo (first run only) and push it onto sys.path."""
+    if not Path(REPO_DIR).exists():
+        print("[setup] Cloning Motion‑Diffusion‑Model repo …")
+        subprocess.run(
+            [
+                "git",
+                "clone",
+                "https://github.com/GuyTevet/motion-diffusion-model.git",
+                REPO_DIR,
+            ],
+            check=True,
+        )
+
+    repo_abs = str(Path(REPO_DIR).resolve())
+    if repo_abs not in sys.path:
+        sys.path.insert(0, repo_abs)
+
+
+# ---------------------------------------------------------------------------
+# Core generation
+# ---------------------------------------------------------------------------
+
+def run_mdm(prompt: str, length: float, seed: int) -> Optional[str]:
+    """Run sample/generate.py and return the produced MP4 (or None)."""
+    ensure_repo_ready()
+
+    ckpt = Path(CHECKPOINT_PATH).resolve()
+    if not ckpt.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {ckpt}")
+
+    # Ensure output dir exists
+    Path(OUTPUT_DIR).mkdir(exist_ok=True)
+
+    # Script path is relative to repo root; we chdir into REPO_DIR later
+    script_path = Path("sample") / "generate.py"
+    if not (Path(REPO_DIR) / script_path).exists():
+        raise FileNotFoundError("sample/generate.py not found in the repo!")
+
+    # Assemble CLI
+    cmd = [
+        "python",
+        str(script_path),
+        "--model_path",
+        str(ckpt),
+        "--text_prompt",
+        prompt,
+        "--motion_length",
+        f"{min(length, MAX_LEN_SEC):.2f}",
+        "--seed",
+        str(seed),
     ]
-    
-    # Animation update function
-    def update(frame):
-        ax.clear()
-        
-        # Set axis limits
-        max_range = max(4, np.max(np.abs(motion)))
-        ax.set_xlim([-max_range/2, max_range/2 + motion[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)')
-        ax.set_ylabel('Y (sideways)')
-        ax.set_zlabel('Z (upward)')
-        
-        # Plot joints
-        ax.scatter(motion[frame, :, 0], 
-                  motion[frame, :, 1], 
-                  motion[frame, :, 2], c='b', marker='o')
-        
-        # Plot connections
-        for start, end in connections:
-            ax.plot([motion[frame, start, 0], motion[frame, end, 0]],
-                   [motion[frame, start, 1], motion[frame, end, 1]],
-                   [motion[frame, start, 2], motion[frame, end, 2]], 'r-')
-        
-        # Add action type to 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(frames, 180), interval=1000/30)
-    
-    # Save animation
-    anim.save(output_path, writer='ffmpeg', fps=30)
-    plt.close()
-    
-    print(f"Animation saved to {output_path}")
-    return output_path
-
-def text_to_motion(text_prompt, motion_length=3.0, seed=0):
-    """Generate motion from text prompt"""
+
+    print("[run]", " ".join(cmd))
     try:
-        # Each call creates a new animation with different parameters
-        return create_motion(text_prompt, motion_length, seed)
-    except Exception as e:
-        import traceback
-        print(f"Error generating motion: {str(e)}")
+        subprocess.run(cmd, cwd=REPO_DIR, check=True)
+    except subprocess.CalledProcessError as exc:
+        print("[error] sample/generate.py failed:", exc)
+        return None
+
+    # Locate newest MP4 in repo (generator writes to repo root/results/…)
+    mp4_files = list(Path(REPO_DIR).rglob("*.mp4"))
+    if not mp4_files:
+        print("[warn] No MP4 produced by generator.")
+        return None
+
+    newest = max(mp4_files, key=lambda p: p.stat().st_mtime)
+    final_path = Path(OUTPUT_DIR) / newest.name
+    newest.replace(final_path)  # move instead of copy to save disk
+
+    print(f"[ok] Motion video saved to {final_path}")
+    return str(final_path)
+
+
+def text_to_motion(prompt: str, length: float = 3.0, seed: int = 0):
+    """Wrapper called by Gradio UI."""
+    try:
+        return run_mdm(prompt, length, seed)
+    except Exception:
         print(traceback.format_exc())
         return None
 
-# Create the Gradio interface
+
+# ---------------------------------------------------------------------------
+# Gradio Interface
+# ---------------------------------------------------------------------------
+
 demo = gr.Interface(
     fn=text_to_motion,
     inputs=[
-        gr.Textbox(label="Text Prompt", placeholder="A person walks forward, then turns left", lines=3, value="A person walking"),
-        gr.Slider(minimum=1.0, maximum=9.8, value=3.0, label="Motion Length (seconds)"),
-        gr.Number(label="Random Seed", value=0)
+        gr.Textbox(
+            label="Text Prompt",
+            lines=3,
+            value="A person walks forward and waves.",
+        ),
+        gr.Slider(
+            minimum=1.0,
+            maximum=MAX_LEN_SEC,
+            step=0.1,
+            value=3.0,
+            label="Motion Length (seconds)",
+        ),
+        gr.Number(label="Random Seed", value=0, precision=0),
     ],
     outputs=gr.Video(label="Generated Motion"),
-    title="Motion Generation Demo",
-    description="Generate human motions from text descriptions. Try prompts with actions like 'walk', 'run', 'jump', 'dance', 'turn', or 'wave'."
+    title="Motion Diffusion Model Demo (HumanML)",
+    description=(
+        "Enter an action description such as 'A person runs in a circle and jumps'.\n"
+        "The HumanML checkpoint renders a skeletal MP4."
+    ),
 )
 
-# Launch the app
+# ---------------------------------------------------------------------------
+# Launch
+# ---------------------------------------------------------------------------
+
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()