app.py

import json
import os
import shutil
import subprocess
import sys
from pathlib import Path

import gradio as gr
import numpy as np
from huggingface_hub import snapshot_download

import spaces


ROOT = Path(__file__).resolve().parent
CODE_DIR = ROOT / "DrawMotion"
MODEL_REPO = "I0u0I/DrawMotion"
GIT_REPO = "https://github.com/InvertedForest/DrawMotion.git"

ASSET_PATTERNS = [
    "logs/human_ml3d/last.ckpt",
    "mid_feat/t2m/mid_feat.pt",
    "stickman/weight/real_init/t2m/stickman_encoder.ckpt",
]

EXAMPLES = {
    "forward line": [[0, 0], [40, 0], [90, 0], [150, 0], [220, 0]],
    "left arc": [[0, 0], [35, -20], [75, -55], [120, -90], [180, -115], [240, -120]],
    "right arc": [[0, 0], [35, 20], [75, 55], [120, 90], [180, 115], [240, 120]],
    "zigzag": [[0, 0], [45, -45], [90, 35], [135, -35], [180, 45], [230, 0]],
    "circle": [[0, 0], [35, -55], [95, -75], [155, -45], [165, 20], [110, 55], [45, 45], [0, 0]],
}

runner = None


def ensure_drawmotion_code():
    if not CODE_DIR.exists():
        subprocess.run(["git", "clone", "--depth", "1", GIT_REPO, str(CODE_DIR)], check=True)
    snapshot_download(
        repo_id=MODEL_REPO,
        repo_type="model",
        allow_patterns=ASSET_PATTERNS,
        local_dir=CODE_DIR,
    )
    if str(CODE_DIR) not in sys.path:
        sys.path.insert(0, str(CODE_DIR))
    blender_dir = CODE_DIR / "blender"
    blender_dir.mkdir(exist_ok=True)
    (blender_dir / "__init__.py").write_text("", encoding="utf-8")
    (blender_dir / "deal_joint.py").write_text(
        "import numpy as np\n\n"
        "def threed2rot(joints):\n"
        "    return np.zeros((len(joints), joints.shape[1], 3), dtype=np.float32)\n",
        encoding="utf-8",
    )
    for rel_path in [
        "data/datasets/human_ml3d/mean.npy",
        "data/datasets/human_ml3d/std.npy",
        "data/datasets/kit_ml/mean.npy",
        "data/datasets/kit_ml/std.npy",
    ]:
        target = CODE_DIR / rel_path
        target.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(ROOT / rel_path, target)
    os.chdir(CODE_DIR)


ensure_drawmotion_code()

from demo.drawmotion_studio.app import validate_generate_payload
from demo.drawmotion_studio.runner import DrawMotionRunner
from mogen.utils.plot_utils import plot_3d_motion, t2m_kinematic_chain


def get_runner():
    global runner
    if runner is None:
        runner = DrawMotionRunner(
            ckpt_path="logs/human_ml3d/last.ckpt",
            gpu="0",
            sample_index=0,
            output_dir=str(ROOT / "runs"),
        )
    return runner


def normalize_custom_points(custom_points):
    points = json.loads(custom_points)
    normalized = []
    for point in points:
        if isinstance(point, dict):
            normalized.append({"x": float(point["x"]), "y": float(point["y"])})
        else:
            normalized.append({"x": float(point[0]), "y": float(point[1])})
    return normalized


def preset_points(name):
    return [{"x": float(x), "y": float(y)} for x, y in EXAMPLES[name]]


def format_result_json(result):
    slim = dict(result)
    slim["pred_joint"] = np.asarray(slim["pred_joint"]).round(5).tolist()
    slim["input_trajectory"] = np.asarray(slim["input_trajectory"]).round(5).tolist()
    slim["pred_trajectory"] = np.asarray(slim["pred_trajectory"]).round(5).tolist()
    return json.dumps(slim, indent=2)


@spaces.GPU(duration=300)
def generate(text, trajectory_mode, custom_trajectory, frames, alpha, trajectory_scale, ifg_repeat, ifg_scale):
    if trajectory_mode == "custom JSON":
        trajectory = normalize_custom_points(custom_trajectory)
    else:
        trajectory = preset_points(trajectory_mode)

    payload = {
        "text": text,
        "trajectory": trajectory,
        "length": int(frames),
        "density": float(alpha),
        "trajectory_scale": float(trajectory_scale),
        "ifg_repeat": int(ifg_repeat),
        "ifg_scale": float(ifg_scale),
        "stickmen": [],
    }
    payload = validate_generate_payload(payload)
    result = get_runner().generate(payload)

    run_dir = sorted((ROOT / "runs").iterdir())[-1]
    video_path = run_dir / "motion.mp4"
    plot_3d_motion(
        str(video_path),
        t2m_kinematic_chain,
        np.asarray(result["pred_joint"], dtype=np.float32),
        title=result["text"],
        fps=20,
    )

    result_json = format_result_json(result)
    result_path = run_dir / "result_for_download.json"
    result_path.write_text(result_json, encoding="utf-8")
    return str(video_path), result_json, str(result_path)


def fill_custom_example(name):
    if name == "custom JSON":
        name = "left arc"
    return json.dumps(EXAMPLES[name], indent=2)


with gr.Blocks(title="DrawMotion") as demo:
    gr.Markdown("# DrawMotion")
    gr.Markdown("Text and trajectory conditioned 3D human motion generation.")
    with gr.Row():
        with gr.Column(scale=1):
            text = gr.Textbox(
                label="Text",
                value="A person walks forward and turns left.",
                lines=2,
            )
            trajectory_mode = gr.Dropdown(
                choices=list(EXAMPLES.keys()) + ["custom JSON"],
                value="left arc",
                label="Trajectory",
            )
            custom_trajectory = gr.Textbox(
                label="Custom trajectory JSON",
                value=fill_custom_example("left arc"),
                lines=8,
            )
            with gr.Row():
                frames = gr.Slider(32, 196, value=120, step=1, label="Frames")
                alpha = gr.Slider(0, 1, value=0.2, step=0.05, label="Alpha")
            with gr.Row():
                trajectory_scale = gr.Slider(20, 200, value=50, step=1, label="Trajectory scale")
                ifg_repeat = gr.Slider(0, 100, value=50, step=1, label="IFG repeat")
            ifg_scale = gr.Slider(0, 200, value=50, step=1, label="IFG scale")
            run_button = gr.Button("Generate", variant="primary")
        with gr.Column(scale=1):
            video = gr.Video(label="Generated motion")
            result_json = gr.Code(label="Result JSON", language="json", lines=18)
            result_file = gr.File(label="Download result.json")

    trajectory_mode.change(
        fn=fill_custom_example,
        inputs=trajectory_mode,
        outputs=custom_trajectory,
        show_progress="hidden",
    )
    run_button.click(
        fn=generate,
        inputs=[text, trajectory_mode, custom_trajectory, frames, alpha, trajectory_scale, ifg_repeat, ifg_scale],
        outputs=[video, result_json, result_file],
        concurrency_limit=1,
    )

demo.queue(max_size=8).launch()