Skip to navigation
ReferenceFull Library ReferenceNemo AutomodelNemo AutomodelComponentsModelsMimo V2 Flashnemo_automodel.components.models.mimo_v2_flash.processing

nemo_automodel.components.models.mimo_v2_flash.processing

View as Markdown

MiMo-V2 image preprocessing aligned with the official serving/training stack.

The stock Qwen2VLImageProcessor that AutoProcessor resolves for MiMo-V2 checkpoints resizes with PIL/torchvision bicubic (antialiased, uint8-round-trip) and normalizes in 0-1 space. That is not what the model sees in production: Xiaomi’s own processor in the reference serving stack (sglang MiMoProcessor, contributed as MiMo day-0 support and used for the RL rollouts of MiMo-V2.6-*-RL) resizes with torch bilinear on 0-255 floats, normalizes in 0-255 space, and flattens RGBA onto white instead of dropping alpha. Grid geometry (smart_resize targets, patch layout, token counts) is identical for images whose short side is at least the patch factor; below that the reference upscales the short side to the factor before rounding while HF clamps each side independently, so a few tiny inputs (e.g. 16x192) get different grids. For every other image only pixel values differ, so the mismatch is a silent distribution shift of up to ~0.2 std per pixel.

Normalization statistics are read from the checkpoint’s own preprocessor_config.json (via the stock image processor built by AutoProcessor), never hard-coded: MiMo-V2.6 checkpoints declare ImageNet values (Xiaomi commit “Align preprocessor config to training values”), while MiMo-V2.5 declares CLIP values, and each is what that family trained with.

Module Contents

Classes

NameDescription
MimoV2ImageProcessorDrop-in image processor for MiMo-V2 checkpoints matching sglang pixels.
MimoV2ImageProcessorKwargsPer-call keyword arguments understood by MimoV2ImageProcessor.

Functions

NameDescription
_register_with_auto_image_processorRegister the class so a saved processor directory reloads through HF auto classes.
_to_rgbConvert to RGB with the sglang semantics: RGBA is flattened onto white.
build_mimo_v2_processorBuild the MiMo-V2 chat processor with sglang-aligned image preprocessing.
smart_resizeCompute resize targets with the sglang MiMoProcessor semantics.

Data

ImageInput

logger

API

class nemo_automodel.components.models.mimo_v2_flash.processing.MimoV2ImageProcessor(
    patch_size: int = 16,
    temporal_patch_size: int = 2,
    merge_size: int = 2,
    min_pixels: int = 8192,
    max_pixels: int = 8388608,
    image_mean: tuple[float, float, float] = (0.485, 0.456, 0.406),
    image_std: tuple[float, float, float] = (0.229, 0.224, 0.225),
    kwargs = {}
)

Bases: BaseImageProcessor

Drop-in image processor for MiMo-V2 checkpoints matching sglang pixels.

Implements the pixel path of the reference MiMoProcessor: load as RGB (RGBA flattened onto white), cast to 0-255 float, resize with F.interpolate(mode="bilinear", align_corners=False), standardize with mean/std expressed in 0-255 scale, duplicate the frame along the temporal axis, and emit patches in the Qwen2VL merge-block order. Geometry (smart_resize targets and patch order) matches Qwen2VLImageProcessor except for the sub-factor short-side branch described in the module docstring; only resampling and normalization arithmetic differ otherwise.

Inheriting BaseImageProcessor provides the save/reload contract (save_pretrained/from_pretrained) that ProcessorMixin relies on when a checkpoint is written, and is required for registration with AutoImageProcessor (see _register_with_auto_image_processor) so a saved processor directory reloads through AutoProcessor.from_pretrained. Only JSON-serializable constructor arguments may be stored as instance attributes so to_dict() stays serializable.

model_input_names
= ['pixel_values', 'image_grid_thw']
size
dict[str, int]

Pixel bounds in the HF size mapping shape used by token estimators.

nemo_automodel.components.models.mimo_v2_flash.processing.MimoV2ImageProcessor.__call__(
    images: nemo_automodel.components.models.mimo_v2_flash.processing.ImageInput | list[nemo_automodel.components.models.mimo_v2_flash.processing.ImageInput] | list[list[nemo_automodel.components.models.mimo_v2_flash.processing.ImageInput]],
    min_pixels: int | None = None,
    max_pixels: int | None = None,
    return_tensors: str | None = 'pt',
    kwargs = {}
) -> transformers.BatchFeature

Preprocess images into MiMo-V2 vision-tower inputs.

Parameters:

images
ImageInput | list[ImageInput] | list[list[ImageInput]]

One PIL image / local path, a flat list of them, or a per-sample nested list ([[img1], [img2], ...]) as produced by ProcessorMixin.apply_chat_template on batched conversations. Nested input is flattened one level, preserving order.

min_pixels
int | NoneDefaults to None

Per-call override of the instance pixel bound.

max_pixels
int | NoneDefaults to None

Per-call override of the instance pixel bound.

return_tensors
str | NoneDefaults to 'pt'

Kept for Qwen2VLImageProcessor interface compatibility; outputs are always torch tensors.

**kwargs
Defaults to {}

Ignored, for the same interface compatibility.

Returns: BatchFeature

