Cosmos3 edge support

src/diffusers/models/transformers/transformer_cosmos3.py

@@ -139,15 +139,35 @@ def forward(self, position_ids, device, dtype):
         return emb.cos().to(dtype=dtype), emb.sin().to(dtype=dtype)  # each: [B,N,head_dim]
 
 
+class Cosmos3NemotronRMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.float()
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+        return (self.weight.float() * hidden_states).to(input_dtype)
+
+
 class Cosmos3VLTextMLP(nn.Module):
-    def __init__(self, hidden_size: int, intermediate_size: int):
+    def __init__(self, hidden_size: int, intermediate_size: int, hidden_act: str = "silu"):
         super().__init__()
-        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        if hidden_act not in ("relu2", "silu"):
+            raise ValueError(f"Cosmos3 only supports `hidden_act` values 'relu2' and 'silu', got {hidden_act!r}.")
+        self.hidden_act = hidden_act
+        if hidden_act == "silu":
+            self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
         self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
         self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
-        self.act_fn = nn.SiLU()
+        self.act_fn = nn.SiLU() if hidden_act == "silu" else None
 
     def forward(self, x):
+        if self.hidden_act == "relu2":
+            return self.down_proj(torch.relu(self.up_proj(x)).square())
         return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
 
 
@@ -183,8 +203,7 @@ def forward(self, x: torch.Tensor, domain_id: torch.Tensor) -> torch.Tensor:
 
 
 class Cosmos3PackedMoTAttention(nn.Module, AttentionModuleMixin):
-    """Dual-pathway packed attention for Qwen3VL MoT — separate projections for
-    understanding (causal) and generation (full) token streams."""
+    """Dual-pathway packed attention with separate projections for the understanding and generation token streams."""
 
     _default_processor_cls = Cosmos3AttnProcessor
     _available_processors = [Cosmos3AttnProcessor]
@@ -197,6 +216,8 @@ def __init__(
         num_key_value_heads: int,
         attention_bias: bool,
         rms_norm_eps: float,
+        qk_norm_for_text: bool = True,
+        norm_type: str = "rms_norm",
         processor=None,
     ):
         super().__init__()
@@ -212,16 +233,27 @@ def __init__(
         self.to_k = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=attention_bias)
         self.to_v = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=attention_bias)
         self.to_out = nn.Linear(num_attention_heads * head_dim, hidden_size, bias=attention_bias)
-        self.norm_q = RMSNorm(head_dim, eps=rms_norm_eps, elementwise_affine=True, bias=False)
-        self.norm_k = RMSNorm(head_dim, eps=rms_norm_eps, elementwise_affine=True, bias=False)
+        if not qk_norm_for_text:
+            self.norm_q = nn.Identity()
+            self.norm_k = nn.Identity()
+        elif norm_type == "nemotron_rms_norm":
+            self.norm_q = Cosmos3NemotronRMSNorm(head_dim, eps=rms_norm_eps)
+            self.norm_k = Cosmos3NemotronRMSNorm(head_dim, eps=rms_norm_eps)
+        else:
+            self.norm_q = RMSNorm(head_dim, eps=rms_norm_eps, elementwise_affine=True, bias=False)
+            self.norm_k = RMSNorm(head_dim, eps=rms_norm_eps, elementwise_affine=True, bias=False)
 
         # Generation pathway
         self.add_q_proj = nn.Linear(hidden_size, num_attention_heads * head_dim, bias=attention_bias)
         self.add_k_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=attention_bias)
         self.add_v_proj = nn.Linear(hidden_size, num_key_value_heads * head_dim, bias=attention_bias)
         self.to_add_out = nn.Linear(num_attention_heads * head_dim, hidden_size, bias=attention_bias)
