"""Gradio API replacing kimodo_demo's Viser entrypoint. Exposes a single endpoint at `/gradio_api/call/kimodo_motion` that accepts: (prompt, num_frames, seed, cfg, num_steps, constraints_json) and returns a JSON envelope: { "status": "ok", "numFrames": int, "fps": 30, "rootTranslation": [[x,y,z], ...], # [N, 3] "jointRotMats": [[[[...]]]], # [N, 30, 3, 3] "footContacts": [[lh, lt, rh, rt]], # [N, 4] (optional) "summary": str } The webapp's src/lib/services/kimodo.ts polls `/gradio_api/call/kimodo_motion/` for the SSE event stream. """ from __future__ import annotations import json import os import sys import traceback import gradio as gr import numpy as np import torch from constraints_schema import parse_constraints # Lazy imports of kimodo so import-time failures (e.g. missing CUDA on the # Space build container) don't kill `python server.py --help`. _model = None _skeleton = None _device = None def _load_model(): global _model, _skeleton, _device if _model is not None: return _model, _skeleton, _device print("[server] loading Kimodo-SOMA-RP-v1.1 ...", file=sys.stderr, flush=True) from kimodo import load_model _device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"[server] device = {_device}", file=sys.stderr, flush=True) model, resolved = load_model( "Kimodo-SOMA-RP-v1.1", device=_device, default_family="Kimodo", return_resolved_name=True, ) print(f"[server] resolved model = {resolved}", file=sys.stderr, flush=True) _model = model _skeleton = model.skeleton return _model, _skeleton, _device def kimodo_motion( prompt: str, num_frames: int, seed: int, cfg: float, num_steps: int, constraints_json: str, progress: gr.Progress = gr.Progress(), # noqa: B008 — Gradio convention ) -> dict: """Generate one SOMA motion sample. Heavy work runs on the GPU; constraint parsing on the CPU. Returns the JSON envelope the webapp expects.""" try: if not prompt or not prompt.strip(): return {"status": "error", "error": "prompt is empty"} n = int(num_frames) if n < 1 or n > 300: return { "status": "error", "error": f"num_frames must be in [1, 300]; got {n}", } # Validate the constraints payload up front so a bad request doesn't # waste GPU time. We accept the same JSON the kimodo CLI accepts — # extra cross-list validation in constraints_schema bounds-checks frame # indices against num_frames. try: raw = json.loads(constraints_json) if constraints_json else [] parse_constraints(raw, n) # validates shape + bounds except (ValueError, json.JSONDecodeError) as e: return {"status": "error", "error": f"constraint validation: {e}"} progress(0.02, desc="Loading model...") model, skeleton, device = _load_model() # Convert the JSON list of dicts into kimodo constraint objects via # the official loader — accepts a list-of-dicts directly. from kimodo.constraints import load_constraints_lst constraint_lst = load_constraints_lst(raw, skeleton, device=device) if seed is not None and int(seed) >= 0: from kimodo.tools import seed_everything seed_everything(int(seed)) progress(0.10, desc=f"Diffusion ({int(num_steps)} steps)...") cfg_kwargs = {"cfg_type": "regular", "cfg_weight": float(cfg)} # Single sample, single prompt. If you want multi-prompt later, this is # where you'd thread it through. output = model( [prompt.strip()], [n], constraint_lst=constraint_lst, num_denoising_steps=int(num_steps), num_samples=1, multi_prompt=True, num_transition_frames=20, return_numpy=True, **cfg_kwargs, ) progress(0.92, desc="Serializing...") # Output keys we know exist (per generate.py): posed_joints, global_rot_mats. # Shapes: posed_joints [n_samples, T, J, 3], global_rot_mats [n_samples, T, J, 3, 3]. if "posed_joints" not in output or "global_rot_mats" not in output: return { "status": "error", "error": f"unexpected model output keys: {list(output.keys())}", } posed_joints = output["posed_joints"] global_rot_mats = output["global_rot_mats"] if posed_joints.ndim != 4 or global_rot_mats.ndim != 5: return { "status": "error", "error": ( f"unexpected shapes: posed_joints={posed_joints.shape}, " f"global_rot_mats={global_rot_mats.shape}" ), } # Convert global rotation matrices → local (parent-relative) so the # client can apply per-bone rotations directly without doing inverse # FK in the browser. from kimodo.skeleton import global_rots_to_local_rots joints_pos_t = torch.from_numpy(posed_joints[0]).to(device) joints_rot_t = torch.from_numpy(global_rot_mats[0]).to(device) local_rot_mats_t = global_rots_to_local_rots(joints_rot_t, skeleton) local_rot_mats = local_rot_mats_t.detach().cpu().numpy().astype(np.float32) # Root translation = posed_joints at the root joint index. root_idx = int(getattr(skeleton, "root_idx", 0)) root_translation = ( joints_pos_t[:, root_idx, :].detach().cpu().numpy().astype(np.float32) ) # Spot-check the SOMA shape: 30 joints expected for SOMA-RP-v1.1. T, J = local_rot_mats.shape[0], local_rot_mats.shape[1] if (T, J) != (n, 30): return { "status": "error", "error": ( f"expected ({n}, 30, 3, 3) for local_rot_mats, got " f"{local_rot_mats.shape}" ), } # Optional foot_contacts if the model emitted them. foot_contacts_out = None if "foot_contacts" in output: fc = output["foot_contacts"] # Drop the leading sample dim if present if fc.ndim == 3: fc = fc[0] foot_contacts_out = np.asarray(fc, dtype=np.float32).tolist() progress(1.0, desc="Done") return { "status": "ok", "numFrames": int(T), "fps": int(getattr(model, "fps", 30)), "rootTranslation": root_translation.tolist(), "jointRotMats": local_rot_mats.tolist(), "footContacts": foot_contacts_out, "summary": prompt.strip(), } except Exception as e: traceback.print_exc() return {"status": "error", "error": f"{type(e).__name__}: {e}"} with gr.Blocks(title="Genga Kimodo") as demo: gr.Markdown( "# Genga × Kimodo\n" "API-only Space. Inference endpoint at `/gradio_api/call/kimodo_motion`.\n\n" "This Space backs the GengaMachines webapp and is not a public sandbox. " "For the official interactive Kimodo demo, see " "[nvidia/Kimodo](https://huggingface.co/spaces/nvidia/Kimodo)." ) in_prompt = gr.Textbox(label="Prompt", value="A person waves hello with their right hand.") in_frames = gr.Slider(30, 300, value=90, step=6, label="num_frames (30 fps)") in_seed = gr.Number(value=42, label="seed (use -1 to skip seeding)", precision=0) in_cfg = gr.Slider(1.0, 10.0, value=5.0, step=0.5, label="cfg_weight") in_steps = gr.Slider(10, 50, value=30, step=1, label="num_denoising_steps") in_constraints = gr.Textbox(label="constraints_json", value="[]", lines=4) btn = gr.Button("Generate") out = gr.JSON(label="result") btn.click( fn=kimodo_motion, inputs=[in_prompt, in_frames, in_seed, in_cfg, in_steps, in_constraints], outputs=out, api_name="kimodo_motion", ) if __name__ == "__main__": demo.queue(max_size=4).launch( server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)), show_api=True, )