Improve MDM integration for better animation quality

app.py

@@ -2,8 +2,14 @@
 """
 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.
+Generate human motion from a text prompt with **Motion‑Diffusion‑Model (MDM)**.
+We keep the user‑supplied layout:
+
+* repo folder        → motion-diffusion-model/
+* checkpoint file    → checkpoints/opt000750000.pt
+
+The script simply changes into the repo and executes `python -m sample.generate`,
+appearing exactly like the original workflow that was already working for you.
 """
 
 from __future__ import annotations
@@ -17,20 +23,19 @@ from typing import Optional
 import gradio as gr
 
 # ---------------------------------------------------------------------------
-# Configuration
+# Config – edit only if your Space layout changes
 # ---------------------------------------------------------------------------
 
-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
+REPO_DIR = "motion-diffusion-model"               # cloned repo folder name
+CHECKPOINT_PATH = "checkpoints/opt000750000.pt"   # user‑supplied ckpt path
+OUTPUT_DIR = "output"                             # final MP4 destination
+MAX_LEN_SEC = 9.8                                  # model’s max sequence length
 
 # ---------------------------------------------------------------------------
-# Setup helpers
+# Helper: make sure repo exists + on sys.path
 # ---------------------------------------------------------------------------
 
 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(
@@ -47,31 +52,24 @@ def ensure_repo_ready() -> None:
     if repo_abs not in sys.path:
         sys.path.insert(0, repo_abs)
 
-
 # ---------------------------------------------------------------------------
-# Core generation
+# Core Generation
 # ---------------------------------------------------------------------------
 
 def run_mdm(prompt: str, length: float, seed: int) -> Optional[str]:
-    """Run sample/generate.py and return the produced MP4 (or None)."""
+    """Call the official generator inside the repo and return path to MP4."""
     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),
+        "-m",
+        "sample.generate",
         "--model_path",
         str(ckpt),
         "--text_prompt",
@@ -86,34 +84,35 @@ def run_mdm(prompt: str, length: float, seed: int) -> Optional[str]:
     try:
         subprocess.run(cmd, cwd=REPO_DIR, check=True)
     except subprocess.CalledProcessError as exc:
-        print("[error] sample/generate.py failed:", exc)
+        print("[error] Generator failed:", exc)
         return None
 
-    # Locate newest MP4 in repo (generator writes to repo root/results/…)
+    # Look for the newest MP4 inside the repo after generation
     mp4_files = list(Path(REPO_DIR).rglob("*.mp4"))
     if not mp4_files:
-        print("[warn] No MP4 produced by generator.")
+        print("[warn] No MP4 produced.")
         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
+    newest.replace(final_path)
 
     print(f"[ok] Motion video saved to {final_path}")
     return str(final_path)
 
+# ---------------------------------------------------------------------------
+# Gradio wrapper
+# ---------------------------------------------------------------------------
 
 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
 
-
 # ---------------------------------------------------------------------------
-# Gradio Interface
+# Interface
 # ---------------------------------------------------------------------------
 
 demo = gr.Interface(
@@ -136,8 +135,8 @@ demo = gr.Interface(
     outputs=gr.Video(label="Generated Motion"),
     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."
+        "Describe an action — e.g. 'A person runs in a circle and jumps'. "
+        "The HumanML checkpoint returns a skeletal MP4."
     ),
 )