mstar.model.cosmos3.config

Configuration for the Cosmos3 omni generator.

A single Cosmos3Config describes every Cosmos3 checkpoint (Nano, Super, Policy-DROID, and the Super task variants). The checkpoints share one architecture; they differ only in the transformer dimensions (num_hidden_layers / hidden_size / num_attention_heads / intermediate_size) and two capability flags (sound_gen, action_gen).

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

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).
Cosmos3SchedulerConfig([scheduler_type, ...])	UniPC multistep flow scheduler settings (scheduler/scheduler_config).
Cosmos3VAEConfig([z_dim, ...])	The Wan2.2-TI2V-5B VAE (AutoencoderKLWan) parameters we need at the serving layer.

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, 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, 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, cuda_graph=True, graph_max_latent_area=2000, compile_denoise=True, attention_backend='dense_gen', vae=<factory>, scheduler=<factory>, 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)

• 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)

• 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)

• cuda_graph (bool)

• graph_max_latent_area (int)

• compile_denoise (bool)

• attention_backend (str)

• vae (Cosmos3VAEConfig)

• scheduler (Cosmos3SchedulerConfig)

• local_dir (str)

action_gen: bool = True

attention_backend: str = 'dense_gen'

attention_bias: bool = False

base_fps: int = 24

compile_denoise: bool = True

cuda_graph: bool = True

enable_fps_modulation: bool = True

enable_sound: bool = True

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

graph_max_latent_area: int = 2000

head_dim: int = 128

hidden_size: int = 4096

intermediate_size: int = 12288

joint_attn_implementation: str = 'two_way'

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

mrope_interleaved: bool = True

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_video: int = 35

num_key_value_heads: int = 8

patch_latent_dim: int = 192

qk_norm_for_diffusion: bool = True

qk_norm_for_text: bool = True

rms_norm_eps: float = 1e-06

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

rope_theta: float = 5000000.0

scheduler: Cosmos3SchedulerConfig

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

vae: Cosmos3VAEConfig

video_temporal_causal: bool = False

vocab_size: int = 151936

class mstar.model.cosmos3.config.Cosmos3SchedulerConfig(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)

Bases: object

UniPC multistep flow scheduler settings (scheduler/scheduler_config).

The denoise loop drives a diffusers UniPCMultistepScheduler configured from these fields; we do not re-implement the bh2 corrector.

Parameters:

• 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)

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_type: str = 'unipc'

sigma_max: float = 200.0

sigma_min: float = 0.147

solver_order: int = 2

solver_type: str = 'bh2'

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