mstar.model.cosmos3.components.vision

mstar.model.cosmos3.components.vision#

The Cosmos3-Edge reasoner’s vision tower and projector.

A packed, variable-resolution SigLIP2-style encoder (visual): a linear patch embedding over patch_size x patch_size RGB patches, a learned square position grid bilinearly resized to every frame’s patch grid, N pre-LayerNorm encoder blocks whose attention stays within one frame, and a post LayerNorm. The projector (projector) merges each 2x2 block of patches into one token (LayerNorm per patch -> concat -> Linear -> GELU -> Linear) in the text hidden size, which the understanding tower consumes in place of the <|image_pad|> / <|video_pad|> tokens.

Parameter names follow vision_encoder/model.safetensors with its leading model. stripped (see loader.vision_encoder_name_remapper), so the checkpoint loads by name. Shapes and the packed patch order come from the Hugging Face Cosmos3EdgeVisionModel / Cosmos3EdgePatchMerger.

Frames of one video share a grid, and an image is a single frame, so the attention runs one batched SDPA per distinct frame size instead of a block-diagonal mask over the whole pack.

Functions

frame_attention(q, k, v, frame_lens)

Non-causal attention within each frame of a packed [N, H, D] batch.

resize_position_grid(grid_embeds, grid_thw, ...)

Resize the learned [S, S, hidden] position grid to every frame's (h, w) patch grid and lay it out in the processor's block-major (2x2-merge) patch order, repeated over the frame's t.

Classes

Cosmos3PatchMerger(config)

projector: per-patch LayerNorm, 2x2 merge into the channel dim, Linear -> GELU -> Linear into the text hidden size.

Cosmos3VisionAttention(config)

Cosmos3VisionEmbeddings(config)

Cosmos3VisionEncoder(config)

Cosmos3VisionEncoderLayer(config)

Cosmos3VisionMLP(config)

Cosmos3VisionModel(config)

Vision tower + projector: packed pixel patches -> text-space tokens.

Cosmos3VisionTransformer(config)

visual: embeddings -> encoder -> post LayerNorm, over packed patches.
class mstar.model.cosmos3.components.vision.Cosmos3PatchMerger(config)[source]#

Bases: Module

projector: per-patch LayerNorm, 2x2 merge into the channel dim, Linear -> GELU -> Linear into the text hidden size.

Parameters:

config (Cosmos3ReasonerConfig)

forward(x)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

x (Tensor)

Return type:

Tensor

class mstar.model.cosmos3.components.vision.Cosmos3VisionAttention(config)[source]#

Bases: Module

Parameters:

config (Cosmos3VisionEncoderConfig)

forward(x, frame_lens)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:
Return type:

Tensor

class mstar.model.cosmos3.components.vision.Cosmos3VisionEmbeddings(config)[source]#

Bases: Module

Parameters:

config (Cosmos3VisionEncoderConfig)

forward(pixel_values, grid_thw)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:
Return type:

Tensor

class mstar.model.cosmos3.components.vision.Cosmos3VisionEncoder(config)[source]#

Bases: Module

Parameters:

config (Cosmos3VisionEncoderConfig)

forward(x, frame_lens)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:
Return type:

Tensor

class mstar.model.cosmos3.components.vision.Cosmos3VisionEncoderLayer(config)[source]#

Bases: Module

Parameters:

config (Cosmos3VisionEncoderConfig)

forward(x, frame_lens)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:
Return type:

Tensor

class mstar.model.cosmos3.components.vision.Cosmos3VisionMLP(config)[source]#

Bases: Module

Parameters:

config (Cosmos3VisionEncoderConfig)

forward(x)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:

x (Tensor)

Return type:

Tensor

class mstar.model.cosmos3.components.vision.Cosmos3VisionModel(config)[source]#

Bases: Module

Vision tower + projector: packed pixel patches -> text-space tokens.

pixel_values is [total_patches, 3 * patch_size**2] in the processor’s block-major order (see components.reasoner), grid_thw one (t, h, w) patch grid per image/video. Returns [total_patches // merge_size**2, text_hidden]: one token per merged 2x2 block, in the order the <|image_pad|> / <|video_pad|> tokens appear in the prompt.

Parameters:

config (Cosmos3ReasonerConfig)

forward(pixel_values, grid_thw)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:
Return type:

Tensor

property tokens_per_patch_block: int#
class mstar.model.cosmos3.components.vision.Cosmos3VisionTransformer(config)[source]#

Bases: Module

visual: embeddings -> encoder -> post LayerNorm, over packed patches.

Parameters:

config (Cosmos3VisionEncoderConfig)

forward(pixel_values, grid_thw)[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Parameters:
Return type:

Tensor

mstar.model.cosmos3.components.vision.frame_attention(q, k, v, frame_lens)[source]#

Non-causal attention within each frame of a packed [N, H, D] batch.

frame_lens lists the consecutive frame token counts. Frames of equal length are batched into one SDPA call; the outputs come back in pack order.

Parameters:
Return type:

Tensor

mstar.model.cosmos3.components.vision.resize_position_grid(grid_embeds, grid_thw, merge_size)[source]#

Resize the learned [S, S, hidden] position grid to every frame’s (h, w) patch grid and lay it out in the processor’s block-major (2x2-merge) patch order, repeated over the frame’s t.

Bilinear, align_corners=False, antialiased — the reference’s F.interpolate call; fp32 on CPU because antialias has no bf16 kernel there. Returns [total_patches, hidden] in pack order.

Parameters:
Return type:

Tensor