Improve MDM integration for better animation quality

app.py

@@ -1,406 +1,156 @@
-import gradio as gr
-import torch
-import os
-import sys
-import numpy as np
-from pathlib import Path
-import traceback
-import subprocess
+# app.py
+"""
+Motion Diffusion Demo on Hugging Face Spaces
+-------------------------------------------
+Generates human motion from a text prompt using the Motion-Diffusion-Model (MDM)
+checkpoint already uploaded to this Space.
+
+Key points
+~~~~~~~~~~
+* **Repo location**          : motion-diffusion-model/
+* **Checkpoint location**    : checkpoints/opt000750000.pt (path kept intact)
+* We call the official `sample.generate` CLI so we inherit every default the
+  authors bundled with the checkpoint (vocab, SMPL params, diffusion schedule …).
+* If anything goes wrong the function falls back to returning `None`, allowing
+  Gradio to show an empty result instead of crashing the Space.
+"""
+
+from __future__ import annotations
 
-def ensure_mdm_repo():
-    """Ensure the MDM repository is cloned and set up"""
-    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"])
-        
-        # Set up the repository
-        print("Setting up the repository...")
-        subprocess.run(["python", "-m", "spacy", "download", "en_core_web_sm"])
-        
-        # Add necessary 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("..")
-    
-    # Add the repository to the Python path
-    if "./motion-diffusion-model" not in sys.path:
-        sys.path.append("./motion-diffusion-model")
-
-def text_to_motion(text_prompt, motion_length=3.0, seed=0):
-    """Generate motion from text prompt using MDM"""
-    try:
-        print(f"Generating motion for: '{text_prompt}', length: {motion_length}s, seed: {seed}")
-        
-        # Ensure the MDM repository is set up
-        ensure_mdm_repo()
-        
-        # Create output directory
-        os.makedirs("output", exist_ok=True)
-        
-        # Get absolute path to the checkpoint
-        checkpoint_path = os.path.abspath("checkpoints/opt000750000.pt")
-        print(f"Using checkpoint: {checkpoint_path}")
-        
-        # Change to the MDM repository directory
-        original_dir = os.getcwd()
-        os.chdir("motion-diffusion-model")
-        
-        # List the sample directory to see what scripts are available
-        print("Available scripts in sample directory:")
-        if os.path.exists("sample"):
-            for file in os.listdir("sample"):
-                print(f"  - {file}")
-        
-        # Find the generate script
-        generate_script = None
-        for root, dirs, files in os.walk("."):
-            for file in files:
-                if file.endswith(".py") and "generate" in file:
-                    generate_script = os.path.join(root, file)
-                    print(f"Found generate script: {generate_script}")
-                    break
-            if generate_script:
-                break
-        
-        if not generate_script:
-            print("Could not find generate script")
-            os.chdir(original_dir)
-            return None
-        
-        # Create a simple Python script that uses our model
-        with open("run_mdm.py", "w") as f:
-            f.write("""
 import os
 import sys
-import torch
-import numpy as np
+import subprocess
+import traceback
 from pathlib import Path
+from typing import Optional
 
-# Add current directory to path
-sys.path.insert(0, os.getcwd())
-
-# Import required modules
-from utils.model_util import create_model_and_diffusion, load_saved_model
-from utils import dist_util
+import gradio as gr
 
-def generate_motion(model_path, text_prompt, motion_length, seed):
-    # Set up model
-    model, diffusion = create_model_and_diffusion(
-        model_path=model_path,
-        dataset='humanml',
-        diffusion_steps=1000,
-        num_frames=motion_length * 20,  # Assuming 20 fps
-    )
-    
-    # Load checkpoint
-    load_saved_model(model, model_path)
-    model.eval()
-    
-    # Set seed
-    torch.manual_seed(seed)
-    
-    # Generate motion
-    with torch.no_grad():
-        # Process text
-        text_emb = model.encode_text(text_prompt)
-        
-        # Generate motion
-        samples = diffusion.p_sample_loop(
-            model.forward_with_text,
-            shape=(1, model.njoints, model.nfeats, int(motion_length * 20)),
-            text_emb=text_emb,
-            clip_denoised=True,
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+
+REPO_DIR = "motion-diffusion-model"               # repo folder (already synced)
+CHECKPOINT_PATH = "checkpoints/opt000750000.pt"   # keep as-is per user request
+OUTPUT_DIR = "output"                             # where final MP4 files live
+MAX_LEN_SEC = 9.8                                  # model’s hard limit
+
+# ---------------------------------------------------------------------------
+# Helper functions
+# ---------------------------------------------------------------------------
+
+def ensure_repo_ready() -> None:
+    """Clone the repo only if it isn’t present 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,
         )
-        
-        # Save to file
-        os.makedirs('output', exist_ok=True)
-        output_path = f'output/motion_{abs(hash(text_prompt) % 10000)}_{int(motion_length)}_{seed}.mp4'
-        
-        # Visualize and save
-        from visualization.visualize import visualize
-        visualize(samples.cpu().numpy(), output_path)
-        
-        return output_path
 