-        self.norm_added_q = RMSNorm(head_dim, eps=rms_norm_eps, elementwise_affine=True, bias=False)
-        self.norm_added_k = RMSNorm(head_dim, eps=rms_norm_eps, elementwise_affine=True, bias=False)
+        if norm_type == "nemotron_rms_norm":
+            self.norm_added_q = Cosmos3NemotronRMSNorm(head_dim, eps=rms_norm_eps)
+            self.norm_added_k = Cosmos3NemotronRMSNorm(head_dim, eps=rms_norm_eps)
+        else:
+            self.norm_added_q = RMSNorm(head_dim, eps=rms_norm_eps, elementwise_affine=True, bias=False)
+            self.norm_added_k = RMSNorm(head_dim, eps=rms_norm_eps, elementwise_affine=True, bias=False)
 
         if processor is None:
             processor = self._default_processor_cls()
@@ -237,12 +269,7 @@ def forward(
 
 
 class Cosmos3VLTextMoTDecoderLayer(nn.Module):
-    """
-    Qwen3VL text MoT (Mixture of Tokens) decoder layer. Features dual-pathway attention for understanding vs
-    generation.
-
-    This is used for both Dense and MoE models.
-    """
+    """Cosmos3 text MoT decoder layer for the Qwen3 and Nemotron dense backbones."""
 
     def __init__(
         self,
@@ -253,28 +280,43 @@ def __init__(
         intermediate_size: int,
         attention_bias: bool,
         rms_norm_eps: float,
+        hidden_act: str = "silu",
+        qk_norm_for_text: bool = True,
     ):
         super().__init__()
         self.hidden_size = hidden_size
+        norm_type = "nemotron_rms_norm" if hidden_act == "relu2" else "rms_norm"
         self.self_attn = Cosmos3PackedMoTAttention(
             hidden_size=hidden_size,
             head_dim=head_dim,
             num_attention_heads=num_attention_heads,
             num_key_value_heads=num_key_value_heads,
             attention_bias=attention_bias,
             rms_norm_eps=rms_norm_eps,
+            qk_norm_for_text=qk_norm_for_text,
+            norm_type=norm_type,
         )
 
-        self.mlp = Cosmos3VLTextMLP(hidden_size=hidden_size, intermediate_size=intermediate_size)
-        self.mlp_moe_gen = Cosmos3VLTextMLP(hidden_size=hidden_size, intermediate_size=intermediate_size)
-
-        self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False)
-        self.input_layernorm_moe_gen = RMSNorm(hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False)
-        self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False)
-        self.post_attention_layernorm_moe_gen = RMSNorm(
-            hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False
+        self.mlp = Cosmos3VLTextMLP(
+            hidden_size=hidden_size, intermediate_size=intermediate_size, hidden_act=hidden_act
+        )
+        self.mlp_moe_gen = Cosmos3VLTextMLP(
+            hidden_size=hidden_size, intermediate_size=intermediate_size, hidden_act=hidden_act
         )
 
+        if norm_type == "nemotron_rms_norm":
+            self.input_layernorm = Cosmos3NemotronRMSNorm(hidden_size, eps=rms_norm_eps)
+            self.input_layernorm_moe_gen = Cosmos3NemotronRMSNorm(hidden_size, eps=rms_norm_eps)
+            self.post_attention_layernorm = Cosmos3NemotronRMSNorm(hidden_size, eps=rms_norm_eps)
+            self.post_attention_layernorm_moe_gen = Cosmos3NemotronRMSNorm(hidden_size, eps=rms_norm_eps)
+        else:
+            self.input_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False)
+            self.input_layernorm_moe_gen = RMSNorm(hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False)
+            self.post_attention_layernorm = RMSNorm(hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False)
+            self.post_attention_layernorm_moe_gen = RMSNorm(
+                hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False
+            )
+
     def forward(
         self,
         und_seq: torch.Tensor,
@@ -335,10 +377,18 @@ def __init__(
         sound_latent_fps: float = 25.0,
         timestep_scale: float = 0.001,
         vocab_size: int = 151936,
+        hidden_act: str = "silu",
+        qk_norm_for_text: bool = True,
+        rope_axes_dim: tuple[int, int, int] | list[int] | None = None,
+        backbone_type: str = "cosmos3_qwen3vl",
+        temporal_compression_factor: int = 4,
     ):
         super().__init__()
 
-        rope_axes_dim = rope_scaling.get("mrope_section", [24, 20, 20]) if rope_scaling is not None else [24, 20, 20]
+        if rope_axes_dim is None:
+            rope_axes_dim = (
+                rope_scaling.get("mrope_section", [24, 20, 20]) if rope_scaling is not None else [24, 20, 20]
+            )
         self.register_to_config(rope_axes_dim=rope_axes_dim)
 
         # Text-model layers live directly on the transformer (flat layout). The published
@@ -355,12 +405,18 @@ def __init__(
                     intermediate_size=intermediate_size,
                     attention_bias=attention_bias,
                     rms_norm_eps=rms_norm_eps,
+                    hidden_act=hidden_act,
+                    qk_norm_for_text=qk_norm_for_text,
                 )
                 for _ in range(num_hidden_layers)
             ]
         )
