mstar.model.cosmos3.components.reasoner

mstar.model.cosmos3.components.reasoner#

Prompt-side plumbing for the Cosmos3-Edge reasoner: media preprocessing, placeholder expansion and the 3D mRoPE positions of a VLM prompt.

Pure tensor / string helpers with no model state, shared by Cosmos3Model.process_prompt (data worker) and the reasoner submodule, and matching the Hugging Face Cosmos3EdgeProcessor / Cosmos3EdgeModel byte-for-byte:

  • preprocess_image / preprocess_video: bicubic antialiased resize to multiples of patch_size * merge_size inside [min_pixels, max_pixels] (on 8-bit pixels through the torchvision v2 functional, whose native uint8 kernel is what the reference calls — the v1 API rounds differently by one 8-bit step on some pixels), (x/255 - mean) / std, then block-major 2x2 patchify so consecutive patches form the merger’s blocks.

  • expand_placeholders: one <|image_pad|> per merged block for images; videos become one timestamped <t seconds><|vision_start|>...<|vision_end|> span per frame.

  • mrope_position_ids: text tokens advance all three axes together; a vision span puts t on the temporal axis and the merged (h, w) grid on the spatial ones, all offset by the current position, and advances the cursor by max(h, w) merged patches.

Functions

expand_placeholders(text, tokenizer, cfg, ...)

Replace the chat template's single-token media placeholders with the tokens the encoder will fill: <|image_pad|> x merged blocks per image; <|vision_start|><|video_pad|><|vision_end|> -> one <t seconds><|vision_start|>{pads}<|vision_end|> span per frame.

mrope_position_ids(input_ids, cfg, ...)

3D mRoPE ids [3, N] (temporal, height, width) of a rendered prompt and the position the first generated token takes.

patchify(frames, patch_size, merge_size)

[T, C, H, W] -> [T * gh * gw, patch_size**2 * C] patches in time-major, block-major order (the merge_size x merge_size patches of one merger block are consecutive), pixel values ordered (ph, pw, C) inside a patch.

preprocess_image(image, cfg)

One [C, H, W] image (8-bit or [0, 1] float) -> packed patches [gh * gw, patch_size**2 * C] and its grid.

preprocess_video(video, cfg, source_fps[, ...])

A decoded [T, C, H, W] clip -> packed patches [T' * gh * gw, ...] over the sampled frames and its grid with per-frame timestamps.

render_chat(tokenizer, parts, cfg[, ...])

Render ordered prompt parts (text / image / video, as written) through the checkpoint's chat template with the generation prompt appended.

sample_frame_indices(total_frames, ...[, ...])

Uniform frame sampling: fps frames per second of source video (default the processor's), clamped to [min_frames, max_frames] and to the clip, or exactly num_frames.

smart_resize(num_frames, height, width, ...)

Target (height, width): both multiples of factor, the (temporal patch count x) pixel count inside [min_pixels, max_pixels], the aspect ratio kept as closely as possible.

Classes

MediaGrid(t, h, w[, timestamps])

The patch grid of one preprocessed image (t == 1) or video, plus a video's per-frame timestamps in seconds.
class mstar.model.cosmos3.components.reasoner.MediaGrid(t, h, w, timestamps=())[source]#

Bases: object

The patch grid of one preprocessed image (t == 1) or video, plus a video’s per-frame timestamps in seconds.

Parameters:
h: int#
t: int#
property thw: tuple[int, int, int]#
timestamps: tuple[float, ...] = ()#
tokens(merge_size)[source]#

LLM tokens this media occupies (one per merged block).

Parameters:

merge_size (int)

Return type:

int

tokens_per_frame(merge_size)[source]#
Parameters:

merge_size (int)

Return type:

int

w: int#
mstar.model.cosmos3.components.reasoner.expand_placeholders(text, tokenizer, cfg, image_grids, video_grids)[source]#

Replace the chat template’s single-token media placeholders with the tokens the encoder will fill: <|image_pad|> x merged blocks per image; <|vision_start|><|video_pad|><|vision_end|> -> one <t seconds><|vision_start|>{pads}<|vision_end|> span per frame.

Parameters:
Return type:

str

mstar.model.cosmos3.components.reasoner.mrope_position_ids(input_ids, cfg, image_grids, video_grids)[source]#

3D mRoPE ids [3, N] (temporal, height, width) of a rendered prompt and the position the first generated token takes.

Text runs put the same increasing ids on all three axes. Each run of <|image_pad|> / <|video_pad|> tokens is one frame of its media (videos are rendered one span per frame): temporal = the cursor, height / width = the cursor plus the merged-grid coordinates; the cursor then advances by max(h, w) merged patches. Decoding continues at max(position) + 1 on all axes.

Parameters:
Return type:

tuple[Tensor, int]

mstar.model.cosmos3.components.reasoner.patchify(frames, patch_size, merge_size)[source]#

[T, C, H, W] -> [T * gh * gw, patch_size**2 * C] patches in time-major, block-major order (the merge_size x merge_size patches of one merger block are consecutive), pixel values ordered (ph, pw, C) inside a patch.

Parameters:
Return type:

Tensor

mstar.model.cosmos3.components.reasoner.preprocess_image(image, cfg)[source]#

One [C, H, W] image (8-bit or [0, 1] float) -> packed patches [gh * gw, patch_size**2 * C] and its grid.

Parameters:
Return type:

tuple[Tensor, MediaGrid]

mstar.model.cosmos3.components.reasoner.preprocess_video(video, cfg, source_fps, num_frames=None, fps=None)[source]#

A decoded [T, C, H, W] clip -> packed patches [T' * gh * gw, ...] over the sampled frames and its grid with per-frame timestamps.

Parameters:
Return type:

tuple[Tensor, MediaGrid]

mstar.model.cosmos3.components.reasoner.render_chat(tokenizer, parts, cfg, enable_thinking=None, system_prompt=None)[source]#

Render ordered prompt parts (text / image / video, as written) through the checkpoint’s chat template with the generation prompt appended. Media parts become the template’s single-token placeholders, which expand_placeholders then grows. system_prompt leads only when the client sent no system message of its own.

Parameters:
Return type:

str

mstar.model.cosmos3.components.reasoner.sample_frame_indices(total_frames, source_fps, cfg, num_frames=None, fps=None)[source]#

Uniform frame sampling: fps frames per second of source video (default the processor’s), clamped to [min_frames, max_frames] and to the clip, or exactly num_frames. Indices are linspace-rounded over the whole clip, like the reference.

Parameters:
Return type:

list[int]

mstar.model.cosmos3.components.reasoner.smart_resize(num_frames, height, width, temporal_factor, factor, min_pixels, max_pixels)[source]#

Target (height, width): both multiples of factor, the (temporal patch count x) pixel count inside [min_pixels, max_pixels], the aspect ratio kept as closely as possible.

Parameters:

  • num_frames (int)

  • height (int)

  • width (int)

  • temporal_factor (int)

  • factor (int)

  • min_pixels (int)

  • max_pixels (int)

Return type:

tuple[int, int]