Deploy ZeroGPU ARDY motion generation API

README.md

@@ -1,13 +1,26 @@
 ---
-title: CozyClay ZeroGPU
-emoji: 🏃
-colorFrom: purple
-colorTo: pink
+title: CozyClay ZeroGPU Motion Engine
+emoji: 🎬
+colorFrom: yellow
+colorTo: gray
 sdk: gradio
-sdk_version: 6.24.0
-python_version: '3.12'
+sdk_version: 6.8.0
+python_version: 3.12.12
 app_file: app.py
 pinned: false
+license: gpl-3.0
+suggested_hardware: zero-a10g
+preload_from_hub:
+  - NousResearch/Meta-Llama-3-8B-Instruct
+  - McGill-NLP/LLM2Vec-Meta-Llama-3-8B-Instruct-mntp-supervised
+  - nvidia/ARDY-Core-RP-20FPS-Horizon40
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# CozyClay ZeroGPU Motion Engine
+
+GPU API for the [CozyClay Studio](https://huggingface.co/spaces/VIDraft/CozyClay).
+It turns a natural-language instruction into a newly sampled ARDY 3D skeletal
+motion and returns the NPZ motion to the browser editor.
+
+CozyClay is GPL-3.0-or-later. NVIDIA ARDY is an independent Apache-2.0 project;
+its released weights have their own NVIDIA Open Model License.

app.py

@@ -0,0 +1,214 @@
+import json
+import os
+import tempfile
+from pathlib import Path
+
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+
+MODEL_NAME = "ARDY-Core-RP-20FPS-Horizon40"
+PUBLIC_LLAMA = "NousResearch/Meta-Llama-3-8B-Instruct"
+_model = None
+
+
+def _load_model():
+    """Load once into ZeroGPU's CUDA-emulated process, then reuse it."""
+    global _model
+    if _model is not None:
+        return _model
+
+    from ardy.model.load_model import TEXT_ENCODER_PRESETS, load_model
+    from ardy.model.llm2vec.llm2vec import LLM2Vec
+
+    # ARDY's released LLM2Vec adapter normally names Meta's gated Llama repo.
+    # This public mirror has identical base weights and requires no visitor token.
+    TEXT_ENCODER_PRESETS["llm2vec"]["kwargs"]["base_model_name_or_path"] = PUBLIC_LLAMA
+
+    # Preserve the instruction wrapper used during ARDY training even though
+    # the public mirror has a different repository name in its config.
+    original_prepare = LLM2Vec.prepare_for_tokenization
+    if not getattr(original_prepare, "_cozyclay_public_llama", False):
+        def prepare_for_tokenization(self, text):
+            model_name = str(getattr(self.model.config, "_name_or_path", ""))
+            if model_name == PUBLIC_LLAMA:
+                return "<|start_header_id|>user<|end_header_id|>\n\n" + text.strip() + "<|eot_id|>"
+            return original_prepare(self, text)
+
+        prepare_for_tokenization._cozyclay_public_llama = True
+        LLM2Vec.prepare_for_tokenization = prepare_for_tokenization
+
+    _model = load_model(
+        MODEL_NAME,
+        device="cuda:0",
+        text_encoder_mode="local",
+    ).eval()
+    return _model
+
+
+def _validate_request(raw):
+    try:
+        body = json.loads(raw)
+    except Exception as exc:
+        raise gr.Error(f"Invalid request JSON: {exc}") from exc
+    if not isinstance(body, dict):
+        raise gr.Error("Request must be a JSON object.")
+    prompt = str(body.get("prompt", "")).strip()
+    if not prompt or len(prompt) > 500:
+        raise gr.Error("Prompt must contain 1 to 500 characters.")
+    duration = float(body.get("duration", 5.0))
+    if not 0.15 <= duration <= 10.0:
+        raise gr.Error("Duration must be between 0.15 and 10 seconds.")
+    seed = body.get("seed")
+    if seed is not None:
+        seed = int(seed)
+        if not 0 <= seed <= 2**31 - 1:
+            raise gr.Error("Seed is outside the supported range.")
+    if body.get("poses") or body.get("motionEdit") or body.get("segments"):
+        raise gr.Error(
+            "This first ZeroGPU release supports prompt and root-path generation. "
+            "Pose pins, prompt schedules, and motion edits are being enabled next."
+        )
+    waypoints = body.get("waypoints") or []
+    if not isinstance(waypoints, list) or len(waypoints) > 32:
+        raise gr.Error("Waypoints must be a list with at most 32 entries.")
+    return prompt, duration, seed, waypoints
+
+
+def _duration(raw):
+    try:
+        seconds = float(json.loads(raw).get("duration", 5.0))
+    except Exception:
+        seconds = 5.0
+    return min(300, max(120, int(90 + seconds * 24)))
+
+
+@spaces.GPU(duration=_duration, size="large")
+def generate_motion(request_json):
+    """Generate a real ARDY motion NPZ for the CozyClay browser client."""
+    prompt, duration, seed, waypoints = _validate_request(request_json)
+    model = _load_model()
+
+    from ardy.model.loading import get_env_var
+    from ardy.model.registry import resolve_model_name
+    from ardy.motion_rep.tools import length_to_mask
+    from ardy.postprocess import post_process_motion
+    from ardy.skeleton import SOMASkeleton30
+    from ardy.tools import seed_everything, to_numpy
+
+    device = "cuda:0"
+    skeleton = model.skeleton
+    fps = model.motion_rep.fps
+    frames = int(duration * fps)
+    if seed is not None:
+        seed_everything(seed)
+
+    constraints = []
+    if waypoints:
+        from ardy.constraints import Root2DConstraintSet
+        ordered = sorted(waypoints, key=lambda item: int(item["frame"]))
+        frame_ids = [int(item["frame"]) for item in ordered]
+        if frame_ids[0] < 0 or frame_ids[-1] >= frames or len(frame_ids) != len(set(frame_ids)):
+            raise gr.Error(f"Waypoint frames must be unique and inside 0..{frames - 1}.")
+        xz = [[float(item["x"]), float(item["z"])] for item in ordered]
+        constraints.append(
+            Root2DConstraintSet(
+                skeleton,
+                frame_indices=torch.tensor(frame_ids),
+                root_2d=torch.tensor(xz, device=device, dtype=torch.float32),
+                global_root_heading=None,
+            )
+        )
+
+    lengths = torch.tensor([frames], device=device)
+    observed, mask = None, None
+    if constraints:
+        observed, mask = model.motion_rep.create_conditions_from_constraints_batched(
+            constraints, lengths, to_normalize=True, device=device
+        )
+
+    patch = model.num_frames_per_token
+    max_window = (int(10 * fps) // patch) * patch
+    history = ((max_window - model.gen_horizon_len) // patch) * patch
+    with torch.inference_mode():
+        motion = model(
+            [prompt],
+            frames,
+            num_denoising_steps=int(model.diffusion.num_base_steps),
+            pad_mask=length_to_mask(lengths),
+            first_heading_angle=torch.zeros(1, device=device),
+            motion_mask=mask,
+            observed_motion=observed,
+            cfg_weight=(2.0, 2.0),
+            progress_bar=lambda values: values,
+            crop_history_length=history,
+        )
+        output = model.motion_rep.inverse(motion, is_normalized=True)
+
+    resolved = resolve_model_name(MODEL_NAME, checkpoints_dir=get_env_var("CHECKPOINTS_DIR"))
+    if "g1" not in resolved.lower():
+        output.update(
+            post_process_motion(
+                output["local_rot_mats"], output["root_positions"],
+                output["foot_contacts"], skeleton,
+                constraint_lst=constraints or None,
+            )
+        )
+    if isinstance(skeleton, SOMASkeleton30):
+        output = skeleton.output_to_SOMASkeleton77(output)
+    output = to_numpy(output)
+    arrays = {
+        key: (value[0] if hasattr(value, "shape") and value.ndim > 0 and value.shape[0] == 1 else value)
+        for key, value in output.items()
+    }
+    arrays["fps"] = np.asarray(fps)
+    arrays["text"] = np.asarray(prompt)
+
+    out_dir = Path(tempfile.mkdtemp(prefix="cozyclay-"))
+    output_path = out_dir / "generated-motion.npz"
+    np.savez(output_path, **arrays)
+    report = {
+        "target_space": "skeleton_joint_center",
+        "frames": frames,
+        "fps": int(fps),
+        "model": resolved,
+        "prompt": prompt,
+        "bytes": output_path.stat().st_size,
+        "output_name": output_path.name,
+        "waypoints": len(waypoints),
+        "device": torch.cuda.get_device_name(0),
+    }
+    return str(output_path), report
+
+
+with gr.Blocks(title="CozyClay ZeroGPU Motion Engine") as demo:
+    gr.Markdown(
+        "# CozyClay · ZeroGPU Motion Engine\n"
+        "Natural-language ARDY motion generation for the browser-based CozyClay studio. "
+        "Use the full editor at **[VIDraft/CozyClay](https://huggingface.co/spaces/VIDraft/CozyClay)**."
+    )
+    with gr.Row():
+        prompt = gr.Textbox(
+            value="A person walks forward, turns left, and waves.",
+            label="Motion prompt",
+        )
+        duration = gr.Slider(1, 10, value=5, step=0.5, label="Seconds")
+    run = gr.Button("Generate 3D motion", variant="primary")
+    motion_file = gr.File(label="Generated ARDY NPZ")
+    report = gr.JSON(label="Generation report")
+
+    def make_request(text, seconds):
+        return json.dumps({"prompt": text, "duration": seconds})
+
+    request = gr.Textbox(visible=False)
+    run.click(make_request, [prompt, duration], request, queue=False).then(
+        generate_motion,
+        request,
+        [motion_file, report],
+        api_name="generate_motion",
+        concurrency_limit=1,
+    )
+
+demo.queue(default_concurrency_limit=1)
+

packages.txt

@@ -0,0 +1,3 @@
+cmake
+g++
+

requirements.txt

@@ -0,0 +1,5 @@
+spaces
+gradio==6.8.0
+torch==2.10.0
+git+https://github.com/nv-tlabs/ardy.git@693f74d13b3d04a0a22ce127ee79c929dd89756b
+