-if __name__ == '__main__':
-    import argparse
-    
-    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()
-    
-    output_path = generate_motion(
-        args.model_path,
-        args.text_prompt,
-        args.motion_length,
-        args.seed
-    )
-    
-    print(f"Generated motion saved to: {output_path}")
-""")
-        
-        # Run our custom script
-        cmd = [
-            "python", 
-            "run_mdm.py",
-            "--model_path", checkpoint_path,
-            "--text_prompt", text_prompt,
-            "--motion_length", str(motion_length),
-            "--seed", str(int(seed))
-        ]
-        
-        print(f"Running command: {' '.join(cmd)}")
-        result = subprocess.run(cmd, capture_output=True, text=True)
-        
-        # Print the output for debugging
-        print("Command output:", result.stdout)
-        if result.stderr:
-            print("Command error:", result.stderr)
-        
-        # Check for output files
-        output_mp4 = None
-        for root, dirs, files in os.walk("."):
-            for file in files:
-                if file.endswith(".mp4"):
-                    output_mp4 = os.path.join(root, file)
-                    print(f"Found output file: {output_mp4}")
-                    break
-            if output_mp4:
-                break
-        
-        # Return to the original directory
-        os.chdir(original_dir)
-        
-        # If we found an output file, copy it to our output directory
-        if output_mp4:
-            output_path = f"output/output_{abs(hash(text_prompt) % 10000)}_{int(motion_length)}_{seed}.mp4"
-            subprocess.run(["cp", os.path.join("motion-diffusion-model", output_mp4), output_path])
-            print(f"Copied output to {output_path}")
-            return output_path
-        
-        # Fall back to simplified motion generation
-        print("MDM generation failed, falling back to simplified motion")
-        return create_simplified_motion(text_prompt, motion_length, seed)
-        
-    except Exception as e:
-        print(f"Error generating motion: {str(e)}")
-        print(traceback.format_exc())
-        
-        # Fall back to simplified motion generation
-        try:
-            return create_simplified_motion(text_prompt, motion_length, seed)
-        except:
-            return None
-
-def create_simplified_motion(text_prompt, motion_length, seed):
-    """Create a simplified motion animation as fallback"""
-    print("Creating simplified motion animation...")
-    
-    # Create output directory
-    os.makedirs("output", exist_ok=True)
-    output_path = f"output/simplified_{abs(hash(text_prompt) % 10000)}_{int(motion_length)}_{seed}.mp4"
-    
-    # Create a standalone script to generate the motion
-    with open("simplified_motion.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
-
-# Set random 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
+    repo_abs = str(Path(REPO_DIR).resolve())
+    if repo_abs not in sys.path:
+        sys.path.insert(0, repo_abs)
+
+
+def run_mdm(prompt: str, length: float, seed: int) -> Optional[str]:
+    """Generate a motion MP4 via the authors’ sample.generate script."""
+    ensure_repo_ready()
+
+    ckpt = Path(CHECKPOINT_PATH).resolve()
+    if not ckpt.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {ckpt}")
+
+    # The script creates its own result folder; we just need somewhere to move
+    # the freshest MP4 afterwards.
+    Path(OUTPUT_DIR).mkdir(exist_ok=True)
+
+    cmd = [
+        "python",
+        "-m",
+        "sample.generate",
+        "--model_path",
+        str(ckpt),
+        "--text_prompt",
+        prompt,
+        "--motion_length",
+        f"{min(length, MAX_LEN_SEC):.2f}",
+        "--seed",
+        str(seed),
+    ]
+
+    print("[run]", " ".join(cmd))
+    try:
+        subprocess.run(cmd, cwd=REPO_DIR, check=True)
+    except subprocess.CalledProcessError as exc:
+        print("[error] sample.generate failed:", exc)
+        return None
 
-# Set parameters
-frames = int({motion_length} * 30)  # 30 fps
-speed = 4.0 if running else 2.0 if walking else 1.0
+    # Grab the newest MP4 produced by the script
+    mp4_files = list(Path(REPO_DIR).rglob("*.mp4"))
+    if not mp4_files:
+        print("[warn] No MP4 file produced by the generator.")
+        return None
 
-# Create motion data - 16 joints with 3D coordinates
-joints = 16
-dims = 3
-motion = np.zeros((frames, joints, dims))
+    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/quota
 
-# 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
+    print(f"[ok] Motion video saved to {final_path}")
+    return str(final_path)
 
-# Create figure for visualization
-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
-]
+def fallback_motion(prompt: str, length: float, seed: int) -> Optional[str]:
+    """Placeholder fallback – returns None so the UI stays clean."""
+    print("[fallback] Returning empty result.")
+    return None
 
-# 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)
+def text_to_motion(prompt: str, length: float = 3.0, seed: int = 0):
+    try:
+        return run_mdm(prompt, length, seed) or fallback_motion(prompt, length, seed)
+    except Exception:
+        print(traceback.format_exc())
+        return fallback_motion(prompt, length, seed)
 
-# 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}")
-""")
-    
-    # Run the script
-    subprocess.run(["python", "simplified_motion.py"])
-    
-    if os.path.exists(output_path):
-        return output_path
-    else:
-        return None
+# ---------------------------------------------------------------------------
+# Gradio UI
+# ---------------------------------------------------------------------------
 
-# Create the 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 Diffusion Model 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 (e.g. 'A person runs in a circle and jumps').\n"
+        "The model returns a skeletal MP4 generated with the HumanML checkpoint."
+    ),
 )
 
-# Launch the app
+# ---------------------------------------------------------------------------
+# Launch
+# ---------------------------------------------------------------------------
+
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()