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()