From 1a0f2d603db6a60d636ac239dd8e9781a0498435 Mon Sep 17 00:00:00 2001 From: Jaime Alfaro Salazar <62313620+jialfaro@users.noreply.github.com> Date: Sat, 5 Sep 2026 01:52:23 -0500 Subject: [PATCH] Gr00t-N1.7: honor processor_config preprocessing flags (use_albumentations, crop_fraction, letter_box_transform) /claim #7081 --- .../gr00t/modeling/processing_gr00t_n1d7.py | 118 ++++++++++++++++-- 1 file changed, 108 insertions(+), 10 deletions(-) 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 ade5eac83b1..1e58c7c0831 100644 --- a/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py +++ b/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py @@ -46,20 +46,108 @@ def __call__(self, img: torch.Tensor) -> torch.Tensor: return transforms.functional.pad(img, padding=[pad_left, pad_top, pad_right, pad_bottom], fill=0) +class FractionalCenterCrop: + """Center-crop an image by a fractional factor of its current dimensions. + + The crop size is computed as int(dim * fraction) for each spatial dim (floor behavior), + matching the behavior used in Isaac-GR00T's FractionalCenterCrop. + """ + + def __init__(self, fraction: float): + assert 0.0 < fraction <= 1.0, "fraction must be in (0, 1]" + self.fraction = fraction + + def __call__(self, img: torch.Tensor) -> torch.Tensor: + import math + + *leading_dims, c, h, w = img.shape + new_h = int(h * self.fraction) + new_w = int(w * self.fraction) + top = (h - new_h) // 2 + left = (w - new_w) // 2 + if leading_dims: + batch_size = math.prod(leading_dims) + img_r = img.reshape(batch_size, c, h, w) + cropped = img_r[:, :, top : top + new_h, left : left + new_w] + return cropped.reshape(leading_dims + [c, new_h, new_w]) + return img[:, top : top + new_h, left : left + new_w] + + +class SmallestMaxSize: + """Resize while preserving aspect ratio so that the smallest side equals `size`. + + Uses the provided interpolation mode. This mirrors albumentations.SmallestMaxSize. + """ + + def __init__(self, size: int, interpolation=transforms.InterpolationMode.BOX): + assert size > 0 + self.size = int(size) + self.interpolation = interpolation + + def __call__(self, img: torch.Tensor) -> torch.Tensor: + *leading_dims, c, h, w = img.shape + if h == 0 or w == 0: + return img + if min(h, w) == self.size: + return img + scale = self.size / float(min(h, w)) + new_h = int(round(h * scale)) + new_w = int(round(w * scale)) + # Use torchvision Resize with explicit (H, W) + resize = transforms.Resize(size=(new_h, new_w), interpolation=self.interpolation) + if leading_dims: + batch_size = int(torch.tensor(leading_dims).prod().item()) if leading_dims else 1 + img_r = img.reshape(batch_size, c, h, w) + resized = resize(img_r) + return resized.reshape(leading_dims + [c, new_h, new_w]) + return resize(img) + + def _build_eval_image_transform( image_target_size: list[int], image_crop_size: list[int], + use_albumentations: bool = False, + letter_box_transform: bool = True, + crop_fraction: float = 0.95, + shortest_image_edge: int = 256, ) -> transforms.Compose: - """Deterministic eval/inference image transform (letterbox → resize → centercrop → resize).""" - return transforms.Compose( - [ - transforms.ToImage(), - LetterBoxTransform(), - transforms.Resize(size=image_target_size), - transforms.CenterCrop(size=image_crop_size), - transforms.Resize(size=image_target_size), - ] - ) + """Deterministic eval/inference image transform. + + Two modes are supported to match upstream Isaac-GR00T checkpoints: + - use_albumentations=True: emulates the albumentations recipe used in Isaac-GR00T + (optional letterbox -> SmallestMaxSize(shortest_image_edge, INTER_AREA) -> + FractionalCenterCrop(crop_fraction) -> Resize(image_target_size, INTER_AREA)). + - use_albumentations=False: preserves the original vLLM-Omni pipeline + (optional letterbox -> Resize(image_target_size) -> CenterCrop(image_crop_size) + -> Resize(image_target_size)). + + The function intentionally avoids importing albumentations to keep the repo + dependency-free and reproduces the deterministic behaviour using torchvision + transforms and small helpers above. + """ + steps = [transforms.ToImage()] + + if use_albumentations: + # Only apply letterbox when explicitly requested in the processor config. + if letter_box_transform: + steps.append(LetterBoxTransform()) + if shortest_image_edge is not None: + # Use BOX (area-like) interpolation to match cv2.INTER_AREA used by Isaac. + steps.append(SmallestMaxSize(shortest_image_edge, interpolation=transforms.InterpolationMode.BOX)) + if crop_fraction is not None: + steps.append(FractionalCenterCrop(crop_fraction)) + # Resize back to target size using BOX/area-like interpolation. + steps.append(transforms.Resize(size=image_target_size, interpolation=transforms.InterpolationMode.BOX)) + else: + # Preserve historical vLLM-Omni behavior for callers that don't opt into + # the albumentations-style pipeline. + if letter_box_transform: + steps.append(LetterBoxTransform()) + steps.append(transforms.Resize(size=image_target_size)) + steps.append(transforms.CenterCrop(size=image_crop_size)) + steps.append(transforms.Resize(size=image_target_size)) + + return transforms.Compose(steps) logger = init_logger(__name__) @@ -176,6 +264,8 @@ def __init__( image_target_size: list[int] = None, shortest_image_edge: int = 256, crop_fraction: float = 0.95, + use_albumentations: bool = False, + letter_box_transform: bool = True, random_rotation_angle: int | None = None, color_jitter_params: dict[str, float] | None = None, formalize_language: bool = True, @@ -237,6 +327,8 @@ def __init__( 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 +336,10 @@ def __init__( self.eval_image_transform = _build_eval_image_transform( image_target_size, image_crop_size, + use_albumentations=self.use_albumentations, + letter_box_transform=self.letter_box_transform, + crop_fraction=self.crop_fraction, + shortest_image_edge=self.shortest_image_edge, ) self._collator = self.data_collator_class( model_name=model_name, @@ -470,6 +566,8 @@ def save_pretrained(self, save_directory: str | Path) -> list[Path]: "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,