Cosmos3 edge support

src/diffusers/models/transformers/transformer_cosmos3.py

@@ -13,19 +13,39 @@
 # limitations under the License.
 
 import math
+from dataclasses import dataclass
 
 import torch
 import torch.nn as nn
 
 from ...configuration_utils import ConfigMixin, register_to_config
 from ...loaders import PeftAdapterMixin
+from ...utils import BaseOutput
 from ..attention import AttentionMixin, AttentionModuleMixin
 from ..attention_dispatch import dispatch_attention_fn
 from ..embeddings import TimestepEmbedding, Timesteps
 from ..modeling_utils import ModelMixin
 from ..normalization import RMSNorm
 
 
+@dataclass
+class Cosmos3OmniTransformerOutput(BaseOutput):
+    """Output of [`Cosmos3OmniTransformer`].
+
+    Args:
+        sample (`list[torch.Tensor]`):
+            Per-item vision velocity predictions.
+        sound (`list[torch.Tensor]`, *optional*):
+            Per-item sound velocity predictions when sound generation is enabled.
+        action (`list[torch.Tensor]`, *optional*):
+            Per-item action velocity predictions when action generation is enabled.
+    """
+
+    sample: list[torch.Tensor]
+    sound: list[torch.Tensor] | None = None
+    action: list[torch.Tensor] | None = None
+
+
 class Cosmos3AttnProcessor:
     """Dual-pathway attention processor for Cosmos3.
 
@@ -139,15 +159,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,11 +223,11 @@ 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]
+    _supports_qkv_fusion = False
 
     def __init__(
         self,
@@ -197,6 +237,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 +254,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 +290,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 +301,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 +398,16 @@ 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,
     ):
         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 +424,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
         )
@@ -577,7 +652,10 @@ def forward(
         action_timesteps: torch.Tensor | None = None,
         action_noisy_frame_indexes: list[torch.Tensor] | None = None,
         action_domain_ids: list[torch.Tensor] | None = None,
-    ) -> tuple[list[torch.Tensor], list[torch.Tensor] | None, list[torch.Tensor] | None]:
+        return_dict: bool = True,
+    ) -> (
+        Cosmos3OmniTransformerOutput | tuple[list[torch.Tensor], list[torch.Tensor] | None, list[torch.Tensor] | None]
+    ):
         """Run a full denoising-step forward pass.
 
         Args:
@@ -605,10 +683,11 @@ def forward(
             action_timesteps: Optional per-token diffusion timesteps for action tokens.
             action_noisy_frame_indexes: Optional noisy frame indices per action item.
             action_domain_ids: Optional per-item domain IDs selecting the action head weights.
+            return_dict: Whether to return a [`Cosmos3OmniTransformerOutput`] instead of a tuple.
 
         Returns:
-            ``(preds_vision, preds_sound, preds_action)`` — lists of per-modality predictions. Optional modalities
-            return ``None`` when their inputs are omitted.
+            A [`Cosmos3OmniTransformerOutput`] or a tuple of per-modality prediction lists. Optional modalities return
+            ``None`` when their inputs are omitted.
         """
         has_sound = sound_tokens is not None and sound_sequence_indexes is not None
         has_action = action_tokens is not None and action_sequence_indexes is not None
@@ -722,4 +801,7 @@ def forward(
                 preds_action_packed, action_token_shapes, action_noisy_frame_indexes
             )
 
-        return preds_vision, preds_sound, preds_action
+        if not return_dict:
+            return preds_vision, preds_sound, preds_action
+
+        return Cosmos3OmniTransformerOutput(sample=preds_vision, sound=preds_sound, action=preds_action)

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/denoise.py

@@ -215,7 +215,7 @@ def __call__(self, components: Cosmos3OmniModularPipeline, block_state: BlockSta
             transformer_kwargs = {
                 name: value for name, value in transformer_kwargs.items() if name in transformer_args
             }
-            preds_vision, preds_sound, preds_action = components.transformer(**transformer_kwargs)
+            preds_vision, preds_sound, preds_action = components.transformer(**transformer_kwargs, return_dict=False)
             velocities[pass_name] = components._mask_velocity_predictions(
                 preds_vision,
                 preds_sound,
@@ -633,6 +633,7 @@ def _forward(components, static, vision_tokens, vision_timesteps):
             vision_mse_loss_indexes=static["vision_mse_loss_indexes"],
             vision_timesteps=vision_timesteps,
             vision_noisy_frame_indexes=static["vision_noisy_frame_indexes"],
+            return_dict=False,
         )
         return preds_vision[-1]