sder111/ardy

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 7import base648import gzip9import json10import random11import sys12import time13import xml.etree.ElementTree as ET14from pathlib import Path15 16import spaces  # must precede torch / CUDA-touching imports17import torch18import numpy as np19import trimesh20import gradio as gr21 22# The vendored ARDY package lives next to this file.23sys.path.insert(0, str(Path(__file__).resolve().parent))24 25from ardy.model import load_model  # noqa: E40226from ardy.model.load_model import load_text_encoder  # noqa: E40227from ardy.motion_rep.tools import length_to_mask  # noqa: E40228from ardy.tools import seed_everything, to_numpy  # noqa: E40229 30# Two rigs are offered in the playground:31#   "human" -> ARDY-Core-RP-20FPS-Horizon40, 27-joint skeleton @ 20 fps32#   "robot" -> ARDY-G1-RP-25FPS-Horizon52, 34-joint Unitree G1 robot @ 25 fps33# The robot rig matches the sibling Space hugging-apps/ardy-g1-motion-generation.34DEFAULT_RIG = "human"35MAX_SEED = 2**31 - 136 37# -----------------------------------------------------------------------------38# Model loading (module scope).39#40# ZeroGPU has no live GPU at startup: it intercepts .to("cuda") *placement* but41# NOT arbitrary CUDA compute. LLM2Vec's PEFT load runs a LoRA `merge_and_unload`42# (real matmuls), so we build everything on CPU first, then .to("cuda") — the43# placement call is what the ZeroGPU hijack packs to disk and streams into VRAM44# on the first @spaces.GPU request.45# -----------------------------------------------------------------------------46# transformers / PEFT / safetensors infer their load device from47# torch.cuda.is_available(), which ZeroGPU reports True at startup even though no48# real GPU is attached — so a plain load tries to place weights on cuda and dies49# with "No CUDA GPUs are available". Force every loader onto CPU by masking50# is_available() during construction, then restore it and .to("cuda") (which the51# ZeroGPU hijack packs to disk + streams into VRAM on the first request).52#53# A single text encoder is built once and shared across both motion models54# (load_text_encoder is explicitly designed for this — see its docstring).55_real_cuda_available = torch.cuda.is_available56torch.cuda.is_available = lambda: False57try:58    print("Loading ARDY text encoder (LLM2Vec-Llama-3-8B) on CPU…", flush=True)59    _text_encoder = load_text_encoder(mode="local", device="cpu")60 61    print("Loading ARDY human (core) motion model on CPU…", flush=True)62    MODEL_HUMAN = load_model("core", device="cpu", text_encoder=_text_encoder)63 64    print("Loading ARDY robot (G1) motion model on CPU…", flush=True)65    MODEL_ROBOT = load_model("g1", device="cpu", text_encoder=_text_encoder)66finally:67    torch.cuda.is_available = _real_cuda_available68 69print("Moving models to CUDA (ZeroGPU-intercepted placement)…", flush=True)70_text_encoder = _text_encoder.to("cuda")71for _m in (MODEL_HUMAN, MODEL_ROBOT):72    _m.to("cuda")73    # self.device was captured at construction (device="cpu"); generation creates74    # many tensors on it, so retarget it to cuda to match the moved weights.75    _m.device = "cuda"76    _m.eval()77 78 79def _rig_params(model):80    fps = float(model.motion_rep.fps)81    skeleton = model.skeleton82    parents = skeleton.joint_parents.cpu().numpy().astype(int).tolist()83    patch = model.num_frames_per_token84    gen_horizon = model.gen_horizon_len85    num_base_steps = int(model.diffusion.num_base_steps)86    # History carried between autoregressive windows. The reference streaming demo87    # keeps this SHORT so a new prompt takes effect within a window.88    hist_crop = max(patch, (4 // patch) * patch)89    root_idx = parents.index(-1) if -1 in parents else 090    return {91        "model": model,92        "fps": fps,93        "skeleton": skeleton,94        "parents": parents,95        "patch": patch,96        "gen_horizon": gen_horizon,97        "num_base_steps": num_base_steps,98        "hist_crop": hist_crop,99        "root_idx": root_idx,100    }101 102 103RIGS = {104    "human": _rig_params(MODEL_HUMAN),105    "robot": _rig_params(MODEL_ROBOT),106}107# The largest base-step count across rigs bounds the diffusion-steps slider.108NUM_BASE_STEPS = max(r["num_base_steps"] for r in RIGS.values())109for _name, _r in RIGS.items():110    print(111        f"[{_name}] rig ready: {_r['skeleton'].nbjoints} joints, {_r['fps']} fps, "112        f"horizon {_r['gen_horizon']}, patch {_r['patch']}, "113        f"hist_crop {_r['hist_crop']}, base_steps {_r['num_base_steps']}",114        flush=True,115    )116 117 118def _normalize_rig(rig) -> str:119    """Map any UI/API rig value onto a valid RIGS key ('human' | 'robot')."""120    if rig is None:121        return DEFAULT_RIG122    key = str(rig).strip().lower()123    if key in RIGS:124        return key125    if key.startswith("hum") or "core" in key or "person" in key:126        return "human"127    if key.startswith("rob") or "g1" in key or "unitree" in key:128        return "robot"129    return DEFAULT_RIG130 131 132# -----------------------------------------------------------------------------133# HUMAN rig: skinned body mesh (ARDY "CoreSkin" linear-blend skinning).134#135# The reference viz (ardy/viz/viser_utils.py) renders a smooth humanoid body by136# skinning a bind mesh with the per-frame *global* joint transforms:137#     verts = CoreSkin.skin(global_rot_mats, posed_joints, rot_is_global=True)138# The browser holds the static skin data (bind vertices / faces / LBS139# indices+weights) and does the per-vertex blend, while the server sends only the140# tiny per-frame joint affine matrices141#     A[f,j] = fk[f,j] @ bind_rig_transform_inv[j]      (fk = [R_global | pos])142# so the payload stays small (~0.1 MB/clip).143# -----------------------------------------------------------------------------144_HUMAN_SKEL = RIGS["human"]["skeleton"]145_SKIN_PATH = Path(_HUMAN_SKEL.folder) / "skin_standard.npz"146_skin = np.load(_SKIN_PATH)147BIND_RIG_INV = np.linalg.inv(148    np.asarray(_skin["bind_rig_transform"], dtype=np.float64)149).astype(np.float32)  # [J, 4, 4]150 151 152def _build_skin_blob():153    """Pack the static skin data into one gzip+base64 blob (loaded once by the154    browser). Layout: bind_vertices f32[V,3] | faces u32[F,3] | lbs_idx u8[V,W]155    | lbs_wt f32[V,W]."""156    bind_v = np.asarray(_skin["bind_vertices"], dtype=np.float32)157    faces = np.asarray(_skin["faces"], dtype=np.uint32)158    idx = np.asarray(_skin["lbs_indices"], dtype=np.uint8)159    wt = np.asarray(_skin["lbs_weights"], dtype=np.float32)160    raw = (161        np.ascontiguousarray(bind_v).tobytes()162        + np.ascontiguousarray(faces).tobytes()163        + np.ascontiguousarray(idx).tobytes()164        + np.ascontiguousarray(wt).tobytes()165    )166    meta = {"V": int(bind_v.shape[0]), "F": int(faces.shape[0]), "W": int(idx.shape[1])}167    return base64.b64encode(gzip.compress(raw, 6)).decode("ascii"), meta168 169 170SKIN_B64, SKIN_META = _build_skin_blob()171print(f"CoreSkin ready: {SKIN_META['V']} verts / {SKIN_META['F']} faces, "172      f"blob {len(SKIN_B64) // 1024} KB", flush=True)173 174 175def _joint_affines_human(global_rot_mats: np.ndarray, posed_joints: np.ndarray) -> str:176    """Per-frame joint affine matrices A = fk @ bind_rig_inv, base64 f32 [T,J,12].177 178    global_rot_mats: [T, J, 3, 3]; posed_joints: [T, J, 3]."""179    T, J = posed_joints.shape[:2]180    fk = np.tile(np.eye(4, dtype=np.float32), (T, J, 1, 1))181    fk[..., :3, :3] = global_rot_mats.astype(np.float32)182    fk[..., :3, 3] = posed_joints.astype(np.float32)183    A = (fk @ BIND_RIG_INV)[..., :3, :]  # [T, J, 3, 4]184    A = np.ascontiguousarray(A.reshape(T, J, 12).astype(np.float32))185    return base64.b64encode(A.tobytes()).decode("ascii")186 187 188# -----------------------------------------------------------------------------189# ROBOT rig: G1 robot mesh rig (rigid per-joint STL meshes).190#191# Unlike the human skeleton (rendered with one skinned body mesh via LBS), the192# Unitree G1 robot is rendered by attaching a rigid STL mesh to each articulated193# joint. This mirrors ardy/viz/g1_rig.py (G1MeshRig): each mesh has a local194# transform (geom_pos, geom_rot) relative to its joint, read from the MuJoCo195# g1.xml, plus a coordinate change from MuJoCo to ARDY axes. We precompute — for196# each mesh — its geometry PRE-TRANSFORMED into the joint-local frame197# (v' = geom_rot @ v + geom_pos), so at render time the browser just applies the198# per-frame joint transform: world_v = joint_pos + joint_rot @ v'.199# -----------------------------------------------------------------------------200# G1 joint -> STL mesh mapping (mirrors ardy/viz/g1_rig.py G1_MESH_JOINT_MAP).201G1_MESH_JOINT_MAP = {202    "pelvis_skel": ["pelvis.STL", "pelvis_contour_link.STL"],203    "left_hip_pitch_skel": ["left_hip_pitch_link.STL"],204    "left_hip_roll_skel": ["left_hip_roll_link.STL"],205    "left_hip_yaw_skel": ["left_hip_yaw_link.STL"],206    "left_knee_skel": ["left_knee_link.STL"],207    "left_ankle_pitch_skel": ["left_ankle_pitch_link.STL"],208    "left_ankle_roll_skel": ["left_ankle_roll_link.STL"],209    "right_hip_pitch_skel": ["right_hip_pitch_link.STL"],210    "right_hip_roll_skel": ["right_hip_roll_link.STL"],211    "right_hip_yaw_skel": ["right_hip_yaw_link.STL"],212    "right_knee_skel": ["right_knee_link.STL"],213    "right_ankle_pitch_skel": ["right_ankle_pitch_link.STL"],214    "right_ankle_roll_skel": ["right_ankle_roll_link.STL"],215    "waist_yaw_skel": ["waist_yaw_link_rev_1_0.STL", "waist_yaw_link.STL"],216    "waist_roll_skel": ["waist_roll_link_rev_1_0.STL", "waist_roll_link.STL"],217    "waist_pitch_skel": [218        "torso_link_rev_1_0.STL",219        "torso_link.STL",220        "logo_link.STL",221        "head_link.STL",222    ],223    "left_shoulder_pitch_skel": ["left_shoulder_pitch_link.STL"],224    "left_shoulder_roll_skel": ["left_shoulder_roll_link.STL"],225    "left_shoulder_yaw_skel": ["left_shoulder_yaw_link.STL"],226    "left_elbow_skel": ["left_elbow_link.STL"],227    "left_wrist_roll_skel": ["left_wrist_roll_link.STL"],228    "left_wrist_pitch_skel": ["left_wrist_pitch_link.STL"],229    "left_wrist_yaw_skel": ["left_wrist_yaw_link.STL", "left_rubber_hand.STL"],230    "right_shoulder_pitch_skel": ["right_shoulder_pitch_link.STL"],231    "right_shoulder_roll_skel": ["right_shoulder_roll_link.STL"],232    "right_shoulder_yaw_skel": ["right_shoulder_yaw_link.STL"],233    "right_elbow_skel": ["right_elbow_link.STL"],234    "right_wrist_roll_skel": ["right_wrist_roll_link.STL"],235    "right_wrist_pitch_skel": ["right_wrist_pitch_link.STL"],236    "right_wrist_yaw_skel": ["right_wrist_yaw_link.STL", "right_rubber_hand.STL"],237}238 239_ROBOT_SKEL = RIGS["robot"]["skeleton"]240_MUJOCO_TO_ARDY = np.array(241    [[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]], dtype=np.float64242)243_G1_SKEL_DIR = Path(_ROBOT_SKEL.folder)244_G1_MESH_DIR = _G1_SKEL_DIR / "meshes" / "g1"245_G1_XML = _G1_SKEL_DIR / "xml" / "g1.xml"246 247 248def _quat_wxyz_to_matrix(wxyz: np.ndarray) -> np.ndarray:249    w, x, y, z = wxyz250    n = np.sqrt(w * w + x * x + y * y + z * z)251    if n < 1e-12:252        return np.eye(3)253    w, x, y, z = w / n, x / n, y / n, z / n254    return np.array(255        [256            [1 - 2 * (y * y + z * z), 2 * (x * y - z * w), 2 * (x * z + y * w)],257            [2 * (x * y + z * w), 1 - 2 * (x * x + z * z), 2 * (y * z - x * w)],258            [2 * (x * z - y * w), 2 * (y * z + x * w), 1 - 2 * (x * x + y * y)],259        ],260        dtype=np.float64,261    )262 263 264def _mesh_local_transforms() -> dict:265    """mesh_file -> (geom_pos[3], geom_rot[3x3]) in ARDY axes, parsed from g1.xml."""266    if not _G1_XML.exists():267        return {}268    root = ET.parse(_G1_XML).getroot()269    file_to_name = {}270    for mesh in root.findall(".//asset/mesh"):271        name, file = mesh.get("name"), mesh.get("file")272        if name and file:273            file_to_name[file] = name274    name_to_tf = {}275    for geom in root.findall(".//geom"):276        name = geom.get("mesh")277        if name is None:278            continue279        pos = geom.get("pos")280        quat = geom.get("quat")281        gp = np.zeros(3) if pos is None else np.array([float(v) for v in pos.split()])282        gr_ = np.eye(3) if quat is None else _quat_wxyz_to_matrix(283            np.array([float(v) for v in quat.split()])284        )285        name_to_tf[name] = (gp, gr_)286    out = {}287    for file, name in file_to_name.items():288        gp, gr_ = name_to_tf.get(name, (np.zeros(3), np.eye(3)))289        gp = _MUJOCO_TO_ARDY @ gp290        gr_ = _MUJOCO_TO_ARDY @ gr_ @ _MUJOCO_TO_ARDY.T291        out[file] = (gp, gr_)292    return out293 294 295def _build_g1_mesh_blob():296    """Pack all rigid G1 meshes, pre-transformed into joint-local frame, into one297    gzip+base64 blob. Returns (b64, meta). meta.parts lists per-mesh298    {joint, v_off, v_cnt}. Layout: all verts f32[Vtot,3] then all faces299    u32[Ftot,3] (face indices are GLOBAL into the concatenated vertex array)."""300    skeleton = _ROBOT_SKEL301    local_tf = _mesh_local_transforms()302    all_v = []303    all_f = []304    parts = []305    v_cursor = 0306    for joint_name, mesh_files in G1_MESH_JOINT_MAP.items():307        if joint_name not in skeleton.bone_index:308            continue309        joint_idx = int(skeleton.bone_index[joint_name])310        for mesh_file in mesh_files:311            mp = _G1_MESH_DIR / mesh_file312            if not mp.exists():313                continue314            mesh = trimesh.load_mesh(str(mp), process=True)315            if isinstance(mesh, trimesh.Scene):316                mesh = trimesh.util.concatenate(mesh.dump())317            verts = np.asarray(mesh.vertices, dtype=np.float64) @ _MUJOCO_TO_ARDY.T318            faces = np.asarray(mesh.faces, dtype=np.int64)319            gp, gr_ = local_tf.get(mesh_file, (np.zeros(3), np.eye(3)))320            # Pre-apply the mesh's joint-local transform: v' = geom_rot @ v + geom_pos.321            verts = (verts @ gr_.T) + gp322            vcnt = verts.shape[0]323            parts.append({"joint": joint_idx, "v_off": v_cursor, "v_cnt": vcnt})324            all_v.append(verts.astype(np.float32))325            all_f.append((faces + v_cursor).astype(np.uint32))  # global vertex indices326            v_cursor += vcnt327    V = np.concatenate(all_v, axis=0) if all_v else np.zeros((0, 3), np.float32)328    F = np.concatenate(all_f, axis=0) if all_f else np.zeros((0, 3), np.uint32)329    raw = np.ascontiguousarray(V).tobytes() + np.ascontiguousarray(F).tobytes()330    meta = {"Vtot": int(V.shape[0]), "Ftot": int(F.shape[0]), "parts": parts}331    b64 = base64.b64encode(gzip.compress(raw, 6)).decode("ascii")332    return b64, meta333 334 335G1_MESH_B64, G1_MESH_META = _build_g1_mesh_blob()336print(337    f"G1 rig ready: {len(G1_MESH_META['parts'])} meshes, "338    f"{G1_MESH_META['Vtot']} verts / {G1_MESH_META['Ftot']} faces, "339    f"blob {len(G1_MESH_B64) // 1024} KB",340    flush=True,341)342 343 344def _joint_transforms_robot(global_rot_mats: np.ndarray, posed_joints: np.ndarray) -> str:345    """Per-frame joint affine matrices [R_global | pos], base64 f32 [T,J,12].346 347    The browser applies world_v = pos + R_global @ v' per mesh (v' already in348    joint-local frame)."""349    T, J = posed_joints.shape[:2]350    A = np.zeros((T, J, 3, 4), dtype=np.float32)351    A[..., :3, :3] = global_rot_mats.astype(np.float32)352    A[..., :3, 3] = posed_joints.astype(np.float32)353    A = np.ascontiguousarray(A.reshape(T, J, 12).astype(np.float32))354    return base64.b64encode(A.tobytes()).decode("ascii")355 356 357# --- Autoregressive generation with a persistable latent state ---------------358# ARDY is autoregressive: it generates one `gen_horizon_len`-frame window at a359# time, conditioned on a history of previous frames. `autoregressive_step` is360# the streaming primitive — it returns the *normalized motion-feature tensor*361# for (history + new window), which can be fed straight back in as the next362# window's history. We thread that tensor to (a) fill a requested clip length363# and (b) CONTINUE a clip with a new prompt, exactly like the reference364# interactive demo. Persisting the tensor in a gr.State lets a second "Continue"365# click resume from where the first clip ended, on the same character.366def _generate_sequence(rig_key, prompt, num_new_frames, steps, cfg_weight, init_tensor):367    """Run the AR loop for one prompt on the selected rig. `init_tensor`:368    normalized feature tensor [1, Th, D] on cuda (prior motion to continue), or369    None to start fresh. Returns the full normalized feature tensor."""370    r = RIGS[rig_key]371    model = r["model"]372    gen_horizon = r["gen_horizon"]373    patch = r["patch"]374    hist_crop = r["hist_crop"]375    text_feat, text_pad_mask = model._encode_text([prompt])376    motion_tensor = init_tensor377    target_new = max(1, int(np.ceil(num_new_frames / gen_horizon))) * gen_horizon378    produced = 0379    while produced < target_new:380        if motion_tensor is None:381            history, hist_len = None, 0382        else:383            hist_len = (min(motion_tensor.shape[1], hist_crop) // patch) * patch384            history = motion_tensor[:, motion_tensor.shape[1] - hist_len:] if hist_len else None385            hist_len = history.shape[1] if history is not None else 0386        samples = model.autoregressive_step(387            num_frames=hist_len + gen_horizon,   # exactly history + one window (no future)388            num_denoising_steps=steps,389            motion_mask=None,390            observed_motion=None,391            cfg_weight=float(cfg_weight),392            texts=None,393            text_feat=text_feat,394            text_pad_mask=text_pad_mask,395            init_history_sequence=history,396            init_global_translation=None,   # first window -> defaults (origin / +Z heading)397            init_first_heading_angle=None,398        )399        new_tail = samples[:, hist_len:]     # the freshly generated window400        motion_tensor = new_tail if motion_tensor is None else torch.cat([motion_tensor, new_tail], dim=1)401        produced += new_tail.shape[1]402    return motion_tensor403 404 405def _pack_payload(rig_key, motion_tensor, prompt, seed):406    """Decode a normalized feature tensor to the browser payload (per-frame joint407    affines + root ground-track), tagged with the rig so the viewer loads the408    correct skeleton/model."""409    r = RIGS[rig_key]410    model = r["model"]411    with torch.no_grad():412        out = to_numpy(model.motion_rep.inverse(motion_tensor, is_normalized=True))413    posed = np.asarray(out["posed_joints"])[0]      # [T, J, 3] global joint positions414    grm = np.asarray(out["global_rot_mats"])[0]      # [T, J, 3, 3] global rotations415    if rig_key == "robot":416        affines = _joint_transforms_robot(grm, posed)417    else:418        affines = _joint_affines_human(grm, posed)419    return {420        "rig": rig_key,421        "fps": r["fps"],422        "num_frames": int(posed.shape[0]),423        "num_joints": int(posed.shape[1]),424        "affines": affines,                          # drives the browser rig425        "root": np.round(posed[:, r["root_idx"], :].astype(np.float32), 4).tolist(),426        "prompt": prompt,427        "seed": seed,428    }429 430 431def _core_generate(rig, prompt, duration, diffusion_steps, cfg_weight, seed, randomize_seed, init_np):432    rig_key = _normalize_rig(rig)433    r = RIGS[rig_key]434    prompt = (prompt or "").strip()435    if not prompt:436        raise gr.Error("Please enter a text prompt describing the motion.")437    if randomize_seed:438        seed = random.randint(0, MAX_SEED)439    seed = int(seed)440    seed_everything(seed)441    steps = max(1, min(int(diffusion_steps), r["num_base_steps"]))442    num_new = max(r["patch"], int(round(float(duration) * r["fps"])))443    init_tensor = None if init_np is None else torch.from_numpy(init_np).to("cuda")444 445    t0 = time.perf_counter()446    with torch.no_grad():447        full = _generate_sequence(rig_key, prompt, num_new, steps, cfg_weight, init_tensor)448    payload = _pack_payload(rig_key, full, prompt, seed)449    tag = "continue" if init_np is not None else "generate"450    print(f"[{tag}:{rig_key}] '{prompt[:50]}' +{num_new}f -> {full.shape[1]}f total "451          f"steps={steps} seed={seed} {time.perf_counter() - t0:.1f}s", flush=True)452    return json.dumps(payload), seed, full.detach().cpu().numpy()453 454 455@spaces.GPU456def ui_generate(prompt, rig=DEFAULT_RIG, duration=5.0, diffusion_steps=NUM_BASE_STEPS,457                cfg_weight=2.0, seed=0, randomize_seed=True):458    """Start a fresh clip (resets the running sequence)."""459    return _core_generate(rig, prompt, duration, diffusion_steps, cfg_weight, seed, randomize_seed, None)460 461 462@spaces.GPU463def ui_continue(prompt, rig=DEFAULT_RIG, duration=5.0, diffusion_steps=NUM_BASE_STEPS,464                cfg_weight=2.0, seed=0, randomize_seed=True, state=None):465    """Append a new action, continuing from the previous clip's final pose."""466    return _core_generate(rig, prompt, duration, diffusion_steps, cfg_weight, seed, randomize_seed, state)467 468 469@spaces.GPU470def generate_motion(prompt: str, rig: str = DEFAULT_RIG, duration: float = 5.0,471                    diffusion_steps: int = NUM_BASE_STEPS, cfg_weight: float = 2.0,472                    seed: int = 0, randomize_seed: bool = True) -> tuple[str, int]:473    """Generate a 3D motion clip from a text prompt with ARDY.474 475    Args:476        prompt: Natural-language description of the motion (e.g. "a person walks in a circle").477        rig: Which character to animate — "human" (27-joint skeleton) or "robot" (Unitree G1).478        duration: Length of the generated motion in seconds.479        diffusion_steps: Number of denoising steps (1..num_base_steps).480        cfg_weight: Classifier-free-guidance weight for the text prompt.481        seed: Random seed for reproducibility.482        randomize_seed: If True, ignore `seed` and draw a fresh random one.483 484    Returns:485        A JSON string with the animated skeleton payload plus the seed used.486    """487    payload_json, seed, _ = _core_generate(488        rig, prompt, duration, diffusion_steps, cfg_weight, seed, randomize_seed, None489    )490    return payload_json, seed491 492 493# -----------------------------------------------------------------------------494# Front-end: a self-contained Three.js playground, delivered as a Gradio-native495# custom HTML component (Gradio 6 `gr.HTML` templates + js_on_load).496#497# The motion JSON is carried as the component's own `value` prop. `js_on_load`498# imports three.js, builds the scene once, wires the controls, then registers a499# `watch('value', ...)` callback that Gradio fires whenever the component is500# updated as the output of a Python event (Generate button / Examples).501#502# Each payload is tagged with its `rig`. The viewer ships the static data for503# BOTH rigs (human skin blob + G1 rigid-mesh blob) and switches at load time:504#  - "human": one skinned body mesh (linear-blend skinning in the browser).505#  - "robot": rigid per-joint G1 STL meshes posed by the joint transforms.506# -----------------------------------------------------------------------------507 508PLAYER_TEMPLATE = """509<div class="ardy-playground">510  <div class="ardy-canvas-wrap">511    <div class="ardy-hint">Generate a motion to load it into the playground.</div>512  </div>513  <div class="ardy-controls">514    <button class="ardy-play ardy-btn" type="button">▶ Play</button>515    <input class="ardy-scrub" type="range" min="0" max="0" value="0" step="1" />516    <span class="ardy-frame">0 / 0</span>517    <label class="ardy-lbl">Speed518      <select class="ardy-speed">519        <option value="0.25">0.25×</option>520        <option value="0.5">0.5×</option>521        <option value="1" selected>1×</option>522        <option value="2">2×</option>523      </select>524    </label>525    <label class="ardy-lbl"><input type="checkbox" class="ardy-loop" checked/> Loop</label>526    <label class="ardy-lbl"><input type="checkbox" class="ardy-trail"/> Root trail</label>527  </div>528  <div class="ardy-caption"></div>529</div>530"""531 532# css_template rules are auto-scoped to this component by Gradio.533PLAYER_CSS_TEMPLATE = """534.ardy-playground { width: 100%; }535.ardy-canvas-wrap {536  position: relative; width: 100%; height: 480px;537  border-radius: 12px; overflow: hidden;538  background: #ffffff;539  border: 1px solid #e5e7eb;540}541.ardy-canvas-wrap canvas { display:block; width:100% !important; height:100% !important; }542.ardy-hint {543  position:absolute; top:50%; left:50%; transform:translate(-50%,-50%);544  color:#98a2b3; font-size:14px; text-align:center; pointer-events:none;545}546.ardy-controls {547  display:flex; align-items:center; gap:12px; flex-wrap:wrap;548  margin-top:10px; padding:8px 4px;549}550.ardy-controls .ardy-btn {551  background:#76B900; color:#fff;552  border:none; border-radius:8px; padding:6px 16px; cursor:pointer; font-weight:600;553}554.ardy-scrub { flex:1; min-width:160px; accent-color:#76B900; }555.ardy-frame { font-variant-numeric: tabular-nums; color: var(--body-text-color); min-width:70px; }556.ardy-lbl { font-size:13px; color: var(--body-text-color); display:flex; align-items:center; gap:4px; }557.ardy-caption { margin-top:6px; font-size:13px; color:#667085; }558"""559 560APP_CSS = """561#col-container { max-width: 1200px; margin: 0 auto; }562.dark .gradio-container { color: var(--body-text-color); }563"""564 565# Runs once, when the component first renders. `element`, `props`, and `watch`566# are injected by Gradio. We import three.js, decode the static skin data (human567# rig) and the static rigid-mesh data (robot rig), build the scene, and subscribe568# to value changes with `watch('value', ...)`. Each generated payload carries the569# per-frame joint affine matrices plus a `rig` tag; the viewer renders whichever570# rig the payload requests, swapping the on-screen mesh as needed.571PLAYER_JS_ON_LOAD = r"""572const root = element;573const q = (sel) => root.querySelector(sel);574 575const state = {576  ready:false,577  THREE:null, OrbitControls:null,578  renderer:null, scene:null, camera:null, controls:null,579  trailLine:null,580  data:null, verts:null, frame:0, playing:false, lastT:0,581  speed:1, loop:true, trail:false, pending:null,582  rig:null,                 // which rig mesh is currently mounted in the scene583  // human rig584  skin:null, humanMesh:null, humanGeom:null,585  // robot rig586  robotMesh:null, robotGeom:null,587  robotBaseVerts:null, robotFaces:null, robotVtot:0, robotParts:null,588};589 590// --- binary helpers ---------------------------------------------------------591function b64ToBytes(b64){592  const bin = atob(b64); const out = new Uint8Array(bin.length);593  for(let i=0;i<bin.length;i++) out[i]=bin.charCodeAt(i);594  return out;595}596async function gunzip(bytes){597  const ds = new DecompressionStream("gzip");598  const buf = await new Response(new Blob([bytes]).stream().pipeThrough(ds)).arrayBuffer();599  return buf;600}601 602// Decode the one-time static human skin blob into typed arrays.603async function decodeSkin(){604  const buf = await gunzip(b64ToBytes(ARDY_SKIN_B64));605  const V = ARDY_SKIN_META.V, F = ARDY_SKIN_META.F, W = ARDY_SKIN_META.W;606  let o = 0;607  const bindV = new Float32Array(buf.slice(o, o+V*3*4)); o += V*3*4;608  const faces = new Uint32Array(buf.slice(o, o+F*3*4));  o += F*3*4;609  const idx   = new Uint8Array(buf.slice(o, o+V*W));      o += V*W;610  const wt    = new Float32Array(buf.slice(o, o+V*W*4));  o += V*W*4;611  return {V, F, W, bindV, faces, idx, wt};612}613 614// Decode the one-time static robot rig blob: pre-transformed mesh vertices615// (joint-local frame) + global-indexed faces + per-mesh part table.616async function decodeRig(){617  const buf = await gunzip(b64ToBytes(G1_MESH_B64));618  const V = G1_MESH_META.Vtot, F = G1_MESH_META.Ftot;619  let o = 0;620  const baseVerts = new Float32Array(buf.slice(o, o+V*3*4)); o += V*3*4;621  const faces     = new Uint32Array(buf.slice(o, o+F*3*4));  o += F*3*4;622  return {V, F, baseVerts, faces, parts: G1_MESH_META.parts};623}624 625// HUMAN: per-vertex linear-blend skinning for every frame (once per clip).626// A[f,j] is a 3x4 affine (row-major, 12 floats); posed vertex =627// sum_k w_k * A[idx_k] @ [bind_x, bind_y, bind_z, 1].628function skinAllFrames(A, T, J){629  const s = state.skin, V = s.V, W = s.W, bindV = s.bindV, idx = s.idx, wt = s.wt;630  const frames = new Array(T);631  for(let f=0; f<T; f++){632    const out = new Float32Array(V*3);633    const Ab = f*J*12;634    for(let v=0; v<V; v++){635      const bx = bindV[v*3], by = bindV[v*3+1], bz = bindV[v*3+2];636      let x=0, y=0, z=0;637      for(let k=0; k<W; k++){638        const w = wt[v*W+k]; if(w===0) continue;639        const a = Ab + idx[v*W+k]*12;640        x += w*(A[a]*bx   + A[a+1]*by  + A[a+2]*bz  + A[a+3]);641        y += w*(A[a+4]*bx + A[a+5]*by  + A[a+6]*bz  + A[a+7]);642        z += w*(A[a+8]*bx + A[a+9]*by  + A[a+10]*bz + A[a+11]);643      }644      out[v*3]=x; out[v*3+1]=y; out[v*3+2]=z;645    }646    frames[f] = out;647  }648  return frames;649}650 651// ROBOT: per-frame rigid transform: for every mesh part, world_v = pos + R @ v'652// where (R,pos) is the driving joint's global transform this frame and v' is653// the vertex already baked into that joint's local frame.654function poseAllFrames(A, T, J){655  const V = state.robotVtot, base = state.robotBaseVerts, parts = state.robotParts;656  const frames = new Array(T);657  for(let f=0; f<T; f++){658    const out = new Float32Array(V*3);659    const Ab = f*J*12;660    for(let p=0; p<parts.length; p++){661      const jp = parts[p];662      const a = Ab + jp.joint*12;663      const r0=A[a],   r1=A[a+1], r2=A[a+2],  px=A[a+3];664      const r3=A[a+4], r4=A[a+5], r5=A[a+6],  py=A[a+7];665      const r6=A[a+8], r7=A[a+9], r8=A[a+10], pz=A[a+11];666      const vs = jp.v_off, ve = jp.v_off + jp.v_cnt;667      for(let v=vs; v<ve; v++){668        const bx=base[v*3], by=base[v*3+1], bz=base[v*3+2];669        out[v*3]   = px + r0*bx + r1*by + r2*bz;670        out[v*3+1] = py + r3*bx + r4*by + r5*bz;671        out[v*3+2] = pz + r6*bx + r7*by + r8*bz;672      }673    }674    frames[f] = out;675  }676  return frames;677}678 679// --- three.js scene ---------------------------------------------------------680function initScene(){681  const THREE = state.THREE, OrbitControls = state.OrbitControls;682  const wrap = q(".ardy-canvas-wrap");683  if(!wrap || state.renderer) return;684  const w = wrap.clientWidth || 800, h = wrap.clientHeight || 480;685  const scene = new THREE.Scene(); scene.background = new THREE.Color(0xffffff);686  const camera = new THREE.PerspectiveCamera(42, w/h, 0.05, 200);687  camera.position.set(3.8, 2.2, 4.7);688  const renderer = new THREE.WebGLRenderer({antialias:true});689  renderer.setSize(w, h); renderer.setPixelRatio(Math.min(window.devicePixelRatio,2));690  renderer.shadowMap.enabled = true; renderer.shadowMap.type = THREE.PCFSoftShadowMap;691  wrap.appendChild(renderer.domElement);692  const controls = new OrbitControls(camera, renderer.domElement);693  controls.target.set(0, 0.9, 0); controls.enableDamping = true;694 695  scene.add(new THREE.HemisphereLight(0xffffff, 0xdfe4ee, 1.4));696  const key = new THREE.DirectionalLight(0xffffff, 1.5);697  key.position.set(3, 6, 4); key.castShadow = true;698  key.shadow.mapSize.set(2048, 2048);699  const c = key.shadow.camera; c.near=0.5; c.far=25; c.left=-3; c.right=3; c.top=3; c.bottom=-3;700  key.shadow.bias = -0.0004;701  scene.add(key);702  scene.add(new THREE.DirectionalLight(0xeef2ff, 0.35).translateX(-4).translateZ(-2));703 704  const ground = new THREE.Mesh(705    new THREE.PlaneGeometry(40, 40),706    new THREE.ShadowMaterial({opacity:0.16})707  );708  ground.rotation.x = -Math.PI/2; ground.position.y = 0; ground.receiveShadow = true;709  scene.add(ground);710 711  const grid = new THREE.GridHelper(10, 20, 0xc4c9d4, 0xe4e7ee);712  grid.position.y = 0.0015; scene.add(grid);713 714  state.renderer=renderer; state.scene=scene; state.camera=camera; state.controls=controls;715 716  new ResizeObserver(()=>{717    const w2 = wrap.clientWidth, h2 = wrap.clientHeight;718    if(w2>0 && h2>0){ camera.aspect=w2/h2; camera.updateProjectionMatrix(); renderer.setSize(w2,h2); }719  }).observe(wrap);720 721  animate();722}723 724// Mount the mesh for the requested rig (lazily built, then shown/hidden). Only725// one rig mesh is visible at a time; both share the scene once created.726function mountRig(rig){727  const THREE = state.THREE;728  if(rig === "robot"){729    if(!state.robotMesh){730      const geom = new THREE.BufferGeometry();731      geom.setIndex(new THREE.BufferAttribute(state.robotFaces, 1));732      geom.setAttribute("position", new THREE.BufferAttribute(new Float32Array(state.robotVtot*3), 3));733      const mat = new THREE.MeshStandardMaterial({color:0xd7dde6, roughness:0.5, metalness:0.55});734      const mesh = new THREE.Mesh(geom, mat);735      mesh.castShadow = true; mesh.frustumCulled = false;736      state.scene.add(mesh); state.robotMesh = mesh; state.robotGeom = geom;737    }738    if(state.humanMesh) state.humanMesh.visible = false;739    state.robotMesh.visible = true;740    return state.robotGeom;741  } else {742    if(!state.humanMesh){743      const s = state.skin;744      const geom = new THREE.BufferGeometry();745      geom.setIndex(new THREE.BufferAttribute(s.faces, 1));746      geom.setAttribute("position", new THREE.BufferAttribute(new Float32Array(s.V*3), 3));747      const mat = new THREE.MeshStandardMaterial({color:0x98bdff, roughness:0.85, metalness:0.0});748      const mesh = new THREE.Mesh(geom, mat);749      mesh.castShadow = true; mesh.frustumCulled = false;750      state.scene.add(mesh); state.humanMesh = mesh; state.humanGeom = geom;751    }752    if(state.robotMesh) state.robotMesh.visible = false;753    state.humanMesh.visible = true;754    return state.humanGeom;755  }756}757 758function activeGeom(){759  return (state.rig === "robot") ? state.robotGeom : state.humanGeom;760}761 762function setFrame(f){763  if(!state.verts) return;764  const T = state.data.num_frames;765  f = Math.max(0, Math.min(T-1, f|0));766  state.frame = f;767  const geom = activeGeom();768  if(!geom) return;769  const pos = geom.getAttribute("position");770  pos.array.set(state.verts[f]);771  pos.needsUpdate = true;772  geom.computeVertexNormals();773  geom.computeBoundingSphere();774  updateTrail();775  const scrub = q(".ardy-scrub"); if(scrub) scrub.value = f;776  const lbl = q(".ardy-frame"); if(lbl) lbl.textContent = (f+1)+" / "+T;777}778 779function updateTrail(){780  const THREE = state.THREE;781  if(!state.data || !state.data.root){ if(state.trailLine) state.trailLine.visible=false; return; }782  if(!state.trail){ if(state.trailLine) state.trailLine.visible=false; return; }783  const T = state.data.num_frames, root = state.data.root;784  if(!state.trailLine){785    const g = new THREE.BufferGeometry();786    g.setAttribute("position", new THREE.BufferAttribute(new Float32Array(T*3),3));787    state.trailLine = new THREE.Line(g, new THREE.LineBasicMaterial({color:0xf59e0b}));788    state.scene.add(state.trailLine);789  }790  state.trailLine.visible = true;791  const attr = state.trailLine.geometry.getAttribute("position");792  for(let t=0;t<T;t++){ attr.setXYZ(t, root[t][0], 0.006, root[t][2]); }793  attr.needsUpdate = true;794  state.trailLine.geometry.setDrawRange(0, Math.max(1, state.frame+1));795}796 797function animate(){798  requestAnimationFrame(animate);799  if(!state.renderer) return;800  const now = performance.now();801  if(state.playing && state.verts){802    const dt = (now - state.lastT)/1000;803    const fps = state.data.fps * state.speed;804    if(dt >= 1/Math.max(1e-3,fps)){805      state.lastT = now;806      let nf = state.frame + 1;807      if(nf >= state.data.num_frames){808        if(state.loop){ nf = 0; } else { nf = state.data.num_frames-1; setPlaying(false); }809      }810      setFrame(nf);811    }812  }813  state.controls.update();814  state.renderer.render(state.scene, state.camera);815}816 817function setPlaying(p){818  state.playing = p;819  const b = q(".ardy-play");820  if(b) b.textContent = p ? "⏸ Pause" : "▶ Play";821  state.lastT = performance.now();822}823 824function loadData(data){825  initScene();826  const rig = (data.rig === "robot") ? "robot" : "human";827  state.rig = rig;828  mountRig(rig);829  state.data = data;830  // Decode per-frame affines and pre-pose every frame (one-time cost per clip).831  const T = data.num_frames, J = data.num_joints;832  const Abytes = b64ToBytes(data.affines);833  const A = new Float32Array(Abytes.buffer, Abytes.byteOffset, Abytes.byteLength/4);834  state.verts = (rig === "robot") ? poseAllFrames(A, T, J) : skinAllFrames(A, T, J);835  if(state.trailLine){ state.scene.remove(state.trailLine); state.trailLine.geometry.dispose(); state.trailLine=null; }836  const scrub = q(".ardy-scrub"); if(scrub){ scrub.max = T-1; scrub.value = 0; }837  const hint = q(".ardy-hint"); if(hint) hint.style.display = "none";838  const cap = q(".ardy-caption");839  if(cap) cap.textContent = '"' + data.prompt + '"  ·  ' + (rig==="robot"?"robot":"human") +840    '  ·  ' + T + ' frames @ ' + data.fps + ' fps  ·  seed ' + data.seed;841  root.dataset.ardyLoaded = "1";842  root.dataset.ardyRig = rig;843  root.dataset.ardyFrames = String(T);844  setFrame(0);845  setPlaying(true);846}847 848function applyPayload(payload){849  if(!payload) return;850  if(!state.ready){ state.pending = payload; return; }   // three.js / rigs still loading851  try { loadData(JSON.parse(payload)); }852  catch(e){ console.error("ARDY playground load error", e); }853}854 855function wireControls(){856  const bind = (sel, ev, fn) => {857    const el = q(sel); if(el && !el.dataset.wired){ el.dataset.wired="1"; el.addEventListener(ev, fn); }858  };859  bind(".ardy-play", "click", ()=> setPlaying(!state.playing));860  bind(".ardy-scrub", "input", (e)=>{ setPlaying(false); setFrame(parseInt(e.target.value)); });861  bind(".ardy-speed", "change", (e)=>{ state.speed = parseFloat(e.target.value); });862  bind(".ardy-loop", "change", (e)=>{ state.loop = e.target.checked; });863  bind(".ardy-trail", "change", (e)=>{ state.trail = e.target.checked; updateTrail(); });864}865 866// Gradio-native hand-off: render whenever the component's value prop updates as867// the output of a Python event (Generate / Continue / Examples).868if (typeof watch === "function") {869  watch("value", () => applyPayload(props.value));870}871 872// Import three.js (esm.sh, not jsDelivr: the OrbitControls addon has an internal873// bare `import ... from "three"` that a browser dynamic import() can't resolve874// without an import map; esm.sh rewrites it and dedupes three) and decode both875// rigs, then flush any value that already arrived.876Promise.all([877  import("https://esm.sh/three@0.160.0"),878  import("https://esm.sh/three@0.160.0/examples/jsm/controls/OrbitControls.js"),879  decodeSkin(),880  decodeRig(),881]).then(([THREE, oc, skin, rig]) => {882  state.THREE = THREE;883  state.OrbitControls = oc.OrbitControls;884  state.skin = skin;885  state.robotBaseVerts = rig.baseVerts;886  state.robotFaces = rig.faces;887  state.robotVtot = rig.V;888  state.robotParts = rig.parts;889  state.ready = true;890  wireControls();891  initScene();892  const start = state.pending || props.value;893  if (start) applyPayload(start);894}).catch((e)=> console.error("ARDY viewer init failed", e));895"""896 897EXAMPLES = [898    ["A person walks forward confidently.", "human", 5.0],899    ["A person walks in a circle.", "human", 6.0],900    ["A person jumps up and down.", "human", 4.0],901    ["The robot walks forward confidently.", "robot", 5.0],902    ["The robot waves with the right hand.", "robot", 4.0],903    ["The robot crouches down and then stands back up.", "robot", 5.0],904]905 906with gr.Blocks() as demo:907    with gr.Column(elem_id="col-container"):908        gr.Markdown(909            """910            # 🕺 ARDY Motion Playground911            Interactive text-to-motion with **[ARDY](https://research.nvidia.com/labs/sil/projects/ardy/)**912            (Autoregressive Diffusion with Hybrid Representation) by NVIDIA.913            Pick a **rig** (human or robot), type a prompt and **Generate** a 3D motion clip,914            then **orbit, scrub, and play** it below.915            Chain actions with **Continue ▸** — the same character keeps going from where it stopped.916            """917        )918 919        # Running latent state (normalized feature tensor, CPU) — lets "Continue"920        # resume the same character from the end of the previous clip.921        seq_state = gr.State(None)922 923        prompt = gr.Textbox(924            label="Motion prompt",925            placeholder="e.g. a person walks in a circle then waves",926            lines=2,927        )928        rig = gr.Radio(929            choices=[("🕺 Human", "human"), ("🤖 Robot (Unitree G1)", "robot")],930            value=DEFAULT_RIG,931            label="Rig",932        )933        with gr.Row():934            run = gr.Button("Generate", variant="primary", scale=2)935            cont = gr.Button("Continue ▸", variant="secondary", scale=1)936 937        # The interactive 3D playground — a Gradio-native custom HTML component.938        # Its `value` (the motion JSON) is set directly by the Generate handler;939        # `watch('value', ...)` in js_on_load renders it. Ship the static data for940        # both rigs (human skin blob + G1 rigid-mesh blob) once, as a header941        # prepended to js_on_load; the per-frame affines ride in each payload.942        _rig_header = (943            f'const ARDY_SKIN_B64="{SKIN_B64}";\n'944            f"const ARDY_SKIN_META={json.dumps(SKIN_META)};\n"945            f'const G1_MESH_B64="{G1_MESH_B64}";\n'946            f"const G1_MESH_META={json.dumps(G1_MESH_META)};\n"947        )948        player = gr.HTML(949            value="",950            html_template=PLAYER_TEMPLATE,951            css_template=PLAYER_CSS_TEMPLATE,952            js_on_load=_rig_header + PLAYER_JS_ON_LOAD,953            elem_id="ardy-player",954        )955 956        with gr.Accordion("Advanced settings", open=False):957            duration = gr.Slider(1.0, 10.0, value=5.0, step=0.5, label="Duration (seconds)")958            diffusion_steps = gr.Slider(959                1, NUM_BASE_STEPS, value=NUM_BASE_STEPS, step=1, label="Diffusion steps"960            )961            cfg_weight = gr.Slider(1.0, 6.0, value=2.0, step=0.5, label="Text guidance (CFG)")962            with gr.Row():963                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)964                seed = gr.Number(label="Seed", value=0, precision=0)965 966        gr.Examples(967            examples=EXAMPLES,968            inputs=[prompt, rig, duration],969            outputs=[player, seed, seq_state],970            fn=ui_generate,971            cache_examples=False,972            run_on_click=True,973        )974 975        gr.Markdown(976            """977            <small>Rigs: **ARDY-Core-RP-20FPS-Horizon40** (human, 27-joint skeleton @ 20 fps)978            and **ARDY-G1-RP-25FPS-Horizon52** (Unitree G1 robot, 34-joint skeleton @ 25 fps).979            Text encoder: LLM2Vec-Llama-3-8B. Post-processing (foot-skate cleanup) is disabled in this demo.980            Motion is generated autoregressively; longer clips take longer.981            **Generate** starts a new clip; **Continue ▸** keeps the same character going,982            transitioning it into the new prompt (like the reference demo's prompt timeline).</small>983            """984        )985 986    _gen_inputs = [prompt, rig, duration, diffusion_steps, cfg_weight, seed, randomize_seed]987    # Generate starts fresh; Continue resumes from the running latent state.988    # The payload is written straight into the player's `value`; its js_on_load989    # `watch('value', ...)` renders it (loading the correct rig from the payload).990    run.click(fn=ui_generate, inputs=_gen_inputs,991              outputs=[player, seed, seq_state], api_name=False)992    cont.click(fn=ui_continue, inputs=_gen_inputs + [seq_state],993               outputs=[player, seed, seq_state], api_name=False)994 995    # Clean single-shot endpoint for the HTTP API / MCP tool (no session state).996    gr.api(generate_motion, api_name="generate")997 998demo.queue()999 1000if __name__ == "__main__":1001    # Gradio 6 moved theme/css from the Blocks constructor to launch(). The1002    # player's JS/CSS now live on the gr.HTML component itself (js_on_load /1003    # css_template), so no global `head=` script is needed.1004    demo.launch(1005        theme=gr.themes.Citrus(),1006        css=APP_CSS,1007        mcp_server=True,1008        ssr_mode=False,1009    )1010