From d7740336ff94d23276d690e3cd74a1ddb25d23cb Mon Sep 17 00:00:00 2001 From: liangmenghuang Date: Tue, 8 Sep 2026 18:59:01 +0800 Subject: [PATCH] [Bugfix][Model] GR00T-N1.7: honor processor_config image pipeline (use_albumentations / crop_fraction / letter_box_transform) to match Isaac-GR00T (#7083) Signed-off-by: liangmengh Co-authored-by: Claude Fable 5.1 --- .../test_processing_eval_image_transform.py | 292 ++++++++++++++++++ .../gr00t/modeling/processing_gr00t_n1d7.py | 125 +++++++- 2 files changed, 410 insertions(+), 7 deletions(-) create mode 100644 tests/diffusion/models/gr00t/test_processing_eval_image_transform.py diff --git a/tests/diffusion/models/gr00t/test_processing_eval_image_transform.py b/tests/diffusion/models/gr00t/test_processing_eval_image_transform.py new file mode 100644 index 000000000..452273185 --- /dev/null +++ b/tests/diffusion/models/gr00t/test_processing_eval_image_transform.py @@ -0,0 +1,292 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""GR00T-N1.7 eval image transform: chosen from ``processor_config.json`` and bit-exact with Isaac-GR00T. + +``GOLDEN`` holds sha256 digests of the ``(C, H, W)`` uint8 output of Isaac-GR00T's own eval transform +(``build_image_transformations_albumentations(...)[1]`` applied through ``apply_with_replay``) for the +synthetic frames built below, at Isaac-GR00T 51d4c89 / albumentations 1.4.18 / opencv 5.0.0.93. The frame +generators are integer-only so they are identical on every platform. + +Regenerating a digest needs Isaac-GR00T installed, whose transformers pin conflicts with this repo's, so it +cannot run here: in an Isaac-GR00T environment push each ``synthetic_frames()`` frame through that transform +with ``letter_box_transform`` False and True and hash the output with ``digest()``. +""" + +from __future__ import annotations + +import hashlib +import json + +import numpy as np +import pytest +import torch +import torchvision.transforms.v2 as transforms + +from vllm_omni.diffusion.models.gr00t.modeling import processing_gr00t_n1d7 as processing +from vllm_omni.diffusion.models.gr00t.modeling.processing_gr00t_n1d7 import ( + AlbumentationsEvalImageTransform, + Gr00tN1d7Processor, + LetterBoxTransform, + _build_eval_image_transform, + _fractional_center_crop, + _smallest_edge_resize, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# nvidia/GR00T-N1.7-3B and nvidia/GR00T-N1.7-DROID both ship exactly these processor_kwargs. +TARGET, CROP, EDGE, FRACTION = [256, 256], [230, 230], 256, 0.95 +MODALITY_CONFIGS = { + "new_embodiment": { + "video": {"delta_indices": [0], "modality_keys": ["cam"]}, + "state": {"delta_indices": [0], "modality_keys": ["joint"]}, + "action": {"delta_indices": [0, 1], "modality_keys": ["joint"]}, + } +} + +# Isaac-GR00T eval transform output for synthetic_frames(); key = "|letterbox=". +GOLDEN: dict[str, tuple[tuple[int, int, int], str]] = { + "gradient_256x256|letterbox=False": ( + (3, 256, 256), + "4e90a1a3c49539b3648aad752bba59af452bd7231e5712e871ab941cbea1d715", + ), + "checker_256x256|letterbox=False": ( + (3, 256, 256), + "d2225b94efea5ab3f1c3e6dbae52b7afa3e0f7172d5de24439a2e729ab56080c", + ), + "noise_256x256|letterbox=False": ( + (3, 256, 256), + "a427ae3b84ebcc6e87470c79c0e45cf684770d6e31158a49f8736685e0da72c6", + ), + "noise_512x768|letterbox=False": ( + (3, 256, 383), + "45bc3a309f3d06c3f31051911bc8a544ff3f067b023b999dd9c2b7e6029ecc93", + ), + "gradient_256x256|letterbox=True": ( + (3, 256, 256), + "4e90a1a3c49539b3648aad752bba59af452bd7231e5712e871ab941cbea1d715", + ), + "checker_256x256|letterbox=True": ( + (3, 256, 256), + "d2225b94efea5ab3f1c3e6dbae52b7afa3e0f7172d5de24439a2e729ab56080c", + ), + "noise_256x256|letterbox=True": ((3, 256, 256), "a427ae3b84ebcc6e87470c79c0e45cf684770d6e31158a49f8736685e0da72c6"), + "noise_512x768|letterbox=True": ((3, 256, 256), "52a3db185e6ee4f00617b87711dd198b8f2e2ad53865713d34e56225a4ec61a2"), +} + + +def _hash_noise(height: int, width: int, seed: int) -> np.ndarray: + """Integer-only pseudo-noise (murmur-style mixing on uint32); no RNG or float, so identical everywhere.""" + ys, xs = np.mgrid[0:height, 0:width].astype(np.uint32) + channels = [] + for channel in range(3): + v = xs * np.uint32(73856093) ^ ys * np.uint32(19349663) ^ np.uint32((seed * 3 + channel) * 83492791) + v ^= v >> np.uint32(13) + v *= np.uint32(0x5BD1E995) + v ^= v >> np.uint32(15) + channels.append((v & np.uint32(0xFF)).astype(np.uint8)) + return np.stack(channels, axis=-1) + + +def synthetic_frames() -> dict[str, np.ndarray]: + """Deterministic HWC uint8 RGB frames: smooth, hard-edged, noisy, and a non-square noisy one. + + The non-square frame is 512x768 so both INTER_AREA downscales (2x without letterbox, 3x with it) + are integer factors, which OpenCV computes with exact integer sums: the goldens do not depend on + the CPU architecture. The 243->256 upscale is a fixed-point bilinear on uint8, also exact. + """ + ys, xs = np.mgrid[0:256, 0:256].astype(np.int64) + frames = { + "gradient_256x256": np.stack([xs % 256, ys % 256, (xs + ys) // 2 % 256], axis=-1).astype(np.uint8), + "checker_256x256": np.stack( + [((ys // 16 + xs // 16) % 2) * 255, ((ys // 8 + xs // 8) % 2) * 255, ((ys // 32) % 2) * 255], + axis=-1, + ).astype(np.uint8), + "noise_256x256": _hash_noise(256, 256, seed=1), + "noise_512x768": _hash_noise(512, 768, seed=2), + } + return frames + + +def digest(tensor: torch.Tensor) -> str: + return hashlib.sha256(np.ascontiguousarray(tensor.numpy()).tobytes()).hexdigest() + + +class _FakeTokenizer: + padding_side = "right" + + +class _FakeVlmProcessor: + """The two attributes ``Gr00tN1d7Processor`` and its collator touch on the Qwen3-VL processor.""" + + def __init__(self) -> None: + self.tokenizer = _FakeTokenizer() + + def apply_chat_template(self, *_args: object, **_kwargs: object) -> str: + return "" + + +@pytest.fixture(autouse=True) +def fake_vlm_processor(monkeypatch: pytest.MonkeyPatch) -> None: + """The real Qwen3-VL processor is a gated download and irrelevant to the image pipeline.""" + monkeypatch.setattr(processing, "build_processor", lambda _model_name, _kwargs: _FakeVlmProcessor()) + + +def _processor(**kwargs) -> Gr00tN1d7Processor: + return Gr00tN1d7Processor( + modality_configs=MODALITY_CONFIGS, image_crop_size=CROP, image_target_size=TARGET, **kwargs + ) + + +def _has_letterbox(compose: transforms.Compose) -> bool: + return any(isinstance(step, LetterBoxTransform) for step in compose.transforms) + + +def test_processor_builds_the_pipeline_named_by_processor_config(): + albumentations = _processor(use_albumentations=True, shortest_image_edge=EDGE, crop_fraction=FRACTION) + transform = albumentations.eval_image_transform + assert isinstance(transform, AlbumentationsEvalImageTransform) + assert (transform.max_size, transform.fraction, transform.letter_box_transform) == (EDGE, FRACTION, False) + assert albumentations.use_albumentations is True + + torchvision = _processor(use_albumentations=False) + assert isinstance(torchvision.eval_image_transform, transforms.Compose) + assert not _has_letterbox(torchvision.eval_image_transform) + + letterboxed = _processor(use_albumentations=False, letter_box_transform=True) + assert _has_letterbox(letterboxed.eval_image_transform) + assert _processor(use_albumentations=True, letter_box_transform=True).eval_image_transform.letter_box_transform + + +def test_legacy_config_without_the_flags_keeps_the_torchvision_recipe(): + """Checkpoints predating the flags behave like Isaac-GR00T's defaults; training-only keys stay ignored.""" + processor = _processor(state_dropout_prob=0.2, extra_augmentation_config=None) + assert processor.use_albumentations is False + assert processor.letter_box_transform is False + assert isinstance(processor.eval_image_transform, transforms.Compose) + assert not _has_letterbox(processor.eval_image_transform) + + +def test_fraction_and_max_size_fall_back_like_isaac_gr00t(): + transform = _build_eval_image_transform( + TARGET, CROP, shortest_image_edge=None, crop_fraction=None, use_albumentations=True + ) + assert isinstance(transform, AlbumentationsEvalImageTransform) + assert transform.fraction == CROP[0] / TARGET[0] + assert transform.max_size == TARGET[0] + with pytest.raises(ValueError, match="crop_fraction is None"): + _build_eval_image_transform(TARGET, None, crop_fraction=None, use_albumentations=True) + + +def test_geometry_follows_isaac_rounding_rules(): + frame = _hash_noise(540, 640, seed=3) + resized = _smallest_edge_resize(frame, EDGE) + assert resized.shape == (256, 303, 3) # round() per dimension, shortest edge -> 256 + cropped = _fractional_center_crop(resized, FRACTION) + assert cropped.shape == (243, 287, 3) # int() truncation: 243.2 -> 243, 287.85 -> 287 + assert np.array_equal(cropped, resized[6:249, 8:295]) # (256-243)//2, (303-287)//2 + square = frame[:256, :256] + assert _smallest_edge_resize(square, EDGE) is square # scale == 1.0 is a no-op, not a resample + + +@pytest.mark.parametrize("case", sorted(GOLDEN)) +def test_albumentations_recipe_is_bit_exact_with_isaac_gr00t(case: str): + name, flag = case.split("|letterbox=") + transform = _build_eval_image_transform( + TARGET, + CROP, + shortest_image_edge=EDGE, + crop_fraction=FRACTION, + use_albumentations=True, + letter_box_transform=(flag == "True"), + ) + out = transform(synthetic_frames()[name]) + expected_shape, expected_digest = GOLDEN[case] + assert out.dtype == torch.uint8 + assert tuple(out.shape) == expected_shape + assert digest(out) == expected_digest + + +def test_letterbox_is_gated_on_both_recipes(): + frame = _hash_noise(540, 640, seed=3) + + def albumentations(letterbox: bool) -> torch.Tensor: + return _build_eval_image_transform( + TARGET, + CROP, + shortest_image_edge=EDGE, + crop_fraction=FRACTION, + use_albumentations=True, + letter_box_transform=letterbox, + )(frame) + + padded, unpadded = albumentations(True), albumentations(False) + assert tuple(padded.shape) == (3, 256, 256) # 540x640 -> 640x640 -> 256x256 -> 243x243 -> 256x256 + assert tuple(unpadded.shape) == (3, 256, 302) # aspect preserved: 256x303 -> 243x287 -> 256x302 + + tv_padded = _build_eval_image_transform(TARGET, CROP, use_albumentations=False, letter_box_transform=True)(frame) + tv_unpadded = _build_eval_image_transform(TARGET, CROP, use_albumentations=False, letter_box_transform=False)(frame) + assert tuple(tv_padded.shape) == tuple(tv_unpadded.shape) == (3, 256, 256) + assert not torch.equal(tv_padded, tv_unpadded) + assert (tv_padded[:, :2, :] == 0).all() and not (tv_unpadded[:, :2, :] == 0).all() # black bars only when gated + + +def test_albumentations_recipe_rejects_frames_outside_the_contract(): + transform = _build_eval_image_transform(TARGET, CROP, use_albumentations=True) + with pytest.raises(ValueError, match="HWC uint8"): + transform(np.zeros((256, 256, 3), dtype=np.float32)) + with pytest.raises(ValueError, match="HWC uint8"): + transform(np.zeros((3, 256, 256), dtype=np.uint8)) + with pytest.raises(ValueError, match="HWC uint8"): + transform(np.zeros((256, 256), dtype=np.uint8)) + + +def test_get_vlm_inputs_runs_the_selected_transform_per_frame(): + processor = _processor(use_albumentations=True, shortest_image_edge=EDGE, crop_fraction=FRACTION) + frames = np.stack([_hash_noise(540, 640, seed=3)] * 2) # (T, H, W, C) like policy.py hands over + vlm_inputs = processor._get_vlm_inputs( + image_keys=["cam"], + images={"cam": frames}, + masks=None, + image_transform=processor.eval_image_transform, + language="x", + ) + images = vlm_inputs["vlm_content"]["images"] + assert len(images) == 2 + assert all(img.size == (302, 256) for img in images) # PIL reports (W, H); no letterbox for this checkpoint + + +def test_flags_round_trip_through_save_pretrained_and_from_pretrained(tmp_path): + saved = _processor( + use_albumentations=True, letter_box_transform=True, shortest_image_edge=EDGE, crop_fraction=FRACTION + ) + saved.save_pretrained(tmp_path) + kwargs = json.loads((tmp_path / "processor_config.json").read_text())["processor_kwargs"] + assert kwargs["use_albumentations"] is True + assert kwargs["letter_box_transform"] is True + + loaded = Gr00tN1d7Processor.from_pretrained(tmp_path) + assert loaded.use_albumentations is True + assert loaded.letter_box_transform is True + assert isinstance(loaded.eval_image_transform, AlbumentationsEvalImageTransform) + assert loaded.eval_image_transform.letter_box_transform is True + + +def test_from_pretrained_defaults_missing_flags_to_isaac_gr00t_defaults(tmp_path): + (tmp_path / "processor_config.json").write_text( + json.dumps( + { + "processor_class": "Gr00tN1d7Processor", + "processor_kwargs": { + "modality_configs": MODALITY_CONFIGS, + "image_crop_size": CROP, + "image_target_size": TARGET, + }, + } + ) + ) + (tmp_path / "statistics.json").write_text("{}") + loaded = Gr00tN1d7Processor.from_pretrained(tmp_path) + assert loaded.use_albumentations is False + assert loaded.letter_box_transform is False + assert isinstance(loaded.eval_image_transform, transforms.Compose) diff --git a/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py b/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py index ade5eac83..d30d1c829 100644 --- a/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py +++ b/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project import json import os import re import warnings +from collections.abc import Callable from copy import deepcopy from pathlib import Path from typing import Any @@ -46,20 +47,119 @@ class LetterBoxTransform: return transforms.functional.pad(img, padding=[pad_left, pad_top, pad_right, pad_bottom], fill=0) +def _smallest_edge_resize(image: np.ndarray, max_size: int) -> np.ndarray: + """albumentations 1.4.18 ``SmallestMaxSize(max_size, interpolation=cv2.INTER_AREA)``. + + Scale off the shortest edge, skip entirely at scale 1.0, round each dimension + independently; ``cv2.resize`` takes ``dsize`` as ``(width, height)``. + """ + import cv2 + + height, width = image.shape[:2] + scale = max_size / float(min(height, width)) + if scale == 1.0: + return image + new_height, new_width = (round(dim * scale) for dim in (height, width)) + if (new_height, new_width) == (height, width): + return image + return cv2.resize(image, dsize=(new_width, new_height), interpolation=cv2.INTER_AREA) + + +def _fractional_center_crop(image: np.ndarray, fraction: float) -> np.ndarray: + """Isaac-GR00T ``FractionalCenterCrop``: ``int()`` truncation (not rounding), floor 1, centred.""" + height, width = image.shape[:2] + crop_height = max(1, int(height * fraction)) + crop_width = max(1, int(width * fraction)) + y_min = (height - crop_height) // 2 + x_min = (width - crop_width) // 2 + return image[y_min : y_min + crop_height, x_min : x_min + crop_width] + + +def _letter_box_pad(image: np.ndarray) -> np.ndarray: + """Isaac-GR00T ``LetterBoxPad``: zero-pad to square, centred, remainder on the far side.""" + import cv2 + + height, width = image.shape[:2] + if height == width: + return image + longest = max(height, width) + pad_h, pad_w = longest - height, longest - width + top, left = pad_h // 2, pad_w // 2 + return cv2.copyMakeBorder(image, top, pad_h - top, left, pad_w - left, cv2.BORDER_CONSTANT, value=0) + + +class AlbumentationsEvalImageTransform: + """Eval transform of the Isaac-GR00T ``use_albumentations=True`` recipe, without albumentations. + + ``[LetterBoxPad] -> SmallestMaxSize(INTER_AREA) -> FractionalCenterCrop -> SmallestMaxSize(INTER_AREA)``, + transcribed from the eval branch of Isaac-GR00T ``build_image_transformations_albumentations`` and + albumentations 1.4.18, so a checkpoint trained with that recipe sees the same pixels at inference. + Per frame: HWC uint8 RGB array in, ``(C, H, W)`` uint8 tensor out (the ``_get_vlm_inputs`` contract). + """ + + def __init__(self, max_size: int, fraction: float, letter_box_transform: bool): + self.max_size = int(max_size) + self.fraction = float(fraction) + self.letter_box_transform = letter_box_transform + + def __call__(self, image: np.ndarray) -> torch.Tensor: + array = np.asarray(image) + if array.dtype != np.uint8 or array.ndim != 3 or array.shape[2] != 3: + raise ValueError( + "GR00T eval image transform expects an HWC uint8 RGB frame, " + f"got shape={tuple(array.shape)} dtype={array.dtype}" + ) + if self.letter_box_transform: + array = _letter_box_pad(array) + array = _smallest_edge_resize(array, self.max_size) + array = _fractional_center_crop(array, self.fraction) + array = _smallest_edge_resize(array, self.max_size) + # The crop is a non-contiguous view; the caller stacks frames right after this. + return torch.from_numpy(np.ascontiguousarray(array)).permute(2, 0, 1) + + def _build_eval_image_transform( image_target_size: list[int], image_crop_size: list[int], -) -> transforms.Compose: - """Deterministic eval/inference image transform (letterbox → resize → centercrop → resize).""" - return transforms.Compose( + *, + shortest_image_edge: int | None = None, + crop_fraction: float | None = None, + use_albumentations: bool = False, + letter_box_transform: bool = False, +) -> Callable[[np.ndarray], torch.Tensor]: + """Deterministic eval/inference image transform, selected the way Isaac-GR00T's processor does. + + ``use_albumentations=True``: ``[letterbox] -> resize (shortest edge) -> fractional centre crop -> resize`` + (:class:`AlbumentationsEvalImageTransform`). Otherwise the torchvision recipe + ``[letterbox] -> resize -> centre crop -> resize``. Letterboxing is gated by ``letter_box_transform`` + on both paths. + """ + if use_albumentations: + if crop_fraction is None: + if image_crop_size is None or image_target_size is None: + raise ValueError("image_crop_size and image_target_size are required when crop_fraction is None") + fraction = image_crop_size[0] / image_target_size[0] + else: + fraction = crop_fraction + if shortest_image_edge is None: + if image_target_size is None: + raise ValueError("image_target_size is required when shortest_image_edge is None") + max_size = image_target_size[0] + else: + max_size = shortest_image_edge + return AlbumentationsEvalImageTransform(max_size, fraction, letter_box_transform) + + steps: list[Callable[..., torch.Tensor]] = [transforms.ToImage()] + if letter_box_transform: + steps.append(LetterBoxTransform()) + steps.extend( [ - transforms.ToImage(), - LetterBoxTransform(), transforms.Resize(size=image_target_size), transforms.CenterCrop(size=image_crop_size), transforms.Resize(size=image_target_size), ] ) + return transforms.Compose(steps) logger = init_logger(__name__) @@ -191,6 +291,9 @@ class Gr00tN1d7Processor(ProcessorMixin): exclude_state: bool = False, # Normalization use_mean_std: bool = False, + # Image pipeline selection; must match the recipe the checkpoint was trained with + use_albumentations: bool = False, + letter_box_transform: bool = False, **kwargs, # absorb deprecated training-only keys from saved processor_config.json ): if transformers_loading_kwargs is None: @@ -237,6 +340,8 @@ class Gr00tN1d7Processor(ProcessorMixin): self.embodiment_id_mapping[k] = v self.shortest_image_edge = shortest_image_edge self.crop_fraction = crop_fraction + self.use_albumentations = use_albumentations + self.letter_box_transform = letter_box_transform self.statistics: dict[str, dict[str, dict[str, dict[str, list[float]]]]] = {} @@ -244,6 +349,10 @@ class Gr00tN1d7Processor(ProcessorMixin): self.eval_image_transform = _build_eval_image_transform( image_target_size, image_crop_size, + shortest_image_edge=shortest_image_edge, + crop_fraction=crop_fraction, + use_albumentations=use_albumentations, + letter_box_transform=letter_box_transform, ) self._collator = self.data_collator_class( model_name=model_name, @@ -428,7 +537,7 @@ class Gr00tN1d7Processor(ProcessorMixin): image_keys: list[str], images: list[Image.Image], masks: dict[str, list[np.ndarray]] | None, - image_transform: transforms.Compose, + image_transform: Callable[[np.ndarray], torch.Tensor], language: str, ): temporal_stacked_images = {} @@ -470,6 +579,8 @@ class Gr00tN1d7Processor(ProcessorMixin): "color_jitter_params": self.color_jitter_params, "shortest_image_edge": self.shortest_image_edge, "crop_fraction": self.crop_fraction, + "use_albumentations": self.use_albumentations, + "letter_box_transform": self.letter_box_transform, "model_name": self.model_name, "model_type": self.model_type, "formalize_language": self.formalize_language,