[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 <liangmengh@nvidia.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
2026-09-08 18:59:01 +08:00
1 parent c9c25b33c3
commit d7740336ff
@@ -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 = "<frame>|letterbox=<flag>".
+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 "<prompt>"
+
+
+@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)
@@ -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,