-        self.norm = RMSNorm(hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False)
-        self.norm_moe_gen = RMSNorm(hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False)
+        if hidden_act == "relu2":
+            self.norm = Cosmos3NemotronRMSNorm(hidden_size, eps=rms_norm_eps)
+            self.norm_moe_gen = Cosmos3NemotronRMSNorm(hidden_size, eps=rms_norm_eps)
+        else:
+            self.norm = RMSNorm(hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False)
+            self.norm_moe_gen = RMSNorm(hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False)
         self.rotary_emb = Cosmos3VLTextRotaryEmbedding(
             head_dim=head_dim, rope_theta=rope_theta, rope_axes_dim=rope_axes_dim
         )

src/diffusers/modular_pipelines/cosmos/before_denoise.py

@@ -1,13 +1,14 @@
 import copy
 
+import numpy as np
 import torch
 
 from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer
 from ...pipelines.cosmos.pipeline_cosmos3_omni import _EMBODIMENT_TO_DOMAIN_ID, CosmosActionCondition
 from ...schedulers import UniPCMultistepScheduler
 from ...utils.torch_utils import randn_tensor
 from ..modular_pipeline import ModularPipelineBlocks, PipelineState
-from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam
 from .modular_pipeline import Cosmos3OmniModularPipeline
 
 
@@ -936,6 +937,10 @@ def description(self) -> str:
     def expected_components(self) -> list[ComponentSpec]:
         return [ComponentSpec("scheduler", UniPCMultistepScheduler)]
 
+    @property
+    def expected_configs(self) -> list[ConfigSpec]:
+        return [ConfigSpec(name="use_native_flow_schedule", default=False)]
+
     @property
     def inputs(self) -> list[InputParam]:
         return [
@@ -953,7 +958,15 @@ def intermediate_outputs(self) -> list[OutputParam]:
     def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
         block_state = self.get_block_state(state)
         device = components._execution_device
-        components.scheduler.set_timesteps(block_state.num_inference_steps, device=device)
+        if components.config.use_native_flow_schedule:
+            sigmas = np.linspace(
+                1.0 - 1.0 / components.scheduler.config.num_train_timesteps,
+                0.0,
+                block_state.num_inference_steps + 1,
+            )[:-1]
+            components.scheduler.set_timesteps(block_state.num_inference_steps, device=device, sigmas=sigmas)
+        else:
+            components.scheduler.set_timesteps(block_state.num_inference_steps, device=device)
         block_state.timesteps = components.scheduler.timesteps
         block_state.num_warmup_steps = (
             len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order

src/diffusers/modular_pipelines/cosmos/encoders.py

@@ -10,7 +10,7 @@
 from ...utils import logging
 from ...video_processor import VideoProcessor
 from ..modular_pipeline import ModularPipelineBlocks, PipelineState
-from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from ..modular_pipeline_utils import ComponentSpec, ConfigSpec, InputParam, OutputParam
 from .modular_pipeline import Cosmos3OmniModularPipeline
 
 
@@ -53,6 +53,13 @@ def expected_components(self) -> list[ComponentSpec]:
             ComponentSpec("text_tokenizer", AutoTokenizer),
         ]
 
+    @property
+    def expected_configs(self) -> list[ConfigSpec]:
+        return [
+            ConfigSpec(name="default_use_system_prompt", default=True),
+            ConfigSpec(name="enable_safety_checker", default=True),
+        ]
+
     @property
     def inputs(self) -> list[InputParam]:
         return [
@@ -76,8 +83,8 @@ def inputs(self) -> list[InputParam]:
             InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."),
             InputParam(
                 name="use_system_prompt",
-                type_hint=bool,
-                default=True,
+                type_hint=bool | None,
+                default=None,
                 description="Whether to prepend the Cosmos3 system prompt.",
             ),
             InputParam(
@@ -115,6 +122,8 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState)
             block_state.height = 720
         if block_state.width is None:
             block_state.width = 1280
+        if block_state.use_system_prompt is None:
+            block_state.use_system_prompt = components.config.default_use_system_prompt
 
         self._check_inputs(block_state)
         if components.requires_safety_checker:
@@ -316,6 +325,13 @@ def expected_components(self) -> list[ComponentSpec]:
             ),
         ]
 
