nemo_automodel.components.models.muse_spark.image_processing

Native image and video preprocessing for Muse Spark Open.

Module Contents

Classes

Name	Description
`MuseSparkImageProcessor`	Resize, normalize, and flatten Muse Spark Open image/video patches.
`MuseSparkVideoFeatures`	Preprocessed flattened video patches and their placeholder metadata.

Functions

Name	Description
`_image_to_uint8`	Convert one PIL/NumPy/Tensor image to contiguous uint8 `[3,H,W]`.
`_least_distance_fit`	Find the closest positive integer grid inside a token-count interval.
`_least_distortion_fit`	Choose an integer grid near the ideal size with minimum aspect error.
`_patchify`	Flatten `[temporal,3,H,W]` into row-major temporal patch vectors.
`_resize_and_normalize`	Lanczos-resize one image and return float32 normalized `[3,H,W]`.
`find_resize_target`	Return an aspect-preserving aligned canvas and visual-token count.

API

class nemo_automodel.components.models.muse_spark.image_processing.MuseSparkImageProcessor(
    patch_size: int = 14,
    patch_temporal: int = 2,
    downsample_factor: int = 2,
    min_image_tokens: int = 1,
    max_image_tokens: int = 4096,
    max_video_tokens: int = 4096,
    max_tokens_per_frame: int = 144,
    num_frames_per_video: int = 96,
    sampling_fps: float | None = 2.0,
    frame_factor: int = 2,
    image_mean: tuple[float, float, float] = (0.5, 0.5, 0.5),
    image_std: tuple[float, float, float] = (0.5, 0.5, 0.5),
    kwargs: typing.Any = {}
)

Bases: BaseImageProcessor

Resize, normalize, and flatten Muse Spark Open image/video patches.

image_mean

= tuple(image_mean)

image_std

= tuple(image_std)

model_input_names

= ['pixel_values', 'image_grid_thw']

token_patch_size

int

Return the pixel side represented by one post-shuffle token.

nemo_automodel.components.models.muse_spark.image_processing.MuseSparkImageProcessor.get_number_of_image_patches(
    height: int,
    width: int,
    images_kwargs: dict | None = None
) -> int

Return the visual placeholder count for an unprocessed image.

nemo_automodel.components.models.muse_spark.image_processing.MuseSparkImageProcessor.preprocess(
    images: typing.Any,
    return_tensors: str | None = None,
    kwargs: typing.Any = {}
) -> transformers.image_processing_utils.BatchFeature

Preprocess a flat image list into patch and grid tensors.

Parameters:

images

Any

One image or a flat list of PIL/NumPy/Tensor images.

return_tensors

str | NoneDefaults to None

Optional Transformers tensor type.

**kwargs

AnyDefaults to {}

Unused standard image-processor overrides.

Returns: BatchFeature

pixel_values has layout [sum(h*w),3*t*p*p] and

nemo_automodel.components.models.muse_spark.image_processing.MuseSparkImageProcessor.preprocess_video(
    video: typing.Any
) -> nemo_automodel.components.models.muse_spark.image_processing.MuseSparkVideoFeatures

Preprocess one video into flattened temporal clips.

Parameters:

video

Any

A [frames,channels,height,width] tensor, a sequence of frames, or (frames, metadata) where metadata may contain fps and frames_indices.

Returns: MuseSparkVideoFeatures

Flattened patches, one [clips,h,w] grid, tokens per clip, and

class nemo_automodel.components.models.muse_spark.image_processing.MuseSparkVideoFeatures(
    patches: torch.Tensor,
    grid_thw: tuple[int, int, int],
    tokens_per_clip: int,
    timestamps: tuple[float, ...]
)

Dataclass

Preprocessed flattened video patches and their placeholder metadata.

grid_thw

tuple[int, int, int]

patches

Tensor

timestamps

tuple[float, ...]

tokens_per_clip

int

nemo_automodel.components.models.muse_spark.image_processing._image_to_uint8(
    image: typing.Any
) -> torch.Tensor

Convert one PIL/NumPy/Tensor image to contiguous uint8 [3,H,W].

nemo_automodel.components.models.muse_spark.image_processing._least_distance_fit(
    ideal_height: float,
    ideal_width: float,
    minimum: int,
    maximum: int
) -> tuple[int, int]

Find the closest positive integer grid inside a token-count interval.

nemo_automodel.components.models.muse_spark.image_processing._least_distortion_fit(
    ideal_height: float,
    ideal_width: float,
    image_height: int,
    image_width: int,
    minimum: int,
    maximum: int
) -> tuple[int, int]

Choose an integer grid near the ideal size with minimum aspect error.

nemo_automodel.components.models.muse_spark.image_processing._patchify(
    clip: torch.Tensor,
    patch_size: int
) -> tuple[torch.Tensor, int, int]

Flatten [temporal,3,H,W] into row-major temporal patch vectors.

nemo_automodel.components.models.muse_spark.image_processing._resize_and_normalize(
    image: typing.Any,
    target_height: int,
    target_width: int,
    image_mean: tuple[float, float, float],
    image_std: tuple[float, float, float]
) -> torch.Tensor

Lanczos-resize one image and return float32 normalized [3,H,W].

nemo_automodel.components.models.muse_spark.image_processing.find_resize_target(
    image_height: int,
    image_width: int,
    token_patch_size: int,
    minimum_tokens: int,
    maximum_tokens: int
) -> tuple[int, int, int]

Return an aspect-preserving aligned canvas and visual-token count.

Parameters:

image_height

int

Source image height in pixels.

image_width

int

Source image width in pixels.

token_patch_size

int

Pixel extent represented by one post-shuffle token.

minimum_tokens

int

Minimum number of visual tokens.

maximum_tokens

int

Maximum number of visual tokens.

Returns: int

(height, width, tokens) with height and width aligned to