cs686/ardy-motion-api
6
1import os2 3os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")4# Load models locally (no external text-encoder service) and prefer the HF cache.5os.environ.setdefault("TEXT_ENCODER_MODE", "local")6# README `preload_from_hub` is mounted read-only at runtime. Keep Xet's mutable7# chunk cache and logs in /tmp, while resolving preloaded model snapshots from8# the normal Hub cache before attempting any network download.9os.environ.setdefault("HF_XET_CACHE", "/tmp/huggingface-xet")10os.environ.setdefault("LOCAL_CACHE", "true")11 12import base6413import gzip14import json15import random16import sys17import tempfile18import time19import xml.etree.ElementTree as ET20from pathlib import Path21 22import spaces # must precede torch / CUDA-touching imports23import torch24import numpy as np25import trimesh26import gradio as gr27 28# The vendored ARDY package lives next to this file.29sys.path.insert(0, str(Path(__file__).resolve().parent))30 31from ardy.model import load_model # noqa: E40232from ardy.exports import write_bvh # noqa: E40233from ardy.model.load_model import load_text_encoder # noqa: E40234from ardy.motion_rep.tools import length_to_mask # noqa: E40235from ardy.tools import seed_everything, to_numpy # noqa: E40236 37# Two rigs are offered in the playground:38# "human" -> ARDY-Core-RP-20FPS-Horizon40, 27-joint skeleton @ 20 fps39# "robot" -> ARDY-G1-RP-25FPS-Horizon52, 34-joint Unitree G1 robot @ 25 fps40# The robot rig matches the sibling Space hugging-apps/ardy-g1-motion-generation.41DEFAULT_RIG = "human"42MAX_SEED = 2**31 - 143 44# -----------------------------------------------------------------------------45# Model loading (module scope).46#47# ZeroGPU has no live GPU at startup: it intercepts .to("cuda") *placement* but48# NOT arbitrary CUDA compute. LLM2Vec's PEFT load runs a LoRA `merge_and_unload`49# (real matmuls), so we build everything on CPU first, then .to("cuda") — the50# placement call is what the ZeroGPU hijack packs to disk and streams into VRAM51# on the first @spaces.GPU request.52# -----------------------------------------------------------------------------53# transformers / PEFT / safetensors infer their load device from54# torch.cuda.is_available(), which ZeroGPU reports True at startup even though no55# real GPU is attached — so a plain load tries to place weights on cuda and dies56# with "No CUDA GPUs are available". Force every loader onto CPU by masking57# is_available() during construction, then restore it and .to("cuda") (which the58# ZeroGPU hijack packs to disk + streams into VRAM on the first request).59#60# A single text encoder is built once and shared across both motion models61# (load_text_encoder is explicitly designed for this — see its docstring).62_real_cuda_available = torch.cuda.is_available63torch.cuda.is_available = lambda: False64try:65 print("Loading ARDY text encoder (LLM2Vec-Llama-3-8B) on CPU…", flush=True)66 _text_encoder = load_text_encoder(mode="local", device="cpu")67 68 print("Loading ARDY human (core) motion model on CPU…", flush=True)69 MODEL_HUMAN = load_model("core", device="cpu", text_encoder=_text_encoder)70 71 print("Loading ARDY robot (G1) motion model on CPU…", flush=True)72 MODEL_ROBOT = load_model("g1", device="cpu", text_encoder=_text_encoder)73finally:74 torch.cuda.is_available = _real_cuda_available75 76print("Moving models to CUDA (ZeroGPU-intercepted placement)…", flush=True)77_text_encoder = _text_encoder.to("cuda")78for _m in (MODEL_HUMAN, MODEL_ROBOT):79 _m.to("cuda")80 # self.device was captured at construction (device="cpu"); generation creates81 # many tensors on it, so retarget it to cuda to match the moved weights.82 _m.device = "cuda"83 _m.eval()84 85 86def _rig_params(model):87 fps = float(model.motion_rep.fps)88 skeleton = model.skeleton89 parents = skeleton.joint_parents.cpu().numpy().astype(int).tolist()90 patch = model.num_frames_per_token91 gen_horizon = model.gen_horizon_len92 num_base_steps = int(model.diffusion.num_base_steps)93 # History carried between autoregressive windows. The reference streaming demo94 # keeps this SHORT so a new prompt takes effect within a window.95 hist_crop = max(patch, (4 // patch) * patch)96 root_idx = parents.index(-1) if -1 in parents else 097 return {98 "model": model,99 "fps": fps,100 "skeleton": skeleton,101 "parents": parents,102 "patch": patch,103 "gen_horizon": gen_horizon,104 "num_base_steps": num_base_steps,105 "hist_crop": hist_crop,106 "root_idx": root_idx,107 }108 109 110RIGS = {111 "human": _rig_params(MODEL_HUMAN),112 "robot": _rig_params(MODEL_ROBOT),113}114# The largest base-step count across rigs bounds the diffusion-steps slider.115NUM_BASE_STEPS = max(r["num_base_steps"] for r in RIGS.values())116for _name, _r in RIGS.items():117 print(118 f"[{_name}] rig ready: {_r['skeleton'].nbjoints} joints, {_r['fps']} fps, "119 f"horizon {_r['gen_horizon']}, patch {_r['patch']}, "120 f"hist_crop {_r['hist_crop']}, base_steps {_r['num_base_steps']}",121 flush=True,122 )123 124 125def _normalize_rig(rig) -> str:126 """Map any UI/API rig value onto a valid RIGS key ('human' | 'robot')."""127 if rig is None:128 return DEFAULT_RIG129 key = str(rig).strip().lower()130 if key in RIGS:131 return key132 if key.startswith("hum") or "core" in key or "person" in key:133 return "human"134 if key.startswith("rob") or "g1" in key or "unitree" in key:135 return "robot"136 return DEFAULT_RIG137 138 139# -----------------------------------------------------------------------------140# HUMAN rig: skinned body mesh (ARDY "CoreSkin" linear-blend skinning).141#142# The reference viz (ardy/viz/viser_utils.py) renders a smooth humanoid body by143# skinning a bind mesh with the per-frame *global* joint transforms:144# verts = CoreSkin.skin(global_rot_mats, posed_joints, rot_is_global=True)145# The browser holds the static skin data (bind vertices / faces / LBS146# indices+weights) and does the per-vertex blend, while the server sends only the147# tiny per-frame joint affine matrices148# A[f,j] = fk[f,j] @ bind_rig_transform_inv[j] (fk = [R_global | pos])149# so the payload stays small (~0.1 MB/clip).150# -----------------------------------------------------------------------------151_HUMAN_SKEL = RIGS["human"]["skeleton"]152_SKIN_PATH = Path(_HUMAN_SKEL.folder) / "skin_standard.npz"153_skin = np.load(_SKIN_PATH)154BIND_RIG_INV = np.linalg.inv(155 np.asarray(_skin["bind_rig_transform"], dtype=np.float64)156).astype(np.float32) # [J, 4, 4]157 158 159def _build_skin_blob():160 """Pack the static skin data into one gzip+base64 blob (loaded once by the161 browser). Layout: bind_vertices f32[V,3] | faces u32[F,3] | lbs_idx u8[V,W]162 | lbs_wt f32[V,W]."""163 bind_v = np.asarray(_skin["bind_vertices"], dtype=np.float32)164 faces = np.asarray(_skin["faces"], dtype=np.uint32)165 idx = np.asarray(_skin["lbs_indices"], dtype=np.uint8)166 wt = np.asarray(_skin["lbs_weights"], dtype=np.float32)167 raw = (168 np.ascontiguousarray(bind_v).tobytes()169 + np.ascontiguousarray(faces).tobytes()170 + np.ascontiguousarray(idx).tobytes()171 + np.ascontiguousarray(wt).tobytes()172 )173 meta = {"V": int(bind_v.shape[0]), "F": int(faces.shape[0]), "W": int(idx.shape[1])}174 return base64.b64encode(gzip.compress(raw, 6)).decode("ascii"), meta175 176 177SKIN_B64, SKIN_META = _build_skin_blob()178print(f"CoreSkin ready: {SKIN_META['V']} verts / {SKIN_META['F']} faces, "179 f"blob {len(SKIN_B64) // 1024} KB", flush=True)180 181 182def _joint_affines_human(global_rot_mats: np.ndarray, posed_joints: np.ndarray) -> str:183 """Per-frame joint affine matrices A = fk @ bind_rig_inv, base64 f32 [T,J,12].184 185 global_rot_mats: [T, J, 3, 3]; posed_joints: [T, J, 3]."""186 T, J = posed_joints.shape[:2]187 fk = np.tile(np.eye(4, dtype=np.float32), (T, J, 1, 1))188 fk[..., :3, :3] = global_rot_mats.astype(np.float32)189 fk[..., :3, 3] = posed_joints.astype(np.float32)190 A = (fk @ BIND_RIG_INV)[..., :3, :] # [T, J, 3, 4]191 A = np.ascontiguousarray(A.reshape(T, J, 12).astype(np.float32))192 return base64.b64encode(A.tobytes()).decode("ascii")193 194 195# -----------------------------------------------------------------------------196# ROBOT rig: G1 robot mesh rig (rigid per-joint STL meshes).197#198# Unlike the human skeleton (rendered with one skinned body mesh via LBS), the199# Unitree G1 robot is rendered by attaching a rigid STL mesh to each articulated200# joint. This mirrors ardy/viz/g1_rig.py (G1MeshRig): each mesh has a local201# transform (geom_pos, geom_rot) relative to its joint, read from the MuJoCo202# g1.xml, plus a coordinate change from MuJoCo to ARDY axes. We precompute — for203# each mesh — its geometry PRE-TRANSFORMED into the joint-local frame204# (v' = geom_rot @ v + geom_pos), so at render time the browser just applies the205# per-frame joint transform: world_v = joint_pos + joint_rot @ v'.206# -----------------------------------------------------------------------------207# G1 joint -> STL mesh mapping (mirrors ardy/viz/g1_rig.py G1_MESH_JOINT_MAP).208G1_MESH_JOINT_MAP = {209 "pelvis_skel": ["pelvis.STL", "pelvis_contour_link.STL"],210 "left_hip_pitch_skel": ["left_hip_pitch_link.STL"],211 "left_hip_roll_skel": ["left_hip_roll_link.STL"],212 "left_hip_yaw_skel": ["left_hip_yaw_link.STL"],213 "left_knee_skel": ["left_knee_link.STL"],214 "left_ankle_pitch_skel": ["left_ankle_pitch_link.STL"],215 "left_ankle_roll_skel": ["left_ankle_roll_link.STL"],216 "right_hip_pitch_skel": ["right_hip_pitch_link.STL"],217 "right_hip_roll_skel": ["right_hip_roll_link.STL"],218 "right_hip_yaw_skel": ["right_hip_yaw_link.STL"],219 "right_knee_skel": ["right_knee_link.STL"],220 "right_ankle_pitch_skel": ["right_ankle_pitch_link.STL"],221 "right_ankle_roll_skel": ["right_ankle_roll_link.STL"],222 "waist_yaw_skel": ["waist_yaw_link_rev_1_0.STL", "waist_yaw_link.STL"],223 "waist_roll_skel": ["waist_roll_link_rev_1_0.STL", "waist_roll_link.STL"],224 "waist_pitch_skel": [225 "torso_link_rev_1_0.STL",226 "torso_link.STL",227 "logo_link.STL",228 "head_link.STL",229 ],230 "left_shoulder_pitch_skel": ["left_shoulder_pitch_link.STL"],231 "left_shoulder_roll_skel": ["left_shoulder_roll_link.STL"],232 "left_shoulder_yaw_skel": ["left_shoulder_yaw_link.STL"],233 "left_elbow_skel": ["left_elbow_link.STL"],234 "left_wrist_roll_skel": ["left_wrist_roll_link.STL"],235 "left_wrist_pitch_skel": ["left_wrist_pitch_link.STL"],236 "left_wrist_yaw_skel": ["left_wrist_yaw_link.STL", "left_rubber_hand.STL"],237 "right_shoulder_pitch_skel": ["right_shoulder_pitch_link.STL"],238 "right_shoulder_roll_skel": ["right_shoulder_roll_link.STL"],239 "right_shoulder_yaw_skel": ["right_shoulder_yaw_link.STL"],240 "right_elbow_skel": ["right_elbow_link.STL"],241 "right_wrist_roll_skel": ["right_wrist_roll_link.STL"],242 "right_wrist_pitch_skel": ["right_wrist_pitch_link.STL"],243 "right_wrist_yaw_skel": ["right_wrist_yaw_link.STL", "right_rubber_hand.STL"],244}245 246_ROBOT_SKEL = RIGS["robot"]["skeleton"]247_MUJOCO_TO_ARDY = np.array(248 [[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]], dtype=np.float64249)250_G1_SKEL_DIR = Path(_ROBOT_SKEL.folder)251_G1_MESH_DIR = _G1_SKEL_DIR / "meshes" / "g1"252_G1_XML = _G1_SKEL_DIR / "xml" / "g1.xml"253 254 255def _quat_wxyz_to_matrix(wxyz: np.ndarray) -> np.ndarray:256 w, x, y, z = wxyz257 n = np.sqrt(w * w + x * x + y * y + z * z)258 if n < 1e-12:259 return np.eye(3)260 w, x, y, z = w / n, x / n, y / n, z / n261 return np.array(262 [263 [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)],264 [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)],265 [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)],266 ],267 dtype=np.float64,268 )269 270 271def _mesh_local_transforms() -> dict:272 """mesh_file -> (geom_pos[3], geom_rot[3x3]) in ARDY axes, parsed from g1.xml."""273 if not _G1_XML.exists():274 return {}275 root = ET.parse(_G1_XML).getroot()276 file_to_name = {}277 for mesh in root.findall(".//asset/mesh"):278 name, file = mesh.get("name"), mesh.get("file")279 if name and file:280 file_to_name[file] = name281 name_to_tf = {}282 for geom in root.findall(".//geom"):283 name = geom.get("mesh")284 if name is None:285 continue286 pos = geom.get("pos")287 quat = geom.get("quat")288 gp = np.zeros(3) if pos is None else np.array([float(v) for v in pos.split()])289 gr_ = np.eye(3) if quat is None else _quat_wxyz_to_matrix(290 np.array([float(v) for v in quat.split()])291 )292 name_to_tf[name] = (gp, gr_)293 out = {}294 for file, name in file_to_name.items():295 gp, gr_ = name_to_tf.get(name, (np.zeros(3), np.eye(3)))296 gp = _MUJOCO_TO_ARDY @ gp297 gr_ = _MUJOCO_TO_ARDY @ gr_ @ _MUJOCO_TO_ARDY.T298 out[file] = (gp, gr_)299 return out300 301 302def _build_g1_mesh_blob():303 """Pack all rigid G1 meshes, pre-transformed into joint-local frame, into one304 gzip+base64 blob. Returns (b64, meta). meta.parts lists per-mesh305 {joint, v_off, v_cnt}. Layout: all verts f32[Vtot,3] then all faces306 u32[Ftot,3] (face indices are GLOBAL into the concatenated vertex array)."""307 skeleton = _ROBOT_SKEL308 local_tf = _mesh_local_transforms()309 all_v = []310 all_f = []311 parts = []312 v_cursor = 0313 for joint_name, mesh_files in G1_MESH_JOINT_MAP.items():314 if joint_name not in skeleton.bone_index:315 continue316 joint_idx = int(skeleton.bone_index[joint_name])317 for mesh_file in mesh_files:318 mp = _G1_MESH_DIR / mesh_file319 if not mp.exists():320 continue321 mesh = trimesh.load_mesh(str(mp), process=True)322 if isinstance(mesh, trimesh.Scene):323 mesh = trimesh.util.concatenate(mesh.dump())324 verts = np.asarray(mesh.vertices, dtype=np.float64) @ _MUJOCO_TO_ARDY.T325 faces = np.asarray(mesh.faces, dtype=np.int64)326 gp, gr_ = local_tf.get(mesh_file, (np.zeros(3), np.eye(3)))327 # Pre-apply the mesh's joint-local transform: v' = geom_rot @ v + geom_pos.328 verts = (verts @ gr_.T) + gp329 vcnt = verts.shape[0]330 parts.append({"joint": joint_idx, "v_off": v_cursor, "v_cnt": vcnt})331 all_v.append(verts.astype(np.float32))332 all_f.append((faces + v_cursor).astype(np.uint32)) # global vertex indices333 v_cursor += vcnt334 V = np.concatenate(all_v, axis=0) if all_v else np.zeros((0, 3), np.float32)335 F = np.concatenate(all_f, axis=0) if all_f else np.zeros((0, 3), np.uint32)336 raw = np.ascontiguousarray(V).tobytes() + np.ascontiguousarray(F).tobytes()337 meta = {"Vtot": int(V.shape[0]), "Ftot": int(F.shape[0]), "parts": parts}338 b64 = base64.b64encode(gzip.compress(raw, 6)).decode("ascii")339 return b64, meta340 341 342G1_MESH_B64, G1_MESH_META = _build_g1_mesh_blob()343print(344 f"G1 rig ready: {len(G1_MESH_META['parts'])} meshes, "345 f"{G1_MESH_META['Vtot']} verts / {G1_MESH_META['Ftot']} faces, "346 f"blob {len(G1_MESH_B64) // 1024} KB",347 flush=True,348)349 350 351def _joint_transforms_robot(global_rot_mats: np.ndarray, posed_joints: np.ndarray) -> str:352 """Per-frame joint affine matrices [R_global | pos], base64 f32 [T,J,12].353 354 The browser applies world_v = pos + R_global @ v' per mesh (v' already in355 joint-local frame)."""356 T, J = posed_joints.shape[:2]357 A = np.zeros((T, J, 3, 4), dtype=np.float32)358 A[..., :3, :3] = global_rot_mats.astype(np.float32)359 A[..., :3, 3] = posed_joints.astype(np.float32)360 A = np.ascontiguousarray(A.reshape(T, J, 12).astype(np.float32))361 return base64.b64encode(A.tobytes()).decode("ascii")362 363 364# --- Autoregressive generation with a persistable latent state ---------------365# ARDY is autoregressive: it generates one `gen_horizon_len`-frame window at a366# time, conditioned on a history of previous frames. `autoregressive_step` is367# the streaming primitive — it returns the *normalized motion-feature tensor*368# for (history + new window), which can be fed straight back in as the next369# window's history. We thread that tensor to (a) fill a requested clip length370# and (b) CONTINUE a clip with a new prompt, exactly like the reference371# interactive demo. Persisting the tensor in a gr.State lets a second "Continue"372# click resume from where the first clip ended, on the same character.373def _generate_sequence(rig_key, prompt, num_new_frames, steps, cfg_weight, init_tensor):374 """Run the AR loop for one prompt on the selected rig. `init_tensor`:375 normalized feature tensor [1, Th, D] on cuda (prior motion to continue), or376 None to start fresh. Returns the full normalized feature tensor."""377 r = RIGS[rig_key]378 model = r["model"]379 gen_horizon = r["gen_horizon"]380 patch = r["patch"]381 hist_crop = r["hist_crop"]382 text_feat, text_pad_mask = model._encode_text([prompt])383 motion_tensor = init_tensor384 target_new = max(1, int(np.ceil(num_new_frames / gen_horizon))) * gen_horizon385 produced = 0386 while produced < target_new:387 if motion_tensor is None:388 history, hist_len = None, 0389 else:390 hist_len = (min(motion_tensor.shape[1], hist_crop) // patch) * patch391 history = motion_tensor[:, motion_tensor.shape[1] - hist_len:] if hist_len else None392 hist_len = history.shape[1] if history is not None else 0393 samples = model.autoregressive_step(394 num_frames=hist_len + gen_horizon, # exactly history + one window (no future)395 num_denoising_steps=steps,396 motion_mask=None,397 observed_motion=None,398 cfg_weight=float(cfg_weight),399 texts=None,400 text_feat=text_feat,401 text_pad_mask=text_pad_mask,402 init_history_sequence=history,403 init_global_translation=None, # first window -> defaults (origin / +Z heading)404 init_first_heading_angle=None,405 )406 new_tail = samples[:, hist_len:] # the freshly generated window407 motion_tensor = new_tail if motion_tensor is None else torch.cat([motion_tensor, new_tail], dim=1)408 produced += new_tail.shape[1]409 return motion_tensor410 411 412def _decode_motion(rig_key, motion_tensor):413 r = RIGS[rig_key]414 model = r["model"]415 with torch.no_grad():416 out = to_numpy(model.motion_rep.inverse(motion_tensor, is_normalized=True))417 return {418 key: np.asarray(value)[0]419 for key, value in out.items()420 }421 422 423def _pack_payload(rig_key, motion_tensor, prompt, seed):424 """Decode a normalized feature tensor to the browser payload (per-frame joint425 affines + root ground-track), tagged with the rig so the viewer loads the426 correct skeleton/model."""427 r = RIGS[rig_key]428 out = _decode_motion(rig_key, motion_tensor)429 posed = out["posed_joints"] # [T, J, 3] global joint positions430 grm = out["global_rot_mats"] # [T, J, 3, 3] global rotations431 if rig_key == "robot":432 affines = _joint_transforms_robot(grm, posed)433 else:434 affines = _joint_affines_human(grm, posed)435 return {436 "rig": rig_key,437 "fps": r["fps"],438 "num_frames": int(posed.shape[0]),439 "num_joints": int(posed.shape[1]),440 "affines": affines, # drives the browser rig441 "root": np.round(posed[:, r["root_idx"], :].astype(np.float32), 4).tolist(),442 "prompt": prompt,443 "seed": seed,444 }445 446 447def _core_generate(rig, prompt, duration, diffusion_steps, cfg_weight, seed, randomize_seed, init_np):448 rig_key = _normalize_rig(rig)449 r = RIGS[rig_key]450 prompt = (prompt or "").strip()451 if not prompt:452 raise gr.Error("Please enter a text prompt describing the motion.")453 if randomize_seed:454 seed = random.randint(0, MAX_SEED)455 seed = int(seed)456 seed_everything(seed)457 steps = max(1, min(int(diffusion_steps), r["num_base_steps"]))458 num_new = max(r["patch"], int(round(float(duration) * r["fps"])))459 init_tensor = None if init_np is None else torch.from_numpy(init_np).to("cuda")460 461 t0 = time.perf_counter()462 with torch.no_grad():463 full = _generate_sequence(rig_key, prompt, num_new, steps, cfg_weight, init_tensor)464 payload = _pack_payload(rig_key, full, prompt, seed)465 tag = "continue" if init_np is not None else "generate"466 print(f"[{tag}:{rig_key}] '{prompt[:50]}' +{num_new}f -> {full.shape[1]}f total "467 f"steps={steps} seed={seed} {time.perf_counter() - t0:.1f}s", flush=True)468 return json.dumps(payload), seed, full.detach().cpu().numpy()469 470 471@spaces.GPU472def ui_generate(prompt, rig=DEFAULT_RIG, duration=5.0, diffusion_steps=NUM_BASE_STEPS,473 cfg_weight=2.0, seed=0, randomize_seed=True):474 """Start a fresh clip (resets the running sequence)."""475 return _core_generate(rig, prompt, duration, diffusion_steps, cfg_weight, seed, randomize_seed, None)476 477 478@spaces.GPU479def ui_continue(prompt, rig=DEFAULT_RIG, duration=5.0, diffusion_steps=NUM_BASE_STEPS,480 cfg_weight=2.0, seed=0, randomize_seed=True, state=None):481 """Append a new action, continuing from the previous clip's final pose."""482 return _core_generate(rig, prompt, duration, diffusion_steps, cfg_weight, seed, randomize_seed, state)483 484 485@spaces.GPU486def generate_motion(prompt: str, rig: str = DEFAULT_RIG, duration: float = 5.0,487 diffusion_steps: int = NUM_BASE_STEPS, cfg_weight: float = 2.0,488 seed: int = 0, randomize_seed: bool = True) -> tuple[str, int]:489 """Generate a 3D motion clip from a text prompt with ARDY.490 491 Args:492 prompt: Natural-language description of the motion (e.g. "a person walks in a circle").493 rig: Which character to animate — "human" (27-joint skeleton) or "robot" (Unitree G1).494 duration: Length of the generated motion in seconds.495 diffusion_steps: Number of denoising steps (1..num_base_steps).496 cfg_weight: Classifier-free-guidance weight for the text prompt.497 seed: Random seed for reproducibility.498 randomize_seed: If True, ignore `seed` and draw a fresh random one.499 500 Returns:501 A JSON string with the animated skeleton payload plus the seed used.502 """503 payload_json, seed, _ = _core_generate(504 rig, prompt, duration, diffusion_steps, cfg_weight, seed, randomize_seed, None505 )506 return payload_json, seed507 508 509@spaces.GPU510def generate_blender_motion(511 prompt: str,512 duration: float = 5.0,513 diffusion_steps: int = NUM_BASE_STEPS,514 cfg_weight: float = 2.0,515 seed: int = 0,516 randomize_seed: bool = True,517) -> tuple[str, str, str, int]:518 """Generate a human motion and return BVH, NPZ, and JSON metadata files."""519 _, seed, motion_np = _core_generate(520 "human",521 prompt,522 duration,523 diffusion_steps,524 cfg_weight,525 seed,526 randomize_seed,527 None,528 )529 return _export_motion_files("human", motion_np, prompt, seed)530 531 532def _export_motion_files(533 rig: str,534 motion_np: np.ndarray,535 prompt: str,536 seed: int,537) -> tuple[str, str, str, int]:538 """Decode an accumulated ARDY sequence and write Blender export files."""539 rig_key = _normalize_rig(rig)540 if motion_np is None:541 raise gr.Error(542 "Generate a motion first, then use Continue to extend it before exporting."543 )544 545 motion_np = np.asarray(motion_np)546 if motion_np.ndim != 3 or motion_np.shape[0] != 1 or motion_np.shape[1] < 1:547 raise gr.Error(f"Invalid ARDY sequence state: shape={motion_np.shape}")548 549 motion_tensor = torch.from_numpy(motion_np).to("cuda")550 out = _decode_motion(rig_key, motion_tensor)551 r = RIGS[rig_key]552 skeleton = r["skeleton"]553 joint_names = list(skeleton.bone_order_names)554 parents = list(r["parents"])555 rest_positions = to_numpy(skeleton.neutral_joints)556 557 output_dir = Path(tempfile.mkdtemp(prefix="ardy_blender_"))558 npz_path = output_dir / f"ardy_{rig_key}_motion.npz"559 bvh_path = output_dir / f"ardy_{rig_key}_motion.bvh"560 metadata_path = output_dir / "metadata.json"561 562 np.savez_compressed(563 npz_path,564 **out,565 fps=np.asarray(r["fps"], dtype=np.float32),566 text=np.asarray(prompt),567 seed=np.asarray(seed, dtype=np.int64),568 rig=np.asarray(rig_key),569 joint_names=np.asarray(joint_names),570 parents=np.asarray(parents, dtype=np.int32),571 rest_positions=np.asarray(rest_positions, dtype=np.float32),572 )573 write_bvh(574 bvh_path,575 joint_names=joint_names,576 parents=parents,577 rest_positions=rest_positions,578 local_rot_mats=out["local_rot_mats"],579 root_positions=out["root_positions"],580 fps=r["fps"],581 )582 metadata_path.write_text(583 json.dumps(584 {585 "model": (586 "nvidia/ARDY-Core-RP-20FPS-Horizon40"587 if rig_key == "human"588 else "nvidia/ARDY-G1-RP-25FPS-Horizon52"589 ),590 "rig": rig_key,591 "latest_prompt": prompt,592 "seed": int(seed),593 "fps": r["fps"],594 "frames": int(out["posed_joints"].shape[0]),595 "duration_seconds": (596 float(out["posed_joints"].shape[0]) / float(r["fps"])597 ),598 "joints": joint_names,599 "parents": parents,600 "coordinate_system": "Y-up, metric",601 "source": "current accumulated Generate/Continue sequence",602 "post_processing": False,603 },604 ensure_ascii=False,605 indent=2,606 ),607 encoding="utf-8",608 )609 return str(bvh_path), str(npz_path), str(metadata_path), int(seed)610 611 612@spaces.GPU613def export_current_sequence(614 state: object,615 rig: str = DEFAULT_RIG,616 prompt: str = "",617 seed: int = 0,618) -> tuple[str, str, str, int]:619 """Export the complete sequence currently held by Generate/Continue."""620 return _export_motion_files(rig, state, prompt, seed)621 622 623# -----------------------------------------------------------------------------624# Front-end: a self-contained Three.js playground, delivered as a Gradio-native625# custom HTML component (Gradio 6 `gr.HTML` templates + js_on_load).626#627# The motion JSON is carried as the component's own `value` prop. `js_on_load`628# imports three.js, builds the scene once, wires the controls, then registers a629# `watch('value', ...)` callback that Gradio fires whenever the component is630# updated as the output of a Python event (Generate button / Examples).631#632# Each payload is tagged with its `rig`. The viewer ships the static data for633# BOTH rigs (human skin blob + G1 rigid-mesh blob) and switches at load time:634# - "human": one skinned body mesh (linear-blend skinning in the browser).635# - "robot": rigid per-joint G1 STL meshes posed by the joint transforms.636# -----------------------------------------------------------------------------637 638PLAYER_TEMPLATE = """639<div class="ardy-playground">640 <div class="ardy-canvas-wrap">641 <div class="ardy-hint">Generate a motion to load it into the playground.</div>642 </div>643 <div class="ardy-controls">644 <button class="ardy-play ardy-btn" type="button">▶ Play</button>645 <input class="ardy-scrub" type="range" min="0" max="0" value="0" step="1" />646 <span class="ardy-frame">0 / 0</span>647 <label class="ardy-lbl">Speed648 <select class="ardy-speed">649 <option value="0.25">0.25×</option>650 <option value="0.5">0.5×</option>651 <option value="1" selected>1×</option>652 <option value="2">2×</option>653 </select>654 </label>655 <label class="ardy-lbl"><input type="checkbox" class="ardy-loop" checked/> Loop</label>656 <label class="ardy-lbl"><input type="checkbox" class="ardy-trail"/> Root trail</label>657 </div>658 <div class="ardy-caption"></div>659</div>660"""661 662# css_template rules are auto-scoped to this component by Gradio.663PLAYER_CSS_TEMPLATE = """664.ardy-playground { width: 100%; }665.ardy-canvas-wrap {666 position: relative; width: 100%; height: 480px;667 border-radius: 12px; overflow: hidden;668 background: #ffffff;669 border: 1px solid #e5e7eb;670}671.ardy-canvas-wrap canvas { display:block; width:100% !important; height:100% !important; }672.ardy-hint {673 position:absolute; top:50%; left:50%; transform:translate(-50%,-50%);674 color:#98a2b3; font-size:14px; text-align:center; pointer-events:none;675}676.ardy-controls {677 display:flex; align-items:center; gap:12px; flex-wrap:wrap;678 margin-top:10px; padding:8px 4px;679}680.ardy-controls .ardy-btn {681 background:#76B900; color:#fff;682 border:none; border-radius:8px; padding:6px 16px; cursor:pointer; font-weight:600;683}684.ardy-scrub { flex:1; min-width:160px; accent-color:#76B900; }685.ardy-frame { font-variant-numeric: tabular-nums; color: var(--body-text-color); min-width:70px; }686.ardy-lbl { font-size:13px; color: var(--body-text-color); display:flex; align-items:center; gap:4px; }687.ardy-caption { margin-top:6px; font-size:13px; color:#667085; }688"""689 690APP_CSS = """691#col-container { max-width: 1200px; margin: 0 auto; }692.dark .gradio-container { color: var(--body-text-color); }693"""694 695# Runs once, when the component first renders. `element`, `props`, and `watch`696# are injected by Gradio. We import three.js, decode the static skin data (human697# rig) and the static rigid-mesh data (robot rig), build the scene, and subscribe698# to value changes with `watch('value', ...)`. Each generated payload carries the699# per-frame joint affine matrices plus a `rig` tag; the viewer renders whichever700# rig the payload requests, swapping the on-screen mesh as needed.701PLAYER_JS_ON_LOAD = r"""702const root = element;703const q = (sel) => root.querySelector(sel);704 705const state = {706 ready:false,707 THREE:null, OrbitControls:null,708 renderer:null, scene:null, camera:null, controls:null,709 trailLine:null,710 data:null, verts:null, frame:0, playing:false, lastT:0,711 speed:1, loop:true, trail:false, pending:null,712 rig:null, // which rig mesh is currently mounted in the scene713 // human rig714 skin:null, humanMesh:null, humanGeom:null,715 // robot rig716 robotMesh:null, robotGeom:null,717 robotBaseVerts:null, robotFaces:null, robotVtot:0, robotParts:null,718};719 720// --- binary helpers ---------------------------------------------------------721function b64ToBytes(b64){722 const bin = atob(b64); const out = new Uint8Array(bin.length);723 for(let i=0;i<bin.length;i++) out[i]=bin.charCodeAt(i);724 return out;725}726async function gunzip(bytes){727 const ds = new DecompressionStream("gzip");728 const buf = await new Response(new Blob([bytes]).stream().pipeThrough(ds)).arrayBuffer();729 return buf;730}731 732// Decode the one-time static human skin blob into typed arrays.733async function decodeSkin(){734 const buf = await gunzip(b64ToBytes(ARDY_SKIN_B64));735 const V = ARDY_SKIN_META.V, F = ARDY_SKIN_META.F, W = ARDY_SKIN_META.W;736 let o = 0;737 const bindV = new Float32Array(buf.slice(o, o+V*3*4)); o += V*3*4;738 const faces = new Uint32Array(buf.slice(o, o+F*3*4)); o += F*3*4;739 const idx = new Uint8Array(buf.slice(o, o+V*W)); o += V*W;740 const wt = new Float32Array(buf.slice(o, o+V*W*4)); o += V*W*4;741 return {V, F, W, bindV, faces, idx, wt};742}743 744// Decode the one-time static robot rig blob: pre-transformed mesh vertices745// (joint-local frame) + global-indexed faces + per-mesh part table.746async function decodeRig(){747 const buf = await gunzip(b64ToBytes(G1_MESH_B64));748 const V = G1_MESH_META.Vtot, F = G1_MESH_META.Ftot;749 let o = 0;750 const baseVerts = new Float32Array(buf.slice(o, o+V*3*4)); o += V*3*4;751 const faces = new Uint32Array(buf.slice(o, o+F*3*4)); o += F*3*4;752 return {V, F, baseVerts, faces, parts: G1_MESH_META.parts};753}754 755// HUMAN: per-vertex linear-blend skinning for every frame (once per clip).756// A[f,j] is a 3x4 affine (row-major, 12 floats); posed vertex =757// sum_k w_k * A[idx_k] @ [bind_x, bind_y, bind_z, 1].758function skinAllFrames(A, T, J){759 const s = state.skin, V = s.V, W = s.W, bindV = s.bindV, idx = s.idx, wt = s.wt;760 const frames = new Array(T);761 for(let f=0; f<T; f++){762 const out = new Float32Array(V*3);763 const Ab = f*J*12;764 for(let v=0; v<V; v++){765 const bx = bindV[v*3], by = bindV[v*3+1], bz = bindV[v*3+2];766 let x=0, y=0, z=0;767 for(let k=0; k<W; k++){768 const w = wt[v*W+k]; if(w===0) continue;769 const a = Ab + idx[v*W+k]*12;770 x += w*(A[a]*bx + A[a+1]*by + A[a+2]*bz + A[a+3]);771 y += w*(A[a+4]*bx + A[a+5]*by + A[a+6]*bz + A[a+7]);772 z += w*(A[a+8]*bx + A[a+9]*by + A[a+10]*bz + A[a+11]);773 }774 out[v*3]=x; out[v*3+1]=y; out[v*3+2]=z;775 }776 frames[f] = out;777 }778 return frames;779}780 781// ROBOT: per-frame rigid transform: for every mesh part, world_v = pos + R @ v'782// where (R,pos) is the driving joint's global transform this frame and v' is783// the vertex already baked into that joint's local frame.784function poseAllFrames(A, T, J){785 const V = state.robotVtot, base = state.robotBaseVerts, parts = state.robotParts;786 const frames = new Array(T);787 for(let f=0; f<T; f++){788 const out = new Float32Array(V*3);789 const Ab = f*J*12;790 for(let p=0; p<parts.length; p++){791 const jp = parts[p];792 const a = Ab + jp.joint*12;793 const r0=A[a], r1=A[a+1], r2=A[a+2], px=A[a+3];794 const r3=A[a+4], r4=A[a+5], r5=A[a+6], py=A[a+7];795 const r6=A[a+8], r7=A[a+9], r8=A[a+10], pz=A[a+11];796 const vs = jp.v_off, ve = jp.v_off + jp.v_cnt;797 for(let v=vs; v<ve; v++){798 const bx=base[v*3], by=base[v*3+1], bz=base[v*3+2];799 out[v*3] = px + r0*bx + r1*by + r2*bz;800 out[v*3+1] = py + r3*bx + r4*by + r5*bz;801 out[v*3+2] = pz + r6*bx + r7*by + r8*bz;802 }803 }804 frames[f] = out;805 }806 return frames;807}808 809// --- three.js scene ---------------------------------------------------------810function initScene(){811 const THREE = state.THREE, OrbitControls = state.OrbitControls;812 const wrap = q(".ardy-canvas-wrap");813 if(!wrap || state.renderer) return;814 const w = wrap.clientWidth || 800, h = wrap.clientHeight || 480;815 const scene = new THREE.Scene(); scene.background = new THREE.Color(0xffffff);816 const camera = new THREE.PerspectiveCamera(42, w/h, 0.05, 200);817 camera.position.set(3.8, 2.2, 4.7);818 const renderer = new THREE.WebGLRenderer({antialias:true});819 renderer.setSize(w, h); renderer.setPixelRatio(Math.min(window.devicePixelRatio,2));820 renderer.shadowMap.enabled = true; renderer.shadowMap.type = THREE.PCFSoftShadowMap;821 wrap.appendChild(renderer.domElement);822 const controls = new OrbitControls(camera, renderer.domElement);823 controls.target.set(0, 0.9, 0); controls.enableDamping = true;824 825 scene.add(new THREE.HemisphereLight(0xffffff, 0xdfe4ee, 1.4));826 const key = new THREE.DirectionalLight(0xffffff, 1.5);827 key.position.set(3, 6, 4); key.castShadow = true;828 key.shadow.mapSize.set(2048, 2048);829 const c = key.shadow.camera; c.near=0.5; c.far=25; c.left=-3; c.right=3; c.top=3; c.bottom=-3;830 key.shadow.bias = -0.0004;831 scene.add(key);832 scene.add(new THREE.DirectionalLight(0xeef2ff, 0.35).translateX(-4).translateZ(-2));833 834 const ground = new THREE.Mesh(835 new THREE.PlaneGeometry(40, 40),836 new THREE.ShadowMaterial({opacity:0.16})837 );838 ground.rotation.x = -Math.PI/2; ground.position.y = 0; ground.receiveShadow = true;839 scene.add(ground);840 841 const grid = new THREE.GridHelper(10, 20, 0xc4c9d4, 0xe4e7ee);842 grid.position.y = 0.0015; scene.add(grid);843 844 state.renderer=renderer; state.scene=scene; state.camera=camera; state.controls=controls;845 846 new ResizeObserver(()=>{847 const w2 = wrap.clientWidth, h2 = wrap.clientHeight;848 if(w2>0 && h2>0){ camera.aspect=w2/h2; camera.updateProjectionMatrix(); renderer.setSize(w2,h2); }849 }).observe(wrap);850 851 animate();852}853 854// Mount the mesh for the requested rig (lazily built, then shown/hidden). Only855// one rig mesh is visible at a time; both share the scene once created.856function mountRig(rig){857 const THREE = state.THREE;858 if(rig === "robot"){859 if(!state.robotMesh){860 const geom = new THREE.BufferGeometry();861 geom.setIndex(new THREE.BufferAttribute(state.robotFaces, 1));862 geom.setAttribute("position", new THREE.BufferAttribute(new Float32Array(state.robotVtot*3), 3));863 const mat = new THREE.MeshStandardMaterial({color:0xd7dde6, roughness:0.5, metalness:0.55});864 const mesh = new THREE.Mesh(geom, mat);865 mesh.castShadow = true; mesh.frustumCulled = false;866 state.scene.add(mesh); state.robotMesh = mesh; state.robotGeom = geom;867 }868 if(state.humanMesh) state.humanMesh.visible = false;869 state.robotMesh.visible = true;870 return state.robotGeom;871 } else {872 if(!state.humanMesh){873 const s = state.skin;874 const geom = new THREE.BufferGeometry();875 geom.setIndex(new THREE.BufferAttribute(s.faces, 1));876 geom.setAttribute("position", new THREE.BufferAttribute(new Float32Array(s.V*3), 3));877 const mat = new THREE.MeshStandardMaterial({color:0x98bdff, roughness:0.85, metalness:0.0});878 const mesh = new THREE.Mesh(geom, mat);879 mesh.castShadow = true; mesh.frustumCulled = false;880 state.scene.add(mesh); state.humanMesh = mesh; state.humanGeom = geom;881 }882 if(state.robotMesh) state.robotMesh.visible = false;883 state.humanMesh.visible = true;884 return state.humanGeom;885 }886}887 888function activeGeom(){889 return (state.rig === "robot") ? state.robotGeom : state.humanGeom;890}891 892function setFrame(f){893 if(!state.verts) return;894 const T = state.data.num_frames;895 f = Math.max(0, Math.min(T-1, f|0));896 state.frame = f;897 const geom = activeGeom();898 if(!geom) return;899 const pos = geom.getAttribute("position");900 pos.array.set(state.verts[f]);901 pos.needsUpdate = true;902 geom.computeVertexNormals();903 geom.computeBoundingSphere();904 updateTrail();905 const scrub = q(".ardy-scrub"); if(scrub) scrub.value = f;906 const lbl = q(".ardy-frame"); if(lbl) lbl.textContent = (f+1)+" / "+T;907}908 909function updateTrail(){910 const THREE = state.THREE;911 if(!state.data || !state.data.root){ if(state.trailLine) state.trailLine.visible=false; return; }912 if(!state.trail){ if(state.trailLine) state.trailLine.visible=false; return; }913 const T = state.data.num_frames, root = state.data.root;914 if(!state.trailLine){915 const g = new THREE.BufferGeometry();916 g.setAttribute("position", new THREE.BufferAttribute(new Float32Array(T*3),3));917 state.trailLine = new THREE.Line(g, new THREE.LineBasicMaterial({color:0xf59e0b}));918 state.scene.add(state.trailLine);919 }920 state.trailLine.visible = true;921 const attr = state.trailLine.geometry.getAttribute("position");922 for(let t=0;t<T;t++){ attr.setXYZ(t, root[t][0], 0.006, root[t][2]); }923 attr.needsUpdate = true;924 state.trailLine.geometry.setDrawRange(0, Math.max(1, state.frame+1));925}926 927function animate(){928 requestAnimationFrame(animate);929 if(!state.renderer) return;930 const now = performance.now();931 if(state.playing && state.verts){932 const dt = (now - state.lastT)/1000;933 const fps = state.data.fps * state.speed;934 if(dt >= 1/Math.max(1e-3,fps)){935 state.lastT = now;936 let nf = state.frame + 1;937 if(nf >= state.data.num_frames){938 if(state.loop){ nf = 0; } else { nf = state.data.num_frames-1; setPlaying(false); }939 }940 setFrame(nf);941 }942 }943 state.controls.update();944 state.renderer.render(state.scene, state.camera);945}946 947function setPlaying(p){948 state.playing = p;949 const b = q(".ardy-play");950 if(b) b.textContent = p ? "⏸ Pause" : "▶ Play";951 state.lastT = performance.now();952}953 954function loadData(data){955 initScene();956 const rig = (data.rig === "robot") ? "robot" : "human";957 state.rig = rig;958 mountRig(rig);959 state.data = data;960 // Decode per-frame affines and pre-pose every frame (one-time cost per clip).961 const T = data.num_frames, J = data.num_joints;962 const Abytes = b64ToBytes(data.affines);963 const A = new Float32Array(Abytes.buffer, Abytes.byteOffset, Abytes.byteLength/4);964 state.verts = (rig === "robot") ? poseAllFrames(A, T, J) : skinAllFrames(A, T, J);965 if(state.trailLine){ state.scene.remove(state.trailLine); state.trailLine.geometry.dispose(); state.trailLine=null; }966 const scrub = q(".ardy-scrub"); if(scrub){ scrub.max = T-1; scrub.value = 0; }967 const hint = q(".ardy-hint"); if(hint) hint.style.display = "none";968 const cap = q(".ardy-caption");969 if(cap) cap.textContent = '"' + data.prompt + '" · ' + (rig==="robot"?"robot":"human") +970 ' · ' + T + ' frames @ ' + data.fps + ' fps · seed ' + data.seed;971 root.dataset.ardyLoaded = "1";972 root.dataset.ardyRig = rig;973 root.dataset.ardyFrames = String(T);974 setFrame(0);975 setPlaying(true);976}977 978function applyPayload(payload){979 if(!payload) return;980 if(!state.ready){ state.pending = payload; return; } // three.js / rigs still loading981 try { loadData(JSON.parse(payload)); }982 catch(e){ console.error("ARDY playground load error", e); }983}984 985function wireControls(){986 const bind = (sel, ev, fn) => {987 const el = q(sel); if(el && !el.dataset.wired){ el.dataset.wired="1"; el.addEventListener(ev, fn); }988 };989 bind(".ardy-play", "click", ()=> setPlaying(!state.playing));990 bind(".ardy-scrub", "input", (e)=>{ setPlaying(false); setFrame(parseInt(e.target.value)); });991 bind(".ardy-speed", "change", (e)=>{ state.speed = parseFloat(e.target.value); });992 bind(".ardy-loop", "change", (e)=>{ state.loop = e.target.checked; });993 bind(".ardy-trail", "change", (e)=>{ state.trail = e.target.checked; updateTrail(); });994}995 996// Gradio-native hand-off: render whenever the component's value prop updates as997// the output of a Python event (Generate / Continue / Examples).998if (typeof watch === "function") {999 watch("value", () => applyPayload(props.value));1000}1001 1002// Import three.js (esm.sh, not jsDelivr: the OrbitControls addon has an internal1003// bare `import ... from "three"` that a browser dynamic import() can't resolve1004// without an import map; esm.sh rewrites it and dedupes three) and decode both1005// rigs, then flush any value that already arrived.1006Promise.all([1007 import("https://esm.sh/three@0.160.0"),1008 import("https://esm.sh/three@0.160.0/examples/jsm/controls/OrbitControls.js"),1009 decodeSkin(),1010 decodeRig(),1011]).then(([THREE, oc, skin, rig]) => {1012 state.THREE = THREE;1013 state.OrbitControls = oc.OrbitControls;1014 state.skin = skin;1015 state.robotBaseVerts = rig.baseVerts;1016 state.robotFaces = rig.faces;1017 state.robotVtot = rig.V;1018 state.robotParts = rig.parts;1019 state.ready = true;1020 wireControls();1021 initScene();1022 const start = state.pending || props.value;1023 if (start) applyPayload(start);1024}).catch((e)=> console.error("ARDY viewer init failed", e));1025"""1026 1027EXAMPLES = [1028 ["A person walks forward confidently.", "human", 5.0],1029 ["A person walks in a circle.", "human", 6.0],1030 ["A person jumps up and down.", "human", 4.0],1031 ["The robot walks forward confidently.", "robot", 5.0],1032 ["The robot waves with the right hand.", "robot", 4.0],1033 ["The robot crouches down and then stands back up.", "robot", 5.0],1034]1035 1036with gr.Blocks() as demo:1037 with gr.Column(elem_id="col-container"):1038 gr.Markdown(1039 """1040 # 🕺 ARDY Motion Playground1041 Interactive text-to-motion with **[ARDY](https://research.nvidia.com/labs/sil/projects/ardy/)**1042 (Autoregressive Diffusion with Hybrid Representation) by NVIDIA.1043 Pick a **rig** (human or robot), type a prompt and **Generate** a 3D motion clip,1044 then **orbit, scrub, and play** it below.1045 Chain actions with **Continue ▸** — the same character keeps going from where it stopped.1046 """1047 )1048 1049 # Running latent state (normalized feature tensor, CPU) — lets "Continue"1050 # resume the same character from the end of the previous clip.1051 seq_state = gr.State(None)1052 1053 prompt = gr.Textbox(1054 label="Motion prompt",1055 placeholder="e.g. a person walks in a circle then waves",1056 lines=2,1057 )1058 rig = gr.Radio(1059 choices=[("🕺 Human", "human"), ("🤖 Robot (Unitree G1)", "robot")],1060 value=DEFAULT_RIG,1061 label="Rig",1062 )1063 with gr.Row():1064 run = gr.Button("Generate", variant="primary", scale=2)1065 cont = gr.Button("Continue ▸", variant="secondary", scale=1)1066 1067 # The interactive 3D playground — a Gradio-native custom HTML component.1068 # Its `value` (the motion JSON) is set directly by the Generate handler;1069 # `watch('value', ...)` in js_on_load renders it. Ship the static data for1070 # both rigs (human skin blob + G1 rigid-mesh blob) once, as a header1071 # prepended to js_on_load; the per-frame affines ride in each payload.1072 _rig_header = (1073 f'const ARDY_SKIN_B64="{SKIN_B64}";\n'1074 f"const ARDY_SKIN_META={json.dumps(SKIN_META)};\n"1075 f'const G1_MESH_B64="{G1_MESH_B64}";\n'1076 f"const G1_MESH_META={json.dumps(G1_MESH_META)};\n"1077 )1078 player = gr.HTML(1079 value="",1080 html_template=PLAYER_TEMPLATE,1081 css_template=PLAYER_CSS_TEMPLATE,1082 js_on_load=_rig_header + PLAYER_JS_ON_LOAD,1083 elem_id="ardy-player",1084 )1085 1086 with gr.Accordion("Advanced settings", open=False):1087 duration = gr.Slider(1.0, 10.0, value=5.0, step=0.5, label="Duration (seconds)")1088 diffusion_steps = gr.Slider(1089 1, NUM_BASE_STEPS, value=NUM_BASE_STEPS, step=1, label="Diffusion steps"1090 )1091 cfg_weight = gr.Slider(1.0, 6.0, value=2.0, step=0.5, label="Text guidance (CFG)")1092 with gr.Row():1093 randomize_seed = gr.Checkbox(label="Randomize seed", value=True)1094 seed = gr.Number(label="Seed", value=0, precision=0)1095 1096 with gr.Accordion("Blender export", open=False):1097 gr.Markdown(1098 "Export the **complete motion currently in the player**, including "1099 "every segment added with **Continue ▸**, as Blender **BVH**, "1100 "full-fidelity **NPZ**, and JSON metadata."1101 )1102 export_blender = gr.Button("Export current sequence", variant="secondary")1103 with gr.Row():1104 bvh_file = gr.File(label="Blender BVH")1105 npz_file = gr.File(label="ARDY NPZ")1106 metadata_file = gr.File(label="Metadata")1107 1108 gr.Examples(1109 examples=EXAMPLES,1110 inputs=[prompt, rig, duration],1111 outputs=[player, seed, seq_state],1112 fn=ui_generate,1113 cache_examples=False,1114 run_on_click=True,1115 )1116 1117 gr.Markdown(1118 """1119 <small>Rigs: **ARDY-Core-RP-20FPS-Horizon40** (human, 27-joint skeleton @ 20 fps)1120 and **ARDY-G1-RP-25FPS-Horizon52** (Unitree G1 robot, 34-joint skeleton @ 25 fps).1121 Text encoder: LLM2Vec-Llama-3-8B. Post-processing (foot-skate cleanup) is disabled in this demo.1122 Motion is generated autoregressively; longer clips take longer.1123 **Generate** starts a new clip; **Continue ▸** keeps the same character going,1124 transitioning it into the new prompt (like the reference demo's prompt timeline).</small>1125 """1126 )1127 1128 _gen_inputs = [prompt, rig, duration, diffusion_steps, cfg_weight, seed, randomize_seed]1129 # Generate starts fresh; Continue resumes from the running latent state.1130 # The payload is written straight into the player's `value`; its js_on_load1131 # `watch('value', ...)` renders it (loading the correct rig from the payload).1132 run.click(fn=ui_generate, inputs=_gen_inputs,1133 outputs=[player, seed, seq_state], api_name=False)1134 cont.click(fn=ui_continue, inputs=_gen_inputs + [seq_state],1135 outputs=[player, seed, seq_state], api_name=False)1136 export_blender.click(1137 fn=export_current_sequence,1138 inputs=[seq_state, rig, prompt, seed],1139 outputs=[bvh_file, npz_file, metadata_file, seed],1140 api_name=False,1141 )1142 1143 # Clean single-shot endpoint for the HTTP API / MCP tool (no session state).1144 gr.api(generate_motion, api_name="generate")1145 gr.api(generate_blender_motion, api_name="generate_blender")1146 1147demo.queue()1148 1149if __name__ == "__main__":1150 # Gradio 6 moved theme/css from the Blocks constructor to launch(). The1151 # player's JS/CSS now live on the gr.HTML component itself (js_on_load /1152 # css_template), so no global `head=` script is needed.1153 demo.launch(1154 theme=gr.themes.Citrus(),1155 css=APP_CSS,1156 mcp_server=True,1157 ssr_mode=False,1158 )1159