mstar.model.cosmos3

Cosmos3 omni generator model package.

class mstar.model.cosmos3.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.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.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

Modules

components	Cosmos3 backbone components.
config	Configuration for the Cosmos3 omni generator.
constants	Graph-walk names shared by the Cosmos3 model graph and its submodules.
loader	Weight loading for the Cosmos3 generator backbone and the Edge reasoner's vision tower.
sessions	Per-session world state for windowed rollouts.
submodules	NodeSubmodule wrappers for the Cosmos3 generator nodes.
tests