mstar.model.cosmos3.loader

Weight loading for the Cosmos3 generator backbone.

The published checkpoint is the diffusers transformer/ layout: flat layers.N.* keys with unfused attention projections (to_q/to_k/to_v for the understanding pathway, add_q_proj/add_k_proj/add_v_proj for the generation pathway) and _moe_gen-suffixed GEN MLP/norms. Our backbone module mirrors that layout one-to-one, so loading needs no key remapping and no stacked-parameter fusion — only the unused text lm_head is dropped.

Functions

cosmos3_name_remapper(name)	Map a checkpoint key to a backbone parameter path, or None to drop.
load_transformer_weights(model, checkpoint_dir)	Stream the transformer/ shards into model and return loaded keys.
read_transformer_weight_keys(checkpoint_dir)	Return every tensor key declared by the transformer/ shard index.
read_transformer_weight_shapes(checkpoint_dir)	Return {key: shape} for every transformer/ tensor by reading only the safetensors headers — no tensor data is materialized.

mstar.model.cosmos3.loader.cosmos3_name_remapper(name)

Map a checkpoint key to a backbone parameter path, or None to drop.

Identity for every key the backbone owns; None for the intentional drop-list. Kept explicit so an unexpected checkpoint key surfaces as a coverage failure rather than being silently ignored.

Parameters:

name (str)

Return type:

str | None

mstar.model.cosmos3.loader.load_transformer_weights(model, checkpoint_dir, device='cpu')

Stream the transformer/ shards into model and return loaded keys.

Mirrors the meta-device + load_hf_weights path the other model packages use, but resolves the shard list from the diffusers diffusion_pytorch_model index (the generic iterator only knows the model.safetensors name). No stacked-parameter rules: the checkpoint’s projections are unfused and match the backbone parameter names directly. Raises if any backbone parameter is left unfilled — the completeness guarantee bagel’s loader also enforces.

Parameters:

• model (Module)

• checkpoint_dir (str | Path)

• device (str)

Return type:

set[str]

mstar.model.cosmos3.loader.read_transformer_weight_keys(checkpoint_dir)

Return every tensor key declared by the transformer/ shard index.

Parameters:

checkpoint_dir (str | Path)

Return type:

set[str]

mstar.model.cosmos3.loader.read_transformer_weight_shapes(checkpoint_dir)

Return {key: shape} for every transformer/ tensor by reading only the safetensors headers — no tensor data is materialized. Enables CPU-side shape verification of the meta-built backbone against the checkpoint.

Parameters:

checkpoint_dir (str | Path)

Return type:

dict[str, tuple[int, …]]