BatchFeature with pixel_values of shape [total_patches,

nemo_automodel.components.models.mimo_v2_flash.processing.MimoV2ImageProcessor._load(
    image: nemo_automodel.components.models.mimo_v2_flash.processing.ImageInput
) -> torch.Tensor

Load one image as a float32 tensor of shape [channels, height, width], 0-255 scale.

nemo_automodel.components.models.mimo_v2_flash.processing.MimoV2ImageProcessor._pixel_bounds(
    images_kwargs: dict | None
) -> tuple[int, int]
nemo_automodel.components.models.mimo_v2_flash.processing.MimoV2ImageProcessor._process_one(
    image: nemo_automodel.components.models.mimo_v2_flash.processing.ImageInput,
    min_pixels: int,
    max_pixels: int
) -> tuple[torch.Tensor, list[int]]

Run one image through the sglang pixel path.

Parameters:

image
ImageInput

PIL image or local path.

Returns: torch.Tensor

Tuple of patches of shape [grid_tgrid_hgrid_w,

nemo_automodel.components.models.mimo_v2_flash.processing.MimoV2ImageProcessor.fetch_images(
    image_url_or_urls: nemo_automodel.components.models.mimo_v2_flash.processing.ImageInput | list
) -> PIL.Image.Image | list

Normalize raw image inputs to PIL images.

Required by ProcessorMixin.prepare_inputs_layout (transformers >= 5.18) before the images reach __call__. Mirrors ImageProcessingMixin.fetch_images for the supported input types; RGB conversion stays in __call__.

nemo_automodel.components.models.mimo_v2_flash.processing.MimoV2ImageProcessor.get_number_of_image_patches(
    height: int,
    width: int,
    images_kwargs: dict | None = None
) -> int

Number of 16x16 patches for one image, before the merge_size**2 reduction.

Called by Qwen2_5_VLProcessor to expand <|image_pad|> placeholders.

Parameters:

height
int

Source image height in pixels.

width
int

Source image width in pixels.

images_kwargs
dict | NoneDefaults to None

Optional per-call min_pixels/max_pixels overrides.

Returns: int

grid_h * grid_w for the resized image.

class nemo_automodel.components.models.mimo_v2_flash.processing.MimoV2ImageProcessorKwargs()

Bases: ImagesKwargs

Per-call keyword arguments understood by MimoV2ImageProcessor.

ProcessorMixin routes a flat kwarg into images_kwargs only when the key is declared in the image processor’s valid_kwargs TypedDict; undeclared keys are dropped with a one-time warning. Declaring the pixel bounds here keeps per-call overrides such as processor(..., max_pixels=...) working, mirroring Qwen2VLImageProcessor.

max_pixels
int
min_pixels
int
nemo_automodel.components.models.mimo_v2_flash.processing._register_with_auto_image_processor() -> None

Register the class so a saved processor directory reloads through HF auto classes.

ProcessorMixin.save_pretrained writes "image_processor_type": "MimoV2ImageProcessor" into the checkpoint’s processor_config.json; without this mapping AutoProcessor.from_pretrained on that directory raises Unrecognized image processor.

nemo_automodel.components.models.mimo_v2_flash.processing._to_rgb(
    pil_image: PIL.Image.Image
) -> PIL.Image.Image

Convert to RGB with the sglang semantics: RGBA is flattened onto white.

EXIF orientation is deliberately not applied, matching the reference serving stack so training and serving see identical pixels.

nemo_automodel.components.models.mimo_v2_flash.processing.build_mimo_v2_processor(
    pretrained_model_name_or_path: str,
    kwargs: typing.Any = {}
) -> transformers.ProcessorMixin

Build the MiMo-V2 chat processor with sglang-aligned image preprocessing.

Loads the checkpoint’s stock processor (tokenizer, chat template, video and audio handling stay exactly as released) and replaces only the image processor with MimoV2ImageProcessor, initialized from the checkpoint’s own preprocessor_config.json values.

Parameters:

pretrained_model_name_or_path
str

HF model id or local checkpoint path.

**kwargs
AnyDefaults to {}

Forwarded to AutoProcessor.from_pretrained, e.g. trust_remote_code or max_pixels (recipes use the latter to cap image token counts).

Returns: ProcessorMixin

The stock MiMo-V2 processor whose image_processor is a

nemo_automodel.components.models.mimo_v2_flash.processing.smart_resize(
    height: int,
    width: int,
    factor: int,
    min_pixels: int,
    max_pixels: int
) -> tuple[int, int]

Compute resize targets with the sglang MiMoProcessor semantics.

Same algorithm as HF smart_resize except: a short side below factor is upscaled to factor first (that branch also skips the aspect-ratio guard), and degenerate shrink results are rejected instead of clamped.

Parameters:

height
int

Source image height in pixels.

width
int

Source image width in pixels.

factor
int

Both output sides are rounded to a multiple of this.

min_pixels
int

Lower bound on output area in pixels.

max_pixels
int

Upper bound on output area in pixels.

Returns: tuple[int, int]

(resized_height, resized_width) in pixels.

nemo_automodel.components.models.mimo_v2_flash.processing.ImageInput = Image.Image | str
 
nemo_automodel.components.models.mimo_v2_flash.processing.logger = logging.getLogger(__name__)