mstar.model.cosmos3.config

Configuration for the Cosmos3 omni generator.

A single Cosmos3Config describes every Cosmos3 checkpoint (Nano, Super, Edge, the Policy-DROID fine-tunes and the Super task variants). The checkpoints share one dual-pathway MoT architecture; they differ in the transformer dimensions (num_hidden_layers / hidden_size / num_attention_heads / intermediate_size), the two capability flags (sound_gen, action_gen) and, for Edge, the backbone family: a dense Nemotron text tower (hidden_act="relu2", Nemotron RMSNorm ordering, no text QK-norm, a k_norm_und_for_gen on the understanding K the generation tower reads) instead of Nano’s Qwen3-VL one.

Edge checkpoints also carry the reasoner (the understanding tower served as a VLM): a top-level config.json with the SigLIP2-style vision tower and patch-merger projector, plus vision_encoder/model.safetensors. That is parsed into Cosmos3Config.reasoner (None for checkpoints without it).

Values load from a local HF checkpoint directory laid out the diffusers way:

<ckpt>/transformer/config.json   -> the DiT (dual-pathway MoT) dimensions
<ckpt>/vae/config.json           -> AutoencoderKLWan factors + latent stats
<ckpt>/scheduler/scheduler_config.json -> UniPC flow scheduler settings (or the
                                    distilled checkpoints' FlowMatchEuler SDE sampler)
<ckpt>/model_index.json          -> pipeline flags (native flow schedule)
<ckpt>/modular_model_index.json  -> distilled sampler (is_distilled, distilled_sigmas)
<ckpt>/config.json               -> reasoner (vision tower + projector), Edge only
<ckpt>/preprocessor_config.json, video_preprocessor_config.json -> reasoner media processors

Dataclass defaults mirror Cosmos3-Nano so a bare Cosmos3Config() is a valid Nano config without any file present.

Classes

Cosmos3Config([hidden_size, ...])	Cosmos3 generator configuration (one architecture, swappable weights).
Cosmos3MediaProcessorConfig([patch_size, ...])	Resize/normalize/patchify settings of the reasoner's image and video processors (preprocessor_config.json / video_preprocessor_config.json).
Cosmos3ReasonerConfig([vision, ...])	Everything the understanding tower needs beyond the DiT weights to be served as a VLM (the Edge config.json): the vision tower, the patch-merger projector (LayerNorm -> 2x2 merge -> Linear -> GELU -> Linear into the text hidden size), the placeholder token ids, and the media processors.
Cosmos3SchedulerConfig([scheduler_class, ...])	Flow scheduler settings (scheduler/scheduler_config).
Cosmos3VAEConfig([z_dim, ...])	The Wan2.2-TI2V-5B VAE (AutoencoderKLWan) parameters we need at the serving layer.
Cosmos3VisionEncoderConfig([hidden_size, ...])	The reasoner's packed SigLIP2-style vision tower (vision_config of the Edge config.json): patch-embedding linear over patch_size**2 * 3 pixel patches, a learned square position grid of num_patches entries that is bilinearly resized to each image's patch grid, num_hidden_layers pre-LayerNorm encoder blocks attending within one frame, and a post LayerNorm.

class mstar.model.cosmos3.config.Cosmos3Config(hidden_size=4096, num_hidden_layers=36, num_attention_heads=32, num_key_value_heads=8, head_dim=128, intermediate_size=12288, vocab_size=151936, rms_norm_eps=1e-06, attention_bias=False, max_position_embeddings=262144, rope_theta=5000000.0, rope_axes_dim=(24, 20, 20), mrope_interleaved=True, unified_3d_mrope_temporal_modality_margin=15000, unified_3d_mrope_reset_spatial_ids=True, base_fps=24, enable_fps_modulation=True, latent_channel=48, latent_patch_size=2, patch_latent_dim=192, timestep_scale=0.001, joint_attn_implementation='two_way', qk_norm_for_diffusion=True, qk_norm_for_text=True, use_moe=True, hidden_act='silu', backbone_type=None, use_und_k_norm_for_gen=False, action_gen=True, max_action_dim=64, num_embodiment_domains=32, sound_gen=True, sound_dim=64, sound_latent_fps=25.0, temporal_compression_factor_sound=1, sound_sample_rate=48000, enable_sound=True, enable_windowed_video=False, max_windows=64, window_frames_default=29, overlap_frames_default=8, kv_window_frames_default=29, context_frames_default=61, windowed_decode_context_latents=8, windowed_decode_warm_cache=True, windowed_decode_max_warm_streams=0, windowed_decode_latents_per_call=1, vae_decode_channels_last=True, vae_decode_cudnn_benchmark=False, session_store_size=8, session_timeout_s=600.0, session_timeout_max_s=3600.0, video_temporal_causal=False, freeze_und=False, num_inference_steps=50, num_inference_steps_video=35, max_inference_steps=100, fps=24.0, num_frames_video=17, num_inference_steps_action=30, guidance_scale_action=1.0, flow_shift_action=5.0, image_size_default=(1024, 1024), video_size_default=None, guidance_scale=6.0, flow_shift_image=3.0, flow_shift_video=None, conditioning_resize='aspect_crop', use_native_flow_schedule=False, is_distilled=False, distilled_sigmas=None, accelerator_graph=True, graph_max_latent_area=2000, gen_capture_video=(), compile_denoise=True, compile_reasoner_decode=True, compile_reasoner_prefill=False, attention_backend='dense_gen', vae=<factory>, scheduler=<factory>, reasoner=None, enable_reasoner=True, reasoner_temperature=0.6, local_dir='')

Bases: object

Cosmos3 generator configuration (one architecture, swappable weights).

Parameters:

• hidden_size (int)

• num_hidden_layers (int)

• num_attention_heads (int)

• num_key_value_heads (int)

• head_dim (int)

• intermediate_size (int)

• vocab_size (int)

• rms_norm_eps (float)

• attention_bias (bool)

• max_position_embeddings (int)

• rope_theta (float)

• rope_axes_dim (tuple[int, int, int])

• mrope_interleaved (bool)

• unified_3d_mrope_temporal_modality_margin (int)

• unified_3d_mrope_reset_spatial_ids (bool)

• base_fps (int)

• enable_fps_modulation (bool)

• latent_channel (int)

• latent_patch_size (int)

• patch_latent_dim (int)

• timestep_scale (float)

• joint_attn_implementation (str)

• qk_norm_for_diffusion (bool)

• qk_norm_for_text (bool)

• use_moe (bool)

• hidden_act (str)

• backbone_type (str | None)

• use_und_k_norm_for_gen (bool)

• action_gen (bool)

• max_action_dim (int)

• num_embodiment_domains (int)

• sound_gen (bool)

• sound_dim (int | None)

• sound_latent_fps (float)

• temporal_compression_factor_sound (int)

• sound_sample_rate (int)

• enable_sound (bool)

• enable_windowed_video (bool)

• max_windows (int)

• window_frames_default (int)

• overlap_frames_default (int)

• kv_window_frames_default (int)

• context_frames_default (int)

• windowed_decode_context_latents (int)

• windowed_decode_warm_cache (bool)

• windowed_decode_max_warm_streams (int)

• windowed_decode_latents_per_call (int)

• vae_decode_channels_last (bool)

• vae_decode_cudnn_benchmark (bool)

• session_store_size (int)

• session_timeout_s (float)

• session_timeout_max_s (float)

• video_temporal_causal (bool)

• freeze_und (bool)

• num_inference_steps (int)

• num_inference_steps_video (int)

• max_inference_steps (int)

• fps (float)

• num_frames_video (int)

• num_inference_steps_action (int)

• guidance_scale_action (float)

• flow_shift_action (float)

• image_size_default (tuple[int, int])

• video_size_default (tuple[int, int] | None)

• guidance_scale (float)

• flow_shift_image (float | None)

• flow_shift_video (float | None)

• conditioning_resize (str)

• use_native_flow_schedule (bool)

• is_distilled (bool)

• distilled_sigmas (tuple[float, ...] | None)

• accelerator_graph (bool)

• graph_max_latent_area (int)

• gen_capture_video (tuple[tuple[int, int, int], ...])

• compile_denoise (bool)

• compile_reasoner_decode (bool)

• compile_reasoner_prefill (bool)

• attention_backend (str)

• vae (Cosmos3VAEConfig)

• scheduler (Cosmos3SchedulerConfig)

• reasoner (Cosmos3ReasonerConfig | None)

• enable_reasoner (bool)

• reasoner_temperature (float)

• local_dir (str)

accelerator_graph: bool = True

action_gen: bool = True

attention_backend: str = 'dense_gen'

attention_bias: bool = False

backbone_type: str | None = None

base_fps: int = 24

check_streaming_decode()

Reject streaming-decoder settings the decoder would otherwise reinterpret silently: a non-boolean switch (a yaml "false" is truthy), a negative warm-stream cap (read as no cap), or more latents per call than the default window holds (read as one call per window).

Return type:

None

compile_denoise: bool = True

compile_reasoner_decode: bool = True

compile_reasoner_prefill: bool = False

conditioning_resize: str = 'aspect_crop'

context_frames_default: int = 61

distilled_sigmas: tuple[float, ...] | None = None

enable_fps_modulation: bool = True

enable_reasoner: bool = True

enable_sound: bool = True

enable_windowed_video: bool = False

flow_shift_action: float = 5.0

flow_shift_image: float | None = 3.0

flow_shift_video: float | None = None

fps: float = 24.0

freeze_und: bool = False

classmethod from_pretrained(local_dir)

Load from a diffusers-layout checkpoint directory.

Parameters:

local_dir (str | Path)

Return type:

Cosmos3Config

classmethod from_transformer_dict(d)

Build from a diffusers transformer/config.json dict alone.

Sub-configs are left at their defaults; use from_pretrained to also populate VAE/scheduler from their sibling folders.

Parameters:

d (dict[str, Any])

Return type:

Cosmos3Config

property gated_mlp: bool

SwiGLU (gate/up/down) MLPs vs the dense two-projection relu2 ones.

gen_capture_video: tuple[tuple[int, int, int], ...] = ()

graph_max_latent_area: int = 2000

guidance_scale: float = 6.0

guidance_scale_action: float = 1.0

head_dim: int = 128

hidden_act: str = 'silu'

hidden_size: int = 4096

image_size_default: tuple[int, int] = (1024, 1024)

intermediate_size: int = 12288

is_distilled: bool = False

joint_attn_implementation: str = 'two_way'

kv_window_frames_default: int = 29

latent_channel: int = 48

latent_patch_size: int = 2

local_dir: str = ''

max_action_dim: int = 64

max_inference_steps: int = 100

max_position_embeddings: int = 262144

max_windows: int = 64

mrope_interleaved: bool = True

property nemotron_norm: bool

Whether every RMSNorm uses the Nemotron ordering (fp32 weight multiply, then one cast) — the dense relu2 backbone family.

num_attention_heads: int = 32

num_embodiment_domains: int = 32

num_frames_video: int = 17

num_hidden_layers: int = 36

num_inference_steps: int = 50

num_inference_steps_action: int = 30

num_inference_steps_video: int = 35

num_key_value_heads: int = 8

overlap_frames_default: int = 8

patch_latent_dim: int = 192

qk_norm_for_diffusion: bool = True

qk_norm_for_text: bool = True

reasoner: Cosmos3ReasonerConfig | None = None

reasoner_temperature: float = 0.6

rms_norm_eps: float = 1e-06

rope_axes_dim: tuple[int, int, int] = (24, 20, 20)

rope_theta: float = 5000000.0

scheduler: Cosmos3SchedulerConfig

property serves_reasoner: bool

session_store_size: int = 8

session_timeout_max_s: float = 3600.0

session_timeout_s: float = 600.0

sound_dim: int | None = 64

sound_gen: bool = True

sound_latent_fps: float = 25.0

sound_sample_rate: int = 48000

temporal_compression_factor_sound: int = 1

timestep_scale: float = 0.001

unified_3d_mrope_reset_spatial_ids: bool = True

unified_3d_mrope_temporal_modality_margin: int = 15000

use_moe: bool = True

use_native_flow_schedule: bool = False

use_und_k_norm_for_gen: bool = False

vae: Cosmos3VAEConfig

vae_decode_channels_last: bool = True

vae_decode_cudnn_benchmark: bool = False

video_size_default: tuple[int, int] | None = None

video_temporal_causal: bool = False

vocab_size: int = 151936

window_frames_default: int = 29

windowed_decode_context_latents: int = 8

windowed_decode_latents_per_call: int = 1

windowed_decode_max_warm_streams: int = 0

windowed_decode_warm_cache: bool = True

class mstar.model.cosmos3.config.Cosmos3MediaProcessorConfig(patch_size=16, merge_size=2, temporal_patch_size=1, min_pixels=65536, max_pixels=16777216, image_mean=(0.5, 0.5, 0.5), image_std=(0.5, 0.5, 0.5), fps=2.0, min_frames=4, max_frames=768)

Bases: object

Resize/normalize/patchify settings of the reasoner’s image and video processors (preprocessor_config.json / video_preprocessor_config.json).

An input is resized (bicubic, antialiased) so both sides are multiples of patch_size * merge_size and the pixel count lands in [min_pixels, max_pixels], normalized with image_mean / image_std, and cut into patch_size patches in block-major 2x2 order. Video inputs are first sampled at fps frames per second, clamped to [min_frames, max_frames].

Parameters:

• patch_size (int)

• merge_size (int)

• temporal_patch_size (int)

• min_pixels (int)

• max_pixels (int)

• image_mean (tuple[float, float, float])

• image_std (tuple[float, float, float])

• fps (float)

• min_frames (int)

• max_frames (int)

fps: float = 2.0

classmethod from_dict(d)

Parameters:

d (dict[str, Any])

Return type:

Cosmos3MediaProcessorConfig

image_mean: tuple[float, float, float] = (0.5, 0.5, 0.5)

image_std: tuple[float, float, float] = (0.5, 0.5, 0.5)

max_frames: int = 768

max_pixels: int = 16777216

merge_size: int = 2

min_frames: int = 4

min_pixels: int = 65536

patch_size: int = 16

temporal_patch_size: int = 1

class mstar.model.cosmos3.config.Cosmos3ReasonerConfig(vision=<factory>, image_processor=<factory>, video_processor=<factory>, projector_input_hidden_size=1152, projector_hidden_size=11520, projector_out_hidden_size=2048, use_postshuffle_norm=False, image_token_id=19, video_token_id=18, vision_start_token_id=20, vision_end_token_id=21, eos_token_id=11, max_position_embeddings=131072, enable_thinking=True)

Bases: object

Everything the understanding tower needs beyond the DiT weights to be served as a VLM (the Edge config.json): the vision tower, the patch-merger projector (LayerNorm -> 2x2 merge -> Linear -> GELU -> Linear into the text hidden size), the placeholder token ids, and the media processors. The text tower itself is the DiT’s understanding pathway (embed_tokens / layers.N.self_attn.to_* / mlp / norm / lm_head).

Parameters:

• vision (Cosmos3VisionEncoderConfig)

• image_processor (Cosmos3MediaProcessorConfig)

• video_processor (Cosmos3MediaProcessorConfig)

• projector_input_hidden_size (int)

• projector_hidden_size (int)

• projector_out_hidden_size (int)

• use_postshuffle_norm (bool)

• image_token_id (int)

• video_token_id (int)

• vision_start_token_id (int)

• vision_end_token_id (int)

• eos_token_id (int)

• max_position_embeddings (int)

• enable_thinking (bool)

enable_thinking: bool = True

eos_token_id: int = 11

classmethod from_dict(d)

Parameters:

d (dict[str, Any])

Return type:

Cosmos3ReasonerConfig

image_processor: Cosmos3MediaProcessorConfig

image_token_id: int = 19

max_position_embeddings: int = 131072

projector_hidden_size: int = 11520

projector_input_hidden_size: int = 1152

projector_out_hidden_size: int = 2048

use_postshuffle_norm: bool = False

video_processor: Cosmos3MediaProcessorConfig

video_token_id: int = 18

vision: Cosmos3VisionEncoderConfig

vision_end_token_id: int = 21

vision_start_token_id: int = 20

class mstar.model.cosmos3.config.Cosmos3SchedulerConfig(scheduler_class='UniPCMultistepScheduler', scheduler_type='unipc', prediction_type='flow_prediction', predict_x0=True, solver_order=2, solver_type='bh2', use_flow_sigmas=True, use_karras_sigmas=True, final_sigmas_type='zero', num_train_timesteps=1000, flow_shift=1.0, sigma_min=0.147, sigma_max=200.0, stochastic_sampling=False)

Bases: object

Flow scheduler settings (scheduler/scheduler_config).

The denoise loop drives a diffusers scheduler configured from these fields — UniPCMultistepScheduler for the base checkpoints (we do not re-implement the bh2 corrector), FlowMatchEulerDiscreteScheduler with stochastic (SDE) steps over the fixed distilled_sigmas for the 4-step distilled ones; scheduler_class records which.

Parameters:

• scheduler_class (str)

• scheduler_type (str)

• prediction_type (str)

• predict_x0 (bool)

• solver_order (int)

• solver_type (str)

• use_flow_sigmas (bool)

• use_karras_sigmas (bool)

• final_sigmas_type (str)

• num_train_timesteps (int)

• flow_shift (float)

• sigma_min (float)

• sigma_max (float)

• stochastic_sampling (bool)

final_sigmas_type: str = 'zero'

flow_shift: float = 1.0

classmethod from_dict(d)

Parameters:

d (dict[str, Any])

Return type:

Cosmos3SchedulerConfig

num_train_timesteps: int = 1000

predict_x0: bool = True

prediction_type: str = 'flow_prediction'

scheduler_class: str = 'UniPCMultistepScheduler'

scheduler_type: str = 'unipc'

sigma_max: float = 200.0

sigma_min: float = 0.147

solver_order: int = 2

solver_type: str = 'bh2'

stochastic_sampling: bool = False

use_flow_sigmas: bool = True

use_karras_sigmas: bool = True

class mstar.model.cosmos3.config.Cosmos3VAEConfig(z_dim=48, scale_factor_spatial=16, scale_factor_temporal=4, latents_mean=<factory>, latents_std=<factory>)

Bases: object

The Wan2.2-TI2V-5B VAE (AutoencoderKLWan) parameters we need at the serving layer. The full VAE module loads from the vae/ subfolder via diffusers; here we only track the latent geometry and the per-channel normalization statistics the pipeline applies to/from latent space.

Parameters:

• z_dim (int)

• scale_factor_spatial (int)

• scale_factor_temporal (int)

• latents_mean (list[float])

• latents_std (list[float])

classmethod from_dict(d)

Parameters:

d (dict[str, Any])

Return type:

Cosmos3VAEConfig

latents_mean: list[float]

latents_std: list[float]

scale_factor_spatial: int = 16

scale_factor_temporal: int = 4

z_dim: int = 48

class mstar.model.cosmos3.config.Cosmos3VisionEncoderConfig(hidden_size=1152, intermediate_size=4304, num_hidden_layers=27, num_attention_heads=16, num_channels=3, patch_size=16, num_patches=256, spatial_merge_size=2, hidden_act='gelu_pytorch_tanh', layer_norm_eps=1e-06)

Bases: object

The reasoner’s packed SigLIP2-style vision tower (vision_config of the Edge config.json): patch-embedding linear over patch_size**2 * 3 pixel patches, a learned square position grid of num_patches entries that is bilinearly resized to each image’s patch grid, num_hidden_layers pre-LayerNorm encoder blocks attending within one frame, and a post LayerNorm. spatial_merge_size is the projector’s 2x2 patch merge.

Parameters:

• hidden_size (int)

• intermediate_size (int)

• num_hidden_layers (int)

• num_attention_heads (int)

• num_channels (int)

• patch_size (int)

• num_patches (int)

• spatial_merge_size (int)

• hidden_act (str)

• layer_norm_eps (float)

classmethod from_dict(d)

Parameters:

d (dict[str, Any])

Return type:

Cosmos3VisionEncoderConfig

hidden_act: str = 'gelu_pytorch_tanh'

hidden_size: int = 1152

intermediate_size: int = 4304

layer_norm_eps: float = 1e-06

num_attention_heads: int = 16

num_channels: int = 3

num_hidden_layers: int = 27

num_patches: int = 256

patch_size: int = 16

spatial_merge_size: int = 2