+    @property
+    def expected_configs(self) -> list[ConfigSpec]:
+        return [
+            ConfigSpec(name="default_use_system_prompt", default=True),
+            ConfigSpec(name="enable_safety_checker", default=True),
+        ]
+
     @property
     def inputs(self) -> list[InputParam]:
         return [
@@ -345,8 +361,8 @@ def inputs(self) -> list[InputParam]:
             InputParam(name="fps", type_hint=float, default=24.0, description="Frame rate of the generated video."),
             InputParam(
                 name="use_system_prompt",
-                type_hint=bool,
-                default=True,
+                type_hint=bool | None,
+                default=None,
                 description="Whether to prepend the Cosmos3 system prompt.",
             ),
             InputParam(
@@ -380,6 +396,8 @@ def intermediate_outputs(self) -> list[OutputParam]:
     def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) -> PipelineState:
         block_state = self.get_block_state(state)
         self._check_inputs(block_state)
+        if block_state.use_system_prompt is None:
+            block_state.use_system_prompt = components.config.default_use_system_prompt
 
         action = block_state.action
         block_state.action_mode = action.mode

src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py

@@ -140,8 +140,9 @@ class Cosmos3AutoTextEncoderStep(AutoPipelineBlocks):
               The text prompt that guides Cosmos3 generation.
           negative_prompt (`str`, *optional*):
               The negative text prompt used for classifier-free guidance.
-          use_system_prompt (`bool`, *optional*, defaults to True):
-              Whether to prepend the Cosmos3 transfer system prompt.
+          use_system_prompt (`bool`, *optional*):
+              Whether to prepend the system prompt. Defaults to the pipeline configuration for standard and action
+              workflows and to True for transfer.
           action (`CosmosActionCondition`, *optional*):
               Action-conditioning metadata and its reference visual input.
           fps (`float`, *optional*, defaults to 24.0):
@@ -1165,8 +1166,9 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks):
               The text prompt that guides Cosmos3 generation.
           negative_prompt (`str`, *optional*):
               The negative text prompt used for classifier-free guidance.
-          use_system_prompt (`bool`, *optional*, defaults to True):
-              Whether to prepend the Cosmos3 transfer system prompt.
+          use_system_prompt (`bool`, *optional*):
+              Whether to prepend the system prompt. Defaults to the pipeline configuration for standard and action
+              workflows and to True for transfer.
           action (`CosmosActionCondition`, *optional*):
               Action-conditioning metadata and its reference visual input.
           fps (`float`, *optional*, defaults to 24.0):

src/diffusers/modular_pipelines/cosmos/modular_pipeline.py

@@ -77,7 +77,7 @@ def disable_safety_checker(self):
 
     @property
     def requires_safety_checker(self):
-        return getattr(self, "_is_safety_checker_enabled", True)
+        return getattr(self, "_is_safety_checker_enabled", self.config.enable_safety_checker)
 
     def _encode_video(self, x):
         return Cosmos3OmniPipeline._encode_video(self, x)