# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. # SPDX-License-Identifier: Apache-2.0 """Host-side observation -> model-input encoding for GR00T N1.5 / N1.6 / N1.7 (CPU, bit-exact vs the goldens). Implements option (c) of ``host-pipeline-parity.md`` §7: a minimal re-implementation using only what the tt-metal venv ships (torch / torchvision CPU, cv2, Pillow, transformers tokenizer + ``Qwen2VLImageProcessor``). Every recipe below is the one that reproduced the golden ``inputs/`` bit-exactly in §4.2 (uint8 frames, fp32 pixel values, int64 ids, fp32 state): * N1.5 -- ``FourierGr1ArmsOnlyDataConfig`` eval transforms (``x/255 -> CenterCrop(int(256*0.95)) -> Resize(224, bilinear, antialias) -> (x*255).uint8``) + Eagle2.5 single-tile normalisation + bracketed ``repr([instruction])`` prompt + sin/cos state (``Isaac-GR00T-n1.5/gr00t/data/transform/video.py``, ``model/transforms.py``). * N1.6 -- albumentations eval chain (``LetterBoxPad -> SmallestMaxSize(256, INTER_AREA) -> FractionalCenterCrop(0.95) -> SmallestMaxSize(256)``) + Eagle3 ``fetch_image`` (``smart_resize`` factor 28 -> PIL **BICUBIC** 252x252) + ``(u8 - 127.5) / 127.5`` + text-then-image prompt (``processing_eagle3_vl.py``, ``processing_gr00t_n1d6.py``). * N1.7 -- the same albumentations chain per frame (2 frames x V cameras, frame-major) + the transformers ``Qwen2VLImageProcessor`` recipe (bicubic+antialias ``smart_resize`` factor 32, T=2 duplication, ``(c,t,ph,pw)`` patch rows in merge-window order) + Qwen3-VL prompt + q01/q99 state (``processing_gr00t_n1d7.py:703-752``). The pixels come from :func:`qwen2vl_pixel_patches`: the processor's own torchvision calls (``tvF.resize`` BICUBIC antialias on the stacked uint8 frames, fused ``tvF.normalize`` with ``mean * 255``, the patch reshape / permute) issued directly with the parameters of the cached processor object, bit-equal to ``ip(images=...)`` on every golden sample and on random frames (``tests/common/test_host_parity.py``), without its per-call validation and grouping overhead (docs/plan/n17-perf.md). The static-shape side (``S_pad``, slot layout, position ids, RoPE tables, N1.7 vision tables) is delegated to :mod:`prompts`; statistics and state normalisation to :mod:`normalization`. Nothing here imports ``ttnn``. Per-call cost: the prompt tokenization and the per-layout tables (mRoPE tables per padded id sequence, the N1.7 vision tables per grid) are cached in :mod:`prompts` (pure functions of their arguments), so ``encode`` returns the **same tensor objects** for a repeated instruction / grid -- treat every ``ModelInputs`` tensor as read-only. """ from __future__ import annotations import json import math from dataclasses import dataclass, field from functools import lru_cache from pathlib import Path from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple import cv2 import numpy as np import torch from PIL import Image from models.experimental.gr00t.common.configs import GR00T_VERSIONS, get_config from models.experimental.gr00t.common.normalization import EmbodimentSpec, encode_state, get_embodiment_spec from models.experimental.gr00t.common.prompts import ( LLM_ROPE_INV_FREQ_DTYPE, ShapeKey, SlotLayout, TokenizedPrompt, build_slot_layout, encode_prompt, grid_key, mrope_tables_cached, rope_cos_sin_1d, select_layout, tokenizer_dir, vision_tables_cached, ) # -------------------------------------------------------------------------------------------------------------------- # Observation # -------------------------------------------------------------------------------------------------------------------- @dataclass class Observation: """One raw robot observation at batch 1. ``images``: camera key -> ``(T, H, W, 3)`` uint8 RGB frames, ``T`` = frames per camera of the embodiment's video ``delta_indices`` (1 for N1.5 / N1.6, 2 for N1.7: ``[t-15, t]`` oldest first). ``state``: joint group -> ``(T_s, D)`` float (``T_s`` = 1); the raw physical state, also the reference frame of RELATIVE actions at decode time. ``instruction``: the raw language instruction (formalisation / bracketing happens in ``encode``). """ images: Dict[str, np.ndarray] state: Dict[str, np.ndarray] instruction: str embodiment_tag: str def validate(self) -> None: for k, v in self.images.items(): if not isinstance(v, np.ndarray) or v.dtype != np.uint8 or v.ndim != 4 or v.shape[-1] != 3: raise ValueError( f"image {k!r} must be a (T, H, W, 3) uint8 array, got {getattr(v, 'shape', None)} {getattr(v, 'dtype', None)}" ) for k, v in self.state.items(): if not isinstance(v, np.ndarray) or v.ndim not in (1, 2) or v.dtype.kind != "f": raise ValueError(f"state {k!r} must be a (T_s, D) float array, got {getattr(v, 'shape', None)}") if not isinstance(self.instruction, str): raise TypeError(f"instruction must be str, got {type(self.instruction).__name__}") def load_golden_observation(version: str, obs_dir: Path, embodiment_tag: Optional[str] = None) -> Observation: """Read the raw observation saved next to a golden run (``obs/observation.npz`` for N1.5, ``obs/raw_observation.npz`` (+ ``observation.json`` for the N1.7 language / tag) for N1.6 / N1.7). N1.6 / N1.7 arrays carry a leading batch axis of 1 that is dropped. The embodiment tag comes from ``observation.json`` when present (N1.7), else from ``embodiment_tag`` / the version's canonical embodiment. """ _check_version(version) obs_dir = Path(obs_dir) meta: Dict[str, Any] = {} meta_path = obs_dir / "observation.json" if meta_path.is_file(): with open(meta_path) as fh: meta = json.load(fh) tag = embodiment_tag or meta.get("embodiment_tag") or get_config(version).canonical_embodiment if version == "n15": z = np.load(obs_dir / "observation.npz", allow_pickle=False) images = {k[len("video.") :]: np.ascontiguousarray(z[k]) for k in z.files if k.startswith("video.")} state = {k[len("state.") :]: np.ascontiguousarray(z[k]) for k in z.files if k.startswith("state.")} lang_key = "annotation.human.action.task_description" if lang_key not in z.files: raise KeyError(f"{obs_dir}: observation.npz lacks {lang_key!r}") arr = z[lang_key].reshape(-1) if arr.shape != (1,): raise ValueError(f"{obs_dir}: expected one instruction, got {z[lang_key].shape}") instruction = str(arr[0]) else: z = np.load(obs_dir / "raw_observation.npz", allow_pickle=False) images, state = {}, {} for k in z.files: if k.startswith("video."): v = z[k] if v.ndim != 5 or v.shape[0] != 1: raise ValueError(f"{obs_dir}: {k} must be (1, T, H, W, 3), got {v.shape}") images[k[len("video.") :]] = np.ascontiguousarray(v[0]) elif k.startswith("state."): v = z[k] if v.ndim != 3 or v.shape[0] != 1: raise ValueError(f"{obs_dir}: {k} must be (1, T_s, D), got {v.shape}") state[k[len("state.") :]] = np.ascontiguousarray(v[0]) if version == "n16": if "language" not in z.files: raise KeyError(f"{obs_dir}: raw_observation.npz lacks 'language'") instruction = str(z["language"].reshape(-1)[0]) else: if "language_raw" in meta: instruction = str(meta["language_raw"]) elif "language" in meta: instruction = str(meta["language"]) else: raise KeyError(f"{obs_dir}: observation.json lacks 'language_raw'") obs = Observation(images=images, state=state, instruction=instruction, embodiment_tag=str(tag)) obs.validate() return obs def _check_version(version: str) -> None: if version not in GR00T_VERSIONS: raise KeyError(f"unknown GR00T version {version!r}; expected one of {GR00T_VERSIONS}") # -------------------------------------------------------------------------------------------------------------------- # Image primitives # -------------------------------------------------------------------------------------------------------------------- def normalize_u8(chw: torch.Tensor) -> torch.Tensor: """transformers fused rescale+normalize with mean = std = 0.5: ``(u8 - 127.5) / 127.5`` in fp32 (``F.normalize(x.float(), mean*255, std*255)``). The naive ``((x/255) - 0.5) / 0.5`` is 5.9e-8 off on 27 % of the values and is NOT bit-exact (host-pipeline-parity §4.2 "Common").""" if chw.dtype != torch.uint8: raise TypeError(f"normalize_u8 expects uint8, got {chw.dtype}") return (chw.to(torch.float32) - 127.5) / 127.5 def smallest_max_size(img: np.ndarray, max_size: int, interpolation: int = cv2.INTER_AREA) -> np.ndarray: """albumentations 1.4.18 ``SmallestMaxSize`` (``geometric/functional.py:427-445`` + ``resize`` ``:377-383``): ``scale = max_size / min(h, w)``; if ``scale != 1`` resize to ``(round(h*scale), round(w*scale))`` with cv2 (Python banker's ``round``), returning the input untouched when the size does not change.""" h, w = img.shape[:2] scale = max_size / float(min(h, w)) if scale != 1.0: nh, nw = round(h * scale), round(w * scale) if (nh, nw) != (h, w): return cv2.resize(img, dsize=(nw, nh), interpolation=interpolation) return img def fractional_center_crop(img: np.ndarray, crop_fraction: float) -> np.ndarray: """``FractionalCenterCrop`` (``gr00t_n1d{6,7}/image_augmentations.py``): ``ch = int(h * f)``, ``cw = int(w * f)`` (at least 1), top-left ``((h - ch) // 2, (w - cw) // 2)``. 256 -> 243 at 0.95; 256x455 -> 243x432 at (6, 11).""" if not 0.0 < crop_fraction <= 1.0: raise ValueError(f"crop_fraction must be in (0, 1], got {crop_fraction}") h, w = img.shape[:2] ch = max(1, int(h * crop_fraction)) cw = max(1, int(w * crop_fraction)) y0 = (h - ch) // 2 x0 = (w - cw) // 2 return img[y0 : y0 + ch, x0 : x0 + cw] def letterbox_pad(img: np.ndarray) -> np.ndarray: """``LetterBoxPad``: zero-pad to a square (``pad // 2`` before, remainder after); identity for square frames.""" h, w = img.shape[:2] if h == w: return img m = max(h, w) ph, pw = m - h, m - w return cv2.copyMakeBorder(img, ph // 2, ph - ph // 2, pw // 2, pw - pw // 2, cv2.BORDER_CONSTANT, value=0) def albumentations_eval_transform( img: np.ndarray, shortest_edge: int = 256, crop_fraction: float = 0.95, letterbox: bool = True ) -> np.ndarray: """Eval branch of ``build_image_transformations_albumentations`` (N1.6 always letterboxes; N1.7 only with ``letter_box_transform=True``): ``[LetterBoxPad] -> SmallestMaxSize -> FractionalCenterCrop -> SmallestMaxSize``, uint8 HWC in / out (host-pipeline-parity §4.2 N1.6 step 1, N1.7 step 1).""" if img.ndim != 3 or img.shape[-1] != 3 or img.dtype != np.uint8: raise ValueError(f"expected a (H, W, 3) uint8 frame, got {img.shape} {img.dtype}") if letterbox: img = letterbox_pad(img) img = smallest_max_size(img, shortest_edge, cv2.INTER_AREA) img = fractional_center_crop(img, crop_fraction) return smallest_max_size(img, shortest_edge, cv2.INTER_AREA) def eagle3_smart_resize( height: int, width: int, factor: int = 28, min_pixels: int = 4 * 28 * 28, max_pixels: int = 4096 * 28 * 28, max_size: int = 500 * 14, max_ratio: int = 200, ) -> Tuple[int, int]: """Eagle3 ``smart_resize`` (``processing_eagle3_vl.py:105-138``): round both sides to a multiple of ``factor`` (clamped to ``[factor, max_size]``), then shrink / grow to the pixel budget. 256x256 -> 252x252 -> 81 tokens.""" if max(height, width) / min(height, width) > max_ratio: raise ValueError(f"aspect ratio {max(height, width) / min(height, width)} exceeds {max_ratio}") hb = min(max(factor, round(height / factor) * factor), max_size) wb = min(max(factor, round(width / factor) * factor), max_size) if hb * wb > max_pixels: beta = math.sqrt((hb * wb) / max_pixels) hb = math.floor(hb / beta / factor) * factor wb = math.floor(wb / beta / factor) * factor elif hb * wb < min_pixels: beta = math.sqrt(min_pixels / (height * width)) hb = math.ceil(height * beta / factor) * factor wb = math.ceil(width * beta / factor) * factor return int(hb), int(wb) def qwen2vl_smart_resize( height: int, width: int, factor: int = 32, min_pixels: int = 65_536, max_pixels: int = 16_777_216 ) -> Tuple[int, int]: """Qwen2-VL ``smart_resize`` (transformers ``image_processing_qwen2_vl``): the resized size the N1.7 image processor will pick -- used only to predict the ViT grid (``h/16, w/16``); the pixels come from the processor. 256x455 -> 256x448 -> grid (1, 16, 28).""" hb = round(height / factor) * factor wb = round(width / factor) * factor if hb * wb > max_pixels: beta = math.sqrt((height * width) / max_pixels) hb = max(factor, math.floor(height / beta / factor) * factor) wb = max(factor, math.floor(width / beta / factor) * factor) elif hb * wb < min_pixels: beta = math.sqrt(min_pixels / (height * width)) hb = math.ceil(height * beta / factor) * factor wb = math.ceil(width * beta / factor) * factor return int(hb), int(wb) def n15_crop_resize(frames: np.ndarray, crop_fraction: float = 0.95, out_size: int = 224) -> torch.Tensor: """N1.5 ``VideoToTensor -> VideoCrop(eval: CenterCrop(int(H*0.95))) -> VideoResize(224, bilinear, antialias) -> VideoToNumpy`` (``video.py:276-277, 316-333, 556-557, 584``): ``(T, H, W, 3)`` uint8 -> ``(T, 224, 224, 3)`` uint8, the final ``(x * 255).to(uint8)`` being a **truncation**. ColorJitter is a no-op in eval.""" from torchvision.transforms import InterpolationMode from torchvision.transforms import functional as TF if frames.ndim != 4 or frames.shape[-1] != 3 or frames.dtype != np.uint8: raise ValueError(f"expected (T, H, W, 3) uint8 frames, got {frames.shape} {frames.dtype}") h, w = frames.shape[1:3] x = torch.from_numpy(np.ascontiguousarray(frames)).to(torch.float32) / 255.0 x = x.permute(0, 3, 1, 2) x = TF.center_crop(x, [int(h * crop_fraction), int(w * crop_fraction)]) x = TF.resize(x, [out_size, out_size], interpolation=InterpolationMode.BILINEAR, antialias=True) return (x.permute(0, 2, 3, 1) * 255).to(torch.uint8) # -------------------------------------------------------------------------------------------------------------------- # Per-version vision encoders # -------------------------------------------------------------------------------------------------------------------- @dataclass class VisionInputs: """Vision-side tensors of one observation in the model's image order (frame-major: ``t0 v0, t0 v1, t1 v0, ...``).""" images_uint8: torch.Tensor # (n_img, 3, H, W) uint8 frames handed to the VLM image processor pixel_values: torch.Tensor # n15/n16 (n_img, 3, H, W) fp32 | n17 (n_img * P, 1536) fp32 patch rows tokens_per_image: Tuple[int, ...] image_sizes: Optional[torch.Tensor] = None # (n_img, 2) int64 [[H, W]] (N1.5 / N1.6) image_grid_thw: Optional[torch.Tensor] = None # (n_img, 3) int64 (N1.7) images_pre_resize_uint8: Optional[torch.Tensor] = None # N1.6: (n_img, 3, 256, 256) after the eval chain @property def n_images(self) -> int: return int(self.images_uint8.shape[0]) @property def grid_thw(self) -> Optional[Tuple[int, int, int]]: if self.image_grid_thw is None: return None return tuple(int(x) for x in self.image_grid_thw[0].tolist()) def _stack_frame_major(per_camera_chw: Sequence[torch.Tensor]) -> torch.Tensor: """``[(T, C, H, W)] * V -> (T * V, C, H, W)`` in the reference order (``processing_gr00t_n1d7.py:745-748``, N1.5 ``rearrange("v t c h w -> (t v) c h w")``).""" shapes = {tuple(t.shape) for t in per_camera_chw} if len(shapes) != 1: raise ValueError(f"all cameras must yield identically shaped frames (static shape key), got {sorted(shapes)}") return torch.stack(list(per_camera_chw), dim=1).flatten(0, 1) def n15_encode_images(frames_per_camera: Sequence[np.ndarray], spec: EmbodimentSpec) -> VisionInputs: """N1.5: crop/resize each camera's frames, Eagle2.5 single 224x224 tile (``F.resize`` identity + normalise), ``eagle_image_sizes = [[224, 224]]`` per image, 256 tokens per image (host-pipeline-parity §4.2 N1.5 1-4).""" per_cam = [] for f in frames_per_camera: u8 = n15_crop_resize(f, spec.crop_fraction, 224) # (T, 224, 224, 3) per_cam.append(u8.permute(0, 3, 1, 2).contiguous()) imgs = _stack_frame_major(per_cam) if tuple(imgs.shape[-2:]) != (224, 224): raise NotImplementedError(f"Eagle2.5 dynamic tiling for {tuple(imgs.shape[-2:])} frames is not modelled") pv = normalize_u8(imgs) n = imgs.shape[0] return VisionInputs( images_uint8=imgs, pixel_values=pv, tokens_per_image=(256,) * n, image_sizes=torch.tensor([[224, 224]] * n, dtype=torch.int64), ) def n16_encode_images(frames_per_camera: Sequence[np.ndarray], spec: EmbodimentSpec) -> VisionInputs: """N1.6: albumentations eval chain -> Eagle3 ``fetch_image`` (``smart_resize`` factor 28 + PIL BICUBIC, the default resample for RGB; BILINEAR / LANCZOS / NEAREST are 21-98 LSB off) -> normalise; ``h*w // 784`` tokens (host-pipeline-parity §4.2 N1.6 1-2).""" per_cam_eval, per_cam_eagle, tokens = [], [], [] for f in frames_per_camera: ev, eg = [], [] for frame in f: tr = albumentations_eval_transform( frame, spec.shortest_image_edge, spec.crop_fraction, spec.letter_box_transform ) pil = Image.fromarray(tr).convert("RGB") rh, rw = eagle3_smart_resize(pil.height, pil.width) pil_r = pil.resize( (rw, rh), Image.Resampling.BICUBIC ) # fetch_image: image.resize((w, h)), default = BICUBIC ev.append(torch.from_numpy(np.ascontiguousarray(tr)).permute(2, 0, 1)) eg.append(torch.from_numpy(np.array(pil_r, dtype=np.uint8)).permute(2, 0, 1)) tokens.append(rh * rw // 784) per_cam_eval.append(torch.stack(ev)) per_cam_eagle.append(torch.stack(eg)) imgs = _stack_frame_major(per_cam_eagle) pre = _stack_frame_major(per_cam_eval) n = imgs.shape[0] h, w = imgs.shape[-2:] return VisionInputs( images_uint8=imgs.contiguous(), pixel_values=normalize_u8(imgs), tokens_per_image=(int(h * w // 784),) * n, image_sizes=torch.tensor([[int(h), int(w)]] * n, dtype=torch.int64), images_pre_resize_uint8=pre.contiguous(), ) @lru_cache(maxsize=None) def get_qwen2vl_image_processor(): """Cached transformers ``Qwen2VLImageProcessor`` from the Cosmos-Reason2-2B shell ``preprocessor_config.json`` (patch 16, temporal 2, merge 2, mean = std = 0.5, bicubic, shortest_edge 65536 px). The torchvision backend is the one verified bit-exact (host-pipeline-parity §4.2 / §9 item 7); anything else raises.""" import transformers from transformers import AutoImageProcessor prev = transformers.logging.get_verbosity() transformers.logging.set_verbosity_error() try: ip = AutoImageProcessor.from_pretrained(str(tokenizer_dir("n17"))) finally: transformers.logging.set_verbosity(prev) name = type(ip).__name__ if name not in ("Qwen2VLImageProcessor", "Qwen2VLImageProcessorFast"): raise RuntimeError(f"expected the torchvision Qwen2VLImageProcessor, got {name} (PIL backend is unverified)") expected = dict( patch_size=16, temporal_patch_size=2, merge_size=2, do_resize=True, do_rescale=True, do_normalize=True ) for k, v in expected.items(): if getattr(ip, k, None) != v: raise RuntimeError(f"Qwen2VLImageProcessor.{k} = {getattr(ip, k, None)}, expected {v}") if tuple(ip.image_mean) != (0.5, 0.5, 0.5) or tuple(ip.image_std) != (0.5, 0.5, 0.5): raise RuntimeError(f"Qwen2VLImageProcessor mean/std {ip.image_mean}/{ip.image_std} != 0.5") return ip @dataclass(frozen=True) class Qwen2VLPatchParams: """The ``Qwen2VLImageProcessor`` parameters :func:`qwen2vl_pixel_patches` replays (read once from the cached processor object): ``smart_resize`` factor / pixel budget, patch geometry and the fused ``normalize`` mean / std (``image_mean * (1 / rescale_factor)``, transformers ``_fuse_mean_std_and_rescale_factor``).""" patch_size: int merge_size: int temporal_patch_size: int min_pixels: int max_pixels: int mean: torch.Tensor std: torch.Tensor @lru_cache(maxsize=None) def qwen2vl_patch_params() -> Qwen2VLPatchParams: ip = get_qwen2vl_image_processor() size = ip.size lo, hi = size["shortest_edge"], size["longest_edge"] if not lo or not hi: raise RuntimeError(f"Qwen2VLImageProcessor.size {size} lacks shortest_edge / longest_edge") from PIL import Image as _Image if int(ip.resample) != int(_Image.Resampling.BICUBIC): raise RuntimeError(f"Qwen2VLImageProcessor.resample {ip.resample} != BICUBIC (the verified recipe)") inv = 1.0 / float(ip.rescale_factor) return Qwen2VLPatchParams( patch_size=int(ip.patch_size), merge_size=int(ip.merge_size), temporal_patch_size=int(ip.temporal_patch_size), min_pixels=int(lo), max_pixels=int(hi), mean=torch.tensor(ip.image_mean) * inv, std=torch.tensor(ip.image_std) * inv, ) def qwen2vl_pixel_patches(images_uint8: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """``(B, 3, H, W)`` uint8 frames of one shape -> ``(pixel_values (B * gh * gw, C * T * p * p) fp32, image_grid_thw (B, 3) int64)`` exactly as ``Qwen2VLImageProcessor._preprocess`` computes them for a same-shape batch: ``tvF.resize`` (BICUBIC, ``antialias=True``) of the stacked uint8 tensor to ``smart_resize(H, W)``, fused ``tvF.normalize`` on the fp32 cast, the ``(c, t, ph, pw)`` merge-window patch reshape with the temporal duplication (module docstring).""" from torchvision.transforms import InterpolationMode from torchvision.transforms.v2 import functional as tvF if not isinstance(images_uint8, torch.Tensor) or images_uint8.dtype != torch.uint8 or images_uint8.ndim != 4: raise TypeError(f"expected a (B, 3, H, W) uint8 tensor, got {getattr(images_uint8, 'shape', None)}") b, c, h, w = (int(d) for d in images_uint8.shape) if c != 3: raise ValueError(f"expected 3 channels, got {c}") p = qwen2vl_patch_params() ps, m, tps = p.patch_size, p.merge_size, p.temporal_patch_size rh, rw = qwen2vl_smart_resize(h, w, ps * m, p.min_pixels, p.max_pixels) x = tvF.resize(images_uint8, [rh, rw], interpolation=InterpolationMode.BICUBIC, antialias=True) x = tvF.normalize(x.to(torch.float32), p.mean, p.std) gh, gw = rh // ps, rw // ps x = x.reshape(b, c, gh // m, m, ps, gw // m, m, ps).permute(0, 2, 5, 3, 6, 1, 4, 7) pv = x.unsqueeze(6).expand(-1, -1, -1, -1, -1, -1, tps, -1, -1).reshape(b * gh * gw, c * tps * ps * ps) grid = torch.tensor([[1, gh, gw]] * b, dtype=torch.int64) return pv, grid def qwen2vl_pixel_patches_transformers(images_uint8: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """The transformers processor itself on the same frames (the oracle :func:`qwen2vl_pixel_patches` is tested against; slower by its per-call validation / grouping).""" enc = get_qwen2vl_image_processor()(images=[im for im in images_uint8], return_tensors="pt") return enc["pixel_values"], enc["image_grid_thw"].to(torch.int64) def n17_encode_images(frames_per_camera: Sequence[np.ndarray], spec: EmbodimentSpec) -> VisionInputs: """N1.7: albumentations eval chain per frame (no letterbox for DROID), frame-major stack ``(T*V, 3, 256, 455)``, then the ``Qwen2VLImageProcessor`` recipe (:func:`qwen2vl_pixel_patches`) -> ``pixel_values (T*V*P, 1536)`` fp32 + ``image_grid_thw``; ``t*h*w // 4`` tokens per image (host-pipeline-parity §4.2 N1.7 1-2). A hand-written bicubic resize is NOT bit-exact (§4.2 last row); the torchvision calls of the processor are.""" per_cam = [] for f in frames_per_camera: per_cam.append( torch.stack( [ torch.from_numpy( np.ascontiguousarray( albumentations_eval_transform( fr, spec.shortest_image_edge, spec.crop_fraction, spec.letter_box_transform ) ) ).permute(2, 0, 1) for fr in f ] ) ) imgs = _stack_frame_major(per_cam).contiguous() pv, grid = qwen2vl_pixel_patches(imgs) if pv.dtype != torch.float32: raise RuntimeError(f"pixel patches are {pv.dtype}, expected fp32") if grid.shape != (imgs.shape[0], 3): raise RuntimeError(f"image_grid_thw {tuple(grid.shape)} for {imgs.shape[0]} images") if len({tuple(r) for r in grid.tolist()}) != 1: raise NotImplementedError(f"mixed ViT grids {grid.tolist()} are not supported by the static layout (§3.2)") n_patches = int(grid.prod(dim=1).sum()) if tuple(pv.shape) != (n_patches, 3 * 2 * 16 * 16): raise RuntimeError(f"pixel_values {tuple(pv.shape)} != ({n_patches}, 1536)") tokens = tuple(int(t * h * w) // 4 for t, h, w in grid.tolist()) return VisionInputs(images_uint8=imgs, pixel_values=pv.contiguous(), tokens_per_image=tokens, image_grid_thw=grid) # -------------------------------------------------------------------------------------------------------------------- # ModelInputs + encode # -------------------------------------------------------------------------------------------------------------------- @dataclass class ModelInputs: """Every tensor the device model needs for one observation, plus the static-shape bookkeeping. Field contract (see ``docs/research/foundation-f2-host.md`` §2 for the per-version table): * identity: ``version``, ``embodiment_tag``, ``embodiment_id (1,) int64``. * prompt (unpadded, == golden ``inputs/``): ``prompt`` text, ``input_ids (1, S)``, ``attention_mask (1, S)`` int64, ``seq_len`` S, ``n_text`` (instruction BPE tokens incl. N1.5's brackets), ``image_slots`` runs, ``text_span``. * vision: ``images_uint8 (n_img, 3, H, W)``, ``pixel_values`` (N1.5 ``(n_img,3,224,224)``, N1.6 ``(n_img,3,252,252)``, N1.7 ``(n_img*P, 1536)`` patch rows), ``image_sizes`` (N1.5/N1.6) or ``image_grid_thw`` (N1.7), ``n_images``, ``tokens_per_image``, ``images_pre_resize_uint8`` (N1.6 only). * state: ``state (1, 1, D_max)`` fp32, ``state_mask (1, 1, D_max)`` bool, ``state_dim`` real dims. * static layout (static-shape-plan §4): ``shape_key``, ``slots`` (``input_ids_padded (1, S_pad)``, ``key_valid``, ``slot_of_token``, padded ``image_slots``, ``pad_slots``, 1-D ``position_ids``), ``position_ids`` (``(1, S_pad)`` for N1.5 / N1.6, ``(3, 1, S_pad)`` mRoPE for N1.7), ``rope_cos`` / ``rope_sin (1, S_pad, 128)`` fp32. * N1.7 vision host tables (None otherwise): ``vision_pos_embed_indices (4, n_img*P)`` int64, ``vision_pos_embed_weights (4, n_img*P)`` fp32, ``vision_rope_cos`` / ``vision_rope_sin (n_img*P, 64)`` fp32, ``vision_cu_seqlens``. """ version: str embodiment_tag: str embodiment_id: torch.Tensor # prompt prompt: str input_ids: torch.Tensor attention_mask: torch.Tensor seq_len: int n_text: int image_slots: Tuple[Tuple[int, int], ...] text_span: Tuple[int, int] # vision n_images: int tokens_per_image: Tuple[int, ...] images_uint8: torch.Tensor pixel_values: torch.Tensor image_sizes: Optional[torch.Tensor] image_grid_thw: Optional[torch.Tensor] images_pre_resize_uint8: Optional[torch.Tensor] # state state: torch.Tensor state_mask: torch.Tensor state_dim: int # static layout shape_key: ShapeKey slots: SlotLayout position_ids: torch.Tensor rope_cos: torch.Tensor rope_sin: torch.Tensor # N1.7 vision tables vision_pos_embed_indices: Optional[torch.Tensor] = None vision_pos_embed_weights: Optional[torch.Tensor] = None vision_rope_cos: Optional[torch.Tensor] = None vision_rope_sin: Optional[torch.Tensor] = None vision_cu_seqlens: Optional[Tuple[int, ...]] = None extras: Dict[str, Any] = field(default_factory=dict) @property def s_pad(self) -> int: return self.shape_key.s_pad @property def input_ids_padded(self) -> torch.Tensor: return self.slots.input_ids_padded @property def key_valid(self) -> torch.Tensor: return self.slots.key_valid def golden_inputs(self) -> Dict[str, torch.Tensor]: """The unpadded tensors under the file names of the golden ``inputs/`` (N1.5 ``eagle_*``, N1.6 ``tensors/`` incl. the list element ``pixel_values.{k}``, N1.7 ``inputs/``).""" if self.version == "n15": return { "eagle_input_ids": self.input_ids, "eagle_attention_mask": self.attention_mask, "eagle_pixel_values": self.pixel_values, "eagle_image_sizes": self.image_sizes, "state": self.state, "state_mask": self.state_mask, "embodiment_id": self.embodiment_id, } if self.version == "n16": out = { "input_ids": self.input_ids, "attention_mask": self.attention_mask, "image_sizes": self.image_sizes, "state": self.state, "embodiment_id": self.embodiment_id, } for k in range(self.n_images): out[f"pixel_values.{k}"] = self.pixel_values[k : k + 1] return out return { "input_ids": self.input_ids, "attention_mask": self.attention_mask, "pixel_values": self.pixel_values, "image_grid_thw": self.image_grid_thw, "state": self.state, "embodiment_id": self.embodiment_id, } def _frames_for_camera(spec: EmbodimentSpec, obs: Observation, camera: str) -> np.ndarray: if camera not in obs.images: raise KeyError( f"{spec.version}/{spec.tag}: camera {camera!r} missing from the observation (have {sorted(obs.images)})" ) f = obs.images[camera] if f.ndim != 4 or f.shape[-1] != 3 or f.dtype != np.uint8: raise ValueError(f"{camera}: frames must be (T, H, W, 3) uint8, got {f.shape} {f.dtype}") if f.shape[0] != spec.frames_per_camera: raise ValueError( f"{spec.version}/{spec.tag}: camera {camera!r} must provide {spec.frames_per_camera} frame(s) " f"(video delta_indices {spec.video_delta_indices}), got {f.shape[0]}" ) return f def encode( version: str, obs: Observation, embodiment: Optional[str] = None, layout_name: Optional[str] = None ) -> ModelInputs: """Raw :class:`Observation` -> :class:`ModelInputs` for ``version`` (``embodiment`` defaults to ``obs.embodiment_tag``; ``layout_name`` forces a static layout, otherwise :func:`prompts.select_layout`). Steps: per-camera frames -> version image encoder (:func:`n15_encode_images` / :func:`n16_encode_images` / :func:`n17_encode_images`) -> prompt build + tokenize + structural check (:func:`prompts.encode_prompt`) -> state normalisation (:func:`normalization.encode_state`) -> static slot layout, position ids and RoPE tables (:mod:`prompts`) -> N1.7 vision host tables. Raises on any unexpected shape / key / drift; never truncates. """ _check_version(version) obs.validate() cfg = get_config(version) tag = embodiment if embodiment is not None else obs.embodiment_tag spec = get_embodiment_spec(version, tag) frames = [_frames_for_camera(spec, obs, cam) for cam in spec.video_keys] if version == "n15": vis = n15_encode_images(frames, spec) elif version == "n16": vis = n16_encode_images(frames, spec) else: vis = n17_encode_images(frames, spec) if vis.n_images != spec.n_images: raise AssertionError(f"{vis.n_images} images encoded for a {spec.n_images}-image embodiment") tp: TokenizedPrompt = encode_prompt(version, obs.instruction, vis.tokens_per_image) state, state_mask = encode_state(spec, obs.state) if tuple(state.shape) != (1, len(spec.state_delta_indices), spec.max_state_dim): raise AssertionError(f"state shape {tuple(state.shape)}") layout = select_layout(cfg, vis.n_images, vis.grid_thw, vis.tokens_per_image[0], tp.n_text, layout_name) slots = build_slot_layout(cfg, tp, layout) head_dim = cfg.llm.head_dim inv_dtype = LLM_ROPE_INV_FREQ_DTYPE[version] vision_tables: Dict[str, Any] = {} if version == "n17": gkey = grid_key(vis.image_grid_thw) pos3, cos, sin = mrope_tables_cached( tuple(slots.input_ids_padded[0].tolist()), cfg.prompt.image_token_id, gkey, cfg.vision.merge_size, head_dim, cfg.llm.rope_theta, tuple(int(x) for x in cfg.llm.mrope_section), inv_dtype, ) position_ids = pos3[:, None, :] idx, wts, vcos, vsin, cu = vision_tables_cached( gkey, cfg.vision.pos_grid, cfg.vision.merge_size, cfg.vision.head_dim, cfg.vision.rope_theta ) vision_tables = dict( vision_pos_embed_indices=idx, vision_pos_embed_weights=wts, vision_rope_cos=vcos, vision_rope_sin=vsin, vision_cu_seqlens=cu, ) else: cos, sin = rope_cos_sin_1d(slots.position_ids[0], head_dim, cfg.llm.rope_theta, inv_dtype) position_ids = slots.position_ids return ModelInputs( version=version, embodiment_tag=tag, embodiment_id=torch.tensor([spec.embodiment_id], dtype=torch.int64), prompt=tp.text, input_ids=tp.input_ids, attention_mask=tp.attention_mask, seq_len=tp.seq_len, n_text=tp.n_text, image_slots=tp.image_slots, text_span=tp.text_span, n_images=vis.n_images, tokens_per_image=vis.tokens_per_image, images_uint8=vis.images_uint8, pixel_values=vis.pixel_values, image_sizes=vis.image_sizes, image_grid_thw=vis.image_grid_thw, images_pre_resize_uint8=vis.images_pre_resize_uint8, state=state, state_mask=state_mask, state_dim=spec.state_dim, shape_key=slots.shape_key, slots=slots, position_ids=position_ids, rope_cos=cos[None], rope_sin=sin[None], **vision_tables, ) __all__ = [ "ModelInputs", "Observation", "Qwen2VLPatchParams", "VisionInputs", "albumentations_eval_transform", "eagle3_smart_resize", "encode", "fractional_center_crop", "get_qwen2vl_image_processor", "letterbox_pad", "load_golden_observation", "n15_crop_resize", "n15_encode_images", "n16_encode_images", "n17_encode_images", "normalize_u8", "qwen2vl_patch_params", "qwen2vl_pixel_patches", "qwen2vl_pixel_patches_transformers", "qwen2vl_smart_resize", "smallest_max_size", ]