diff --git a/CHANGELOG.md b/CHANGELOG.md index 94526581..5743416a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,14 @@ Notable changes to vla.cpp. Format loosely follows [Keep a Changelog](https://ke ### Added +- **New model: Cosmos3 Policy** (`lerobot/Cosmos3-Edge-Policy-LIBERO`, a world-action model: a 28-layer + Mixture-of-Transformers denoises 4 latent video frames and a 16-step action chunk together, 30 UniPC + steps). `src/models/cosmos3_policy.cpp` follows lerobot's `policies/cosmos3`: the Wan2.2 VAE encoder and the + two-stream backbone as ggml graphs, the UniPC solver on the host, the canvas, the prompt (JSON caption, chat + template, byte-level BPE from the GGUF) and the rot6d decode for `vla-cli --text ... --libero`. + `scripts/convert_cosmos3_policy_to_gguf.py` converts the checkpoint. Agrees with lerobot's policy in bf16 to + 2e-3 to 2e-2 on the action chunk; the client preset `cosmos3_policy` and `eval/run_libero.sh -m cosmos3_policy` + serve it. CPU and CUDA only. - **Hexagon NPU on the Qualcomm Dragonwing IQ-10 (Linux).** `scripts/build_linux_hexagon.sh` cross-builds vla.cpp with `GGML_HEXAGON=ON` in Qualcomm's toolchain container, `vla-server` included. On the IQ-10 (v81, diff --git a/CMakeLists.txt b/CMakeLists.txt index 3d78ea00..9310d400 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -259,6 +259,7 @@ add_library(vla_core ${VLA_CORE_LIB_TYPE} src/models/turbovla.cpp src/models/act.cpp src/models/octo.cpp + src/models/cosmos3_policy.cpp src/tokenizer.cpp ) target_include_directories(vla_core diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index a2cdd710..d46e00cc 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -29,7 +29,7 @@ case say so in the commit message and back it with a LIBERO sweep. `VLA_IMG_SIZE` must match the model or `predict` returns empty: 512 for SmolVLA, 448 for Evo-1, 256 for GR00T N1.7 and VLA-JEPA, 256 with 2 views for -TurboVLA, 224 for the rest. Other knobs: `VLA_BENCH_ITERS` (timing), +TurboVLA and Cosmos3 Policy, 224 for the rest. Other knobs: `VLA_BENCH_ITERS` (timing), `VLA_TIMING=phase`, `VLA_EXTRA_TOKEN` / `VLA_EXTRA_COUNT` (VLA-JEPA needs its `` tokens; Octo needs exactly 16 tokens, so `VLA_EXTRA_TOKEN=0 VLA_EXTRA_COUNT=10`), `VLA_OCTO_UNNORM_DATASET` (Octo, e.g. diff --git a/README.md b/README.md index 7b3e6873..b38eaec7 100644 --- a/README.md +++ b/README.md @@ -101,7 +101,8 @@ Use [ggml-compiler](https://github.com/khanhnd61-vr/ggml-compiler) and follow [t | [VLA-JEPA](https://hf.co/vrfai/vla-jepa-libero) | Y | Y | Y | Y | Y | Y | Y | | [Octo-Small](https://hf.co/vrfai/octo-small-libero-gguf) | Y | Y | Y | Y | Y | - | Y | | [TurboVLA](https://hf.co/vrfai/turbovla-libero-gguf) | Y | Y | Y | Y | Y | Y | Y | -| [ACT](https://hf.co/ravediamond/act-aloha-sim-transfer-cube-gguf) | Y | Y | Y | ~ | Y | ~ | Y | +| [ACT](https://hf.co/ravediamond/act-aloha-sim-transfer-cube-gguf) | Y | Y | Y | ~ | Y | ~ | ~ | +| [Cosmos3 Policy](https://hf.co/vrfai/cosmos3-policy-libero-gguf) | ~ | ~ | - | - | - | - | - | Legend: `Y` = supported (released and benchmarked), `~` = in progress, `-` = planned. diff --git a/eval/client/run_sim_client_direct.py b/eval/client/run_sim_client_direct.py index 7f3f8786..6adbd857 100644 --- a/eval/client/run_sim_client_direct.py +++ b/eval/client/run_sim_client_direct.py @@ -95,6 +95,11 @@ "it, gr00t_n1_7 returns the raw normalized [40, 132] chunk and gr00t_n1_6 " "returns the raw [50, 128] chunk (won't drive LIBERO).", ) + parser.add_argument( + "--gguf", type=str, default=None, + help="[cosmos3_policy] the GGUF vla-server runs: the caption settings (view description, fps, canvas, ...) " + "are read from it, so the prompt is the one the server's model was converted for.", + ) parser.add_argument( "--bitvla-unnorm-key", type=str, default=None, help="[bitvla/gr00t_n1_6/gr00t_n1_7] embodiment key inside dataset_statistics.json " @@ -120,6 +125,7 @@ n_action_steps=args.n_action_steps, stats_json=args.stats_json, bitvla_unnorm_key=args.bitvla_unnorm_key, + cosmos3_gguf=args.gguf, ) if args.arch == "evo1": client = Evo1PipelineAdapter(client=client) diff --git a/eval/client/vla_cpp_client.py b/eval/client/vla_cpp_client.py index 501176e8..28ceedbd 100644 --- a/eval/client/vla_cpp_client.py +++ b/eval/client/vla_cpp_client.py @@ -58,6 +58,12 @@ }, "gr00t_n1_7": {"image_size": 256, "tokenizer": "nvidia/Cosmos-Reason2-2B", "max_state_dim": 132}, + # Cosmos3 Policy (LIBERO): two 256x256 views, no state, a 16x10 rot6d chunk that the client turns into 16 LIBERO + # commands; run it with --n-action-steps 16 (the policy executes the whole chunk). The tokenizer is the one + # the policy was trained with (config.json tokenizer_name); a local copy works too (`tokenizer/` of the + # checkpoint) with --tokenizer. + "cosmos3_policy": {"image_size": 256, "tokenizer": "nvidia/Cosmos3 Policy", "max_state_dim": 0, "max_length": 192}, + "gr00t_n1_5": {"image_size": 224, "tokenizer": "lerobot/eagle2hg-processor-groot-n1p5", "max_state_dim": 64, "trust_remote_code": True}, @@ -91,6 +97,90 @@ + "." ) +# Cosmos3 Policy: the prompt is lerobot's Cosmos3PromptProcessorStep caption in the chat template of +# Cosmos3TokenizerProcessorStep. The caption's settings (view description, fps, chunk, canvas, aspect ratio, +# use_state) are the ones the GGUF was converted with (`cosmos3_policy.caption.*`, what src/models/cosmos3_policy_prompt.cpp +# reads for vla-cli --text), so the client reads them from the GGUF vla-server runs (--gguf). Without one it assumes +# the values of the published LIBERO checkpoint, and says so. +COSMOS3_POLICY_N_VIEWS = 2 +COSMOS3_POLICY_CONCAT_VIEW = "This video contains concatenated views from multiple camera perspectives." +COSMOS3_POLICY_START_OF_GENERATION = "<|vision_start|>" +COSMOS3_POLICY_LIBERO_SETTINGS = { # lerobot/Cosmos3-Edge-Policy-LIBERO, compiled for a 192x320 canvas + "view_description": "The left half shows the third-person view; the right half shows the wrist-mounted camera.", + "aspect_ratio": "16,9", "fps": 20.0, "use_state": False, "chunk_size": 16, + "canvas_h": 192, "canvas_w": 320, "prompt_len": 192, +} + + +def _cosmos3_policy_settings(gguf_path: str | Path) -> dict: + """The caption settings of a Cosmos3 Policy GGUF (written by scripts/convert_cosmos3_policy_to_gguf.py --policy-dir).""" + try: + import gguf # llama.cpp's gguf-py + except ImportError as e: + raise ImportError("reading the GGUF needs the gguf package (pip install gguf)") from e + reader = gguf.GGUFReader(str(gguf_path)) + + def field(key: str, cast): + f = reader.fields.get(key) + if f is None: + raise KeyError(f"{gguf_path} has no {key}; convert it with --policy-dir") + v = f.parts[-1] + return bytes(v).decode() if cast is str else cast(v[0]) + + return { + "view_description": field("cosmos3_policy.caption.view_description", str), + "aspect_ratio": field("cosmos3_policy.caption.aspect_ratio", str), + "fps": field("cosmos3_policy.caption.fps", float), + "use_state": bool(field("cosmos3_policy.caption.use_state", int)), + "chunk_size": field("cosmos3_policy.chunk_size", int), + "canvas_h": field("cosmos3_policy.canvas_h", int), + "canvas_w": field("cosmos3_policy.canvas_w", int), + "prompt_len": field("cosmos3_policy.prompt_len", int), + } + + +def _cosmos3_policy_caption(task: str, s: dict) -> str: + """Cosmos3PromptProcessorStep.caption with 0 idle frames for the settings `s`, as the C++ builds it.""" + task = task.strip() + if not task: + return "" + if not task.endswith((".", "!", "?")): + task = f"{task}." + duration = (s["chunk_size"] + 1) / s["fps"] + rounded = round(duration) + framing = " ".join(t for t in (COSMOS3_POLICY_CONCAT_VIEW, s["view_description"].strip()) if t) + return json.dumps({ + "cinematography": {"framing": framing}, + "actions": [{"time": f"0:00-{rounded // 60}:{rounded % 60:02d}", "description": task, + "idle_frame": f"0 out of {s['chunk_size'] + (1 if s['use_state'] else 0)}."}], + "duration": f"{int(duration)}s", + "fps": float(int(s["fps"])), + "resolution": {"H": s["canvas_h"], "W": s["canvas_w"]}, + "aspect_ratio": s["aspect_ratio"], + }) + + +def _cosmos3_policy_libero(chunk: np.ndarray) -> np.ndarray: + """[16, 10] rot6d rows -> [16, 7] LIBERO commands, as lerobot's Cosmos3LiberoActionProcessorStep does: + the two columns and their cross product make a matrix, projected onto SO(3) by an SVD, then converted to an + axis-angle vector (all in float64); the gripper goes from 1 open / 0 closed to -1 open / +1 closed.""" + a = np.asarray(chunk, dtype=np.float64) + first, second = a[:, 3:6], a[:, 6:9] + m = np.stack([first, second, np.cross(first, second)], axis=-1) # columns + u, _, vh = np.linalg.svd(m) + reflection = np.linalg.det(u @ vh) < 0 + u = np.where(reflection[:, None, None], u * np.array([1.0, 1.0, -1.0]), u) + r = u @ vh + skew = np.stack([r[:, 2, 1] - r[:, 1, 2], r[:, 0, 2] - r[:, 2, 0], r[:, 1, 0] - r[:, 0, 1]], -1) + sin_a = np.linalg.norm(skew, axis=-1) / 2 + cos_a = (np.trace(r, axis1=-2, axis2=-1) - 1) / 2 + angle = np.arctan2(sin_a, cos_a) + scale = np.where(sin_a > 1e-8, angle / np.maximum(sin_a, 1e-8), 1.0) + axis_angle = skew * scale[:, None] / 2 + gripper = np.clip(1.0 - 2.0 * a[:, 9:], -1.0, 1.0) + return np.concatenate([a[:, :3], axis_angle, gripper], axis=-1).astype(np.float32) + + def _load_pb(): if "vla_pb2" in sys.modules: @@ -159,6 +249,7 @@ def __init__( stats_json: str | Path | None = None, bitvla_unnorm_key: str | None = None, rel_stats_json: str | Path | None = None, + cosmos3_gguf: str | Path | None = None, ): if arch not in ARCH_PRESETS: raise ValueError(f"unknown arch {arch!r}; expected one of {sorted(ARCH_PRESETS)}") @@ -166,6 +257,17 @@ def __init__( tokenizer_name = tokenizer_name if tokenizer_name is not None else preset["tokenizer"] image_size = image_size if image_size is not None else preset["image_size"] max_state_dim = max_state_dim if max_state_dim is not None else preset.get("max_state_dim", 32) + # cosmos3_policy: the caption settings and the prompt length come from the GGUF the server runs. + self._cosmos3 = None + if arch == "cosmos3_policy": + if cosmos3_gguf: + self._cosmos3 = _cosmos3_policy_settings(cosmos3_gguf) + if max_length is None: + max_length = self._cosmos3["prompt_len"] + else: + self._cosmos3 = dict(COSMOS3_POLICY_LIBERO_SETTINGS) + print("vla-cpp-direct: no --gguf given; the caption is built with the settings of " + "lerobot/Cosmos3-Edge-Policy-LIBERO (pass the server's GGUF to read them from it)", flush=True) max_length = max_length if max_length is not None else preset.get("max_length", 48) if tokenizer_name is None: raise ValueError( @@ -670,6 +772,8 @@ def _predict_chunk(self, observations: dict[str, Any]) -> np.ndarray: if self.arch == "turbovla": return self._predict_chunk_turbovla(observations) + if self.arch == "cosmos3_policy": + return self._predict_chunk_cosmos3_policy(observations) if self.arch == "evo1": return self._predict_chunk_evo1(observations) if self.arch == "gr00t_n1_7": @@ -820,6 +924,63 @@ def _predict_chunk_turbovla(self, observations: dict[str, Any]) -> np.ndarray: ) return self._turbovla_action_unnorm(chunk) + def _cosmos3_policy_prompt_ids(self, task: str) -> list[int]: + """Token ids of the prompt of `task`, as Cosmos3TokenizerProcessorStep.encode builds them: one user turn + in the chat template, the end-of-turn token and <|vision_start|>, no special tokens added.""" + messages = [{"role": "user", "content": [{"type": "text", "text": _cosmos3_policy_caption(task, self._cosmos3)}]}] + text = self.tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + text += self.tok.eos_token + COSMOS3_POLICY_START_OF_GENERATION + return self.tok(text, add_special_tokens=False)["input_ids"] + + def _predict_chunk_cosmos3_policy(self, observations: dict[str, Any]) -> np.ndarray: + """Two views, the instruction as tokens, no state; the server returns a [16, 10] rot6d chunk in dataset + units (the engine de-normalizes), decoded here to [16, 7] LIBERO commands.""" + if len(self.image_keys) != COSMOS3_POLICY_N_VIEWS: + raise ValueError(f"Cosmos3 Policy takes {COSMOS3_POLICY_N_VIEWS} views [third person, wrist], " + f"got image keys {self.image_keys}") + images_f32: list[np.ndarray] = [] + for key in self.image_keys: + if key not in observations: + raise KeyError(f"Cosmos3 Policy image key {key!r} missing; got {list(observations)}") + image = observations[key] + if isinstance(image, torch.Tensor): + image = image.numpy() + image = np.asarray(image, dtype=np.float32) + if image.ndim != 3 or image.shape[0] != 3: + raise ValueError(f"{key}: expected CHW float32 [3,H,W] in [0,1], got {image.shape}") + images_f32.append(np.ascontiguousarray(np.transpose(image, (1, 2, 0)), dtype=np.float32)) + + task = observations.get("task", "") + if isinstance(task, bytes): + task = task.decode() + if not isinstance(task, str) or not task.strip(): + raise ValueError("Cosmos3 Policy requires a non-empty task instruction") + lang = self._cosmos3_policy_prompt_ids(task) + if len(lang) > self.max_length: + raise ValueError(f"the prompt is {len(lang)} tokens; the model was compiled for {self.max_length}") + + req = self.pb.PredictRequest() + req.request_id = self._step + self._step += 1 + for image in images_f32: + proto_image = req.images.add() + proto_image.encoding = self.pb.Image.F32_RGB_01 + proto_image.height, proto_image.width = image.shape[:2] + proto_image.data = image.tobytes() + req.lang_tokens.extend(lang) + # no state, and no noise: the server samples it + + self.sock.send(req.SerializeToString()) + resp = self.pb.PredictResponse() + resp.ParseFromString(self.sock.recv()) + if resp.error: + raise RuntimeError(f"vla-server error: {resp.error}") + self._last_response = resp + chunk = np.array(resp.action_chunk, dtype=np.float32).reshape(resp.chunk_size, resp.action_dim) + if chunk.shape != (self._cosmos3["chunk_size"], 10): + raise RuntimeError(f"expected a {self._cosmos3['chunk_size']}x10 chunk from the Cosmos3 Policy server, got {chunk.shape}") + return _cosmos3_policy_libero(chunk) + _VJ_PS, _VJ_TPS, _VJ_MERGE, _VJ_SIDE = 16, 2, 2, 256 @classmethod diff --git a/eval/run_libero.sh b/eval/run_libero.sh index 6b90d6ca..1fc8f44f 100644 --- a/eval/run_libero.sh +++ b/eval/run_libero.sh @@ -38,7 +38,7 @@ Usage: $(basename "$0") -i [-o ] [-n ] [- -m MODEL which model to run: smol | pi0 | pi05 | bit | evo1 | vla_adapter | openvla_oft | gr00t_n1_5 | gr00t_n1_6 | gr00t_n1_7 | - octo | turbovla | vla_jepa | all + octo | turbovla | vla_jepa | cosmos3_policy | all (default: all) -h show this help @@ -49,6 +49,10 @@ Env overrides: BIND_ADDR, CLIENT_ADDR, BITVLA_TOKENIZER, GR00T_N1_6_TOKENIZER, SERVER_BIN (prebuilt vla-server; setting it skips the build), VLA_JEPA_PYTHON (python for the vla_jepa client; it needs LIBERO and transformers>=5.4, default is the LIBERO venv), + COSMOS3_POLICY_PYTHON (python for the cosmos3_policy client, which needs LIBERO and a + transformers that reads the policy's tokenizer, default is the LIBERO venv), + COSMOS3_POLICY_TOKENIZER (the policy's tokenizer/ directory, default nvidia/Cosmos3 Policy + from the Hub), COSMOS3_POLICY_RECV_TIMEOUT_MS (client wait per chunk, default 300000), VLA_FIXED_NOISE_SEED (client-side noise, for paired A/B runs) EOF } @@ -74,9 +78,9 @@ done shift $((OPTIND - 1)) case "${MODEL}" in - smol|pi0|pi05|bit|evo1|vla_adapter|openvla_oft|gr00t_n1_5|gr00t_n1_6|gr00t_n1_7|octo|turbovla|vla_jepa|all) ;; + smol|pi0|pi05|bit|evo1|vla_adapter|openvla_oft|gr00t_n1_5|gr00t_n1_6|gr00t_n1_7|octo|turbovla|vla_jepa|cosmos3_policy|all) ;; *) - echo "ERROR: -m must be one of: smol | pi0 | pi05 | bit | evo1 | vla_adapter | openvla_oft | gr00t_n1_5 | gr00t_n1_6 | gr00t_n1_7 | octo | turbovla | vla_jepa | all (got '${MODEL}')" >&2 + echo "ERROR: -m must be one of: smol | pi0 | pi05 | bit | evo1 | vla_adapter | openvla_oft | gr00t_n1_5 | gr00t_n1_6 | gr00t_n1_7 | octo | turbovla | vla_jepa | cosmos3_policy | all (got '${MODEL}')" >&2 exit 1 ;; esac @@ -113,6 +117,9 @@ TASK_IDS="${TASK_IDS:-0 1 2 3 4 5 6 7 8 9}" PALIGEMMA_TOKENIZER="${PALIGEMMA_TOKENIZER:-}" VENV_PY="${REPO_ROOT}/eval/sim/libero/libero_uv/.venv/bin/python" VLA_JEPA_PYTHON="${VLA_JEPA_PYTHON:-${VENV_PY}}" +COSMOS3_POLICY_PYTHON="${COSMOS3_POLICY_PYTHON:-${VENV_PY}}" +COSMOS3_POLICY_TOKENIZER="${COSMOS3_POLICY_TOKENIZER:-}" +COSMOS3_POLICY_RECV_TIMEOUT_MS="${COSMOS3_POLICY_RECV_TIMEOUT_MS:-300000}" CLIENT="${REPO_ROOT}/eval/client/run_sim_client_direct.py" BIND_ADDR="${BIND_ADDR:-tcp://*:5555}" CLIENT_ADDR="${CLIENT_ADDR:-tcp://localhost:5555}" @@ -147,6 +154,7 @@ N_ACTION_STEPS_GR00T_N1_7="${N_ACTION_STEPS_GR00T_N1_7:-16}" # N1.7 H4 closeout N_ACTION_STEPS_OCTO="${N_ACTION_STEPS_OCTO:-4}" N_ACTION_STEPS_TURBOVLA="${N_ACTION_STEPS_TURBOVLA:-12}" N_ACTION_STEPS_VLA_JEPA="${N_ACTION_STEPS_VLA_JEPA:-7}" +N_ACTION_STEPS_COSMOS3_POLICY="${N_ACTION_STEPS_COSMOS3_POLICY:-16}" # Cosmos3 Policy executes the whole 16-action chunk mkdir -p "${OUTPUT_ROOT}" OUTPUT_ROOT="$(cd "${OUTPUT_ROOT}" && pwd)" @@ -175,7 +183,12 @@ if [[ ! -x "${SERVER_BIN}" ]]; then echo "ERROR: ${SERVER_BIN} not found after build." >&2 exit 1 fi -if [[ ! -x "${VENV_PY}" ]]; then +# The cosmos3_policy client may run in another python that has LIBERO (COSMOS3_POLICY_PYTHON), as vla_jepa's may. +need_libero_venv=1 +if [[ "${MODEL}" == cosmos3_policy && "${COSMOS3_POLICY_PYTHON}" != "${VENV_PY}" && -x "${COSMOS3_POLICY_PYTHON}" ]]; then + need_libero_venv=0 +fi +if (( need_libero_venv )) && [[ ! -x "${VENV_PY}" ]]; then echo "ERROR: LIBERO venv not found at ${VENV_PY}." >&2 echo " Run: bash eval/sim/libero/setup_libero.sh" >&2 exit 1 @@ -316,6 +329,17 @@ run_model() { if [[ "${arch}" == vla_jepa ]]; then client_py="${VLA_JEPA_PYTHON}" fi + # cosmos3_policy: the tokenizer is the policy's (nvidia/Cosmos3 Policy on the Hub unless a local copy is given), the + # caption settings are read from the GGUF the server runs, and a chunk takes seconds on a GPU but about a minute + # on a CPU, so the client waits longer than its 30 s default. + if [[ "${arch}" == cosmos3_policy ]]; then + client_py="${COSMOS3_POLICY_PYTHON}" + client_extra+=(--gguf "${server_args[0]}") + if [[ -n "${COSMOS3_POLICY_TOKENIZER}" ]]; then + client_extra+=(--tokenizer "${COSMOS3_POLICY_TOKENIZER}") + fi + client_extra+=(--recv-timeout-ms "${COSMOS3_POLICY_RECV_TIMEOUT_MS}") + fi # bitvla auto-loads tokenizer + dataset_statistics.json from the GGUF repo on # the Hub; only pass --tokenizer when BITVLA_TOKENIZER overrides with a local dir. @@ -552,5 +576,20 @@ if should_run vla_jepa; then fi fi +# cosmos3_policy: the GGUF of scripts/convert_cosmos3_policy_to_gguf.py, not +# published yet, so a missing one is skipped instead of failing a sweep over all models. +if should_run cosmos3_policy; then + cosmos3_gguf="${MODELS_ROOT}/cosmos3-policy-libero-gguf/cosmos3-policy-libero-bf16.gguf" + if [[ ! -f "${cosmos3_gguf}" ]]; then + echo "[skip] cosmos3_policy: ${cosmos3_gguf} not found" + else + run_model cosmos3_policy \ + "${MODELS_ROOT}/cosmos3-policy-libero-gguf" \ + "${N_ACTION_STEPS_COSMOS3_POLICY}" \ + "" \ + "${cosmos3_gguf}" + fi +fi + echo "====================" echo "Done. Results under ${OUTPUT_ROOT}" diff --git a/scripts/convert_cosmos3_policy_to_gguf.py b/scripts/convert_cosmos3_policy_to_gguf.py new file mode 100644 index 00000000..516120c4 --- /dev/null +++ b/scripts/convert_cosmos3_policy_to_gguf.py @@ -0,0 +1,369 @@ +# Copyright 2026 VinRobotics +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Cosmos3 Policy (lerobot/Cosmos3-Edge-Policy-LIBERO) -> vla.cpp GGUF. + +Reads the checkpoint directory: config.json, model.safetensors, the postprocessor's unnormalizer stats and +tokenizer/. Only what inference needs is written: + +- the Wan2.2 VAE encoder. Its causal 3-D convolutions see one frame behind two zero frames, so only the last + temporal tap of each kernel reaches the output: that tap is written as a 2-D kernel. The decoder and the + temporal convolutions of the downsamplers (skipped on a first frame) are left out; +- the 28 MoT decoder layers, both streams. The understanding stream's final norm, the language head, the + SigLIP vision tower and its projector play no part in action inference and are left out; +- the projections between the backbone and the latents: vae2llm, llm2vae, and action2llm / llm2action of the + config's embodiment (`domain_id`), cut to `action_dim` since the other channels are zero; +- the timestep embedding, in f32 as the policy keeps it; +- the latent statistics of the VAE, the action quantiles of the unnormalizer, the canvas layout of the policy's + cameras, the caption settings of the prompt step and the byte-level BPE tokenizer (llama.cpp's + `tokenizer.ggml.*` keys), so `vla-cli --text` builds the prompt itself (src/models/cosmos3_policy.cpp). + +GEMM weights keep the checkpoint's bf16 (the engine casts them to its matmul type at load), norms and biases +are f32. + + python scripts/convert_cosmos3_policy_to_gguf.py --ckpt --out cosmos3-policy-libero-bf16.gguf +""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import numpy as np +import torch +from safetensors import safe_open + +from gguf_common import ( + add, + add_array, + add_f32, + arg_parser, + finish, + kv_f32, + kv_prefix, + kv_u32, + open_writer, + read_json, + require, + resolve_out, +) + +ARCH = "cosmos3_policy" +KV = kv_prefix(ARCH) +NET = "model.network." +LM = NET + "mot_backbone.model.language_model." + +# Video canvases of the 256 tier by aspect ratio name, (width, height): lerobot's VIDEO_CANVAS_SIZES. +CANVAS_256 = {"1,1": (256, 256), "4,3": (320, 256), "3,4": (256, 320), "16,9": (320, 192), "9,16": (192, 320)} + + +def canvas_layout(config: dict) -> dict: + """The canvas of the policy's cameras, as lerobot's Cosmos3VideoCanvasProcessorStep lays it out: every view + at `view_size`, two side by side, fit into the closest canvas of the 256 tier without enlarging, then the + latent grid of the content. One or two cameras are supported (three use another layout).""" + cameras = [k for k, f in config["input_features"].items() if f["type"] == "VISUAL"] + if len(cameras) not in (1, 2): + raise SystemExit(f"{len(cameras)} cameras; the canvas layout is written for one or two") + view_h, view_w = (int(x) for x in config["view_size"]) + if min(view_h, view_w) > 256: + raise SystemExit(f"a {view_h}x{view_w} view needs a higher resolution tier than the 256 one written here") + frame_h, frame_w = view_h, view_w * len(cameras) + canvas_w, canvas_h = min(CANVAS_256.values(), key=lambda s: abs(frame_h / frame_w - s[1] / s[0])) + scale = min(canvas_w / frame_w, canvas_h / frame_h, 1.0) + content_h, content_w = int(scale * frame_h + 0.5), int(scale * frame_w + 0.5) + spatial = int(config["vae"]["scale_factor_spatial"]) + ratio = next(name for name, size in CANVAS_256.items() if size == (canvas_w, canvas_h)) + return { + "n_views": len(cameras), "view_h": view_h, "view_w": view_w, + "canvas_h": canvas_h, "canvas_w": canvas_w, "content_h": content_h, "content_w": content_w, + "latent_h": content_h // spatial, "latent_w": content_w // spatial, "aspect_ratio": ratio, + "cameras": cameras, + } + + +def conv_tap(w: torch.Tensor) -> torch.Tensor: + """The 2-D kernel a causal 3-D convolution applies to a first frame: the input is padded with KT-1 zero frames + in front, so of the kernel [OC, IC, KT, KH, KW] only tap KT-1 reaches that frame's output.""" + if w.dim() != 5: + raise SystemExit(f"expected a 3-D kernel, got shape {tuple(w.shape)}") + return w[:, :, -1].contiguous() + + +def add_tokenizer(writer, tok_dir: Path) -> None: + """The BPE of tokenizer.json as llama.cpp's GGUF vocabulary keys (gpt2 model, "tekken" pre-tokenizer).""" + d = json.loads((tok_dir / "tokenizer.json").read_text()) + cfg = json.loads((tok_dir / "tokenizer_config.json").read_text()) + m = d["model"] + if m["type"] != "BPE" or not m.get("ignore_merges") or m.get("byte_fallback"): + raise SystemExit("expected a byte-level BPE tokenizer with ignore_merges, as the Cosmos3 tokenizer is") + vocab = m["vocab"] + tokens = [None] * len(vocab) + for tok, i in vocab.items(): + tokens[i] = tok + if any(t is None for t in tokens): + raise SystemExit("the vocabulary has gaps") + types = [1] * len(tokens) # llama.cpp: 1 normal, 2 unknown, 3 control, 4 user defined + for a in d["added_tokens"]: + if a["id"] >= len(tokens): # the C++ reads the vocabulary's spellings; a token outside it has none + raise SystemExit(f"added token {a['content']!r} (id {a['id']}) is not in the vocabulary of {len(tokens)}") + types[a["id"]] = 3 if a["special"] else 4 + # The specials of tokenizer_config.json: a string, an AddedToken dict or null (a byte-level BPE needs no + # unknown token); one that is not a token is left out rather than guessed. + ids = dict(vocab) + ids.update((a["content"], a["id"]) for a in d["added_tokens"]) + + def special_id(name: str) -> int | None: + t = cfg.get(name) + if isinstance(t, dict): + t = t.get("content") + return ids.get(t) if isinstance(t, str) else None + + unk = special_id("unk_token") + if unk is not None: + types[unk] = 2 + merges = [" ".join(x) if isinstance(x, list) else x for x in m["merges"]] + writer.add_string("tokenizer.ggml.model", "gpt2") + writer.add_string("tokenizer.ggml.pre", "tekken") + writer.add_array("tokenizer.ggml.tokens", tokens) + writer.add_array("tokenizer.ggml.token_type", types) + writer.add_array("tokenizer.ggml.merges", merges) + for key, name in (("bos_token_id", "bos_token"), ("eos_token_id", "eos_token"), ("unknown_token_id", "unk_token"), + ("padding_token_id", "pad_token")): + i = special_id(name) + if i is not None: + writer.add_uint32(f"tokenizer.ggml.{key}", i) + writer.add_bool("tokenizer.ggml.add_bos_token", False) + print(f"tokenizer: {len(tokens)} tokens, {len(merges)} merges, {sum(t in (3, 4) for t in types)} special") + + +def load_quantiles(ckpt: Path, action_dim: int) -> tuple[np.ndarray, np.ndarray]: + """q01 and q99 of the action from the postprocessor's unnormalizer (QUANTILES normalization).""" + meta = read_json(ckpt / "policy_postprocessor.json") + state_file = next((s.get("state_file") for s in meta.get("steps", []) + if s.get("registry_name") == "unnormalizer_processor"), None) + if not state_file or not (ckpt / state_file).is_file(): + raise SystemExit("policy_postprocessor.json names no unnormalizer state file") + with safe_open(str(ckpt / state_file), framework="pt") as f: + keys = set(f.keys()) + if "action.q01" not in keys or "action.q99" not in keys: + raise SystemExit(f"{state_file} has no action.q01/q99 (the policy normalizes actions by quantiles)") + q01 = f.get_tensor("action.q01").float().numpy().reshape(-1) + q99 = f.get_tensor("action.q99").float().numpy().reshape(-1) + if q01.size != action_dim or q99.size != action_dim: + raise SystemExit(f"action quantiles have {q01.size} values, the action has {action_dim}") + return q01.astype(np.float32), q99.astype(np.float32) + + +class Weights: + """The checkpoint's tensors by name, read one at a time; remembers which were used.""" + + def __init__(self, path: Path): + self.f = safe_open(str(path), framework="pt") + self.keys = set(self.f.keys()) + self.used: set[str] = set() + + def __call__(self, name: str) -> torch.Tensor: + if name not in self.keys: + raise SystemExit(f"missing tensor {name}") + self.used.add(name) + return self.f.get_tensor(name) + + def has(self, name: str) -> bool: + return name in self.keys + + +def check_config(config: dict) -> None: + """The engine is written for the LIBERO policy's setup; refuse what it does not do.""" + if config.get("type") != "cosmos3_edge": # lerobot's policy type of the Edge backbone + raise SystemExit(f"config.json is a {config.get('type')!r} policy, not cosmos3_edge") + text = config["mot_backbone"]["text_config"] + if text["hidden_act"] != "relu2" or text["rope_parameters"]["rope_type"] != "default": + raise SystemExit("the engine is written for Nemotron 3's relu^2 MLP and default M-RoPE") + if text["joint_attn_implementation"] != "two_way": + raise SystemExit("the engine is written for two_way joint attention") + if config["mode"] != "wam" or config.get("use_state") or config["guidance_scale"] != 1.0 or config["n_obs_steps"] != 1: + raise SystemExit("the engine is written for mode wam, no state, no classifier-free guidance and one observation step") + if config["sampler"] != "unipc": + raise SystemExit(f"sampler {config['sampler']!r}: the engine runs UniPC") + if config["latent_patch_size"] != 2 or config["vae"]["patch_size"] != 2: + raise SystemExit("the engine is written for a latent patch of 2 and a VAE pixel patch of 2") + if config["vae"].get("is_residual") is not True: + raise SystemExit("the engine is written for the Wan2.2 VAE (residual down blocks)") + + +def main() -> int: + ap = arg_parser(ARCH, "checkpoint directory (config.json, model.safetensors, policy_postprocessor*, tokenizer/)", + description=__doc__) + args = ap.parse_args() + ckpt = args.ckpt.resolve() + require(ckpt / "config.json", ckpt / "model.safetensors", ckpt / "tokenizer" / "tokenizer.json") + out = args.out.resolve() if args.out else resolve_out(args, ckpt, ARCH) + + config = read_json(ckpt / "config.json") + check_config(config) + text = config["mot_backbone"]["text_config"] + vae = config["vae"] + action_dim = int(config["output_features"]["action"]["shape"][0]) + chunk = int(config["chunk_size"]) + layout = canvas_layout(config) + n_layers = int(text["num_hidden_layers"]) + hidden = int(text["hidden_size"]) + n_blocks = len(vae["dim_mult"]) + + W = Weights(ckpt / "model.safetensors") + writer = open_writer(out, ARCH) + + # --- metadata ------------------------------------------------------------------------------------------- + kv_u32(writer, KV, { + "n_layers": n_layers, "hidden": hidden, "intermediate": text["intermediate_size"], + "n_heads": text["num_attention_heads"], "n_kv_heads": text["num_key_value_heads"], "head_dim": text["head_dim"], + "vocab_size": text["vocab_size"], + "mrope_section_t": text["rope_parameters"]["mrope_section"][0], + "mrope_section_h": text["rope_parameters"]["mrope_section"][1], + "mrope_section_w": text["rope_parameters"]["mrope_section"][2], + "time_embed_dim": config["timestep_embedding_dim"], "temporal_modality_margin": config["temporal_modality_margin"], + "num_steps": config["num_inference_steps"], "num_train_timesteps": config["num_train_timesteps"], + "chunk_size": chunk, "action_dim": action_dim, "max_action_dim": config["max_action_dim"], + "domain_id": config["domain_id"], "latent_patch_size": config["latent_patch_size"], + "vae.z_dim": vae["z_dim"], "vae.base_dim": vae["base_dim"], "vae.patch_size": vae["patch_size"], + "vae.scale_factor_spatial": vae["scale_factor_spatial"], "vae.scale_factor_temporal": vae["scale_factor_temporal"], + "vae.n_blocks": n_blocks, "vae.num_res_blocks": vae["num_res_blocks"], + "n_views": layout["n_views"], "view_h": layout["view_h"], "view_w": layout["view_w"], + "canvas_h": layout["canvas_h"], "canvas_w": layout["canvas_w"], + "content_h": layout["content_h"], "content_w": layout["content_w"], + "latent_h": layout["latent_h"], "latent_w": layout["latent_w"], + "prompt_len": 512, # the longest prompt the engine takes + }) + # Each down block's shortcut averages 2x2 pixels (factor_s) and, where the block also halves time, pairs of + # frames (factor_t); the last block keeps its size. + temporal = list(vae["temperal_downsample"]) + [False] * n_blocks + for i, mult in enumerate(vae["dim_mult"]): + last = i == n_blocks - 1 + kv_u32(writer, KV, {f"vae.down.{i}.dim_mult": mult, + f"vae.down.{i}.factor_t": 2 if (temporal[i] and not last) else 1, + f"vae.down.{i}.factor_s": 1 if last else 2}) + kv_f32(writer, KV, { + "rms_eps": text["rms_norm_eps"], "rope_theta": text["rope_parameters"]["rope_theta"], + "timestep_scale": config["timestep_scale"], "base_fps": config["base_fps"], "fps": config["fps"], + "sampling_shift": config["sampling_shift"], + }) + writer.add_string(KV("caption.view_description"), str(config["view_description"])) + writer.add_string(KV("caption.aspect_ratio"), layout["aspect_ratio"]) + writer.add_float32(KV("caption.fps"), float(config["fps"])) + writer.add_uint32(KV("caption.use_state"), 1 if config.get("use_state") else 0) + writer.add_string(KV("cameras"), ", ".join(layout["cameras"])) + add_tokenizer(writer, ckpt / "tokenizer") + + add_array(writer, KV("latents_mean"), np.asarray(vae["latents_mean"], dtype=np.float32)) + add_array(writer, KV("latents_std"), np.asarray(vae["latents_std"], dtype=np.float32)) + q01, q99 = load_quantiles(ckpt, action_dim) + add_array(writer, KV("action_q01"), q01) + add_array(writer, KV("action_q99"), q99) + + # --- the backbone ---------------------------------------------------------------------------------------- + n = 0 + + def put(name: str, t: torch.Tensor, f32: bool = False) -> None: + nonlocal n + (add_f32 if f32 else add)(writer, name, t) + n += 1 + + put("token_embd.weight", W(LM + "embed_tokens.weight")) + for i in range(n_layers): + p = f"{LM}layers.{i}." + for stream, suffix in (("", ""), ("gen.", "_gen")): + b = f"blk.{i}.{stream}" + put(b + "attn_norm.weight", W(f"{p}input_layernorm{suffix}.weight"), f32=True) + put(b + "ffn_norm.weight", W(f"{p}post_attention_layernorm{suffix}.weight"), f32=True) + for ours, theirs in (("attn_q", "q_proj"), ("attn_k", "k_proj"), ("attn_v", "v_proj"), ("attn_output", "o_proj")): + put(f"{b}{ours}.weight", W(f"{p}self_attn.{theirs}{suffix}.weight")) + put(b + "ffn_up.weight", W(f"{p}mlp{suffix}.up_proj.weight")) + put(b + "ffn_down.weight", W(f"{p}mlp{suffix}.down_proj.weight")) + # The generation stream's QK norms, and the norm of the understanding keys its queries read. + put(f"blk.{i}.gen.attn_q_norm.weight", W(f"{p}self_attn.q_norm_gen.weight"), f32=True) + put(f"blk.{i}.gen.attn_k_norm.weight", W(f"{p}self_attn.k_norm_gen.weight"), f32=True) + put(f"blk.{i}.attn_k_norm_gen.weight", W(f"{p}self_attn.k_norm_und_for_gen.weight"), f32=True) + put("gen.output_norm.weight", W(LM + "norm_gen.weight"), f32=True) + + # --- projections and the timestep embedding ---------------------------------------------------------- + put("vae2llm.weight", W(NET + "vae2llm.weight")) + put("vae2llm.bias", W(NET + "vae2llm.bias"), f32=True) + put("llm2vae.weight", W(NET + "llm2vae.weight")) + put("llm2vae.bias", W(NET + "llm2vae.bias"), f32=True) + domain = int(config["domain_id"]) + max_ad = int(config["max_action_dim"]) + # DomainAwareLinear computes x @ W + b per domain with W [in, out]; nn.Linear's [out, in] is its transpose. + a2l = W(NET + "action2llm.fc.weight")[domain].view(max_ad, hidden)[:action_dim].t().contiguous() # [hidden, A] + put("action2llm.weight", a2l) + put("action2llm.bias", W(NET + "action2llm.bias.weight")[domain], f32=True) + l2a = W(NET + "llm2action.fc.weight")[domain].view(hidden, max_ad)[:, :action_dim].t().contiguous() # [A, hidden] + put("llm2action.weight", l2a) + put("llm2action.bias", W(NET + "llm2action.bias.weight")[domain][:action_dim], f32=True) + put("action_modality_embed.weight", W(NET + "action_modality_embed.weight"), f32=True) + for k in ("linear_1", "linear_2"): + put(f"time_embedder.{k}.weight", W(f"{NET}time_embedder.{k}.weight"), f32=True) + put(f"time_embedder.{k}.bias", W(f"{NET}time_embedder.{k}.bias"), f32=True) + + # --- the VAE encoder --------------------------------------------------------------------------------- + V = NET + "vae." + + def conv3(ours: str, theirs: str) -> None: + put(ours + ".weight", conv_tap(W(theirs + ".weight"))) + put(ours + ".bias", W(theirs + ".bias"), f32=True) + + def norm(ours: str, theirs: str) -> None: + put(ours + ".gamma", W(theirs + ".gamma").flatten(), f32=True) + + def conv1x1(ours: str, theirs: str) -> None: + w = W(theirs + ".weight") # [OC, IC, (1,) 1, 1]: a 1x1 convolution is a linear map + put(ours + ".weight", w.reshape(w.shape[0], w.shape[1]).contiguous()) + put(ours + ".bias", W(theirs + ".bias"), f32=True) + + def resnet(ours: str, theirs: str) -> None: + norm(ours + ".norm1", theirs + ".norm1") + conv3(ours + ".conv1", theirs + ".conv1") + norm(ours + ".norm2", theirs + ".norm2") + conv3(ours + ".conv2", theirs + ".conv2") + if W.has(theirs + ".conv_shortcut.weight"): + conv1x1(ours + ".shortcut", theirs + ".conv_shortcut") + + conv3("vae.conv_in", V + "encoder.conv_in") + for b in range(n_blocks): + for r in range(int(vae["num_res_blocks"])): + resnet(f"vae.down.{b}.res.{r}", f"{V}encoder.down_blocks.{b}.resnets.{r}") + if W.has(f"{V}encoder.down_blocks.{b}.downsampler.resample.1.weight"): + # ZeroPad2d((0, 1, 0, 1)) then Conv2d(dim, dim, 3, stride 2) + put(f"vae.down.{b}.down.weight", W(f"{V}encoder.down_blocks.{b}.downsampler.resample.1.weight")) + put(f"vae.down.{b}.down.bias", W(f"{V}encoder.down_blocks.{b}.downsampler.resample.1.bias"), f32=True) + resnet("vae.mid.res.0", V + "encoder.mid_block.resnets.0") + norm("vae.mid.attn.norm", V + "encoder.mid_block.attentions.0.norm") + conv1x1("vae.mid.attn.qkv", V + "encoder.mid_block.attentions.0.to_qkv") + conv1x1("vae.mid.attn.proj", V + "encoder.mid_block.attentions.0.proj") + resnet("vae.mid.res.1", V + "encoder.mid_block.resnets.1") + norm("vae.norm_out", V + "encoder.norm_out") + conv3("vae.conv_out", V + "encoder.conv_out") + conv1x1("vae.quant_conv", V + "quant_conv") + + skipped = sorted(k for k in W.keys if k not in W.used) + heads = sorted({".".join(k.split(".")[:5]) for k in skipped}) + print(f"{n} tensors written; {len(skipped)} of the checkpoint's left out: {', '.join(heads)}") + print(f"canvas {layout['canvas_h']}x{layout['canvas_w']}, content {layout['content_h']}x{layout['content_w']} of " + f"{layout['n_views']} {layout['view_h']}x{layout['view_w']} views, latent grid {layout['latent_h']}x{layout['latent_w']}; " + f"{config['num_inference_steps']} UniPC steps, chunk {chunk}x{action_dim}, domain {domain}") + return finish(writer, out) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/arch.h b/src/arch.h index ddbfcb17..80d04d64 100644 --- a/src/arch.h +++ b/src/arch.h @@ -73,6 +73,7 @@ enum class Arch { VLA_JEPA, // LeRobot Qwen3-VL-2B-Instruct+V-JEPÀ+DiT-B FM. TURBOVLA, // TurboVLA (DINOv3 + BERT + VL Fusion + ACT decoder). ACT, // LeRobot ACT (ResNet + transformer encoder-decoder, no language). + COSMOS3_POLICY,// NVIDIA Cosmos3 Policy world-action model (Nemotron MoT + Wan2.2 VAE, UniPC). }; /** @@ -251,6 +252,15 @@ std::unique_ptr act_create(const std::string& mmproj_path, const std::string& config_path, const Options& opts); +/** + * @brief Build a Cosmos3 Policy model (LIBERO). The VAE is baked into @p ckpt_path; the graphs are generated. + * @copydetails smolvla_create + */ +std::unique_ptr cosmos3_policy_create(const std::string& mmproj_path, + const std::string& ckpt_path, + const std::string& config_path, + const Options& opts); + /** * @brief Inspect a GGUF and identify the architecture tag. * diff --git a/src/model.cpp b/src/model.cpp index 8d22246e..5e69a617 100644 --- a/src/model.cpp +++ b/src/model.cpp @@ -81,6 +81,7 @@ bool detect_arch_gguf(const std::string& path, Arch* out) { try_str("vla_jepa.architecture", arch_str) || try_str("vla_adapter.architecture", arch_str) || try_str("turbovla.architecture", arch_str) || + try_str("cosmos3_policy.architecture", arch_str) || try_str("act.architecture", arch_str)) { if (arch_str == "smolvla") { *out = Arch::SMOLVLA; @@ -138,6 +139,10 @@ bool detect_arch_gguf(const std::string& path, Arch* out) { *out = Arch::ACT; ok = true; } + else if (arch_str == "cosmos3_policy") { + *out = Arch::COSMOS3_POLICY; + ok = true; + } } gguf_free(gctx); @@ -322,6 +327,10 @@ Model* model_load(const std::string& mmproj_path, const std::string& ckpt_path, std::printf("vla: arch = act\n"); impl = act_create(mmproj_path, ckpt_path, config_path, opts); break; + case Arch::COSMOS3_POLICY: + std::printf("vla: arch = cosmos3_policy\n"); + impl = cosmos3_policy_create(mmproj_path, ckpt_path, config_path, opts); + break; } if (!impl) return nullptr; diff --git a/src/models/cosmos3_policy.cpp b/src/models/cosmos3_policy.cpp new file mode 100644 index 00000000..1fad5fd0 --- /dev/null +++ b/src/models/cosmos3_policy.cpp @@ -0,0 +1,1819 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Cosmos3 Policy (LIBERO): a world-action model. One 28-layer Mixture-of-Transformers backbone (Nemotron 3 +// Dense, two streams) denoises 4 latent video frames and a 16-step action chunk together, conditioned on the +// Wan2.2 VAE latent of the current frame and on the prompt, with 30 UniPC steps. lerobot's +// policies/cosmos3/modeling_cosmos3.py is the reference; this file follows its inference path: +// +// VAE encoder canvas (one frame, pixel-unshuffled to 12 channels) -> normalized latent, 50 patches of 192 +// prefix the understanding stream over the prompt, once: its keys (normed and rotated for the +// generation queries) and values of every layer are kept for the denoising loop +// step x 30 the generation stream over [video patches | actions]: timestep embedding, 28 layers whose +// attention reads [prompt keys | its own keys], then the velocities of latents and actions +// solver UniPC (diffusers' UniPCMultistepScheduler: flow prediction, order 2, bh2, corrector) on the +// host, in f32 like the reference +// output the 16x10 chunk unnormalized with the dataset quantiles: [dpos(3), rot6d(6), gripper] +// +// Input contract (docs/MODELS.md): two camera views in the order [third person, wrist], composed into the +// 192x320 canvas here exactly as lerobot's Cosmos3VideoCanvasProcessorStep does (or one canvas the caller +// composed); the prompt as token ids, which Cosmos3Prompt (below) builds from the instruction +// (vla-cli --text); `Inputs::noise` the 16x10 action noise (the video noise is drawn here from a seed that +// follows it). cosmos3_libero_action decodes the chunk to LIBERO commands. The prompt builder, the tokenizer and +// the decode are the second half of this file. + +#include "arch.h" +#include "backend.h" +#include "cosmos3_policy.h" +#include "gguf_reader.h" +#include "loader.h" +#include "model.h" +#include "options.h" +#include "scratch_ctx.h" +#include "layers/attn.h" +#include "layers/conv.h" +#include "layers/norm.h" +#include "layers/rope.h" + +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "gguf.h" + +#if defined(GGML_USE_CUDA) && !defined(GGML_USE_HIP) +#include +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace vla { +namespace { + +constexpr const char * kTag = "vla(cosmos3_policy)"; + +double now_ms() { + using clock = std::chrono::steady_clock; + return std::chrono::duration(clock::now().time_since_epoch()).count(); +} + +// ------------------------------------------------------------------------------------------------ +// The canvas: lerobot's Cosmos3VideoCanvasProcessorStep +// ------------------------------------------------------------------------------------------------ + +// How the policy's cameras fill the canvas, from the GGUF (scripts/convert_cosmos3_policy_to_gguf.py). +struct Layout { + int n_views = 2, view_h = 256, view_w = 256; // every view resized to view_h x view_w, composed side by side + int canvas_h = 192, canvas_w = 320; // the canvas of the resolution tier + int content_h = 160, content_w = 320; // the composed frame fit into it (never enlarged), top left +}; + +// torch's antialiased resampling (aten UpSampleKernel, align_corners=False): per output index the input +// range and its normalized filter weights. Float arithmetic as in aten, so uint8 results agree. +struct AaTaps { + std::vector xmin, xsize; + std::vector w; // [out][maxk] + int maxk = 0; +}; + +float triangle(float x) { + x = std::fabs(x); + return x < 1.f ? 1.f - x : 0.f; +} + +// Keys' cubic, a = -0.5, written as aten's cubic_convolution1/2 (the same value as the textbook form, but not +// the same float rounding). +float cubic(float x) { + const float a = -0.5f; + x = std::fabs(x); + if (x < 1.f) + return ((a + 2.f) * x - (a + 3.f)) * x * x + 1.f; + if (x < 2.f) + return ((a * x - 5.f * a) * x + 8.f * a) * x - 4.f * a; + return 0.f; +} + +// The types are aten's (_compute_indices_min_size_weights_aa): float values, but sums with the double literals +// 0.5 and 1.0 are done in double, then narrowed. +AaTaps aa_taps(int in, int out, bool bicubic) { + AaTaps t; + const float scale = (float) in / (float) out; + const double interp = bicubic ? 4.0 : 2.0; + const float support = scale >= 1.f ? (float) ((interp * 0.5) * (double) scale) : (float) (interp * 0.5); + const float invscal = scale >= 1.f ? (float) (1.0 / (double) scale) : 1.f; + t.maxk = (int) std::ceil(support) * 2 + 1; + t.xmin.resize((size_t) out); + t.xsize.resize((size_t) out); + t.w.assign((size_t) out * (size_t) t.maxk, 0.f); + for (int i = 0; i < out; i++) { + const float center = (float) ((double) scale * ((double) i + 0.5)); + const int lo = std::max((int) ((double) (center - support) + 0.5), 0); + const int hi = std::min((int) ((double) (center + support) + 0.5), in); + const int n = hi - lo; + float * w = &t.w[(size_t) i * (size_t) t.maxk]; + float total = 0.f; + for (int j = 0; j < n; j++) { + const float d = (float) (j + lo) - center; + const float x = (float) (((double) d + 0.5) * (double) invscal); + w[j] = bicubic ? cubic(x) : triangle(x); + total += w[j]; + } + if (total != 0.f) + for (int j = 0; j < n; j++) + w[j] /= total; + t.xmin[(size_t) i] = lo; + t.xsize[(size_t) i] = n; + } + return t; +} + +// planes: [n_planes][h][w] floats -> [n_planes][oh][ow]; width first, then height, as aten does. +std::vector resize_aa(const std::vector & src, int planes, int h, int w, int oh, int ow, bool bicubic) { + const AaTaps tw = aa_taps(w, ow, bicubic); + const AaTaps th = aa_taps(h, oh, bicubic); + std::vector tmp((size_t) planes * (size_t) h * (size_t) ow); + for (int p = 0; p < planes; p++) + for (int y = 0; y < h; y++) { + const float * row = &src[((size_t) p * h + y) * (size_t) w]; + float * dst = &tmp[((size_t) p * h + y) * (size_t) ow]; + for (int x = 0; x < ow; x++) { + const float * k = &tw.w[(size_t) x * (size_t) tw.maxk]; + float s = 0.f; + for (int j = 0; j < tw.xsize[(size_t) x]; j++) + s = std::fmaf(row[tw.xmin[(size_t) x] + j], k[j], s); + dst[x] = s; + } + } + std::vector out((size_t) planes * (size_t) oh * (size_t) ow); + for (int p = 0; p < planes; p++) + for (int y = 0; y < oh; y++) { + const float * k = &th.w[(size_t) y * (size_t) th.maxk]; + for (int x = 0; x < ow; x++) { + float s = 0.f; + for (int j = 0; j < th.xsize[(size_t) y]; j++) + s = std::fmaf(tmp[((size_t) p * h + (size_t) (th.xmin[(size_t) y] + j)) * (size_t) ow + (size_t) x], k[j], s); + out[((size_t) p * oh + y) * (size_t) ow + (size_t) x] = s; + } + } + return out; +} + +// One camera view as CHW floats in [0, 1] at the view size (bilinear with antialias and uint8 rounding when it +// has another size, as the processor's prepare_view does). +bool view_chw(const ImageView & v, int view_h, int view_w, std::vector & chw) { + if (!v.data || v.w <= 0 || v.h <= 0) + return false; + const size_t n = (size_t) v.w * (size_t) v.h; + chw.resize(3 * n); + for (size_t i = 0; i < n; i++) + for (int c = 0; c < 3; c++) { + const float x = v.format == PixelFormat::U8 ? (float) ((const uint8_t *) v.data)[i * 3 + (size_t) c] / 255.f + : ((const float *) v.data)[i * 3 + (size_t) c]; + chw[(size_t) c * n + i] = x; + } + if (v.w != view_w || v.h != view_h) { + std::vector scaled(chw.size()); + for (size_t i = 0; i < chw.size(); i++) + scaled[i] = chw[i] * 255.f; + std::vector r = resize_aa(scaled, 3, v.h, v.w, view_h, view_w, false); + for (float & x : r) + x = std::min(std::max(std::nearbyint(x), 0.f), 255.f) / 255.f; + chw = std::move(r); + } + return true; +} + +// The views -> the float canvas [3][canvas_h][canvas_w] in uint8 levels. +bool make_canvas(const ImageView * images, int n_images, const Layout & L, std::vector & canvas) { + // A caller that composes the canvas itself (lerobot's processor in Python) sends it as the one image. + if (n_images == 1 && images[0].data && images[0].w == L.canvas_w && images[0].h == L.canvas_h) { + const size_t n = (size_t) L.canvas_w * (size_t) L.canvas_h; + canvas.resize(3 * n); + for (size_t i = 0; i < n; i++) + for (size_t c = 0; c < 3; c++) { + const size_t k = i * 3 + c; + canvas[c * n + i] = images[0].format == PixelFormat::U8 + ? (float) ((const uint8_t *) images[0].data)[k] + : std::nearbyint(((const float *) images[0].data)[k] * 255.f); + } + return true; + } + if (n_images != L.n_views) { + std::fprintf(stderr, "%s: expected %d views [third person, wrist] or one %dx%d canvas, got %d images\n", kTag, + L.n_views, L.canvas_w, L.canvas_h, n_images); + return false; + } + std::vector> views((size_t) L.n_views); + for (int i = 0; i < L.n_views; i++) + if (!view_chw(images[i], L.view_h, L.view_w, views[(size_t) i])) { + std::fprintf(stderr, "%s: view %d is empty\n", kTag, i); + return false; + } + // Side by side, to uint8 by truncation (the processor does (frame * 255).clamp(0, 255).to(uint8)). + const int fw = L.view_w * L.n_views, fh = L.view_h; + std::vector frame((size_t) 3 * (size_t) fh * (size_t) fw); + for (int c = 0; c < 3; c++) + for (int y = 0; y < fh; y++) + for (int v = 0; v < L.n_views; v++) + for (int x = 0; x < L.view_w; x++) { + const float f = views[(size_t) v][((size_t) c * fh + y) * (size_t) L.view_w + (size_t) x] * 255.f; + frame[((size_t) c * fh + y) * (size_t) fw + (size_t) (v * L.view_w + x)] = + std::trunc(std::min(std::max(f, 0.f), 255.f)); + } + // Bicubic with antialias into the content area (rounded to even, as torch.round), then the canvas filled + // to the right and below by reflection (edge replication when the padding is wider than the content). + std::vector content = frame; + if (L.content_h != fh || L.content_w != fw) { + content = resize_aa(frame, 3, fh, fw, L.content_h, L.content_w, true); + for (float & x : content) + x = std::nearbyint(std::min(std::max(x, 0.f), 255.f)); + } + const bool replicate = L.canvas_w - L.content_w >= L.content_w || L.canvas_h - L.content_h >= L.content_h; + auto src = [&](int i, int n) { // the content index a canvas index reads + if (i < n) + return i; + return replicate ? n - 1 : 2 * (n - 1) - i; // reflect without repeating the edge + }; + canvas.assign((size_t) 3 * (size_t) L.canvas_h * (size_t) L.canvas_w, 0.f); + for (int c = 0; c < 3; c++) + for (int y = 0; y < L.canvas_h; y++) { + const int sy = src(y, L.content_h); + for (int x = 0; x < L.canvas_w; x++) + canvas[((size_t) c * L.canvas_h + y) * (size_t) L.canvas_w + (size_t) x] = + content[((size_t) c * L.content_h + (size_t) sy) * (size_t) L.content_w + (size_t) src(x, L.content_w)]; + } + return true; +} + +// ------------------------------------------------------------------------------------------------ +// UniPC: diffusers' UniPCMultistepScheduler as the policy configures it (flow prediction and flow sigmas, bh2, +// solver order 2, predict_x0, lower_order_final, the corrector on). f32 data, f64 coefficients, on the host. +// ------------------------------------------------------------------------------------------------ + +struct UniPC { + std::vector sigmas; // [n + 1], the last one 0 + std::vector timesteps; // [n], the integer timesteps the network is told + + // The native sigma schedule: num_steps noise levels evenly spaced from 1 - 1/num_train_timesteps (as a + // float32, like the native sampler) down to 0, shifted to shift*s / (1 + (shift-1)*s); the timesteps are + // 1000*s truncated to integers (the scheduler stores them as int64). + UniPC(int n, int train_steps, double shift) { + const double sigma_max = (double) (float) (1.0 - 1.0 / (double) train_steps); + sigmas.resize((size_t) n + 1); + timesteps.resize((size_t) n); + for (int i = 0; i < n; i++) { + const double s = sigma_max + (0.0 - sigma_max) * ((double) i / (double) n); // np.linspace(sigma_max, 0, n + 1)[i] + const double ss = shift * s / (1.0 + (shift - 1.0) * s); + sigmas[(size_t) i] = (float) ss; + timesteps[(size_t) i] = (int64_t) (ss * (double) train_steps); + } + sigmas[(size_t) n] = 0.f; + } + + // alpha = 1 - sigma (flow), lambda = log(alpha) - log(sigma). + double lambda(int i) const { + const double s = sigmas[(size_t) i]; + return std::log(1.0 - s) - std::log(s); + } + + struct Stream { + std::vector m0, m1; // model_outputs[-1], [-2]: the x0 predictions of the last two steps + std::vector last_sample; // the sample the last predictor started from + int lower_order_nums = 0, this_order = 1; + bool have_last = false; + }; + + // One scheduler.step() on a stream: v is the network's velocity at x, the sample of step i; x becomes the + // sample of step i + 1. UniC corrects x with the new model output first (from step 1 on), then UniP predicts. + void step(int i, Stream & S, const std::vector & v, std::vector & x) const { + const size_t N = x.size(); + const int n = (int) timesteps.size(); + std::vector x0(N); // convert_model_output: x0 = x - sigma_t * v + for (size_t k = 0; k < N; k++) + x0[k] = x[k] - sigmas[(size_t) i] * v[k]; + + if (i > 0 && S.have_last) { // multistep_uni_c_bh_update(x0, last_sample, x, this_order) + const double sigma_t = sigmas[(size_t) i], sigma_s0 = sigmas[(size_t) i - 1]; + const double alpha_t = 1.0 - sigma_t; + const double h = lambda(i) - lambda(i - 1), hh = -h; + const double h_phi_1 = std::expm1(hh), B_h = std::expm1(hh); + // rhos_c: R rhos = b for the order; order 1 is the fixed 1/2. + double r_m1 = 0.0, r_t = 0.5, rk = 1.0; + if (S.this_order == 2) { + rk = (lambda(i - 2) - lambda(i - 1)) / h; + double h_phi_k = h_phi_1 / hh - 1.0; + const double b0 = h_phi_k / B_h; // factorial 1 + h_phi_k = h_phi_k / hh - 1.0 / 2.0; + const double b1 = h_phi_k * 2.0 / B_h; // factorial 2 + // R = [[1, 1], [rk, 1]] + r_m1 = (b0 - b1) / (1.0 - rk); + r_t = b0 - r_m1; + } + const double c_last = sigma_t / sigma_s0; + for (size_t k = 0; k < N; k++) { + const double m0 = S.m0[k]; + double corr = r_t * ((double) x0[k] - m0); + if (S.this_order == 2) + corr += r_m1 * (((double) S.m1[k] - m0) / rk); + x[k] = (float) (c_last * (double) S.last_sample[k] - alpha_t * h_phi_1 * m0 - alpha_t * B_h * corr); + } + } + // The model-output history, then the order of this step's predictor (lower_order_final, warm-up). + S.m1 = std::move(S.m0); + S.m0 = std::move(x0); + S.this_order = std::min(std::min(2, n - i), S.lower_order_nums + 1); + S.last_sample = x; + S.have_last = true; + + // multistep_uni_p_bh_update(x, this_order): from sigma[i] to sigma[i + 1]. + const double sigma_t = sigmas[(size_t) i + 1], sigma_s0 = sigmas[(size_t) i]; + if (sigma_t <= 0.0) { + // The last step lands on sigma 0: its lambda is infinite, sigma_t/sigma_s0 is 0, h_phi_1 and B_h are -1 + // and alpha_t is 1, so the update is x = m0, the x0 prediction (lower_order_final makes this step + // order 1 whenever there are two steps or more). + for (size_t k = 0; k < N; k++) + x[k] = S.m0[k]; + } else { + const double alpha_t = 1.0 - sigma_t; + const double lambda_t = std::log(1.0 - sigma_t) - std::log(sigma_t); + const double h = lambda_t - lambda(i), hh = -h; + const double h_phi_1 = std::expm1(hh), B_h = std::expm1(hh); + const double c_x = sigma_t / sigma_s0; + double rk = 1.0; + if (S.this_order == 2) + rk = (lambda(i - 1) - lambda(i)) / h; + for (size_t k = 0; k < N; k++) { + const double m0 = S.m0[k]; + double xt = c_x * (double) x[k] - alpha_t * h_phi_1 * m0; + if (S.this_order == 2) + xt -= alpha_t * B_h * 0.5 * (((double) S.m1[k] - m0) / rk); + x[k] = (float) xt; + } + } + if (S.lower_order_nums < 2) + S.lower_order_nums++; + } +}; + +// ------------------------------------------------------------------------------------------------ +// Weights +// ------------------------------------------------------------------------------------------------ + +struct ConvW { + ggml_tensor *w = nullptr, *b = nullptr; +}; + +struct ResW { + ggml_tensor *norm1 = nullptr, *norm2 = nullptr; + ConvW conv1, conv2, shortcut; // shortcut.w null when in_dim == out_dim +}; + +struct DownW { + std::vector res; + ConvW down; // ZeroPad2d((0, 1, 0, 1)) + Conv2d(dim, dim, 3, stride 2); none on the last block + ggml_tensor * avg = nullptr; // AvgDown3D of the shortcut as a strided constant kernel; null when identity + int factor_t = 1, factor_s = 1; +}; + +struct VaeW { + ConvW conv_in; + std::vector down; + ResW mid0, mid1; + ggml_tensor * attn_norm = nullptr; + ConvW qkv, proj; // 1x1 + ggml_tensor * norm_out = nullptr; + ConvW conv_out, quant; // quant 1x1 +}; + +struct StreamW { + ggml_tensor *attn_norm = nullptr, *q = nullptr, *k = nullptr, *v = nullptr, *o = nullptr; + ggml_tensor *ffn_norm = nullptr, *up = nullptr, *down = nullptr; +}; + +struct LayerW { + StreamW und, gen; + ggml_tensor * k_norm_for_gen = nullptr; // the understanding keys' norm for the generation queries + ggml_tensor * q_norm = nullptr, * k_norm = nullptr; // the generation stream's QK norms +}; + +bool cuda_conv_has_mma() { +#if defined(GGML_USE_CUDA) && !defined(GGML_USE_HIP) + int major = 0, minor = 0; + const int dev = backend_device_index(); + return cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, dev) == cudaSuccess && + cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, dev) == cudaSuccess && + major * 10 + minor >= 75; +#else + return false; +#endif +} + +// ------------------------------------------------------------------------------------------------ +// The architecture +// ------------------------------------------------------------------------------------------------ + +struct Cosmos3PolicyModelArch : public ModelArchBase { + Cosmos3PolicyModelArch() : ModelArchBase(Arch::COSMOS3_POLICY) {} + ~Cosmos3PolicyModelArch() override; + + ggml_backend_t backend = nullptr; + int n_threads = default_cpu_threads(); + ggml_type mt = GGML_TYPE_BF16; // GEMM weights + // CUDA (Turing+) runs the VAE as ggml's direct convolution on F16 kernels (implicit GEMM on tensor cores, + // F32 accumulation); elsewhere the kernels stay F32 and go through im2col, as act.cpp does. + bool conv_direct = false; + + // --- sizes, from the GGUF --- + int64_t hidden = 2048, inter = 9216, n_layers = 28, n_heads = 16, n_kv_heads = 8, head_dim = 128; + float rms_eps = 1e-5f, rope_theta = 1e8f; + int mrope[3] = { 24, 20, 20 }; + int64_t time_dim = 256; + float timestep_scale = 0.001f; + int64_t margin = 15000; // temporal_modality_margin: the generation stream starts this far after the text + double base_fps = 24.0, fps = 20.0; + int steps = 30, train_steps = 1000; + double shift = 5.0; + int64_t chunk = 16, adim = 10; + int64_t z_dim = 48, vae_patch = 2, spatial = 16, temporal = 4, latent_patch = 2; + int64_t latent_h = 10, latent_w = 20; // the content's latent grid (the canvas's is cropped to it) + int64_t frames = 5, ppf = 50, n_video = 250, patch_dim = 192; // latent frames, patches per frame + int64_t max_prompt = 512; + Layout layout; + std::vector lat_mean, lat_std, q01, q99; + + // --- weights --- + ggml_context * ctx_w = nullptr; + ggml_backend_buffer_t buf_w = nullptr; + ggml_context * ctx_aux = nullptr; // the AvgDown kernels, made here + ggml_backend_buffer_t buf_aux = nullptr; + ggml_tensor * tok_embd = nullptr; + std::vector layers; + ggml_tensor * gen_out_norm = nullptr; + ggml_tensor *vae2llm_w = nullptr, *vae2llm_b = nullptr, *llm2vae_w = nullptr, *llm2vae_b = nullptr; + ggml_tensor *a2l_w = nullptr, *a2l_b = nullptr, *l2a_w = nullptr, *l2a_b = nullptr, *act_embed = nullptr; + ggml_tensor *te1_w = nullptr, *te1_b = nullptr, *te2_w = nullptr, *te2_b = nullptr; + VaeW vae; + + // --- graphs --- + struct VaeIO { + ggml_tensor *px = nullptr, *lat = nullptr; + }; + struct PrefixIO { + ggml_tensor *ids = nullptr, *cos = nullptr, *sin = nullptr, *mask = nullptr; + }; + struct StepIO { + ggml_tensor *xv = nullptr, *xa = nullptr, *tsin = nullptr; + ggml_tensor *vv = nullptr, *va = nullptr; + }; + graph_cache vae_graph; + graph_cache prefix_graph; // keyed by the prompt length + graph_cache step_graph; + // What every step reads and the prefix sets up, in a buffer of its own so the graph allocator never reuses it + // (a graph input's memory is free after its last use): the understanding stream's keys and values of every + // layer, [head_dim, n_text, n_kv_heads], and the generation stream's rotation tables [head_dim, n_gen], whose + // positions start after the prompt. Re-made when the prompt length changes. + struct KV { + ggml_context * ctx = nullptr; + ggml_backend_buffer_t buf = nullptr; + int64_t n = 0; + std::vector k, v; + ggml_tensor * cos = nullptr, *sin = nullptr; + } kv; + + uint64_t calls = 0; + + bool load_config(gguf_reader & g); + bool load_weights(gguf_reader & g); + bool make_avg_kernels(); + bool ensure_kv(int64_t n); + + ggml_cgraph * build_vae(ggml_context * C, VaeIO & io) const; + ggml_cgraph * build_prefix(ggml_context * C, PrefixIO & io, int64_t n) const; + ggml_cgraph * build_step(ggml_context * C, StepIO & io, int64_t n) const; + + // Graph pieces. + ggml_tensor * lin(ggml_context * C, ggml_tensor * W, ggml_tensor * b, ggml_tensor * x) const; + ggml_tensor * conv(ggml_context * C, const ConvW & cw, ggml_tensor * x, int stride, int pad) const; + ggml_tensor * conv1x1(ggml_context * C, const ConvW & cw, ggml_tensor * x) const; + ggml_tensor * chan_rms(ggml_context * C, ggml_tensor * x, ggml_tensor * gamma) const; + ggml_tensor * resnet(ggml_context * C, const ResW & r, ggml_tensor * x) const; + ggml_tensor * vae_attention(ggml_context * C, ggml_tensor * x) const; + ggml_tensor * attend(ggml_context * C, ggml_tensor * q, ggml_tensor * k, ggml_tensor * v, ggml_tensor * mask, + int64_t n_q) const; + + // Host tables. + void rope_text(int64_t n, std::vector & cs, std::vector & sn) const; + void rope_gen(int64_t n_text, std::vector & cs, std::vector & sn) const; + void time_sinusoid(int64_t t, std::vector & out) const; + + std::vector predict(const Inputs & in) override; +}; + +Cosmos3PolicyModelArch::~Cosmos3PolicyModelArch() { + vae_graph.release(); + prefix_graph.release(); + step_graph.release(); + if (kv.buf) + ggml_backend_buffer_free(kv.buf); + if (kv.ctx) + ggml_free(kv.ctx); + if (buf_aux) + ggml_backend_buffer_free(buf_aux); + if (ctx_aux) + ggml_free(ctx_aux); + if (buf_w) + ggml_backend_buffer_free(buf_w); + if (ctx_w) + ggml_free(ctx_w); + if (backend) + ggml_backend_free(backend); +} + +bool Cosmos3PolicyModelArch::load_config(gguf_reader & g) { + auto u = [&](const std::string & k) { return (int64_t) g.u32(("cosmos3_policy." + k).c_str()); }; + auto f = [&](const std::string & k) { return g.f32(("cosmos3_policy." + k).c_str()); }; + n_layers = u("n_layers"); hidden = u("hidden"); inter = u("intermediate"); + n_heads = u("n_heads"); n_kv_heads = u("n_kv_heads"); head_dim = u("head_dim"); + mrope[0] = (int) u("mrope_section_t"); mrope[1] = (int) u("mrope_section_h"); mrope[2] = (int) u("mrope_section_w"); + time_dim = u("time_embed_dim"); margin = u("temporal_modality_margin"); + steps = (int) u("num_steps"); train_steps = (int) u("num_train_timesteps"); + chunk = u("chunk_size"); adim = u("action_dim"); latent_patch = u("latent_patch_size"); + z_dim = u("vae.z_dim"); vae_patch = u("vae.patch_size"); + spatial = u("vae.scale_factor_spatial"); temporal = u("vae.scale_factor_temporal"); + latent_h = u("latent_h"); latent_w = u("latent_w"); max_prompt = u("prompt_len"); + layout.n_views = (int) u("n_views"); layout.view_h = (int) u("view_h"); layout.view_w = (int) u("view_w"); + layout.canvas_h = (int) u("canvas_h"); layout.canvas_w = (int) u("canvas_w"); + layout.content_h = (int) u("content_h"); layout.content_w = (int) u("content_w"); + rms_eps = f("rms_eps"); rope_theta = f("rope_theta"); timestep_scale = f("timestep_scale"); + base_fps = f("base_fps"); fps = f("fps"); shift = f("sampling_shift"); + + const int64_t n_blocks = u("vae.n_blocks"), n_res = u("vae.num_res_blocks"); + if (n_layers <= 0 || hidden <= 0 || inter <= 0 || n_heads <= 0 || n_kv_heads <= 0 || head_dim <= 0 || + n_heads % n_kv_heads != 0 || head_dim % 2 != 0 || time_dim <= 0 || time_dim % 2 != 0 || steps <= 0 || + train_steps <= 0 || chunk <= 0 || adim <= 0 || z_dim <= 0 || vae_patch != 2 || latent_patch != 2 || + spatial <= 0 || temporal <= 0 || latent_h <= 0 || latent_w <= 0 || n_blocks <= 0 || n_res <= 0 || + max_prompt <= 0 || layout.n_views <= 0 || layout.view_h <= 0 || layout.view_w <= 0 || layout.canvas_h <= 0 || + layout.canvas_w <= 0 || layout.content_h > layout.canvas_h || layout.content_w > layout.canvas_w || + !(rms_eps > 0.f) || !(rope_theta > 0.f) || !(fps > 0.0) || !(shift > 0.0) || + mrope[0] + mrope[1] + mrope[2] != head_dim / 2) { + std::fprintf(stderr, "%s: the GGUF has no valid sizes (converted with an older script?)\n", kTag); + return false; + } + if (layout.canvas_h % spatial != 0 || layout.canvas_w % spatial != 0 || layout.content_h / spatial != latent_h || + layout.content_w / spatial != latent_w || latent_h % latent_patch != 0 || latent_w % latent_patch != 0 || + chunk % temporal != 0) { + std::fprintf(stderr, "%s: canvas %dx%d / content %dx%d do not fit the VAE's grid\n", kTag, layout.canvas_h, + layout.canvas_w, layout.content_h, layout.content_w); + return false; + } + frames = 1 + chunk / temporal; + ppf = (latent_h / latent_patch) * (latent_w / latent_patch); + n_video = frames * ppf; + patch_dim = latent_patch * latent_patch * z_dim; + + vae.down.resize((size_t) n_blocks); + for (int64_t b = 0; b < n_blocks; b++) { + DownW & d = vae.down[(size_t) b]; + d.res.resize((size_t) n_res); + d.factor_t = (int) u("vae.down." + std::to_string(b) + ".factor_t"); + d.factor_s = (int) u("vae.down." + std::to_string(b) + ".factor_s"); + if (d.factor_t < 1 || d.factor_s < 1) { + std::fprintf(stderr, "%s: VAE block %lld has no factors\n", kTag, (long long) b); + return false; + } + } + + lat_mean = g.read_f32("cosmos3_policy.latents_mean"); + lat_std = g.read_f32("cosmos3_policy.latents_std"); + q01 = g.read_f32("cosmos3_policy.action_q01"); + q99 = g.read_f32("cosmos3_policy.action_q99"); + if ((int64_t) lat_mean.size() != z_dim || (int64_t) lat_std.size() != z_dim || (int64_t) q01.size() != adim || + (int64_t) q99.size() != adim) { + std::fprintf(stderr, "%s: latent statistics or action quantiles disagree with the sizes\n", kTag); + return false; + } + for (float s : lat_std) + if (!(s > 0.f)) { + std::fprintf(stderr, "%s: a latent std is not positive\n", kTag); + return false; + } + return true; +} + +bool Cosmos3PolicyModelArch::load_weights(gguf_reader & g) { + const int64_t n_tensors = gguf_get_n_tensors(g.gctx); + ggml_init_params p = { (size_t) (n_tensors + 16) * ggml_tensor_overhead(), nullptr, true }; + ctx_w = ggml_init(p); + if (!ctx_w) + return false; + WeightLoader wl("cosmos3_policy", g, ctx_w, mt); + const ggml_type ct = conv_direct ? GGML_TYPE_F16 : GGML_TYPE_F32; + + tok_embd = wl.gemm("token_embd.weight"); + layers.resize((size_t) n_layers); + for (int64_t i = 0; i < n_layers; i++) { + LayerW & L = layers[(size_t) i]; + auto stream = [&](StreamW & s, const char * pfx) { + s.attn_norm = wl.f32("blk.%lld.%sattn_norm.weight", (long long) i, pfx); + s.q = wl.gemm("blk.%lld.%sattn_q.weight", (long long) i, pfx); + s.k = wl.gemm("blk.%lld.%sattn_k.weight", (long long) i, pfx); + s.v = wl.gemm("blk.%lld.%sattn_v.weight", (long long) i, pfx); + s.o = wl.gemm("blk.%lld.%sattn_output.weight", (long long) i, pfx); + s.ffn_norm = wl.f32("blk.%lld.%sffn_norm.weight", (long long) i, pfx); + s.up = wl.gemm("blk.%lld.%sffn_up.weight", (long long) i, pfx); + s.down = wl.gemm("blk.%lld.%sffn_down.weight", (long long) i, pfx); + }; + stream(L.und, ""); + stream(L.gen, "gen."); + L.k_norm_for_gen = wl.f32("blk.%lld.attn_k_norm_gen.weight", (long long) i); + L.q_norm = wl.f32("blk.%lld.gen.attn_q_norm.weight", (long long) i); + L.k_norm = wl.f32("blk.%lld.gen.attn_k_norm.weight", (long long) i); + } + gen_out_norm = wl.f32("gen.output_norm.weight"); + vae2llm_w = wl.gemm("vae2llm.weight"); vae2llm_b = wl.f32("vae2llm.bias"); + llm2vae_w = wl.gemm("llm2vae.weight"); llm2vae_b = wl.f32("llm2vae.bias"); + a2l_w = wl.gemm("action2llm.weight"); a2l_b = wl.f32("action2llm.bias"); + l2a_w = wl.gemm("llm2action.weight"); l2a_b = wl.f32("llm2action.bias"); + act_embed = wl.f32("action_modality_embed.weight"); + te1_w = wl.f32("time_embedder.linear_1.weight"); te1_b = wl.f32("time_embedder.linear_1.bias"); + te2_w = wl.f32("time_embedder.linear_2.weight"); te2_b = wl.f32("time_embedder.linear_2.bias"); + + auto cw = [&](ConvW & c, const std::string & name, bool required = true) { + c.w = required ? wl.typed(ct, "%s.weight", name.c_str()) : wl.opt_typed(ct, "%s.weight", name.c_str()); + c.b = c.w ? wl.f32("%s.bias", name.c_str()) : nullptr; + }; + auto res = [&](ResW & r, const std::string & name) { + r.norm1 = wl.f32("%s.norm1.gamma", name.c_str()); + r.norm2 = wl.f32("%s.norm2.gamma", name.c_str()); + cw(r.conv1, name + ".conv1"); + cw(r.conv2, name + ".conv2"); + cw(r.shortcut, name + ".shortcut", false); + }; + cw(vae.conv_in, "vae.conv_in"); + for (size_t b = 0; b < vae.down.size(); b++) { + for (size_t r = 0; r < vae.down[b].res.size(); r++) + res(vae.down[b].res[r], "vae.down." + std::to_string(b) + ".res." + std::to_string(r)); + cw(vae.down[b].down, "vae.down." + std::to_string(b) + ".down", false); + } + res(vae.mid0, "vae.mid.res.0"); + vae.attn_norm = wl.f32("vae.mid.attn.norm.gamma"); + cw(vae.qkv, "vae.mid.attn.qkv"); + cw(vae.proj, "vae.mid.attn.proj"); + res(vae.mid1, "vae.mid.res.1"); + vae.norm_out = wl.f32("vae.norm_out.gamma"); + cw(vae.conv_out, "vae.conv_out"); + cw(vae.quant, "vae.quant_conv"); + + if (!wl.ok() || !wl.upload(backend, &buf_w)) + return false; + + // Shape checks that the graphs rely on. + if (tok_embd->ne[0] != hidden || layers[0].und.q->ne[1] != n_heads * head_dim || + layers[0].und.k->ne[1] != n_kv_heads * head_dim || layers[0].und.up->ne[1] != inter || + vae2llm_w->ne[0] != patch_dim || vae2llm_w->ne[1] != hidden || llm2vae_w->ne[1] != patch_dim || + a2l_w->ne[0] != adim || a2l_w->ne[1] != hidden || l2a_w->ne[1] != adim || te1_w->ne[0] != time_dim || + vae.conv_in.w->ne[2] != 3 * vae_patch * vae_patch || vae.quant.w->ne[1] != 2 * z_dim) { + std::fprintf(stderr, "%s: tensor shapes disagree with the GGUF metadata\n", kTag); + return false; + } + for (size_t b = 0; b < vae.down.size(); b++) { + const DownW & d = vae.down[b]; + if ((d.factor_s > 1) != (d.down.w != nullptr)) { + std::fprintf(stderr, "%s: VAE block %zu: the downsampler and its factor disagree\n", kTag, b); + return false; + } + } + return make_avg_kernels(); +} + +// The shortcut of a Wan2.2 residual down block is AvgDown3D: the clip is frame- and pixel-unshuffled (a first +// frame has zero frames in front of it) and groups of the resulting channels are averaged into the block's +// output channels. On one frame that is a strided convolution over the block's input with a constant kernel of +// 0 and 1/g, built here. A block that keeps its size and width (the last) is the identity and gets no kernel. +bool Cosmos3PolicyModelArch::make_avg_kernels() { + struct Pending { + ggml_tensor * t; + std::vector k; + }; + std::vector pending; + const ggml_type ct = conv_direct ? GGML_TYPE_F16 : GGML_TYPE_F32; + ggml_init_params p = { (vae.down.size() + 2) * ggml_tensor_overhead(), nullptr, true }; + ctx_aux = ggml_init(p); + if (!ctx_aux) + return false; + for (size_t b = 0; b < vae.down.size(); b++) { + DownW & d = vae.down[b]; + const int64_t c_in = d.res[0].conv1.w->ne[2]; + const int64_t c_out = d.res.back().conv2.w->ne[3]; + const int ft = d.factor_t, fs = d.factor_s; + const int64_t factor = (int64_t) ft * fs * fs; + if (fs == 1 && ft == 1 && c_in == c_out) + continue; // identity + if ((c_in * factor) % c_out != 0) { + std::fprintf(stderr, "%s: AvgDown3D of %lld x %lld channels into %lld does not divide\n", kTag, + (long long) c_in, (long long) factor, (long long) c_out); + return false; + } + const int64_t g = c_in * factor / c_out; + Pending pend; + pend.t = ggml_new_tensor_4d(ctx_aux, ct, fs, fs, c_in, c_out); + ggml_format_name(pend.t, "vae.avg_down.%zu", b); + pend.k.assign((size_t) (fs * fs * c_in * c_out), 0.f); + for (int64_t o = 0; o < c_out; o++) + for (int64_t c = 0; c < c_in; c++) + for (int hy = 0; hy < fs; hy++) + for (int wx = 0; wx < fs; wx++) { + // The unshuffled channel of (c, the frame itself = tap ft-1, hy, wx), as + // x.view(B, C, T', ft, H', fs, W', fs).permute(0, 1, 3, 5, 7, 2, 4, 6) orders them. + const int64_t e = ((c * ft + (ft - 1)) * fs + hy) * fs + wx; + if (e / g == o) + pend.k[(size_t) (((o * c_in + c) * fs + hy) * fs + wx)] = 1.f / (float) g; + } + d.avg = pend.t; + pending.push_back(std::move(pend)); + } + if (pending.empty()) + return true; + buf_aux = alloc_weights(ctx_aux, backend); + if (!buf_aux) + return false; + for (Pending & pend : pending) { + if (ct == GGML_TYPE_F16) { + std::vector h(pend.k.size()); + ggml_fp32_to_fp16_row(pend.k.data(), h.data(), (int64_t) h.size()); + ggml_backend_tensor_set(pend.t, h.data(), 0, ggml_nbytes(pend.t)); + } else { + ggml_backend_tensor_set(pend.t, pend.k.data(), 0, ggml_nbytes(pend.t)); + } + } + return true; +} + +bool Cosmos3PolicyModelArch::ensure_kv(int64_t n) { + if (kv.ctx && kv.n == n) + return true; + if (kv.buf) + ggml_backend_buffer_free(kv.buf); + if (kv.ctx) + ggml_free(kv.ctx); + kv = KV{}; + ggml_init_params p = { (size_t) (2 * n_layers + 4) * ggml_tensor_overhead(), nullptr, true }; + kv.ctx = ggml_init(p); + if (!kv.ctx) + return false; + kv.cos = ggml_new_tensor_2d(kv.ctx, GGML_TYPE_F32, head_dim, n_video + chunk); + kv.sin = ggml_new_tensor_2d(kv.ctx, GGML_TYPE_F32, head_dim, n_video + chunk); + ggml_set_name(kv.cos, "kv.cos_gen"); + ggml_set_name(kv.sin, "kv.sin_gen"); + kv.k.resize((size_t) n_layers); + kv.v.resize((size_t) n_layers); + for (int64_t i = 0; i < n_layers; i++) { + kv.k[(size_t) i] = ggml_new_tensor_3d(kv.ctx, GGML_TYPE_F32, head_dim, n, n_kv_heads); + kv.v[(size_t) i] = ggml_new_tensor_3d(kv.ctx, GGML_TYPE_F32, head_dim, n, n_kv_heads); + ggml_format_name(kv.k[(size_t) i], "kv.k.%lld", (long long) i); + ggml_format_name(kv.v[(size_t) i], "kv.v.%lld", (long long) i); + } + kv.buf = ggml_backend_alloc_ctx_tensors(kv.ctx, backend); + if (!kv.buf) + return false; + std::vector cs, sn; + rope_gen(n, cs, sn); + ggml_backend_tensor_set(kv.cos, cs.data(), 0, ggml_nbytes(kv.cos)); + ggml_backend_tensor_set(kv.sin, sn.data(), 0, ggml_nbytes(kv.sin)); + kv.n = n; + return true; +} + +// ------------------------------------------------------------------------------------------------ +// Graph pieces +// ------------------------------------------------------------------------------------------------ + +ggml_tensor * Cosmos3PolicyModelArch::lin(ggml_context * C, ggml_tensor * W, ggml_tensor * b, ggml_tensor * x) const { + ggml_tensor * y = ggml_mul_mat(C, W, x); + if (mm_prec_f32_enabled()) + ggml_prec_set_acc(y, GGML_PREC_F32); + return b ? ggml_add(C, y, b) : y; +} + +// x [W, H, C, 1] -> [OW, OH, OC, 1] with the bias. +ggml_tensor * Cosmos3PolicyModelArch::conv(ggml_context * C, const ConvW & cw, ggml_tensor * x, int stride, int pad) const { + ggml_tensor * y = cw.w->type == GGML_TYPE_F16 ? ggml_conv_2d_direct(C, cw.w, x, stride, stride, pad, pad, 1, 1) + : conv_2d_f32(C, cw.w, x, stride, pad); + return cw.b ? ggml_add(C, y, ggml_reshape_4d(C, cw.b, 1, 1, cw.b->ne[0], 1)) : y; +} + +// A 1x1 convolution stored as a linear map [IC, OC]. +ggml_tensor * Cosmos3PolicyModelArch::conv1x1(ggml_context * C, const ConvW & cw, ggml_tensor * x) const { + const ConvW k4 = { ggml_reshape_4d(C, cw.w, 1, 1, cw.w->ne[0], cw.w->ne[1]), cw.b }; + return conv(C, k4, x, 1, 0); +} + +// WanRMS_norm over the channels of x [W, H, C, 1]: F.normalize(x, dim=channels) * sqrt(C) * gamma, which is +// x / rms(x) * gamma. (F.normalize clamps the norm at 1e-12; a channel vector of zeros does not occur.) +ggml_tensor * Cosmos3PolicyModelArch::chan_rms(ggml_context * C, ggml_tensor * x, ggml_tensor * gamma) const { + ggml_tensor * t = ggml_cont(C, ggml_permute(C, x, 1, 2, 0, 3)); // [C, W, H] + t = ggml_mul(C, ggml_rms_norm(C, t, 1e-24f), gamma); + return ggml_cont(C, ggml_permute(C, t, 2, 0, 1, 3)); // [W, H, C] +} + +// WanResidualBlock on one frame: norm, SiLU, conv, norm, SiLU, conv, plus the (1x1) shortcut. +ggml_tensor * Cosmos3PolicyModelArch::resnet(ggml_context * C, const ResW & r, ggml_tensor * x) const { + ggml_tensor * h = r.shortcut.w ? conv1x1(C, r.shortcut, x) : x; + ggml_tensor * y = conv(C, r.conv1, ggml_silu(C, chan_rms(C, x, r.norm1)), 1, 1); + y = conv(C, r.conv2, ggml_silu(C, chan_rms(C, y, r.norm2)), 1, 1); + return ggml_add(C, y, h); +} + +// WanAttentionBlock: single-head attention over the frame's positions, the channels as the head. +ggml_tensor * Cosmos3PolicyModelArch::vae_attention(ggml_context * C, ggml_tensor * x) const { + const int64_t W = x->ne[0], H = x->ne[1], Ch = x->ne[2], N = W * H; + ggml_tensor * qkv = conv1x1(C, vae.qkv, chan_rms(C, x, vae.attn_norm)); // [W, H, 3 Ch] + auto part = [&](int i) { // channel block i as [Ch, N]: one row per position + ggml_tensor * v = ggml_view_3d(C, qkv, W, H, Ch, qkv->nb[1], qkv->nb[2], (size_t) i * (size_t) Ch * qkv->nb[2]); + return ggml_cont(C, ggml_transpose(C, ggml_reshape_2d(C, ggml_cont(C, v), N, Ch))); + }; + ggml_tensor * q = part(0), *k = part(1), *v = part(2); + // Always the explicit form: a head of 640 is outside the flash-attention kernels' head sizes. + ggml_tensor * kq = ggml_mul_mat(C, k, q); // [N_k, N_q] + ggml_prec_set_acc(kq, GGML_PREC_F32); + kq = ggml_soft_max_ext(C, kq, nullptr, 1.0f / std::sqrt((float) Ch), 0.0f); + ggml_tensor * o = ggml_mul_mat(C, ggml_cont(C, ggml_transpose(C, v)), kq); // [Ch, N_q] + ggml_prec_set_acc(o, GGML_PREC_F32); + o = ggml_reshape_3d(C, ggml_cont(C, ggml_transpose(C, ggml_reshape_2d(C, o, Ch, N))), W, H, Ch); + return ggml_add(C, conv1x1(C, vae.proj, o), x); +} + +// q [head_dim, n_q, n_heads], k and v [head_dim, n_k, n_kv_heads] (the query heads of a key head read it, as +// grouped-query attention does) -> [hidden, n_q]. mask [n_k, n_q] F32 or null. +ggml_tensor * Cosmos3PolicyModelArch::attend(ggml_context * C, ggml_tensor * q, ggml_tensor * k, ggml_tensor * v, + ggml_tensor * mask, int64_t n_q) const { + const float scale = 1.0f / std::sqrt((float) head_dim); + if (flash_attn_enabled()) { + ggml_tensor * m16 = mask ? ggml_cast(C, mask, GGML_TYPE_F16) : nullptr; + return ggml_reshape_2d(C, flash_attention(C, q, fa_kv(C, k), fa_kv(C, v), m16, scale), hidden, n_q); + } + ggml_tensor * vt = ggml_cont(C, ggml_permute(C, v, 1, 0, 2, 3)); // [n_k, head_dim, n_kv_heads] + return ggml_reshape_2d(C, attention(C, q, k, vt, mask, scale, hidden, n_q), hidden, n_q); +} + +// ------------------------------------------------------------------------------------------------ +// Graphs +// ------------------------------------------------------------------------------------------------ + +// The Wan2.2 encoder on one frame. px: the canvas in [-1, 1], pixel-unshuffled to [W/2, H/2, 12] (host). +// lat: the mean of the posterior, [W/16, H/16, z_dim], before normalization and cropping (host). +ggml_cgraph * Cosmos3PolicyModelArch::build_vae(ggml_context * C, VaeIO & io) const { + const int64_t W0 = layout.canvas_w / vae_patch, H0 = layout.canvas_h / vae_patch; + io.px = ggml_new_tensor_4d(C, GGML_TYPE_F32, W0, H0, 3 * vae_patch * vae_patch, 1); + ggml_set_input(io.px); + + ggml_tensor * x = conv(C, vae.conv_in, io.px, 1, 1); + for (const DownW & d : vae.down) { + ggml_tensor * x_in = x; + for (const ResW & r : d.res) + x = resnet(C, r, x); + if (d.down.w) + x = conv(C, d.down, ggml_pad(C, x, 1, 1, 0, 0), 2, 0); + ggml_tensor * sc = d.avg ? conv(C, ConvW{ d.avg, nullptr }, x_in, d.factor_s, 0) : x_in; + x = ggml_add(C, x, sc); + } + x = resnet(C, vae.mid0, x); + x = vae_attention(C, x); + x = resnet(C, vae.mid1, x); + x = conv(C, vae.conv_out, ggml_silu(C, chan_rms(C, x, vae.norm_out)), 1, 1); // [w, h, 2 z] + x = conv1x1(C, vae.quant, x); + // DiagonalGaussianDistribution.mode(): the first z_dim channels. + io.lat = ggml_cont(C, ggml_view_3d(C, x, x->ne[0], x->ne[1], z_dim, x->nb[1], x->nb[2], 0)); + ggml_set_output(io.lat); + + ggml_cgraph * gf = ggml_new_graph_custom(C, 4096, false); + ggml_build_forward_expand(gf, io.lat); + return gf; +} + +// The understanding stream over the n prompt tokens (causal), writing every layer's keys as the generation +// queries read them (their own norm, then RoPE) and values into kv. +ggml_cgraph * Cosmos3PolicyModelArch::build_prefix(ggml_context * C, PrefixIO & io, int64_t n) const { + io.ids = ggml_new_tensor_1d(C, GGML_TYPE_I32, n); + io.cos = ggml_new_tensor_2d(C, GGML_TYPE_F32, head_dim, n); + io.sin = ggml_new_tensor_2d(C, GGML_TYPE_F32, head_dim, n); + io.mask = ggml_new_tensor_2d(C, GGML_TYPE_F32, n, n); + for (ggml_tensor * t : { io.ids, io.cos, io.sin, io.mask }) + ggml_set_input(t); + + ggml_cgraph * gf = ggml_new_graph_custom(C, 8192, false); + ggml_tensor * h = ggml_get_rows(C, tok_embd, io.ids); // [hidden, n] f32 + for (int64_t i = 0; i < n_layers; i++) { + const LayerW & L = layers[(size_t) i]; + ggml_tensor * x = rms_norm(C, h, L.und.attn_norm, rms_eps); + ggml_tensor * q = to_heads(C, lin(C, L.und.q, nullptr, x), head_dim, n_heads, n); + ggml_tensor * k = to_heads(C, lin(C, L.und.k, nullptr, x), head_dim, n_kv_heads, n); + ggml_tensor * v = to_heads(C, lin(C, L.und.v, nullptr, x), head_dim, n_kv_heads, n); + // What the generation stream reads of this layer. + ggml_tensor * k_gen = rope_2d(C, rms_norm(C, k, L.k_norm_for_gen, rms_eps), io.cos, io.sin); + ggml_build_forward_expand(gf, ggml_cpy(C, k_gen, kv.k[(size_t) i])); + ggml_build_forward_expand(gf, ggml_cpy(C, v, kv.v[(size_t) i])); + // Nemotron's understanding attention has no QK norm. + q = rope_2d(C, q, io.cos, io.sin); + k = rope_2d(C, k, io.cos, io.sin); + h = ggml_add(C, h, lin(C, L.und.o, nullptr, attend(C, q, k, v, io.mask, n))); + x = rms_norm(C, h, L.und.ffn_norm, rms_eps); + h = ggml_add(C, h, lin(C, L.und.down, nullptr, ggml_sqr(C, ggml_relu(C, lin(C, L.und.up, nullptr, x))))); + } + return gf; +} + +// One network call of the generation stream: [video patches | actions] at the timestep, attending to +// [prompt | itself] in every layer, to the velocities of the video patches and the actions. +ggml_cgraph * Cosmos3PolicyModelArch::build_step(ggml_context * C, StepIO & io, int64_t n) const { + (void) n; // the prompt length enters through kv + const int64_t n_gen = n_video + chunk; + io.xv = ggml_new_tensor_2d(C, GGML_TYPE_F32, patch_dim, n_video); + io.xa = ggml_new_tensor_2d(C, GGML_TYPE_F32, adim, chunk); + io.tsin = ggml_new_tensor_1d(C, GGML_TYPE_F32, time_dim); + for (ggml_tensor * t : { io.xv, io.xa, io.tsin }) + ggml_set_input(t); + + // TimestepEmbedding of the sinusoid (f32): linear, SiLU, linear. + ggml_tensor * temb = lin(C, te2_w, te2_b, ggml_silu(C, lin(C, te1_w, te1_b, io.tsin))); // [hidden] + // Video patches: vae2llm, plus the timestep embedding on the noisy frames (not on the condition frame). + ggml_tensor * ve = lin(C, vae2llm_w, vae2llm_b, io.xv); // [hidden, n_video] + ggml_tensor * cond = ggml_view_2d(C, ve, hidden, ppf, ve->nb[1], 0); + ggml_tensor * rest = ggml_add(C, ggml_view_2d(C, ve, hidden, n_video - ppf, ve->nb[1], (size_t) ppf * ve->nb[1]), temb); + ve = ggml_concat(C, ggml_cont(C, cond), rest, 1); + // Actions: action2llm of the embodiment, the modality embedding, the timestep embedding. + ggml_tensor * ae = ggml_add(C, ggml_add(C, lin(C, a2l_w, a2l_b, io.xa), act_embed), temb); // [hidden, chunk] + ggml_tensor * h = ggml_concat(C, ve, ae, 1); // [hidden, n_gen] + + for (int64_t i = 0; i < n_layers; i++) { + const LayerW & L = layers[(size_t) i]; + ggml_tensor * x = rms_norm(C, h, L.gen.attn_norm, rms_eps); + // QK norm per head, then heads, then RoPE. + auto heads = [&](ggml_tensor * p, ggml_tensor * norm, int64_t nh) { + ggml_tensor * t = ggml_reshape_3d(C, p, head_dim, nh, n_gen); + if (norm) + t = rms_norm(C, t, norm, rms_eps); + return ggml_cont(C, ggml_permute(C, t, 0, 2, 1, 3)); // [head_dim, n_gen, nh] + }; + ggml_tensor * q = rope_2d(C, heads(lin(C, L.gen.q, nullptr, x), L.q_norm, n_heads), kv.cos, kv.sin); + ggml_tensor * k = rope_2d(C, heads(lin(C, L.gen.k, nullptr, x), L.k_norm, n_kv_heads), kv.cos, kv.sin); + ggml_tensor * v = heads(lin(C, L.gen.v, nullptr, x), nullptr, n_kv_heads); + ggml_tensor * K = ggml_concat(C, kv.k[(size_t) i], k, 1); // [head_dim, n + n_gen, n_kv_heads] + ggml_tensor * V = ggml_concat(C, kv.v[(size_t) i], v, 1); + h = ggml_add(C, h, lin(C, L.gen.o, nullptr, attend(C, q, K, V, nullptr, n_gen))); + x = rms_norm(C, h, L.gen.ffn_norm, rms_eps); + h = ggml_add(C, h, lin(C, L.gen.down, nullptr, ggml_sqr(C, ggml_relu(C, lin(C, L.gen.up, nullptr, x))))); + } + h = rms_norm(C, h, gen_out_norm, rms_eps); + ggml_tensor * hv = ggml_view_2d(C, h, hidden, n_video, h->nb[1], 0); + ggml_tensor * ha = ggml_view_2d(C, h, hidden, chunk, h->nb[1], (size_t) n_video * h->nb[1]); + io.vv = lin(C, llm2vae_w, llm2vae_b, hv); // [patch_dim, n_video] + io.va = lin(C, l2a_w, l2a_b, ha); // [adim, chunk] + ggml_set_output(io.vv); + ggml_set_output(io.va); + + ggml_cgraph * gf = ggml_new_graph_custom(C, 8192, false); + ggml_build_forward_expand(gf, io.vv); + ggml_build_forward_expand(gf, io.va); + return gf; +} + +// ------------------------------------------------------------------------------------------------ +// Host tables +// ------------------------------------------------------------------------------------------------ + +// The understanding stream's positions are 0..n-1 on all three M-RoPE axes, so its rotation is the plain one. +// Float arithmetic as the reference: inv_freq = 1 / theta^(2j/d) in f32, angle = position * inv_freq in f32. +void Cosmos3PolicyModelArch::rope_text(int64_t n, std::vector & cs, std::vector & sn) const { + const int64_t half = head_dim / 2; + cs.resize((size_t) (n * head_dim)); + sn.resize((size_t) (n * head_dim)); + for (int64_t s = 0; s < n; s++) + for (int64_t j = 0; j < half; j++) { + const float inv = 1.0f / std::pow(rope_theta, (float) (2 * j) / (float) head_dim); + const float a = (float) s * inv; + cs[(size_t) (s * head_dim + j)] = cs[(size_t) (s * head_dim + half + j)] = (float) std::cos((double) a); + sn[(size_t) (s * head_dim + j)] = sn[(size_t) (s * head_dim + half + j)] = (float) std::sin((double) a); + } +} + +// The generation stream's M-RoPE positions (Cosmos3VFMNetwork.get_rope_index): it starts +// temporal_modality_margin after the last text position; latent frame j sits at j * (base_fps / temporal) * +// (temporal / fps) on the time axis and on its patch grid on the height and width axes; action i at +// (i + 1) * frame_step / temporal with zero height and width. The three axes interleave over the frequencies +// as Qwen3-VL's rotary embedding does ([T H W T H W ... T T] over mrope_section). +void Cosmos3PolicyModelArch::rope_gen(int64_t n_text, std::vector & cs, std::vector & sn) const { + const int64_t half = head_dim / 2, n_gen = n_video + chunk; + const int64_t grid_w = latent_w / latent_patch; + const float gen_start = (float) (n_text - 1) + 1.f + (float) margin; + const float frame_step = (float) (base_fps / (double) temporal) * (float) ((double) temporal / fps); + std::vector pos(3 * (size_t) n_gen); + for (int64_t t = 0; t < n_video; t++) { + const int64_t fr = t / ppf, p = t % ppf; + pos[3 * (size_t) t + 0] = gen_start + (float) fr * frame_step; + pos[3 * (size_t) t + 1] = (float) (p / grid_w); + pos[3 * (size_t) t + 2] = (float) (p % grid_w); + } + for (int64_t i = 0; i < chunk; i++) { + const int64_t t = n_video + i; + pos[3 * (size_t) t + 0] = gen_start + (float) (i + 1) * (frame_step / (float) temporal); + pos[3 * (size_t) t + 1] = pos[3 * (size_t) t + 2] = 0.f; + } + cs.resize((size_t) (n_gen * head_dim)); + sn.resize((size_t) (n_gen * head_dim)); + for (int64_t s = 0; s < n_gen; s++) + for (int64_t j = 0; j < half; j++) { + int axis = 0; + if (j % 3 == 1 && j < 3 * mrope[1]) + axis = 1; + else if (j % 3 == 2 && j < 3 * mrope[2]) + axis = 2; + const float inv = 1.0f / std::pow(rope_theta, (float) (2 * j) / (float) head_dim); + const float a = pos[3 * (size_t) s + (size_t) axis] * inv; + cs[(size_t) (s * head_dim + j)] = cs[(size_t) (s * head_dim + half + j)] = (float) std::cos((double) a); + sn[(size_t) (s * head_dim + j)] = sn[(size_t) (s * head_dim + half + j)] = (float) std::sin((double) a); + } +} + +// SinusoidalTimesteps of t * timestep_scale: cos features then sin, frequencies exp(-ln(10000) j / half). +void Cosmos3PolicyModelArch::time_sinusoid(int64_t t, std::vector & out) const { + const int64_t half = time_dim / 2; + const float ts = (float) t * timestep_scale; + out.resize((size_t) time_dim); + for (int64_t j = 0; j < half; j++) { + const float freq = std::exp(-std::log(10000.f) * (float) j / (float) half); + const float a = ts * freq; + out[(size_t) j] = (float) std::cos((double) a); + out[(size_t) (half + j)] = (float) std::sin((double) a); + } +} + +// ------------------------------------------------------------------------------------------------ +// predict +// ------------------------------------------------------------------------------------------------ + +std::vector Cosmos3PolicyModelArch::predict(const Inputs & in) { + const double t0 = now_ms(); + stats = Stats{}; + if (in.precomputed_img_emb) { + std::fprintf(stderr, "%s: precomputed_img_emb is not supported; pass the camera views\n", kTag); + return {}; + } + + // --- inputs ----------------------------------------------------------------------------------- + std::vector canvas; + if (!make_canvas(in.images, in.n_images, layout, canvas)) + return {}; + if (!in.lang_tokens || in.n_lang <= 0) { + std::fprintf(stderr, "%s: no prompt tokens (build them with the lerobot prompt and tokenizer steps)\n", kTag); + return {}; + } + std::vector ids; + for (int i = 0; i < in.n_lang; i++) { + if (in.attention_mask && i < in.attention_mask_n && in.attention_mask[i] == 0) + continue; // padding: only the real tokens go in, in order + ids.push_back(in.lang_tokens[i]); + } + const int64_t n = (int64_t) ids.size(); + if (n == 0 || n > max_prompt) { + std::fprintf(stderr, "%s: prompt of %lld tokens; 1 to %lld are taken\n", kTag, (long long) n, (long long) max_prompt); + return {}; + } + + // Action noise from the caller; the video noise is drawn here from a seed that follows it. + std::vector an((size_t) (chunk * adim), 0.f); + uint64_t seed = 1469598103934665603ull; // FNV-1a over the action noise, else a fixed seed + if (in.noise) { + std::memcpy(an.data(), in.noise, sizeof(float) * an.size()); + const auto * bytes = (const uint8_t *) in.noise; + for (size_t i = 0; i < sizeof(float) * an.size(); i++) + seed = (seed ^ bytes[i]) * 1099511628211ull; + } else { + // No noise given: sample it like the other archs do, from VLA_COSMOS3_SEED when set (reproducible). + std::random_device rd; + const char * env = std::getenv("VLA_COSMOS3_SEED"); + std::mt19937_64 rng(env ? (uint64_t) std::strtoull(env, nullptr, 10) + calls : ((uint64_t) rd() << 32) ^ rd()); + std::normal_distribution nd(0.f, 1.f); + for (float & x : an) + x = nd(rng); + seed = rng(); // the video noise follows the action noise just drawn + } + std::vector vn((size_t) ((frames - 1) * ppf * patch_dim)); + { + std::mt19937_64 rng(seed); + std::normal_distribution nd(0.f, 1.f); + for (float & x : vn) + x = nd(rng); + } + calls++; + + // --- the VAE: canvas -> latent of the condition frame ----------------------------------------------- + const double t1 = now_ms(); + const int64_t W0 = layout.canvas_w / vae_patch, H0 = layout.canvas_h / vae_patch; + { + const size_t arena = ggml_tensor_overhead() * 4096 + ggml_graph_overhead_custom(4096, false); + if (!vae_graph.ensure(backend, 0, arena, [&](ggml_context * C, VaeIO & io) { return build_vae(C, io); })) { + std::fprintf(stderr, "%s: VAE graph build/alloc failed\n", kTag); + return {}; + } + // [-1, 1], then Wan's patchify: channel (c * p + r) * p + q holds pixel (2h + q, 2w + r). + std::vector px((size_t) (W0 * H0 * 3 * vae_patch * vae_patch)); + for (int64_t c = 0; c < 3; c++) + for (int64_t q = 0; q < vae_patch; q++) + for (int64_t r = 0; r < vae_patch; r++) { + const int64_t ch = (c * vae_patch + r) * vae_patch + q; + for (int64_t y = 0; y < H0; y++) + for (int64_t x = 0; x < W0; x++) + px[(size_t) ((ch * H0 + y) * W0 + x)] = + canvas[(size_t) ((c * layout.canvas_h + (vae_patch * y + q)) * layout.canvas_w + vae_patch * x + r)] / 127.5f - 1.f; + } + VaeIO & io = vae_graph.io(); + ggml_backend_tensor_set(io.px, px.data(), 0, ggml_nbytes(io.px)); + graph_unique_names(vae_graph.graph()); + if (ggml_backend_graph_compute(backend, vae_graph.graph()) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "%s: VAE compute failed\n", kTag); + return {}; + } + } + // The latent [w, h, z], normalized with the VAE's statistics, cropped to the content's grid, patchified + // ("cthpwq->thwpqc"): patch (h, w) holds (p, q, c) = latent[c, 2h + p, 2w + q]. + std::vector xv((size_t) (n_video * patch_dim)); + { + const VaeIO & io = vae_graph.io(); + const int64_t lw = io.lat->ne[0], lh = io.lat->ne[1]; + if (lw < latent_w || lh < latent_h || io.lat->ne[2] != z_dim) { + std::fprintf(stderr, "%s: the VAE produced a %lldx%lld latent, the content needs %lldx%lld\n", kTag, + (long long) lh, (long long) lw, (long long) latent_h, (long long) latent_w); + return {}; + } + std::vector lat((size_t) (lw * lh * z_dim)); + ggml_backend_tensor_get(io.lat, lat.data(), 0, ggml_nbytes(io.lat)); + const int64_t grid_w = latent_w / latent_patch; + for (int64_t c = 0; c < z_dim; c++) + for (int64_t y = 0; y < latent_h; y++) + for (int64_t x = 0; x < latent_w; x++) { + const float z = (lat[(size_t) ((c * lh + y) * lw + x)] - lat_mean[(size_t) c]) / lat_std[(size_t) c]; + const int64_t patch = (y / latent_patch) * grid_w + x / latent_patch; + const int64_t d = ((y % latent_patch) * latent_patch + x % latent_patch) * z_dim + c; + xv[(size_t) (patch * patch_dim + d)] = z; + } + } + // The 4 predicted frames start as pure noise (noise level 1), the condition frame stays clean. + std::memcpy(&xv[(size_t) (ppf * patch_dim)], vn.data(), sizeof(float) * vn.size()); + const double t2 = now_ms(); + + // --- the prompt: the understanding stream's keys and values ------------------------------------------- + { + if (!ensure_kv(n)) { + std::fprintf(stderr, "%s: cannot allocate the prompt cache\n", kTag); + return {}; + } + const size_t arena = ggml_tensor_overhead() * 8192 + ggml_graph_overhead_custom(8192, false); + bool fresh = false; + if (!prefix_graph.ensure(backend, n, arena, [&](ggml_context * C, PrefixIO & io) { fresh = true; return build_prefix(C, io, n); })) { + std::fprintf(stderr, "%s: prefix graph build/alloc failed\n", kTag); + return {}; + } + if (fresh) // the step graph reads kv, which was re-made with the prefix + step_graph.release(); + PrefixIO & io = prefix_graph.io(); + std::vector cs, sn, mask((size_t) (n * n), 0.f); + rope_text(n, cs, sn); + for (int64_t q = 0; q < n; q++) + for (int64_t k = q + 1; k < n; k++) + mask[(size_t) (q * n + k)] = -INFINITY; + ggml_backend_tensor_set(io.ids, ids.data(), 0, ggml_nbytes(io.ids)); + ggml_backend_tensor_set(io.cos, cs.data(), 0, ggml_nbytes(io.cos)); + ggml_backend_tensor_set(io.sin, sn.data(), 0, ggml_nbytes(io.sin)); + ggml_backend_tensor_set(io.mask, mask.data(), 0, ggml_nbytes(io.mask)); + graph_unique_names(prefix_graph.graph()); + if (ggml_backend_graph_compute(backend, prefix_graph.graph()) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "%s: prefix compute failed\n", kTag); + return {}; + } + } + const double t3 = now_ms(); + + // --- the denoising loop ------------------------------------------------------------------------------ + { + const size_t arena = ggml_tensor_overhead() * 8192 + ggml_graph_overhead_custom(8192, false); + if (!step_graph.ensure(backend, n, arena, [&](ggml_context * C, StepIO & io) { return build_step(C, io, n); })) { + std::fprintf(stderr, "%s: step graph build/alloc failed\n", kTag); + return {}; + } + } + StepIO & io = step_graph.io(); + const UniPC solver(steps, train_steps, shift); + UniPC::Stream sv, sa; + std::vector xa = an, vv((size_t) (n_video * patch_dim)), va((size_t) (chunk * adim)), tsin; + for (int k = 0; k < steps; k++) { + time_sinusoid(solver.timesteps[(size_t) k], tsin); + ggml_backend_tensor_set(io.xv, xv.data(), 0, ggml_nbytes(io.xv)); + ggml_backend_tensor_set(io.xa, xa.data(), 0, ggml_nbytes(io.xa)); + ggml_backend_tensor_set(io.tsin, tsin.data(), 0, ggml_nbytes(io.tsin)); + graph_unique_names(step_graph.graph()); + if (ggml_backend_graph_compute(backend, step_graph.graph()) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "%s: step %d compute failed\n", kTag, k); + return {}; + } + ggml_backend_tensor_get(io.vv, vv.data(), 0, ggml_nbytes(io.vv)); + ggml_backend_tensor_get(io.va, va.data(), 0, ggml_nbytes(io.va)); + // mask_velocity: the condition frame has no velocity. + std::fill(vv.begin(), vv.begin() + (ptrdiff_t) (ppf * patch_dim), 0.f); + solver.step(k, sv, vv, xv); + solver.step(k, sa, va, xa); + } + const double t4 = now_ms(); + + // The chunk in dataset units: the unnormalizer's (x + 1) / 2 * (q99 - q01) + q01. + std::vector actions((size_t) (chunk * adim)); + for (int64_t t = 0; t < chunk; t++) + for (int64_t j = 0; j < adim; j++) { + float denom = q99[(size_t) j] - q01[(size_t) j]; + if (denom == 0.f) + denom = 1e-8f; + actions[(size_t) (t * adim + j)] = (xa[(size_t) (t * adim + j)] + 1.f) * denom / 2.f + q01[(size_t) j]; + } + + stats.ms_vision = (float) (t2 - t1); + stats.ms_prefill = (float) (t3 - t2); + stats.ms_denoise = (float) (t4 - t3); + stats.ms_inference = stats.ms_prefill + stats.ms_denoise; + stats.ms_total = (float) (now_ms() - t0); + return actions; +} + +} // namespace + +std::unique_ptr cosmos3_policy_create(const std::string & mmproj_path, const std::string & ckpt_path, + const std::string &, const Options & opts) { + if (!mmproj_path.empty()) + std::printf("%s: note - mmproj '%s' is ignored (the VAE is baked into the GGUF)\n", kTag, mmproj_path.c_str()); + if (opts.act_dtype.value_or(GGML_TYPE_F32) != GGML_TYPE_F32) { + std::fprintf(stderr, "%s: --act-dtype bf16 is not supported; the activations stay F32\n", kTag); + return nullptr; + } + + auto m = std::make_unique(); + m->mt = opts.weight_dtype.value_or(default_weight_dtype(GGML_TYPE_BF16)); + gguf_reader g("cosmos3_policy"); + if (!g.open(ckpt_path)) + return nullptr; + if (!g.has("cosmos3_policy.architecture")) { + std::fprintf(stderr, "%s: %s is not a Cosmos3 Policy GGUF\n", kTag, ckpt_path.c_str()); + return nullptr; + } + if (!m->load_config(g) || !resolve_num_steps("cosmos3_policy", opts, m->steps)) + return nullptr; + + const Backend b = backend_init(kTag, m->n_threads); + if (!b.handle) + return nullptr; + m->backend = b.handle; + // An explicit --weight-dtype f32 keeps the convolutions F32 too. + m->conv_direct = b.is_cuda && cuda_conv_has_mma() && opts.weight_dtype.value_or(GGML_TYPE_F16) != GGML_TYPE_F32; + if (opts.flash_attn.has_value()) + set_flash_attn(*opts.flash_attn); + + if (!m->load_weights(g)) + return nullptr; + + Config & c = m->cfg; + c.n_img = m->n_video; // the video tokens of the generation stream + c.n_lang = m->max_prompt; + c.n_state = 0; + c.n_prefix = m->max_prompt; + c.n_suffix = m->chunk; + c.n_full = c.n_prefix + c.n_img + c.n_suffix; + c.hidden = m->hidden; + c.expert_h = m->hidden; + c.intermediate = m->inter; + c.expert_inter = m->inter; + c.n_q_heads = m->n_heads; + c.n_kv_heads = m->n_kv_heads; + c.head_dim = m->head_dim; + c.q_full_dim = m->n_heads * m->head_dim; + c.kv_full_dim = m->n_kv_heads * m->head_dim; + c.n_layers = m->n_layers; + c.max_state_dim = 0; + c.real_state_dim = 0; + c.max_action_dim = m->adim; // the model's 64 padded channels are dropped in the GGUF + c.real_action_dim = m->adim; + c.num_steps = m->steps; + c.rms_eps = m->rms_eps; + c.norm_eps = m->rms_eps; + c.rope_n_dims = (int) m->head_dim; + c.rope_freq_base = m->rope_theta; + c.denormalized = true; // dataset units: [dpos(3), rot6d(6), gripper] + + std::printf("%s: weights resident %.1f MiB (%s GEMMs, %s convs) - %lld MoT layers, hidden %lld, canvas %dx%d, " + "prompt <= %lld tokens, %d UniPC steps, chunk %lldx%lld\n", + kTag, ggml_backend_buffer_get_size(m->buf_w) / (1024.0 * 1024.0), dtype_name(m->mt), + m->conv_direct ? "f16 direct" : "f32 im2col", (long long) m->n_layers, (long long) m->hidden, + m->layout.canvas_h, m->layout.canvas_w, (long long) m->max_prompt, m->steps, (long long) m->chunk, + (long long) m->adim); + std::printf("%s: send %d views in this order: third person, wrist (or one %dx%d canvas); the instruction as " + "--text (ASCII; needs the tokenizer in the GGUF) or as token ids; noise = %lldx%lld floats\n", + kTag, m->layout.n_views, m->layout.canvas_w, m->layout.canvas_h, (long long) m->chunk, (long long) m->adim); + return m; +} + +// ------------------------------------------------------------------------------------------------ +// The prompt (lerobot's Cosmos3PromptProcessorStep and Cosmos3TokenizerProcessorStep) and the LIBERO action +// decode (Cosmos3LiberoActionProcessorStep), outside the graphs. See cosmos3_policy.h. +// ------------------------------------------------------------------------------------------------ + +namespace { + +constexpr const char * kConcatViewTemplate = "This video contains concatenated views from multiple camera perspectives."; + +// Chat template of the understanding stream (the Cosmos3 Edge template: an empty system turn, one user turn, +// the generation prompt with an empty think block) followed by the end-of-turn token and the token that starts +// the generation stream, as Cosmos3TokenizerProcessorStep.encode builds it. The pieces around the caption. +constexpr const char * kPromptHead = "\n<|im_start|>system\n<|im_end|>\n<|im_start|>user\n"; +constexpr const char * kPromptTail = "<|im_end|>\n<|im_start|>assistant\n\n<|im_end|><|vision_start|>"; + +// ---- ASCII classes of the pre-tokenizer ------------------------------------------------------------ +inline bool is_upper(unsigned char c) { return c >= 'A' && c <= 'Z'; } +inline bool is_lower(unsigned char c) { return c >= 'a' && c <= 'z'; } +inline bool is_digit(unsigned char c) { return c >= '0' && c <= '9'; } +inline bool is_letter(unsigned char c) { return is_upper(c) || is_lower(c); } +inline bool is_nl(unsigned char c) { return c == '\r' || c == '\n'; } +inline bool is_space(unsigned char c) { return c == ' ' || (c >= '\t' && c <= '\r'); } + +// The length of the pre-token that starts at `i`, for the Split regex of the tokenizer: +// [^\r\n\p{L}\p{N}]?[\p{Lu}...]*[\p{Ll}...]+ | [^\r\n\p{L}\p{N}]?[\p{Lu}...]+[\p{Ll}...]* | \p{N} +// | ?[^\s\p{L}\p{N}]+[\r\n/]* | \s*[\r\n]+ | \s+(?!\S) | \s+ +// restricted to ASCII, where the letter classes are A-Z and a-z. The first alternative that matches wins. +size_t pretoken_len(const std::string & t, size_t i) { + const size_t n = t.size(); + auto at = [&](size_t k) { return (unsigned char) t[k]; }; + + // 1 and 2: an optional leading symbol, then capitals and lower case letters. + for (int alt = 1; alt <= 2; alt++) { + for (int with_prefix = 1; with_prefix >= 0; with_prefix--) { + size_t p = i; + if (with_prefix) { + if (p >= n || is_nl(at(p)) || is_letter(at(p)) || is_digit(at(p))) + continue; + p++; + } + size_t u = p; + while (u < n && is_upper(at(u))) + u++; + size_t l = u; + while (l < n && is_lower(at(l))) + l++; + if (alt == 1 && l > u) + return l - i; // upper* lower+ + if (alt == 2 && u > p) + return l - i; // upper+ lower* + } + } + // 3: one digit. + if (is_digit(at(i))) + return 1; + // 4: an optional space, symbols, then newlines and slashes. + { + size_t p = i; + if (at(p) == ' ') + p++; + size_t s = p; + while (s < n && !is_space(at(s)) && !is_letter(at(s)) && !is_digit(at(s))) + s++; + if (s > p) { + while (s < n && (is_nl(at(s)) || at(s) == '/')) + s++; + return s - i; + } + } + // 5, 6 and 7: white space. \s*[\r\n]+ ends at the last line break of the run. + size_t run = i; + while (run < n && is_space(at(run))) + run++; + if (run == i) + return 1; // not reachable for ASCII; keeps the loop moving + for (size_t k = run; k > i; k--) + if (is_nl(at(k - 1))) + return k - i; + if (run == n) + return run - i; // \s+(?!\S) at the end of the text + if (run - i >= 2) + return run - i - 1; // \s+(?!\S) leaves the last space for the next token + return run - i; // \s+ +} + +// GPT-2's bytes_to_unicode: printable bytes stand for themselves, the others take 256 + n in order. +void byte_level_table(std::string out[256], int cp_of_byte[256]) { + int n = 0; + for (int b = 0; b < 256; b++) { + const bool printable = (b >= 33 && b <= 126) || (b >= 161 && b <= 172) || (b >= 174 && b <= 255); + const int cp = printable ? b : 256 + n++; + cp_of_byte[b] = cp; + std::string s; + if (cp < 0x80) { + s += (char) cp; + } else { // cp < 0x800 here + s += (char) (0xC0 | (cp >> 6)); + s += (char) (0x80 | (cp & 0x3F)); + } + out[b] = s; + } +} + +std::string json_string(const std::string & s) { // json.dumps(s) for ASCII s + std::string o = "\""; + for (unsigned char c : s) { + switch (c) { + case '"': o += "\\\""; break; + case '\\': o += "\\\\"; break; + case '\n': o += "\\n"; break; + case '\r': o += "\\r"; break; + case '\t': o += "\\t"; break; + case '\b': o += "\\b"; break; + case '\f': o += "\\f"; break; + default: + if (c >= 0x20 && c <= 0x7e) { + o += (char) c; + } else { + char b[8]; + std::snprintf(b, sizeof(b), "\\u%04x", c); + o += b; + } + } + } + return o + "\""; +} + +std::string strip(const std::string & s) { // str.strip() for ASCII + size_t a = 0, b = s.size(); + while (a < b && is_space((unsigned char) s[a])) + a++; + while (b > a && is_space((unsigned char) s[b - 1])) + b--; + return s.substr(a, b - a); +} + +bool ascii(const std::string & s) { + return std::all_of(s.begin(), s.end(), [](unsigned char c) { return c < 0x80; }); +} + +} // namespace + +// --------------------------------------------------------------------------------------------------- +// Loading +// --------------------------------------------------------------------------------------------------- + +bool Cosmos3Prompt::load(const gguf_reader & g) { + gguf_context * ctx = g.gctx; + const int64_t kt = gguf_find_key(ctx, "tokenizer.ggml.tokens"); + const int64_t kg = gguf_find_key(ctx, "tokenizer.ggml.token_type"); + const int64_t km = gguf_find_key(ctx, "tokenizer.ggml.merges"); + if (kt < 0 || kg < 0 || km < 0 || !g.has("cosmos3_policy.caption.view_description")) { + std::fprintf(stderr, "vla(cosmos3_policy): this GGUF has no tokenizer; convert with --policy-dir, or pass --tokens\n"); + return false; + } + if (gguf_get_arr_type(ctx, kt) != GGUF_TYPE_STRING || gguf_get_arr_type(ctx, km) != GGUF_TYPE_STRING || + gguf_get_arr_type(ctx, kg) != GGUF_TYPE_INT32) { + std::fprintf(stderr, "vla(cosmos3_policy): the tokenizer arrays have unexpected types\n"); + return false; + } + const size_t n = gguf_get_arr_n(ctx, kt); + if (n == 0 || gguf_get_arr_n(ctx, kg) != n) { + std::fprintf(stderr, "vla(cosmos3_policy): the token and token-type arrays disagree\n"); + return false; + } + int cp_of_byte[256]; + byte_level_table(byte_str_, cp_of_byte); + std::unordered_map byte_of_cp; + for (int b = 0; b < 256; b++) + byte_of_cp[cp_of_byte[b]] = b; + + const int32_t * types = (const int32_t *) gguf_get_arr_data(ctx, kg); + tokens_.resize(n); + vocab_.reserve(n * 2); + for (size_t i = 0; i < n; i++) { + tokens_[i] = gguf_get_arr_str(ctx, kt, i); + vocab_.emplace(tokens_[i], (int32_t) i); + // llama.cpp's token types: 3 control, 4 user defined. Both are matched as such in the text. + if (types[i] == 3 || types[i] == 4) { + std::string raw; // back from the byte-level spelling to bytes + bool ok = true; + for (size_t k = 0; k < tokens_[i].size();) { + const unsigned char c = (unsigned char) tokens_[i][k]; + int cp, len; + if (c < 0x80) { + cp = c; len = 1; + } else if ((c & 0xE0) == 0xC0 && k + 1 < tokens_[i].size()) { + cp = ((c & 0x1F) << 6) | ((unsigned char) tokens_[i][k + 1] & 0x3F); len = 2; + } else { + ok = false; + break; + } + const auto it = byte_of_cp.find(cp); + if (it == byte_of_cp.end()) { ok = false; break; } + raw += (char) it->second; + k += (size_t) len; + } + if (ok && !raw.empty()) { + specials_.emplace_back(raw, (int32_t) i); + special_first_[(unsigned char) raw[0]] = true; + } + } + } + std::stable_sort(specials_.begin(), specials_.end(), + [](const auto & a, const auto & b) { return a.first.size() > b.first.size(); }); + + const size_t nm = gguf_get_arr_n(ctx, km); + rank_.reserve(nm * 2); + for (size_t i = 0; i < nm; i++) + rank_.emplace(gguf_get_arr_str(ctx, km, i), (int32_t) i); + + view_description_ = g.str("cosmos3_policy.caption.view_description"); + aspect_ratio_ = g.str("cosmos3_policy.caption.aspect_ratio"); + fps_ = g.f32("cosmos3_policy.caption.fps"); + use_state_ = g.u32("cosmos3_policy.caption.use_state") != 0; + chunk_ = (int) g.u32("cosmos3_policy.chunk_size"); + canvas_h_ = (int) g.u32("cosmos3_policy.canvas_h"); + canvas_w_ = (int) g.u32("cosmos3_policy.canvas_w"); + if (fps_ <= 0.0 || chunk_ <= 0 || aspect_ratio_.empty()) { + std::fprintf(stderr, "vla(cosmos3_policy): the caption settings of the GGUF are incomplete\n"); + return false; + } + return true; +} + +// --------------------------------------------------------------------------------------------------- +// The caption and the chat text +// --------------------------------------------------------------------------------------------------- + +std::string Cosmos3Prompt::caption(const std::string & task_in) const { + std::string task = strip(task_in); + if (task.empty()) + return ""; + if (task.back() != '.' && task.back() != '!' && task.back() != '?') + task += "."; + const double duration = (double) (chunk_ + 1) / fps_; + const long rounded = (long) std::nearbyint(duration); // Python's round(): half to even + char time[48]; + std::snprintf(time, sizeof(time), "0:00-%ld:%02ld", rounded / 60, rounded % 60); + + std::string framing = kConcatViewTemplate; + const std::string vd = strip(view_description_); + if (!vd.empty()) + framing += " " + vd; + + std::string o = "{\"cinematography\": {\"framing\": " + json_string(framing) + "}, \"actions\": [{\"time\": " + + json_string(time) + ", \"description\": " + json_string(task) + ", \"idle_frame\": " + + json_string("0 out of " + std::to_string(chunk_ + (use_state_ ? 1 : 0)) + ".") + "}], \"duration\": " + + json_string(std::to_string((long) duration) + "s") + ", \"fps\": " + std::to_string((long) fps_) + + ".0, \"resolution\": {\"H\": " + std::to_string(canvas_h_) + ", \"W\": " + std::to_string(canvas_w_) + + "}, \"aspect_ratio\": " + json_string(aspect_ratio_) + "}"; + return o; +} + +std::string Cosmos3Prompt::chat_text(const std::string & task) const { + return std::string(kPromptHead) + caption(task) + kPromptTail; +} + +// --------------------------------------------------------------------------------------------------- +// BPE +// --------------------------------------------------------------------------------------------------- + +bool Cosmos3Prompt::bpe_segment(const std::string & text, std::vector & ids, std::string & err) const { + if (!ascii(text)) { + err = "non-ASCII text: this tokenizer implements the ASCII letter classes only"; + return false; + } + for (size_t i = 0; i < text.size();) { + const size_t len = pretoken_len(text, i); + std::string word; // the pre-token in the byte-level alphabet + for (size_t k = i; k < i + len; k++) + word += byte_str_[(unsigned char) text[k]]; + i += len; + + // ignore_merges: a pre-token that is a token is not merged. + const auto whole = vocab_.find(word); + if (whole != vocab_.end()) { + ids.push_back(whole->second); + continue; + } + // Merges by rank, the leftmost pair of the lowest rank first (the order of a priority queue). + std::vector sym; + for (size_t k = 0; k < word.size();) { + const size_t cl = ((unsigned char) word[k] < 0x80) ? 1 : 2; // byte-level characters are 1 or 2 bytes + sym.push_back(word.substr(k, cl)); + k += cl; + } + for (;;) { + int32_t best = INT32_MAX; + size_t at = 0; + for (size_t k = 0; k + 1 < sym.size(); k++) { + const auto it = rank_.find(sym[k] + " " + sym[k + 1]); + if (it != rank_.end() && it->second < best) { + best = it->second; + at = k; + } + } + if (best == INT32_MAX) + break; + sym[at] += sym[at + 1]; + sym.erase(sym.begin() + (long) at + 1); + } + for (const std::string & s : sym) { + const auto it = vocab_.find(s); + if (it == vocab_.end()) { + err = "a symbol is not in the vocabulary: " + s; + return false; + } + ids.push_back(it->second); + } + } + return true; +} + +bool Cosmos3Prompt::encode(const std::string & text, std::vector & ids, std::string & err) const { + size_t start = 0, i = 0; + while (i < text.size()) { + // A special token can start here only with a byte some special starts with (what the GGUF declares, + // whatever its spelling: '<' and '[' for the Cosmos3 tokenizer). + if (special_first_[(unsigned char) text[i]]) { + const std::pair * hit = nullptr; + for (const auto & sp : specials_) // longest first + if (text.compare(i, sp.first.size(), sp.first) == 0) { + hit = &sp; + break; + } + if (hit) { + if (i > start && !bpe_segment(text.substr(start, i - start), ids, err)) + return false; + ids.push_back(hit->second); + i += hit->first.size(); + start = i; + continue; + } + } + i++; + } + return start >= text.size() || bpe_segment(text.substr(start), ids, err); +} + +bool Cosmos3Prompt::tokens(const std::string & task, std::vector & ids, std::string & err) const { + ids.clear(); + // The caption is json.dumps of the task, which escapes non-ASCII as UTF-16 \u sequences: refuse it here + // instead of escaping bytes wrongly. + if (!ascii(task)) { + err = "non-ASCII instruction: only ASCII text is supported"; + return false; + } + return encode(chat_text(task), ids, err); +} + +// --------------------------------------------------------------------------------------------------- +// rot6d -> axis-angle +// --------------------------------------------------------------------------------------------------- + +namespace { + +// Eigen-decomposition of a symmetric 3x3 matrix by cyclic Jacobi rotations; eigenvalues descending, the +// eigenvectors as the columns of v. +void jacobi3(double a[3][3], double w[3], double v[3][3]) { + for (int i = 0; i < 3; i++) + for (int j = 0; j < 3; j++) + v[i][j] = i == j ? 1.0 : 0.0; + for (int sweep = 0; sweep < 50; sweep++) { + const double off = std::fabs(a[0][1]) + std::fabs(a[0][2]) + std::fabs(a[1][2]); + if (off < 1e-300) + break; + for (int p = 0; p < 2; p++) + for (int q = p + 1; q < 3; q++) { + if (std::fabs(a[p][q]) < 1e-300) + continue; + const double theta = (a[q][q] - a[p][p]) / (2.0 * a[p][q]); + const double t = (theta >= 0 ? 1.0 : -1.0) / (std::fabs(theta) + std::sqrt(theta * theta + 1.0)); + const double c = 1.0 / std::sqrt(t * t + 1.0), s = t * c; + for (int k = 0; k < 3; k++) { // columns p and q of a + const double akp = a[k][p], akq = a[k][q]; + a[k][p] = c * akp - s * akq; + a[k][q] = s * akp + c * akq; + } + for (int k = 0; k < 3; k++) { // rows p and q + const double apk = a[p][k], aqk = a[q][k]; + a[p][k] = c * apk - s * aqk; + a[q][k] = s * apk + c * aqk; + } + for (int k = 0; k < 3; k++) { + const double vkp = v[k][p], vkq = v[k][q]; + v[k][p] = c * vkp - s * vkq; + v[k][q] = s * vkp + c * vkq; + } + } + } + int order[3] = { 0, 1, 2 }; + std::sort(order, order + 3, [&](int x, int y) { return a[x][x] > a[y][y]; }); + double v2[3][3]; + for (int j = 0; j < 3; j++) { + w[j] = a[order[j]][order[j]]; + for (int i = 0; i < 3; i++) + v2[i][j] = v[i][order[j]]; + } + std::memcpy(v, v2, sizeof(v2)); +} + +} // namespace + +void cosmos3_libero_action(const float * a, float * out) { + const double f[3] = { a[3], a[4], a[5] }, s[3] = { a[6], a[7], a[8] }; + const double c[3] = { f[1] * s[2] - f[2] * s[1], f[2] * s[0] - f[0] * s[2], f[0] * s[1] - f[1] * s[0] }; + double m[3][3]; // columns: first, second, their cross product + for (int r = 0; r < 3; r++) { + m[r][0] = f[r]; + m[r][1] = s[r]; + m[r][2] = c[r]; + } + // The nearest rotation (the SVD projection onto SO(3)): U V^T, with the last singular direction flipped + // if that is a reflection. Singular values and V from the eigenproblem of M^T M. + double mtm[3][3]; + for (int i = 0; i < 3; i++) + for (int j = 0; j < 3; j++) + mtm[i][j] = m[0][i] * m[0][j] + m[1][i] * m[1][j] + m[2][i] * m[2][j]; + double w[3], v[3][3]; + jacobi3(mtm, w, v); + + double u[3][3] = {}; + const double smax = std::sqrt(std::max(w[0], 0.0)); + bool ok = smax > 0.0; + for (int j = 0; j < 3 && ok; j++) { + const double sj = std::sqrt(std::max(w[j], 0.0)); + if (sj <= 1e-12 * smax) { + if (j == 0) { ok = false; break; } + // A vanishing singular value: complete U to an orthonormal frame. + if (j == 1) { // pick any vector orthogonal to u0 + const double ax = std::fabs(u[0][0]), ay = std::fabs(u[1][0]), az = std::fabs(u[2][0]); + double t[3] = { 0, 0, 0 }; + t[ax <= ay && ax <= az ? 0 : (ay <= az ? 1 : 2)] = 1.0; + const double d = t[0] * u[0][0] + t[1] * u[1][0] + t[2] * u[2][0]; + double x[3] = { t[0] - d * u[0][0], t[1] - d * u[1][0], t[2] - d * u[2][0] }; + const double nx = std::sqrt(x[0] * x[0] + x[1] * x[1] + x[2] * x[2]); + for (int r = 0; r < 3; r++) u[r][1] = x[r] / nx; + } else { + u[0][2] = u[1][0] * u[2][1] - u[2][0] * u[1][1]; + u[1][2] = u[2][0] * u[0][1] - u[0][0] * u[2][1]; + u[2][2] = u[0][0] * u[1][1] - u[1][0] * u[0][1]; + } + continue; + } + for (int r = 0; r < 3; r++) + u[r][j] = (m[r][0] * v[0][j] + m[r][1] * v[1][j] + m[r][2] * v[2][j]) / sj; + } + double R[3][3] = { { 1, 0, 0 }, { 0, 1, 0 }, { 0, 0, 1 } }; + if (ok) { + auto compose = [&]() { + for (int i = 0; i < 3; i++) + for (int j = 0; j < 3; j++) + R[i][j] = u[i][0] * v[j][0] + u[i][1] * v[j][1] + u[i][2] * v[j][2]; + }; + compose(); + const double det = R[0][0] * (R[1][1] * R[2][2] - R[1][2] * R[2][1]) - R[0][1] * (R[1][0] * R[2][2] - R[1][2] * R[2][0]) + + R[0][2] * (R[1][0] * R[2][1] - R[1][1] * R[2][0]); + if (det < 0) { + for (int r = 0; r < 3; r++) u[r][2] = -u[r][2]; + compose(); + } + } + // matrix_to_axis_angle (angles below pi), in double as lerobot does. + const double skew[3] = { R[2][1] - R[1][2], R[0][2] - R[2][0], R[1][0] - R[0][1] }; + const double sin_a = std::sqrt(skew[0] * skew[0] + skew[1] * skew[1] + skew[2] * skew[2]) / 2.0; + const double cos_a = (R[0][0] + R[1][1] + R[2][2] - 1.0) / 2.0; + const double angle = std::atan2(sin_a, cos_a); + const double scale = sin_a > 1e-8 ? angle / std::max(sin_a, 1e-8) : 1.0; + + out[0] = a[0]; + out[1] = a[1]; + out[2] = a[2]; + for (int r = 0; r < 3; r++) + out[3 + r] = (float) (skew[r] * scale / 2.0); + out[6] = std::min(1.0f, std::max(-1.0f, 1.0f - 2.0f * a[9])); +} + +} // namespace vla diff --git a/src/models/cosmos3_policy.h b/src/models/cosmos3_policy.h new file mode 100644 index 00000000..5380f378 --- /dev/null +++ b/src/models/cosmos3_policy.h @@ -0,0 +1,77 @@ +// Copyright 2026 VinRobotics +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Cosmos3 Policy outside the graphs (src/models/cosmos3_policy.cpp): the prompt (lerobot's Cosmos3PromptProcessorStep and +// Cosmos3TokenizerProcessorStep) and the LIBERO action decode (Cosmos3LiberoActionProcessorStep). +// +// * The prompt is the JSON caption of the task in the chat template of the understanding stream, then the +// end-of-turn token and <|vision_start|>. The tokenizer is the model's byte-level BPE (131,072 tokens), +// read from the GGUF (`tokenizer.ggml.*`, the keys llama.cpp uses). It handles ASCII text only: its +// pre-tokenizer splits on Unicode letter classes, which are not implemented here, so other text is refused. +// * rot6d -> axis-angle follows the lerobot step: the two columns and their cross product form a matrix that +// is projected onto SO(3), then converted to a rotation vector, in double precision. + +#pragma once + +#include "gguf_reader.h" + +#include +#include +#include +#include + +namespace vla { + +struct Cosmos3Prompt { + // Loads the tokenizer and the caption settings from a Cosmos3 Policy GGUF written with + // `--policy-dir`. False, with a reason on stderr, when the file has none. + bool load(const gguf_reader & g); + + // The ids of the prompt of `task` (what the policy feeds the understanding stream, unpadded). False, + // with a reason in `err`, for text this tokenizer does not handle. + bool tokens(const std::string & task, std::vector & ids, std::string & err) const; + + // The JSON caption (Cosmos3PromptProcessorStep.caption with 0 idle frames) and the full text the + // tokenizer sees (chat template, end-of-turn and <|vision_start|> included); for tests and debugging. + std::string caption(const std::string & task) const; + std::string chat_text(const std::string & task) const; + + // Token ids of a text with special tokens recognised in it, as Hugging Face's tokenizer does. + bool encode(const std::string & text, std::vector & ids, std::string & err) const; + +private: + bool bpe_segment(const std::string & text, std::vector & ids, std::string & err) const; + + std::vector tokens_; // id -> byte-level spelling + std::unordered_map vocab_; // byte-level spelling -> id + std::unordered_map rank_; // "left right" -> merge priority + std::vector> specials_; // raw text -> id, longest first + bool special_first_[256] = {}; // bytes a special can start with + std::string byte_str_[256]; // byte -> its byte-level character (UTF-8) + + // Caption settings. + std::string view_description_; + std::string aspect_ratio_; + double fps_ = 20.0; + int chunk_ = 16; + int canvas_h_ = 192; + int canvas_w_ = 320; + bool use_state_ = false; +}; + +// The LIBERO command of one 10-wide model action [dpos(3), rot6d(6), gripper]: [dpos(3), axis-angle(3), gripper] +// with the gripper mapped from 1 open / 0 closed to -1 open / +1 closed. +void cosmos3_libero_action(const float * action10, float * out7); + +} // namespace vla diff --git a/src/serving/vla-cli.cpp b/src/serving/vla-cli.cpp index b9074ad8..aa55f76e 100644 --- a/src/serving/vla-cli.cpp +++ b/src/serving/vla-cli.cpp @@ -27,6 +27,7 @@ #include "model.h" #include "options.h" #include "serving/hf_fetch.h" +#include "models/cosmos3_policy.h" #include "tokenizer.h" #define STB_IMAGE_IMPLEMENTATION @@ -132,6 +133,7 @@ const char * arch_slug(Arch a) { case Arch::OCTO: return "octo"; case Arch::TURBOVLA: return "turbovla"; case Arch::ACT: return "act"; + case Arch::COSMOS3_POLICY: return "cosmos3_policy"; } return ""; } @@ -256,6 +258,7 @@ void usage(const char * prog) { " --state proprioception floats, comma-separated (default zeros); pi05\n" " --text needs it, since the state is part of the prompt\n" " --pretty print one action row (max_action_dim values) per line\n" + " --libero Cosmos3 Policy only: print LIBERO commands [dpos(3), axis-angle(3), gripper] instead of rot6d rows\n" " --config policy config.json; its \"runtime\" block sets the flags below,\n" " and flags given here win\n" "%s", @@ -268,6 +271,7 @@ int main(int argc, char ** argv) { std::string mmproj, ckpt, hf, tokens_s, state_s, text_s, config_path; std::vector image_paths; bool pretty = false; + bool libero = false; Options opts; std::string opt_err; @@ -289,6 +293,7 @@ int main(int argc, char ** argv) { else if (a == "--text") text_s = need("--text"); else if (a == "--state") state_s = need("--state"); else if (a == "--pretty") pretty = true; + else if (a == "--libero") libero = true; else if (a == "--config") config_path = need("--config"); else if (opts.parse_arg(argc, argv, i, opt_err)) continue; else if (!opt_err.empty()) { @@ -342,6 +347,11 @@ int main(int argc, char ** argv) { std::fprintf(stderr, "vla-cli: %s takes no instruction; drop --text/--tokens\n", arch_slug(arch)); return 1; } + if (libero && arch != Arch::COSMOS3_POLICY) { + std::fprintf(stderr, "vla-cli: --libero decodes the 10-wide rot6d action of Cosmos3 Policy only; %s is %s\n", + ckpt.c_str(), arch_slug(arch)); + return 1; + } // Validate the cheap args before loading the model. std::vector lang; std::vector attn; // Octo only; empty leaves Inputs::attention_mask null. @@ -350,7 +360,7 @@ int main(int argc, char ** argv) { if (!parse_floats(state_s, state)) return 1; if (!text_s.empty()) { - if (has_spm_tokenizer(ckpt, arch_slug(arch))) { + if (arch == Arch::COSMOS3_POLICY || has_spm_tokenizer(ckpt, arch_slug(arch))) { if (!tokenize_prompt(ckpt, arch_slug(arch), text_s, state, lang, attn)) return 1; for (size_t i=0; i 0 ? cfg.max_action_dim : 1; + int64_t cols = cfg.max_action_dim > 0 ? cfg.max_action_dim : 1; + if (libero) { // checked above: Cosmos3 Policy, 10 columns + std::vector cmd; + for (size_t r = 0; r + 10 <= act.size(); r += 10) { + float row[7]; + cosmos3_libero_action(&act[r], row); + cmd.insert(cmd.end(), row, row + 7); + } + act = std::move(cmd); + cols = 7; + } if (pretty) { for (size_t i=0; i are its last tokens. + const int prompt_len = (int) g.u32("cosmos3_policy.prompt_len"); + if (prompt_len > 0 && (int) ids.size() > prompt_len) { + std::fprintf(stderr, "vla(cosmos3_policy): the prompt is %zu tokens, the model takes at most %d; shorten the instruction\n", + ids.size(), prompt_len); + return false; + } + attention_mask.clear(); + return true; + } + if (a == "octo") { if (!spm_encode(g, arch, text, ids)) return false;