From 278ed5ef4dff9e6b896fe312f3a8869f7681294e Mon Sep 17 00:00:00 2001 From: Atharva Joshi Date: Mon, 13 Jul 2026 13:44:38 -0700 Subject: [PATCH 1/8] Add Cosmos3 Edge modular pipeline support --- .../transformers/transformer_cosmos3.py | 94 +++++++++++----- .../cosmos/before_denoise.py | 17 ++- .../modular_pipelines/cosmos/encoders.py | 28 ++++- .../cosmos/modular_blocks_cosmos3.py | 8 +- .../cosmos/modular_pipeline.py | 2 +- .../pipelines/cosmos/pipeline_cosmos3_omni.py | 34 ++++-- .../test_models_transformer_cosmos3.py | 101 ++++++++++++++++++ .../cosmos/test_modular_pipeline_cosmos3.py | 73 ++++++++++++- tests/pipelines/cosmos/test_cosmos3.py | 96 +++++++++++++++++ 9 files changed, 407 insertions(+), 46 deletions(-) create mode 100644 tests/models/transformers/test_models_transformer_cosmos3.py create mode 100644 tests/pipelines/cosmos/test_cosmos3.py diff --git a/src/diffusers/models/transformers/transformer_cosmos3.py b/src/diffusers/models/transformers/transformer_cosmos3.py index 3e331ba74628..2a150fc879bb 100644 --- a/src/diffusers/models/transformers/transformer_cosmos3.py +++ b/src/diffusers/models/transformers/transformer_cosmos3.py @@ -139,15 +139,37 @@ 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, elementwise_affine: bool = True, bias: bool = False): + super().__init__() + if not elementwise_affine or bias: + raise ValueError("Cosmos3NemotronRMSNorm requires an affine weight without a bias.") + 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 +205,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 +218,8 @@ def __init__( num_key_value_heads: int, attention_bias: bool, rms_norm_eps: float, + qk_norm_for_text: bool = True, + norm_cls: type[nn.Module] = RMSNorm, processor=None, ): super().__init__() @@ -212,16 +235,24 @@ 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) + self.norm_q = ( + norm_cls(head_dim, eps=rms_norm_eps, elementwise_affine=True, bias=False) + if qk_norm_for_text + else nn.Identity() + ) + self.norm_k = ( + norm_cls(head_dim, eps=rms_norm_eps, elementwise_affine=True, bias=False) + if qk_norm_for_text + else nn.Identity() + ) # 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) + self.norm_added_q = norm_cls(head_dim, eps=rms_norm_eps, elementwise_affine=True, bias=False) + self.norm_added_k = norm_cls(head_dim, eps=rms_norm_eps, elementwise_affine=True, bias=False) if processor is None: processor = self._default_processor_cls() @@ -237,12 +268,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,9 +279,12 @@ 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_cls = Cosmos3NemotronRMSNorm if hidden_act == "relu2" else RMSNorm self.self_attn = Cosmos3PackedMoTAttention( hidden_size=hidden_size, head_dim=head_dim, @@ -263,15 +292,21 @@ def __init__( 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_cls=norm_cls, ) - 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.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 + ) - 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( + self.input_layernorm = norm_cls(hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False) + self.input_layernorm_moe_gen = norm_cls(hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False) + self.post_attention_layernorm = norm_cls(hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False) + self.post_attention_layernorm_moe_gen = norm_cls( hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False ) @@ -335,10 +370,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 +398,15 @@ 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) + norm_cls = Cosmos3NemotronRMSNorm if hidden_act == "relu2" else RMSNorm + self.norm = norm_cls(hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False) + self.norm_moe_gen = norm_cls(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 ) diff --git a/src/diffusers/modular_pipelines/cosmos/before_denoise.py b/src/diffusers/modular_pipelines/cosmos/before_denoise.py index 6c1b852419c7..16a18a76ff4a 100644 --- a/src/diffusers/modular_pipelines/cosmos/before_denoise.py +++ b/src/diffusers/modular_pipelines/cosmos/before_denoise.py @@ -1,5 +1,6 @@ import copy +import numpy as np import torch from ...models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer @@ -7,7 +8,7 @@ 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 diff --git a/src/diffusers/modular_pipelines/cosmos/encoders.py b/src/diffusers/modular_pipelines/cosmos/encoders.py index 5d999b0121ba..afb8bab4df84 100644 --- a/src/diffusers/modular_pipelines/cosmos/encoders.py +++ b/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 @@ -45,6 +45,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 [ @@ -68,8 +75,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( @@ -107,6 +114,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: @@ -189,6 +198,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 [ @@ -218,8 +234,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( @@ -253,6 +269,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 diff --git a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py index 191b940768d9..6534c2ba5e9f 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py +++ b/src/diffusers/modular_pipelines/cosmos/modular_blocks_cosmos3.py @@ -57,8 +57,8 @@ class Cosmos3AutoTextEncoderStep(AutoPipelineBlocks): Width of the generated video or image in pixels. fps (`float`, *optional*, defaults to 24.0): Frame rate of the generated video. - use_system_prompt (`bool`, *optional*, defaults to True): - Whether to prepend the Cosmos3 system prompt. + use_system_prompt (`bool`, *optional*): + Whether to prepend the Cosmos3 system prompt. Defaults to the pipeline configuration. add_resolution_template (`bool`, *optional*, defaults to True): Whether to add resolution metadata to the prompt. add_duration_template (`bool`, *optional*, defaults to True): @@ -708,8 +708,8 @@ class Cosmos3OmniBlocks(SequentialPipelineBlocks): Width of the generated video or image in pixels. fps (`float`, *optional*, defaults to 24.0): Frame rate of the generated video. - use_system_prompt (`bool`, *optional*, defaults to True): - Whether to prepend the Cosmos3 system prompt. + use_system_prompt (`bool`, *optional*): + Whether to prepend the Cosmos3 system prompt. Defaults to the pipeline configuration. add_resolution_template (`bool`, *optional*, defaults to True): Whether to add resolution metadata to the prompt. add_duration_template (`bool`, *optional*, defaults to True): diff --git a/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py b/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py index 4628f6bec600..d27fe8310a10 100644 --- a/src/diffusers/modular_pipelines/cosmos/modular_pipeline.py +++ b/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) diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py index 12405755309f..b971c17773fe 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py @@ -377,8 +377,15 @@ def __init__( sound_tokenizer: Cosmos3AVAEAudioTokenizer | None = None, safety_checker: CosmosSafetyChecker | None = None, enable_safety_checker: bool = True, + default_use_system_prompt: bool = True, + use_native_flow_schedule: bool = False, ): super().__init__() + self.register_to_config( + enable_safety_checker=enable_safety_checker, + default_use_system_prompt=default_use_system_prompt, + use_native_flow_schedule=use_native_flow_schedule, + ) if enable_safety_checker: if safety_checker is None: safety_checker = CosmosSafetyChecker() @@ -1083,15 +1090,15 @@ def tokenize_prompt( height: int = 720, width: int = 1280, fps: float = 24.0, - use_system_prompt: bool = True, + use_system_prompt: bool | None = None, add_resolution_template: bool = True, add_duration_template: bool = True, action_mode: str | None = None, action_view_point: str | None = None, ) -> tuple[list[int], list[int]]: - """Apply prompt-augmentation templates and tokenize cond/uncond prompts via the Qwen2 chat template. + """Apply prompt-augmentation templates and tokenize cond/uncond prompts via the configured chat template. - This pipeline does not run a separate text encoder: the joint Cosmos3 transformer consumes raw Qwen2 token IDs + This pipeline does not run a separate text encoder: the joint Cosmos3 transformer consumes raw token IDs alongside vision (and optionally sound) tokens. When ``negative_prompt`` is ``None``, an empty string is used; the Cosmos3 docs page documents recommended @@ -1105,6 +1112,9 @@ def tokenize_prompt( Returns: ``(cond_input_ids, uncond_input_ids)`` — token-id lists for this sample. """ + if use_system_prompt is None: + use_system_prompt = self.config.default_use_system_prompt + is_image = num_frames == 1 if negative_prompt is None: @@ -1272,7 +1282,7 @@ def __call__( action: CosmosActionCondition | None = None, output_type: str = "pil", return_dict: bool = True, - use_system_prompt: bool = True, + use_system_prompt: bool | None = None, callback_on_step_end: Callable[[int, int, dict[str, Any]], None] | PipelineCallback | MultiPipelineCallbacks @@ -1355,9 +1365,9 @@ def __call__( W, C]`), `"pt"` (`torch.Tensor`, `[T, C, H, W]`), or `"latent"` (raw vision latents). return_dict (`bool`, *optional*, defaults to `True`): When `True`, returns a [`Cosmos3OmniPipelineOutput`]; otherwise a plain tuple `(video, sound)`. - use_system_prompt (`bool`, *optional*, defaults to `True`): - When `True`, prepends the mode-specific Cosmos 3 system prompt to the chat template before - tokenization. + use_system_prompt (`bool`, *optional*): + Whether to prepend the mode-specific Cosmos 3 system prompt to the chat template before tokenization. + Defaults to the pipeline's `default_use_system_prompt` configuration. callback_on_step_end (`Callable`, `PipelineCallback`, or `MultiPipelineCallbacks`, *optional*): A callback invoked at the end of each denoising step. Receives `(step_index, timestep, kwargs)` where `kwargs` is keyed by `callback_on_step_end_tensor_inputs`. @@ -1609,7 +1619,15 @@ def __call__( # 6. Set timesteps. UniPCMultistepScheduler keeps per-step state (_step_index, # model_outputs history) on the instance, so sound/action each get their own copy. - self.scheduler.set_timesteps(num_inference_steps, device=device) + if self.config.use_native_flow_schedule: + sigmas = np.linspace( + 1.0 - 1.0 / self.scheduler.config.num_train_timesteps, + 0.0, + num_inference_steps + 1, + )[:-1] + self.scheduler.set_timesteps(num_inference_steps, device=device, sigmas=sigmas) + else: + self.scheduler.set_timesteps(num_inference_steps, device=device) timesteps = self.scheduler.timesteps sound_scheduler = copy.deepcopy(self.scheduler) if sound_latents is not None else None action_scheduler = copy.deepcopy(self.scheduler) if action_latents is not None else None diff --git a/tests/models/transformers/test_models_transformer_cosmos3.py b/tests/models/transformers/test_models_transformer_cosmos3.py new file mode 100644 index 000000000000..633a125d50a0 --- /dev/null +++ b/tests/models/transformers/test_models_transformer_cosmos3.py @@ -0,0 +1,101 @@ +# coding=utf-8 +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from diffusers import Cosmos3OmniTransformer +from diffusers.models.transformers.transformer_cosmos3 import Cosmos3NemotronRMSNorm + + +def get_edge_transformer(action_gen: bool = False): + return Cosmos3OmniTransformer( + action_dim=3 if action_gen else None, + action_gen=action_gen, + backbone_type="cosmos3_edge_nemotron_dense", + head_dim=4, + hidden_act="relu2", + hidden_size=16, + intermediate_size=32, + latent_channel=1, + latent_patch_size=1, + num_attention_heads=4, + num_embodiment_domains=2, + num_hidden_layers=2, + num_key_value_heads=2, + patch_latent_dim=1, + qk_norm_for_text=False, + rms_norm_eps=1e-5, + rope_theta=1e8, + temporal_compression_factor=4, + vocab_size=32, + ) + + +def test_cosmos3_edge_uses_nemotron_parameter_layout(): + transformer = get_edge_transformer(action_gen=True) + state_dict = transformer.state_dict() + + assert transformer.config.backbone_type == "cosmos3_edge_nemotron_dense" + assert transformer.config.temporal_compression_factor == 4 + assert not any("gate_proj" in key for key in state_dict) + assert not any(".norm_q." in key or ".norm_k." in key for key in state_dict) + assert "layers.0.self_attn.norm_added_q.weight" in state_dict + assert "layers.0.self_attn.norm_added_k.weight" in state_dict + assert "layers.0.mlp.up_proj.weight" in state_dict + assert "layers.0.mlp.down_proj.weight" in state_dict + assert "action_proj_in.fc.weight" in state_dict + assert "action_proj_out.fc.weight" in state_dict + + +def test_cosmos3_edge_transformer_runs_action_workflow(): + transformer = get_edge_transformer(action_gen=True).eval() + + with torch.no_grad(): + prediction, sound_prediction, action_prediction = transformer( + input_ids=torch.tensor([1, 2]), + text_indexes=torch.tensor([0, 1]), + position_ids=torch.tensor([[0, 0, 0, 0], [0, 0, 0, 0], [0, 1, 0, 0]]), + und_len=2, + sequence_length=4, + vision_tokens=[torch.randn(1, 1, 1, 1, 1)], + vision_token_shapes=[(1, 1, 1)], + vision_sequence_indexes=torch.tensor([2]), + vision_mse_loss_indexes=torch.tensor([2]), + vision_timesteps=torch.tensor([1]), + vision_noisy_frame_indexes=[torch.tensor([0])], + action_tokens=[torch.randn(1, 3)], + action_token_shapes=[(1, 1, 1)], + action_sequence_indexes=torch.tensor([3]), + action_mse_loss_indexes=torch.tensor([3]), + action_timesteps=torch.tensor([1]), + action_noisy_frame_indexes=[torch.tensor([0])], + action_domain_ids=[torch.tensor(0)], + ) + + assert prediction[0].shape == (1, 1, 1, 1, 1) + assert sound_prediction is None + assert action_prediction[0].shape == (1, 3) + + +def test_cosmos3_nemotron_rms_norm_multiplies_in_float32(): + hidden_states = torch.randn(2, 3, 8, dtype=torch.bfloat16) + norm = Cosmos3NemotronRMSNorm(8, eps=1e-5).bfloat16() + norm.weight.data.copy_(torch.randn(8, dtype=torch.bfloat16)) + + expected = hidden_states.float() + expected = expected * torch.rsqrt(expected.pow(2).mean(-1, keepdim=True) + 1e-5) + expected = (norm.weight.float() * expected).to(hidden_states.dtype) + + torch.testing.assert_close(norm(hidden_states), expected, rtol=0, atol=0) diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py index 5297a441cebf..c4cbbb50ca10 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py @@ -19,15 +19,17 @@ import torch from PIL import Image -from diffusers import Cosmos3OmniTransformer, ModularPipeline +from diffusers import Cosmos3OmniTransformer, ModularPipeline, UniPCMultistepScheduler from diffusers.modular_pipelines import Cosmos3OmniBlocks, Cosmos3OmniModularPipeline from diffusers.modular_pipelines.cosmos.before_denoise import ( Cosmos3ActionDenoiseInputStep, Cosmos3ActionPackSequenceStep, + Cosmos3SetTimestepsStep, Cosmos3SoundDenoiseInputStep, Cosmos3VisionDenoiseInputStep, Cosmos3VisionPackSequenceStep, ) +from diffusers.modular_pipelines.cosmos.encoders import Cosmos3TextEncoderStep from diffusers.modular_pipelines.modular_pipeline import PipelineState from ...testing_utils import torch_device @@ -317,6 +319,60 @@ def test_cosmos3_denoise_input_steps_assemble_modality_segments(): assert state.get("uncond_packed_static") is None +def test_cosmos3_text_step_uses_pipeline_system_prompt_and_safety_configs(monkeypatch): + text_encoder = Cosmos3TextEncoderStep() + pipe = Cosmos3OmniModularPipeline( + blocks=text_encoder, + config_dict={"default_use_system_prompt": False, "enable_safety_checker": False}, + ) + captured = {} + + def tokenize_prompt(*args, **kwargs): + captured["use_system_prompt"] = kwargs["use_system_prompt"] + return torch.zeros((1, 1), dtype=torch.long), torch.zeros((1, 1), dtype=torch.long) + + monkeypatch.setattr(pipe, "tokenize_prompt", tokenize_prompt) + + state = PipelineState() + state.set("prompt", "A small robot moves across a table.") + state.set("negative_prompt", "") + state.set("num_frames", 5) + state.set("height", 32) + state.set("width", 32) + state.set("fps", 24.0) + state.set("add_resolution_template", True) + state.set("add_duration_template", True) + _, state = text_encoder(pipe, state) + + assert captured["use_system_prompt"] is False + assert not pipe.requires_safety_checker + assert state.get("cond_input_ids").shape == (1, 1) + + +def test_cosmos3_native_flow_schedule_uses_edge_sigma_grid(monkeypatch): + set_timesteps = Cosmos3SetTimestepsStep() + pipe = Cosmos3OmniModularPipeline(blocks=set_timesteps, config_dict={"use_native_flow_schedule": True}) + scheduler = UniPCMultistepScheduler(num_train_timesteps=100, use_flow_sigmas=True) + pipe.update_components(scheduler=scheduler) + captured = {} + + def capture_set_timesteps(num_inference_steps, device=None, sigmas=None): + captured["num_inference_steps"] = num_inference_steps + captured["device"] = device + captured["sigmas"] = sigmas + scheduler.timesteps = torch.arange(num_inference_steps) + + monkeypatch.setattr(scheduler, "set_timesteps", capture_set_timesteps) + + state = PipelineState() + state.set("num_inference_steps", 4) + _, state = set_timesteps(pipe, state) + + assert captured["num_inference_steps"] == 4 + assert captured["sigmas"] == pytest.approx([0.99, 0.7425, 0.495, 0.2475]) + assert state.get("timesteps").tolist() == [0, 1, 2, 3] + + def test_cosmos3_modular_model_index_takes_precedence(tmp_path): (tmp_path / "model_index.json").write_text(json.dumps({"_class_name": "Cosmos3OmniDiffusersPipeline"})) (tmp_path / "modular_model_index.json").write_text( @@ -338,9 +394,22 @@ def test_cosmos3_modular_model_index_takes_precedence(tmp_path): def test_cosmos3_model_index_fallback_resolves_modular_pipeline(tmp_path): - (tmp_path / "model_index.json").write_text(json.dumps({"_class_name": "Cosmos3OmniPipeline"})) + (tmp_path / "model_index.json").write_text( + json.dumps( + { + "_class_name": "Cosmos3OmniPipeline", + "default_use_system_prompt": False, + "enable_safety_checker": False, + "use_native_flow_schedule": True, + } + ) + ) pipe = ModularPipeline.from_pretrained(str(tmp_path)) assert isinstance(pipe, Cosmos3OmniModularPipeline) assert isinstance(pipe.blocks, Cosmos3OmniBlocks) + assert not pipe.config.default_use_system_prompt + assert not pipe.config.enable_safety_checker + assert pipe.config.use_native_flow_schedule + assert not pipe.requires_safety_checker diff --git a/tests/pipelines/cosmos/test_cosmos3.py b/tests/pipelines/cosmos/test_cosmos3.py new file mode 100644 index 000000000000..1456a8dbc145 --- /dev/null +++ b/tests/pipelines/cosmos/test_cosmos3.py @@ -0,0 +1,96 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +from types import SimpleNamespace + +import torch + +from diffusers import Cosmos3OmniPipeline, Cosmos3OmniTransformer, UniPCMultistepScheduler + + +class DummyTokenizer: + eos_token_id = 11 + + def __init__(self): + self.conversations = [] + + def convert_tokens_to_ids(self, token): + return {"<|vision_start|>": 20}.get(token, 0) + + def apply_chat_template(self, conversations, **kwargs): + self.conversations.append(conversations) + return SimpleNamespace(input_ids=[1, 2]) + + +class DummyVAE(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.zeros(())) + self.config = SimpleNamespace( + latents_mean=[0.0], + latents_std=[1.0], + scale_factor_spatial=16, + scale_factor_temporal=4, + ) + + @property + def dtype(self): + return self.weight.dtype + + +def get_dummy_pipeline(**kwargs): + transformer = Cosmos3OmniTransformer( + hidden_size=16, + intermediate_size=32, + head_dim=4, + num_attention_heads=4, + num_key_value_heads=2, + num_hidden_layers=1, + latent_channel=1, + latent_patch_size=1, + patch_latent_dim=1, + vocab_size=32, + ) + return Cosmos3OmniPipeline( + transformer=transformer, + text_tokenizer=DummyTokenizer(), + vae=DummyVAE(), + scheduler=UniPCMultistepScheduler(), + enable_safety_checker=False, + **kwargs, + ) + + +def test_cosmos3_pipeline_saves_edge_configuration(tmp_path): + pipeline = get_dummy_pipeline(default_use_system_prompt=False, use_native_flow_schedule=True) + + assert not pipeline.config.enable_safety_checker + assert not pipeline.config.default_use_system_prompt + assert pipeline.config.use_native_flow_schedule + assert pipeline.safety_checker is None + + pipeline.save_config(tmp_path) + model_index = json.loads((tmp_path / "model_index.json").read_text()) + assert model_index["enable_safety_checker"] is False + assert model_index["default_use_system_prompt"] is False + assert model_index["use_native_flow_schedule"] is True + + +def test_cosmos3_tokenize_prompt_uses_checkpoint_system_prompt_default(): + pipeline = get_dummy_pipeline(default_use_system_prompt=False) + + pipeline.tokenize_prompt("A prompt", num_frames=1, add_resolution_template=False) + + assert all(conversation[0]["role"] == "user" for conversation in pipeline.text_tokenizer.conversations) From c811ca7c39fdfb82476932eb4594e7056b6253de Mon Sep 17 00:00:00 2001 From: Atharva Joshi Date: Mon, 13 Jul 2026 13:44:52 -0700 Subject: [PATCH 2/8] Remove Cosmos3 conversion script --- scripts/convert_cosmos3_to_diffusers.py | 628 ------------------ .../test_models_autoencoder_cosmos3_audio.py | 54 -- 2 files changed, 682 deletions(-) delete mode 100644 scripts/convert_cosmos3_to_diffusers.py diff --git a/scripts/convert_cosmos3_to_diffusers.py b/scripts/convert_cosmos3_to_diffusers.py deleted file mode 100644 index 42f230c4c388..000000000000 --- a/scripts/convert_cosmos3_to_diffusers.py +++ /dev/null @@ -1,628 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -""" -Convert a Cosmos3 DCP checkpoint to diffusers format. - -Example: -CUDA_VISIBLE_DEVICES=0 python scripts/convert_cosmos3_to_diffusers.py \ - --checkpoint-path Cosmos3-Nano \ - --output converted/cosmos3-nano-pipeline \ - --save-pipeline -""" - -import argparse -import contextlib -import json -import pathlib -import re - -import torch - - -DEFAULT_SOUND_TOKENIZER_CONFIG = { - "model_type": "autoencoder_v2", - "sampling_rate": 48000, - "stereo": True, - "use_wav_as_input": True, - "normalize_volume": True, - "hop_size": 1920, - "input_channels": 1, - "enc_type": "spec_convnext", - "enc_dim": 192, - "enc_intermediate_dim": 768, - "enc_num_layers": 12, - "enc_num_blocks": 2, - "enc_n_fft": 64, - "enc_hop_length": 16, - "enc_latent_dim": 128, - "enc_c_mults": [1, 2, 4], - "enc_strides": [4, 5, 6], - "enc_identity_init": False, - "enc_use_snake": True, - "dec_type": "oobleck", - "vocoder_input_dim": 64, - "dec_dim": 320, - "dec_c_mults": [1, 2, 4, 8, 16], - "dec_strides": [2, 4, 5, 6, 8], - "dec_use_snake": True, - "dec_final_tanh": False, - "dec_out_channels": 2, - "dec_anti_aliasing": False, - "dec_use_nearest_upsample": False, - "dec_use_tanh_at_final": False, - "bottleneck_type": "vae", - "bottleneck": {"type": "vae"}, - "activation": "snakebeta", - "snake_logscale": True, - "anti_aliasing": False, - "use_cuda_kernel": False, - "causal": False, - "padding_mode": "zeros", - "latent_mean": None, - "latent_std": None, -} - - -def _get_config_value(*configs, name, default=None): - for config in configs: - if config is None: - continue - if hasattr(config, name): - value = getattr(config, name) - if value is not None: - return value - if isinstance(config, dict) and config.get(name) is not None: - return config[name] - return default - - -def _load_sound_tokenizer_state_dict(checkpoint_path: pathlib.Path) -> dict[str, torch.Tensor]: - if checkpoint_path.suffix == ".safetensors": - try: - from safetensors.torch import load_file - except ImportError as exc: - raise ImportError("Loading AVAE .safetensors checkpoints requires safetensors.") from exc - checkpoint = load_file(str(checkpoint_path), device="cpu") - else: - checkpoint = torch.load(checkpoint_path, map_location="cpu") - - if not isinstance(checkpoint, dict): - raise TypeError(f"AVAE checkpoint must be a dict, got {type(checkpoint)!r}.") - - for key in ("generator", "state_dict", "model"): - value = checkpoint.get(key) - if isinstance(value, dict): - checkpoint = value - break - - state_dict = { - key: value.detach().cpu().contiguous() for key, value in checkpoint.items() if isinstance(value, torch.Tensor) - } - if not state_dict: - raise RuntimeError(f"No tensor state dict found in AVAE checkpoint keys: {list(checkpoint.keys())[:16]}") - return state_dict - - -def _load_sound_tokenizer_config(config_path: pathlib.Path | None, fallback_config_path: pathlib.Path) -> dict: - selected_config_path = config_path - if selected_config_path is None and fallback_config_path.exists(): - selected_config_path = fallback_config_path - if selected_config_path is None: - return dict(DEFAULT_SOUND_TOKENIZER_CONFIG) - with open(selected_config_path, encoding="utf-8") as f: - return json.load(f) - - -_SOUND_TOKENIZER_PER_KEY_PREFIXES = ("module.", "generator.", "model.", "state_dict.") -_SOUND_TOKENIZER_RES_UNIT_INNER_NAMES = {0: "snake1", 1: "conv1", 2: "snake2", 3: "conv2"} - - -def _sound_tokenizer_strip_per_key_prefixes(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - out = dict(state_dict) - changed = True - while changed: - changed = False - for prefix in _SOUND_TOKENIZER_PER_KEY_PREFIXES: - if any(key.startswith(prefix) for key in out): - out = {(key[len(prefix) :] if key.startswith(prefix) else key): value for key, value in out.items()} - changed = True - break - if any(key.startswith(("decoder.", "encoder.", "bottleneck.")) for key in out): - break - return out - - -def _sound_tokenizer_filter_supported_modules(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - return { - key: value for key, value in state_dict.items() if key.startswith("encoder.") or key.startswith("decoder.") - } - - -def _sound_tokenizer_infer_num_blocks(state_dict: dict[str, torch.Tensor]) -> int: - block_indices: set[int] = set() - for key in state_dict: - match = re.match(r"decoder\.layers\.(\d+)\.layers\.\d+\.", key) - if match: - block_indices.add(int(match.group(1))) - return len(block_indices) - - -def _sound_tokenizer_remap_flat_layout(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """Convert legacy AVAE `decoder.layers.*` keys to OobleckDecoder attribute keys.""" - if not any(re.match(r"decoder\.layers\.\d+\.", key) for key in state_dict): - return state_dict - - num_blocks = _sound_tokenizer_infer_num_blocks(state_dict) - if num_blocks == 0: - raise RuntimeError("Detected flat `decoder.layers.*` layout but no decoder blocks were found; cannot remap.") - snake1_idx = num_blocks + 1 - conv2_idx = num_blocks + 2 - - def _remap(key: str) -> str: - match = re.fullmatch(r"decoder\.layers\.(\d+)\.layers\.(\d+)\.layers\.(\d+)\.(.+)", key) - if match: - block_n, res_n, inner_n, rest = ( - int(match.group(1)), - int(match.group(2)), - int(match.group(3)), - match.group(4), - ) - if res_n not in (2, 3, 4): - raise RuntimeError(f"Unexpected residual position {res_n} in {key!r}.") - inner_name = _SOUND_TOKENIZER_RES_UNIT_INNER_NAMES.get(inner_n) - if inner_name is None: - raise RuntimeError(f"Unexpected residual inner index {inner_n} in {key!r}.") - return f"decoder.block.{block_n - 1}.res_unit{res_n - 1}.{inner_name}.{rest}" - - match = re.fullmatch(r"decoder\.layers\.(\d+)\.layers\.(\d+)\.(.+)", key) - if match: - block_n, sub_n, rest = int(match.group(1)), int(match.group(2)), match.group(3) - block_idx = block_n - 1 - if sub_n == 0: - return f"decoder.block.{block_idx}.snake1.{rest}" - if sub_n == 1: - return f"decoder.block.{block_idx}.conv_t1.{rest}" - raise RuntimeError(f"Unexpected decoder block sub-index {sub_n} in {key!r}.") - - match = re.fullmatch(r"decoder\.layers\.(\d+)\.(.+)", key) - if match: - layer_n, rest = int(match.group(1)), match.group(2) - if layer_n == 0: - return f"decoder.conv1.{rest}" - if layer_n == snake1_idx: - return f"decoder.snake1.{rest}" - if layer_n == conv2_idx: - return f"decoder.conv2.{rest}" - raise RuntimeError( - f"Unexpected decoder leaf layer index {layer_n} (expected 0, {snake1_idx}, or {conv2_idx}) in {key!r}." - ) - - return key - - return {_remap(key): value for key, value in state_dict.items()} - - -def _sound_tokenizer_reshape_snake_params(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - out: dict[str, torch.Tensor] = {} - for key, value in state_dict.items(): - if ( - key.startswith(("encoder.", "decoder.")) - and (key.endswith(".alpha") or key.endswith(".beta")) - and value.ndim == 1 - ): - value = value.unsqueeze(0).unsqueeze(-1).contiguous() - out[key] = value - return out - - -def _sound_tokenizer_reapply_weight_norm(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """Reconstruct weight-norm parameters if the source checkpoint has folded conv weights.""" - out = dict(state_dict) - candidate_keys = [ - key - for key in state_dict - if key.endswith(".weight") - and ( - any(f".{layer}." in key for layer in ("conv1", "conv2", "conv_t1")) - or re.fullmatch(r"encoder\.layers\.\d+\.weight", key) - ) - ] - for key in candidate_keys: - stem = key[: -len(".weight")] - weight_g_key = f"{stem}.weight_g" - weight_v_key = f"{stem}.weight_v" - if weight_g_key in state_dict or weight_v_key in state_dict: - continue - weight = state_dict[key] - norm_dims = tuple(range(1, weight.ndim)) - out.pop(key) - out[weight_g_key] = weight.norm(p=2, dim=norm_dims, keepdim=True).contiguous() - out[weight_v_key] = weight.contiguous() - return out - - -def _remap_avae_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: - """Convert a legacy AVAE state dict into the Cosmos3AVAEAudioTokenizer state dict.""" - state_dict = _sound_tokenizer_strip_per_key_prefixes(state_dict) - state_dict = _sound_tokenizer_filter_supported_modules(state_dict) - if not state_dict: - raise RuntimeError("Sound tokenizer state dict has no `encoder.*` or `decoder.*` keys after prefix stripping.") - if not any(key.startswith("decoder.") for key in state_dict): - raise RuntimeError("Sound tokenizer state dict has no `decoder.*` keys after prefix stripping.") - state_dict = _sound_tokenizer_remap_flat_layout(state_dict) - state_dict = _sound_tokenizer_reshape_snake_params(state_dict) - state_dict = _sound_tokenizer_reapply_weight_norm(state_dict) - if any(re.match(r"decoder\.layers\.\d+", key) for key in state_dict): - raise RuntimeError("Flat `decoder.layers.*` keys remain after remap; conversion is incomplete.") - return state_dict - - -def _build_sound_tokenizer( - checkpoint_path: pathlib.Path, - config_path: pathlib.Path | None, -): - from diffusers.models.autoencoders.autoencoder_cosmos3_audio import Cosmos3AVAEAudioTokenizer - - config = _load_sound_tokenizer_config(config_path, fallback_config_path=pathlib.Path()) - print(f"Loading AVAE sound tokenizer weights from {checkpoint_path} …") - raw_state_dict = _load_sound_tokenizer_state_dict(checkpoint_path) - state_dict = _remap_avae_state_dict(raw_state_dict) - has_encoder = any(key.startswith("encoder.") for key in state_dict) - print( - f" Remapped {len(raw_state_dict)} → {len(state_dict)} tokenizer keys " - f"({'encoder+decoder' if has_encoder else 'decoder-only'})." - ) - - sound_tokenizer = Cosmos3AVAEAudioTokenizer( - model_type=config.get("model_type", DEFAULT_SOUND_TOKENIZER_CONFIG["model_type"]), - sampling_rate=config.get("sampling_rate", DEFAULT_SOUND_TOKENIZER_CONFIG["sampling_rate"]), - stereo=config.get("stereo", DEFAULT_SOUND_TOKENIZER_CONFIG["stereo"]), - use_wav_as_input=config.get("use_wav_as_input", DEFAULT_SOUND_TOKENIZER_CONFIG["use_wav_as_input"]), - normalize_volume=config.get("normalize_volume", DEFAULT_SOUND_TOKENIZER_CONFIG["normalize_volume"]), - hop_size=config.get("hop_size", DEFAULT_SOUND_TOKENIZER_CONFIG["hop_size"]), - input_channels=config.get("input_channels", DEFAULT_SOUND_TOKENIZER_CONFIG["input_channels"]), - enc_type=config.get("enc_type", DEFAULT_SOUND_TOKENIZER_CONFIG["enc_type"]), - enc_dim=config.get("enc_dim", DEFAULT_SOUND_TOKENIZER_CONFIG["enc_dim"]), - enc_intermediate_dim=config.get( - "enc_intermediate_dim", DEFAULT_SOUND_TOKENIZER_CONFIG["enc_intermediate_dim"] - ), - enc_num_layers=config.get("enc_num_layers", DEFAULT_SOUND_TOKENIZER_CONFIG["enc_num_layers"]), - enc_num_blocks=config.get("enc_num_blocks", DEFAULT_SOUND_TOKENIZER_CONFIG["enc_num_blocks"]), - enc_n_fft=config.get("enc_n_fft", DEFAULT_SOUND_TOKENIZER_CONFIG["enc_n_fft"]), - enc_hop_length=config.get("enc_hop_length", DEFAULT_SOUND_TOKENIZER_CONFIG["enc_hop_length"]), - enc_latent_dim=config.get("enc_latent_dim", DEFAULT_SOUND_TOKENIZER_CONFIG["enc_latent_dim"]), - enc_c_mults=tuple(config.get("enc_c_mults", DEFAULT_SOUND_TOKENIZER_CONFIG["enc_c_mults"])), - enc_strides=tuple(config.get("enc_strides", DEFAULT_SOUND_TOKENIZER_CONFIG["enc_strides"])), - enc_identity_init=config.get("enc_identity_init", DEFAULT_SOUND_TOKENIZER_CONFIG["enc_identity_init"]), - enc_use_snake=config.get("enc_use_snake", DEFAULT_SOUND_TOKENIZER_CONFIG["enc_use_snake"]), - dec_type=config.get("dec_type", DEFAULT_SOUND_TOKENIZER_CONFIG["dec_type"]), - vocoder_input_dim=config.get("vocoder_input_dim", DEFAULT_SOUND_TOKENIZER_CONFIG["vocoder_input_dim"]), - dec_dim=config.get("dec_dim", DEFAULT_SOUND_TOKENIZER_CONFIG["dec_dim"]), - dec_c_mults=tuple(config.get("dec_c_mults", DEFAULT_SOUND_TOKENIZER_CONFIG["dec_c_mults"])), - dec_strides=tuple(config.get("dec_strides", DEFAULT_SOUND_TOKENIZER_CONFIG["dec_strides"])), - dec_use_snake=config.get("dec_use_snake", DEFAULT_SOUND_TOKENIZER_CONFIG["dec_use_snake"]), - dec_final_tanh=config.get("dec_final_tanh", False), - dec_out_channels=config.get("dec_out_channels", DEFAULT_SOUND_TOKENIZER_CONFIG["dec_out_channels"]), - dec_anti_aliasing=config.get("dec_anti_aliasing", DEFAULT_SOUND_TOKENIZER_CONFIG["dec_anti_aliasing"]), - dec_use_nearest_upsample=config.get( - "dec_use_nearest_upsample", DEFAULT_SOUND_TOKENIZER_CONFIG["dec_use_nearest_upsample"] - ), - dec_use_tanh_at_final=config.get( - "dec_use_tanh_at_final", DEFAULT_SOUND_TOKENIZER_CONFIG["dec_use_tanh_at_final"] - ), - bottleneck_type=config.get("bottleneck_type", DEFAULT_SOUND_TOKENIZER_CONFIG["bottleneck_type"]), - bottleneck=config.get("bottleneck", DEFAULT_SOUND_TOKENIZER_CONFIG["bottleneck"]), - activation=config.get("activation", DEFAULT_SOUND_TOKENIZER_CONFIG["activation"]), - snake_logscale=config.get("snake_logscale", DEFAULT_SOUND_TOKENIZER_CONFIG["snake_logscale"]), - anti_aliasing=config.get("anti_aliasing", DEFAULT_SOUND_TOKENIZER_CONFIG["anti_aliasing"]), - use_cuda_kernel=config.get("use_cuda_kernel", DEFAULT_SOUND_TOKENIZER_CONFIG["use_cuda_kernel"]), - causal=config.get("causal", DEFAULT_SOUND_TOKENIZER_CONFIG["causal"]), - padding_mode=config.get("padding_mode", DEFAULT_SOUND_TOKENIZER_CONFIG["padding_mode"]), - latent_mean=config.get("latent_mean", DEFAULT_SOUND_TOKENIZER_CONFIG["latent_mean"]), - latent_std=config.get("latent_std", DEFAULT_SOUND_TOKENIZER_CONFIG["latent_std"]), - encoder_enabled=has_encoder, - ) - load_result = sound_tokenizer.load_state_dict(state_dict, strict=True) - if load_result.missing_keys or load_result.unexpected_keys: - raise RuntimeError( - "Cosmos3 AVAE sound tokenizer load did not match strictly: " - f"missing={load_result.missing_keys}, unexpected={load_result.unexpected_keys}." - ) - return sound_tokenizer - - -@contextlib.contextmanager -def _skip_source_sound_tokenizer_load(omni_mot_model_cls): - original_set_up_tokenizers = omni_mot_model_cls.set_up_tokenizers - - def set_up_tokenizers_without_sound(self): - if not getattr(self.config, "sound_gen", False): - return original_set_up_tokenizers(self) - - sound_gen = self.config.sound_gen - self.config.sound_gen = False - try: - return original_set_up_tokenizers(self) - finally: - self.config.sound_gen = sound_gen - - omni_mot_model_cls.set_up_tokenizers = set_up_tokenizers_without_sound - try: - yield - finally: - omni_mot_model_cls.set_up_tokenizers = original_set_up_tokenizers - - -def main(): - from cosmos3.common.init import init_script - - init_script() - - from accelerate import init_empty_weights - from cosmos3.args import _CHECKPOINTS - from cosmos3.model import Cosmos3OmniModel - from projects.cosmos3.vfm.models.omni_mot_model import OmniMoTModel - from transformers import AutoTokenizer - - from diffusers import AutoencoderKLWan, UniPCMultistepScheduler - from diffusers.models.transformers.transformer_cosmos3 import Cosmos3OmniTransformer - from diffusers.pipelines.cosmos.pipeline_cosmos3_omni import Cosmos3OmniPipeline - - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--checkpoint-path", - default="Cosmos3-Nano", - help="Named checkpoint (e.g. 'Cosmos3-Nano') or path to DCP checkpoint dir.", - ) - parser.add_argument("--output", required=True, help="Directory to save the converted diffusers model.") - parser.add_argument( - "--save-pipeline", - action="store_true", - help="Save the full pipeline (transformer + VAE + tokenizer + scheduler).", - ) - parser.add_argument( - "--dtype", default="bf16", choices=["fp32", "fp16", "bf16"], help="Dtype to save the transformer in." - ) - parser.add_argument( - "--sound-tokenizer-path", help="Optional AVAE sound tokenizer checkpoint to save under sound_tokenizer/." - ) - parser.add_argument( - "--sound-tokenizer-config-path", help="Optional AVAE config JSON to save under sound_tokenizer/config.json." - ) - parser.add_argument( - "--include-sound-tokenizer", - action="store_true", - help="Require saving sound_tokenizer/ even if the source transformer is video-only.", - ) - args = parser.parse_args() - - dtype = {"fp32": torch.float32, "fp16": torch.float16, "bf16": torch.bfloat16}[args.dtype] - sound_tokenizer_path = ( - pathlib.Path(args.sound_tokenizer_path).expanduser().absolute() if args.sound_tokenizer_path else None - ) - sound_tokenizer_config_path = ( - pathlib.Path(args.sound_tokenizer_config_path).expanduser().absolute() - if args.sound_tokenizer_config_path - else None - ) - if args.include_sound_tokenizer and sound_tokenizer_path is None: - raise ValueError("Sound tokenizer output was requested, but --sound-tokenizer-path was not provided.") - if sound_tokenizer_path is not None and not sound_tokenizer_path.exists(): - raise FileNotFoundError(f"Sound tokenizer checkpoint not found: {sound_tokenizer_path}") - if sound_tokenizer_config_path is not None and not sound_tokenizer_config_path.exists(): - raise FileNotFoundError(f"Sound tokenizer config not found: {sound_tokenizer_config_path}") - - checkpoint_name = args.checkpoint_path - if checkpoint_name in _CHECKPOINTS: - checkpoint_path = pathlib.Path(_CHECKPOINTS[checkpoint_name].download()) - else: - checkpoint_path = pathlib.Path(checkpoint_name).expanduser().absolute() - print(f"Resolved checkpoint path: {checkpoint_path}") - - print("Instantiating model and loading weights from DCP checkpoint …") - print("Skipping source AVAE tokenizer instantiation during converter-only model load …") - with _skip_source_sound_tokenizer_load(OmniMoTModel): - _tmp = Cosmos3OmniModel.from_pretrained_dcp(checkpoint_path).model - - # Extract network components and architecture config from DCP model - language_model = _tmp.net.language_model - vae2llm = _tmp.net.vae2llm - llm2vae = _tmp.net.llm2vae - time_embedder = _tmp.net.time_embedder - lm_cfg = _tmp.net.language_model.config - net_cfg = _tmp.net.config - model_cfg = _tmp.config - patch_latent_dim = _tmp.net.patch_latent_dim - hidden_size = _tmp.net.hidden_size - num_attention_heads = _tmp.net.num_heads - num_key_value_heads = _tmp.net.num_kv_heads - head_dim = _tmp.net.head_dim - num_hidden_layers = _tmp.net.num_hidden_layers - latent_patch_size = _tmp.net.latent_patch_size - latent_channel = _tmp.net.latent_channel - timestep_scale = _tmp.net.timestep_scale - base_fps = int(net_cfg.base_fps) - enable_fps_modulation = net_cfg.enable_fps_modulation - unified_3d_mrope_reset_spatial_ids = _tmp.config.diffusion_expert_config.unified_3d_mrope_reset_spatial_ids - unified_3d_mrope_temporal_modality_margin = ( - _tmp.config.diffusion_expert_config.unified_3d_mrope_temporal_modality_margin - ) - sound2llm = getattr(_tmp.net, "sound2llm", None) - llm2sound = getattr(_tmp.net, "llm2sound", None) - sound_modality_embed = getattr(_tmp.net, "sound_modality_embed", None) - has_sound_projection_weights = any(module is not None for module in (sound2llm, llm2sound, sound_modality_embed)) - sound_gen = bool( - _get_config_value(net_cfg, model_cfg, name="sound_gen", default=False) or has_sound_projection_weights - ) - sound_dim = _get_config_value(net_cfg, model_cfg, name="sound_dim", default=None) - if sound_dim is None and sound2llm is not None: - sound_dim = sound2llm.in_features - sound_latent_fps = _get_config_value(net_cfg, model_cfg, name="sound_latent_fps", default=25.0) - if sound_gen: - missing_sound_modules = [ - name - for name, module in ( - ("sound2llm", sound2llm), - ("llm2sound", llm2sound), - ("sound_modality_embed", sound_modality_embed), - ) - if module is None - ] - if missing_sound_modules: - raise RuntimeError( - "Source checkpoint is configured for sound generation but is missing " - f"sound projection weights: {missing_sound_modules}." - ) - if sound_dim is None: - raise RuntimeError("Source checkpoint is configured for sound generation but sound_dim is missing.") - del _tmp - torch.cuda.empty_cache() - - # Init diffusers Cosmos3OmniTransformer with full architecture config from DCP - with init_empty_weights(): - transformer = Cosmos3OmniTransformer( - attention_bias=lm_cfg.attention_bias, - attention_dropout=lm_cfg.attention_dropout, - base_fps=base_fps, - enable_fps_modulation=enable_fps_modulation, - head_dim=head_dim, - hidden_size=hidden_size, - intermediate_size=lm_cfg.intermediate_size, - latent_channel=latent_channel, - latent_patch_size=latent_patch_size, - num_attention_heads=num_attention_heads, - num_hidden_layers=num_hidden_layers, - num_key_value_heads=num_key_value_heads, - patch_latent_dim=patch_latent_dim, - rms_norm_eps=lm_cfg.rms_norm_eps, - rope_scaling=lm_cfg.rope_scaling, - rope_theta=lm_cfg.rope_theta, - sound_dim=sound_dim, - sound_gen=sound_gen, - sound_latent_fps=sound_latent_fps, - timestep_scale=timestep_scale, - unified_3d_mrope_reset_spatial_ids=unified_3d_mrope_reset_spatial_ids, - unified_3d_mrope_temporal_modality_margin=unified_3d_mrope_temporal_modality_margin, - vocab_size=lm_cfg.vocab_size, - ) - # The source language_model nests its transformer stack under a `model.` attribute - # (HF Qwen-style). Diffusers Cosmos3OmniTransformer holds those layers flat, so - # strip the leading `model.` prefix from the language-model state-dict keys. - state_dict = { - (k[len("model.") :] if k.startswith("model.") else k): v for k, v in language_model.state_dict().items() - } - # Remap PackedAttentionMoT attribute names from the source (Qwen-style q_proj/k_proj/... - # plus cosmos-specific *_moe_gen) to the diffusers AttentionModuleMixin canonical names. - # Order matters: the *_moe_gen substrings must be substituted before the plain ones. - _ATTN_KEY_REMAP = [ - (".q_proj_moe_gen.", ".add_q_proj."), - (".k_proj_moe_gen.", ".add_k_proj."), - (".v_proj_moe_gen.", ".add_v_proj."), - (".o_proj_moe_gen.", ".to_add_out."), - (".q_norm_moe_gen.", ".norm_added_q."), - (".k_norm_moe_gen.", ".norm_added_k."), - (".q_proj.", ".to_q."), - (".k_proj.", ".to_k."), - (".v_proj.", ".to_v."), - (".o_proj.", ".to_out."), - (".q_norm.", ".norm_q."), - (".k_norm.", ".norm_k."), - ] - remapped_state_dict: dict[str, torch.Tensor] = {} - for k, v in state_dict.items(): - for old, new in _ATTN_KEY_REMAP: - if old in k: - k = k.replace(old, new) - break - remapped_state_dict[k] = v - state_dict = remapped_state_dict - for k, v in vae2llm.state_dict().items(): - state_dict[f"proj_in.{k}"] = v - for k, v in llm2vae.state_dict().items(): - state_dict[f"proj_out.{k}"] = v - _TIME_EMBEDDER_REMAP = { - "mlp.0.weight": "linear_1.weight", - "mlp.0.bias": "linear_1.bias", - "mlp.2.weight": "linear_2.weight", - "mlp.2.bias": "linear_2.bias", - } - for k, v in time_embedder.state_dict().items(): - state_dict[f"time_embedder.{_TIME_EMBEDDER_REMAP[k]}"] = v - if sound_gen: - for k, v in sound2llm.state_dict().items(): - state_dict[f"audio_proj_in.{k}"] = v - for k, v in llm2sound.state_dict().items(): - state_dict[f"audio_proj_out.{k}"] = v - state_dict["audio_modality_embed"] = sound_modality_embed - transformer.load_state_dict(state_dict, strict=True, assign=True) - del ( - language_model, - vae2llm, - llm2vae, - time_embedder, - sound2llm, - llm2sound, - sound_modality_embed, - state_dict, - ) - torch.cuda.empty_cache() - - transformer = transformer.to(dtype=dtype) - - output_dir = pathlib.Path(args.output) - output_dir.mkdir(parents=True, exist_ok=True) - include_sound_tokenizer = ( - args.include_sound_tokenizer or sound_tokenizer_path is not None or (sound_gen and args.save_pipeline) - ) - if include_sound_tokenizer and sound_tokenizer_path is None: - raise ValueError( - "The source checkpoint is configured for sound generation, so --sound-tokenizer-path " - "is required when saving a full pipeline." - ) - - if args.save_pipeline: - text_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-VL-8B-Instruct") - - diffusers_vae = AutoencoderKLWan.from_pretrained( - "Wan-AI/Wan2.2-TI2V-5B-Diffusers", subfolder="vae", torch_dtype=torch.bfloat16 - ) - sound_tokenizer = None - if include_sound_tokenizer: - assert sound_tokenizer_path is not None - sound_tokenizer = _build_sound_tokenizer(sound_tokenizer_path, sound_tokenizer_config_path) - - # Karras schedule approximating FlowUniPCMultistepScheduler with shift=5, 35 steps. - # Measured from that schedule: first flow-sigma=0.9998, last flow-sigma=0.1281. - # EDM sigma = flow_sigma / (1 - flow_sigma), so: - # sigma_max = 0.9998 / 0.0002 = 4999 (but capped at 200 to avoid duplicate - # integer timesteps from Karras clustering near the top) - # sigma_min = 0.1281 / (1 - 0.1281) = 0.1281 / 0.8719 ≈ 0.147 - scheduler = UniPCMultistepScheduler( - use_karras_sigmas=True, - use_flow_sigmas=True, - prediction_type="flow_prediction", - sigma_max=200.0, - sigma_min=0.147, - ) - - pipeline = Cosmos3OmniPipeline( - transformer=transformer, - text_tokenizer=text_tokenizer, - vae=diffusers_vae, - scheduler=scheduler, - sound_tokenizer=sound_tokenizer, - ) - print(f"Saving full pipeline to {output_dir} …") - pipeline.save_pretrained(str(output_dir), safe_serialization=True, max_shard_size="5GB") - else: - print(f"Saving transformer to {output_dir} …") - transformer.save_pretrained(str(output_dir), safe_serialization=True, max_shard_size="5GB") - if include_sound_tokenizer: - print("Skipping sound_tokenizer/ save because --save-pipeline was not set.") - - print("Done.") - - -if __name__ == "__main__": - main() diff --git a/tests/models/autoencoders/test_models_autoencoder_cosmos3_audio.py b/tests/models/autoencoders/test_models_autoencoder_cosmos3_audio.py index ac7c62322478..c4677ecc67b5 100644 --- a/tests/models/autoencoders/test_models_autoencoder_cosmos3_audio.py +++ b/tests/models/autoencoders/test_models_autoencoder_cosmos3_audio.py @@ -13,9 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import importlib.util -from pathlib import Path - import pytest import torch @@ -144,54 +141,3 @@ def test_cosmos3_audio_tokenizer_decoder_only_state_disables_encode(): assert decoder_only_model.encoder is None with pytest.raises(ValueError, match="decoder-only weights"): decoder_only_model.encode(torch.randn(1, 2, 16)) - - -def _load_converter_module(): - repo_root = Path(__file__).resolve().parents[3] - script_path = repo_root / "scripts" / "convert_cosmos3_to_diffusers.py" - spec = importlib.util.spec_from_file_location("convert_cosmos3_to_diffusers", script_path) - module = importlib.util.module_from_spec(spec) - spec.loader.exec_module(module) - return module - - -def test_cosmos3_audio_converter_keeps_encoder_and_remaps_decoder(): - converter = _load_converter_module() - state_dict = { - "generator.encoder.layers.0.weight": torch.ones(4, 20, 1), - "generator.encoder.layers.1.act.alpha": torch.zeros(16), - "generator.encoder.layers.1.act.beta": torch.zeros(16), - "generator.decoder.layers.0.weight": torch.ones(8, 4, 7), - "generator.decoder.layers.1.layers.0.alpha": torch.zeros(8), - "generator.decoder.layers.1.layers.1.weight": torch.ones(8, 4, 4), - "generator.decoder.layers.1.layers.2.layers.0.alpha": torch.zeros(4), - "generator.decoder.layers.1.layers.2.layers.1.weight": torch.ones(4, 4, 7), - "generator.decoder.layers.2.alpha": torch.zeros(4), - "generator.decoder.layers.3.weight": torch.ones(2, 4, 7), - } - - remapped = converter._remap_avae_state_dict(state_dict) - - assert not any(key.startswith("decoder.layers.") for key in remapped) - assert "encoder.layers.0.weight" not in remapped - assert "encoder.layers.0.weight_g" in remapped - assert "encoder.layers.0.weight_v" in remapped - assert remapped["encoder.layers.1.act.alpha"].shape == (1, 16, 1) - assert remapped["decoder.conv1.weight_g"].shape == (8, 1, 1) - assert remapped["decoder.block.0.snake1.alpha"].shape == (1, 8, 1) - assert remapped["decoder.block.0.res_unit1.snake1.alpha"].shape == (1, 4, 1) - assert remapped["decoder.snake1.alpha"].shape == (1, 4, 1) - - -def test_cosmos3_audio_converter_allows_decoder_only_state_dict(): - converter = _load_converter_module() - state_dict = { - "decoder.conv1.weight": torch.ones(8, 4, 7), - "decoder.snake1.alpha": torch.zeros(4), - } - - remapped = converter._remap_avae_state_dict(state_dict) - - assert not any(key.startswith("encoder.") for key in remapped) - assert "decoder.conv1.weight_g" in remapped - assert remapped["decoder.snake1.alpha"].shape == (1, 4, 1) From f7e27b693e1a7d814ae50ae2f85e00d145bf5711 Mon Sep 17 00:00:00 2001 From: Atharva Joshi Date: Tue, 14 Jul 2026 10:10:40 -0700 Subject: [PATCH 3/8] Refine Cosmos3 Edge norm selection --- .../transformers/transformer_cosmos3.py | 64 +++++++++++-------- .../test_models_transformer_cosmos3.py | 8 +++ 2 files changed, 45 insertions(+), 27 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_cosmos3.py b/src/diffusers/models/transformers/transformer_cosmos3.py index 2a150fc879bb..5121c7f7f3cb 100644 --- a/src/diffusers/models/transformers/transformer_cosmos3.py +++ b/src/diffusers/models/transformers/transformer_cosmos3.py @@ -140,10 +140,8 @@ def forward(self, position_ids, device, dtype): class Cosmos3NemotronRMSNorm(nn.Module): - def __init__(self, dim: int, eps: float, elementwise_affine: bool = True, bias: bool = False): + def __init__(self, dim: int, eps: float): super().__init__() - if not elementwise_affine or bias: - raise ValueError("Cosmos3NemotronRMSNorm requires an affine weight without a bias.") self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) @@ -219,7 +217,7 @@ def __init__( attention_bias: bool, rms_norm_eps: float, qk_norm_for_text: bool = True, - norm_cls: type[nn.Module] = RMSNorm, + norm_type: str = "rms_norm", processor=None, ): super().__init__() @@ -235,24 +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 = ( - norm_cls(head_dim, eps=rms_norm_eps, elementwise_affine=True, bias=False) - if qk_norm_for_text - else nn.Identity() - ) - self.norm_k = ( - norm_cls(head_dim, eps=rms_norm_eps, elementwise_affine=True, bias=False) - if qk_norm_for_text - else nn.Identity() - ) + 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 = norm_cls(head_dim, eps=rms_norm_eps, elementwise_affine=True, bias=False) - self.norm_added_k = norm_cls(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() @@ -284,7 +285,7 @@ def __init__( ): super().__init__() self.hidden_size = hidden_size - norm_cls = Cosmos3NemotronRMSNorm if hidden_act == "relu2" else RMSNorm + norm_type = "nemotron_rms_norm" if hidden_act == "relu2" else "rms_norm" self.self_attn = Cosmos3PackedMoTAttention( hidden_size=hidden_size, head_dim=head_dim, @@ -293,7 +294,7 @@ def __init__( attention_bias=attention_bias, rms_norm_eps=rms_norm_eps, qk_norm_for_text=qk_norm_for_text, - norm_cls=norm_cls, + norm_type=norm_type, ) self.mlp = Cosmos3VLTextMLP( @@ -303,12 +304,18 @@ def __init__( hidden_size=hidden_size, intermediate_size=intermediate_size, hidden_act=hidden_act ) - self.input_layernorm = norm_cls(hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False) - self.input_layernorm_moe_gen = norm_cls(hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False) - self.post_attention_layernorm = norm_cls(hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False) - self.post_attention_layernorm_moe_gen = norm_cls( - hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False - ) + 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, @@ -404,9 +411,12 @@ def __init__( for _ in range(num_hidden_layers) ] ) - norm_cls = Cosmos3NemotronRMSNorm if hidden_act == "relu2" else RMSNorm - self.norm = norm_cls(hidden_size, eps=rms_norm_eps, elementwise_affine=True, bias=False) - self.norm_moe_gen = norm_cls(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 ) diff --git a/tests/models/transformers/test_models_transformer_cosmos3.py b/tests/models/transformers/test_models_transformer_cosmos3.py index 633a125d50a0..b36608fbaa5b 100644 --- a/tests/models/transformers/test_models_transformer_cosmos3.py +++ b/tests/models/transformers/test_models_transformer_cosmos3.py @@ -46,9 +46,17 @@ def get_edge_transformer(action_gen: bool = False): def test_cosmos3_edge_uses_nemotron_parameter_layout(): transformer = get_edge_transformer(action_gen=True) state_dict = transformer.state_dict() + layer = transformer.layers[0] assert transformer.config.backbone_type == "cosmos3_edge_nemotron_dense" assert transformer.config.temporal_compression_factor == 4 + assert isinstance(layer.self_attn.norm_q, torch.nn.Identity) + assert isinstance(layer.self_attn.norm_k, torch.nn.Identity) + assert isinstance(layer.self_attn.norm_added_q, Cosmos3NemotronRMSNorm) + assert isinstance(layer.self_attn.norm_added_k, Cosmos3NemotronRMSNorm) + assert isinstance(layer.input_layernorm, Cosmos3NemotronRMSNorm) + assert isinstance(layer.post_attention_layernorm, Cosmos3NemotronRMSNorm) + assert isinstance(transformer.norm, Cosmos3NemotronRMSNorm) assert not any("gate_proj" in key for key in state_dict) assert not any(".norm_q." in key or ".norm_k." in key for key in state_dict) assert "layers.0.self_attn.norm_added_q.weight" in state_dict From fc9d0ec4b888b2598a6c4fce313e794cfb8af8f8 Mon Sep 17 00:00:00 2001 From: Atharva Joshi Date: Wed, 15 Jul 2026 11:05:31 -0700 Subject: [PATCH 4/8] Use internal Cosmos3 modular test fixture --- tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py index c4cbbb50ca10..4b5589199fed 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py @@ -116,7 +116,7 @@ class TestCosmos3OmniModularPipelineFast(ModularPipelineTesterMixin): pipeline_class = Cosmos3OmniModularPipeline pipeline_blocks_class = Cosmos3OmniBlocks - pretrained_model_name_or_path = "atharvajoshi10/tiny-cosmos3-modular-pipe" + pretrained_model_name_or_path = "hf-internal-testing/tiny-cosmos3-modular-pipe" params = frozenset(["prompt", "height", "width", "num_frames", "guidance_scale"]) batch_params = frozenset() From 32199868324d93658146ea4443102b51c42221cd Mon Sep 17 00:00:00 2001 From: Atharva Joshi Date: Wed, 15 Jul 2026 11:12:15 -0700 Subject: [PATCH 5/8] Use real VAE in Cosmos3 pipeline tests --- tests/pipelines/cosmos/test_cosmos3.py | 30 +++++++++++--------------- 1 file changed, 12 insertions(+), 18 deletions(-) diff --git a/tests/pipelines/cosmos/test_cosmos3.py b/tests/pipelines/cosmos/test_cosmos3.py index 1456a8dbc145..2c8e702c27bb 100644 --- a/tests/pipelines/cosmos/test_cosmos3.py +++ b/tests/pipelines/cosmos/test_cosmos3.py @@ -17,7 +17,7 @@ import torch -from diffusers import Cosmos3OmniPipeline, Cosmos3OmniTransformer, UniPCMultistepScheduler +from diffusers import AutoencoderKLWan, Cosmos3OmniPipeline, Cosmos3OmniTransformer, UniPCMultistepScheduler class DummyTokenizer: @@ -34,22 +34,6 @@ def apply_chat_template(self, conversations, **kwargs): return SimpleNamespace(input_ids=[1, 2]) -class DummyVAE(torch.nn.Module): - def __init__(self): - super().__init__() - self.weight = torch.nn.Parameter(torch.zeros(())) - self.config = SimpleNamespace( - latents_mean=[0.0], - latents_std=[1.0], - scale_factor_spatial=16, - scale_factor_temporal=4, - ) - - @property - def dtype(self): - return self.weight.dtype - - def get_dummy_pipeline(**kwargs): transformer = Cosmos3OmniTransformer( hidden_size=16, @@ -63,10 +47,20 @@ def get_dummy_pipeline(**kwargs): patch_latent_dim=1, vocab_size=32, ) + + torch.manual_seed(0) + vae = AutoencoderKLWan( + base_dim=3, + z_dim=16, + dim_mult=[1, 1, 1, 1], + num_res_blocks=1, + temperal_downsample=[False, True, True], + ) + return Cosmos3OmniPipeline( transformer=transformer, text_tokenizer=DummyTokenizer(), - vae=DummyVAE(), + vae=vae, scheduler=UniPCMultistepScheduler(), enable_safety_checker=False, **kwargs, From bd40cb846854cf6f8f2c6533e73aaf0e2834e7a6 Mon Sep 17 00:00:00 2001 From: Atharva Joshi Date: Wed, 15 Jul 2026 16:19:48 -0700 Subject: [PATCH 6/8] Align Cosmos3 Edge tests with conventions --- .../transformers/transformer_cosmos3.py | 38 ++- .../modular_pipelines/cosmos/denoise.py | 3 +- .../pipelines/cosmos/pipeline_cosmos3_omni.py | 10 + .../test_models_transformer_cosmos3.py | 295 +++++++++++++----- tests/pipelines/cosmos/test_cosmos3.py | 214 +++++++++---- 5 files changed, 404 insertions(+), 156 deletions(-) diff --git a/src/diffusers/models/transformers/transformer_cosmos3.py b/src/diffusers/models/transformers/transformer_cosmos3.py index 5121c7f7f3cb..ea9c6f2f0c5e 100644 --- a/src/diffusers/models/transformers/transformer_cosmos3.py +++ b/src/diffusers/models/transformers/transformer_cosmos3.py @@ -13,12 +13,14 @@ # 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 @@ -26,6 +28,24 @@ 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. @@ -207,6 +227,7 @@ class Cosmos3PackedMoTAttention(nn.Module, AttentionModuleMixin): _default_processor_cls = Cosmos3AttnProcessor _available_processors = [Cosmos3AttnProcessor] + _supports_qkv_fusion = False def __init__( self, @@ -380,8 +401,6 @@ def __init__( 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__() @@ -633,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: @@ -661,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 @@ -778,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) diff --git a/src/diffusers/modular_pipelines/cosmos/denoise.py b/src/diffusers/modular_pipelines/cosmos/denoise.py index f3cbdfa2b655..7d023c7f641c 100644 --- a/src/diffusers/modular_pipelines/cosmos/denoise.py +++ b/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] diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py index b971c17773fe..059424b27e41 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py @@ -1251,6 +1251,14 @@ def _apply_video_safety_check(self, video: Any, output_type: str, device: torch. def current_timestep(self): return self._current_timestep + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def num_timesteps(self): + return self._num_timesteps + @property def interrupt(self): return self._interrupt @@ -1684,6 +1692,7 @@ def __call__( action_timesteps=action_timesteps, action_noisy_frame_indexes=cond_packed_static.get("action_noisy_frame_indexes"), action_domain_ids=[action_domain_id] if action_domain_id is not None else None, + return_dict=False, ) cond_v_vision, cond_v_sound, cond_v_action = self._mask_velocity_predictions( preds_vision, @@ -1723,6 +1732,7 @@ def __call__( action_timesteps=action_timesteps, action_noisy_frame_indexes=uncond_packed_static.get("action_noisy_frame_indexes"), action_domain_ids=[action_domain_id] if action_domain_id is not None else None, + return_dict=False, ) uncond_v_vision, uncond_v_sound, uncond_v_action = self._mask_velocity_predictions( preds_vision, diff --git a/tests/models/transformers/test_models_transformer_cosmos3.py b/tests/models/transformers/test_models_transformer_cosmos3.py index b36608fbaa5b..9204c7e38453 100644 --- a/tests/models/transformers/test_models_transformer_cosmos3.py +++ b/tests/models/transformers/test_models_transformer_cosmos3.py @@ -1,5 +1,5 @@ # coding=utf-8 -# Copyright 2026 The HuggingFace Team. All rights reserved. +# Copyright 2026 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,97 +13,220 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pytest import torch from diffusers import Cosmos3OmniTransformer -from diffusers.models.transformers.transformer_cosmos3 import Cosmos3NemotronRMSNorm - - -def get_edge_transformer(action_gen: bool = False): - return Cosmos3OmniTransformer( - action_dim=3 if action_gen else None, - action_gen=action_gen, - backbone_type="cosmos3_edge_nemotron_dense", - head_dim=4, - hidden_act="relu2", - hidden_size=16, - intermediate_size=32, - latent_channel=1, - latent_patch_size=1, - num_attention_heads=4, - num_embodiment_domains=2, - num_hidden_layers=2, - num_key_value_heads=2, - patch_latent_dim=1, - qk_norm_for_text=False, - rms_norm_eps=1e-5, - rope_theta=1e8, - temporal_compression_factor=4, - vocab_size=32, - ) - - -def test_cosmos3_edge_uses_nemotron_parameter_layout(): - transformer = get_edge_transformer(action_gen=True) - state_dict = transformer.state_dict() - layer = transformer.layers[0] - - assert transformer.config.backbone_type == "cosmos3_edge_nemotron_dense" - assert transformer.config.temporal_compression_factor == 4 - assert isinstance(layer.self_attn.norm_q, torch.nn.Identity) - assert isinstance(layer.self_attn.norm_k, torch.nn.Identity) - assert isinstance(layer.self_attn.norm_added_q, Cosmos3NemotronRMSNorm) - assert isinstance(layer.self_attn.norm_added_k, Cosmos3NemotronRMSNorm) - assert isinstance(layer.input_layernorm, Cosmos3NemotronRMSNorm) - assert isinstance(layer.post_attention_layernorm, Cosmos3NemotronRMSNorm) - assert isinstance(transformer.norm, Cosmos3NemotronRMSNorm) - assert not any("gate_proj" in key for key in state_dict) - assert not any(".norm_q." in key or ".norm_k." in key for key in state_dict) - assert "layers.0.self_attn.norm_added_q.weight" in state_dict - assert "layers.0.self_attn.norm_added_k.weight" in state_dict - assert "layers.0.mlp.up_proj.weight" in state_dict - assert "layers.0.mlp.down_proj.weight" in state_dict - assert "action_proj_in.fc.weight" in state_dict - assert "action_proj_out.fc.weight" in state_dict - - -def test_cosmos3_edge_transformer_runs_action_workflow(): - transformer = get_edge_transformer(action_gen=True).eval() - - with torch.no_grad(): - prediction, sound_prediction, action_prediction = transformer( - input_ids=torch.tensor([1, 2]), - text_indexes=torch.tensor([0, 1]), - position_ids=torch.tensor([[0, 0, 0, 0], [0, 0, 0, 0], [0, 1, 0, 0]]), - und_len=2, - sequence_length=4, - vision_tokens=[torch.randn(1, 1, 1, 1, 1)], - vision_token_shapes=[(1, 1, 1)], - vision_sequence_indexes=torch.tensor([2]), - vision_mse_loss_indexes=torch.tensor([2]), - vision_timesteps=torch.tensor([1]), - vision_noisy_frame_indexes=[torch.tensor([0])], - action_tokens=[torch.randn(1, 3)], - action_token_shapes=[(1, 1, 1)], - action_sequence_indexes=torch.tensor([3]), - action_mse_loss_indexes=torch.tensor([3]), - action_timesteps=torch.tensor([1]), - action_noisy_frame_indexes=[torch.tensor([0])], - action_domain_ids=[torch.tensor(0)], +from diffusers.models.transformers.transformer_cosmos3 import ( + Cosmos3NemotronRMSNorm, + Cosmos3OmniTransformerOutput, +) +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, +) + + +enable_full_determinism() + + +class Cosmos3OmniTransformerTesterConfig(BaseModelTesterConfig): + @property + def model_class(self): + return Cosmos3OmniTransformer + + @property + def main_input_name(self) -> str: + return "vision_tokens" + + @property + def uses_custom_attn_processor(self) -> bool: + return True + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict: + return { + "head_dim": 6, + "hidden_act": "relu2", + "hidden_size": 12, + "intermediate_size": 24, + "latent_channel": 2, + "latent_patch_size": 1, + "num_attention_heads": 2, + "num_hidden_layers": 2, + "num_key_value_heads": 1, + "patch_latent_dim": 2, + "qk_norm_for_text": False, + "rms_norm_eps": 1e-5, + "rope_axes_dim": [1, 1, 1], + "rope_theta": 1e8, + "vocab_size": 32, + } + + def get_dummy_inputs(self, height: int = 1, width: int = 1) -> dict: + num_vision_tokens = height * width + sequence_length = 2 + num_vision_tokens + vision_indexes = torch.arange(2, sequence_length, device=torch_device) + + return { + "input_ids": torch.tensor([1, 2], device=torch_device), + "text_indexes": torch.tensor([0, 1], device=torch_device), + "position_ids": torch.zeros((3, sequence_length), dtype=torch.long, device=torch_device), + "und_len": 2, + "sequence_length": sequence_length, + "vision_tokens": [randn_tensor((1, 2, 1, height, width), generator=self.generator, device=torch_device)], + "vision_token_shapes": [(1, height, width)], + "vision_sequence_indexes": vision_indexes, + "vision_mse_loss_indexes": vision_indexes, + "vision_timesteps": torch.ones(num_vision_tokens, device=torch_device), + "vision_noisy_frame_indexes": [torch.tensor([0], device=torch_device)], + } + + @property + def input_shape(self) -> tuple[int, ...]: + return (1, 2, 1, 1, 1) + + @property + def output_shape(self) -> tuple[int, ...]: + return (1, 2, 1, 1, 1) + + +class TestCosmos3OmniTransformerModel(Cosmos3OmniTransformerTesterConfig, ModelTesterMixin): + def test_output_format(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + + with torch.no_grad(): + output = model(**self.get_dummy_inputs()) + output_tuple = model(**self.get_dummy_inputs(), return_dict=False) + + assert isinstance(output, Cosmos3OmniTransformerOutput) + assert output.sample[0].shape == self.output_shape + assert output.sound is None + assert output.action is None + torch.testing.assert_close(output.sample[0], output_tuple[0][0]) + + def test_determinism(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + + with torch.no_grad(): + first = model(**self.get_dummy_inputs()).sample[0] + second = model(**self.get_dummy_inputs()).sample[0] + + torch.testing.assert_close(first, second) + + def test_outputs_equivalence(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + + with torch.no_grad(): + output = model(**self.get_dummy_inputs()) + output_tuple = model(**self.get_dummy_inputs(), return_dict=False) + + torch.testing.assert_close(output.sample[0], output_tuple[0][0]) + assert output.sound is output_tuple[1] is None + assert output.action is output_tuple[2] is None + + def test_cosmos3_edge_uses_nemotron_parameter_layout(self): + transformer = self.model_class(**self.get_init_dict(), action_dim=3, action_gen=True, num_embodiment_domains=2) + state_dict = transformer.state_dict() + layer = transformer.layers[0] + + assert isinstance(layer.self_attn.norm_q, torch.nn.Identity) + assert isinstance(layer.self_attn.norm_k, torch.nn.Identity) + assert isinstance(layer.self_attn.norm_added_q, Cosmos3NemotronRMSNorm) + assert isinstance(layer.self_attn.norm_added_k, Cosmos3NemotronRMSNorm) + assert isinstance(layer.input_layernorm, Cosmos3NemotronRMSNorm) + assert isinstance(layer.post_attention_layernorm, Cosmos3NemotronRMSNorm) + assert isinstance(transformer.norm, Cosmos3NemotronRMSNorm) + assert not any("gate_proj" in key for key in state_dict) + assert not any(".norm_q." in key or ".norm_k." in key for key in state_dict) + assert "layers.0.self_attn.norm_added_q.weight" in state_dict + assert "layers.0.self_attn.norm_added_k.weight" in state_dict + assert "layers.0.mlp.up_proj.weight" in state_dict + assert "layers.0.mlp.down_proj.weight" in state_dict + assert "action_proj_in.fc.weight" in state_dict + assert "action_proj_out.fc.weight" in state_dict + + def test_cosmos3_edge_transformer_runs_action_workflow(self): + transformer = self.model_class( + **self.get_init_dict(), action_dim=3, action_gen=True, num_embodiment_domains=2 + ).eval() + inputs = self.get_dummy_inputs() + inputs["position_ids"] = torch.zeros((3, 4), dtype=torch.long, device=torch_device) + inputs["sequence_length"] = 4 + inputs.update( + { + "action_tokens": [randn_tensor((1, 3), generator=self.generator, device=torch_device)], + "action_token_shapes": [(1, 1, 1)], + "action_sequence_indexes": torch.tensor([3], device=torch_device), + "action_mse_loss_indexes": torch.tensor([3], device=torch_device), + "action_timesteps": torch.tensor([1], device=torch_device), + "action_noisy_frame_indexes": [torch.tensor([0], device=torch_device)], + "action_domain_ids": [torch.tensor(0, device=torch_device)], + } ) - assert prediction[0].shape == (1, 1, 1, 1, 1) - assert sound_prediction is None - assert action_prediction[0].shape == (1, 3) + with torch.no_grad(): + prediction, sound_prediction, action_prediction = transformer(**inputs, return_dict=False) + + assert prediction[0].shape == self.output_shape + assert sound_prediction is None + assert action_prediction[0].shape == (1, 3) + + def test_cosmos3_nemotron_rms_norm_multiplies_in_float32(self): + hidden_states = torch.randn(2, 3, 8, dtype=torch.bfloat16) + norm = Cosmos3NemotronRMSNorm(8, eps=1e-5).bfloat16() + norm.weight.data.copy_(torch.randn(8, dtype=torch.bfloat16)) + + expected = hidden_states.float() + expected = expected * torch.rsqrt(expected.pow(2).mean(-1, keepdim=True) + 1e-5) + expected = (norm.weight.float() * expected).to(hidden_states.dtype) + + torch.testing.assert_close(norm(hidden_states), expected, rtol=0, atol=0) + + +class TestCosmos3OmniTransformerMemory(Cosmos3OmniTransformerTesterConfig, MemoryTesterMixin): + pass + + +class TestCosmos3OmniTransformerTorchCompile(Cosmos3OmniTransformerTesterConfig, TorchCompileTesterMixin): + @property + def different_shapes_for_compilation(self): + return [(4, 4), (4, 8), (8, 8)] + + def get_dummy_inputs(self, height: int = 4, width: int = 4) -> dict[str, torch.Tensor]: + return super().get_dummy_inputs(height=height, width=width) + + +class TestCosmos3OmniTransformerTraining(Cosmos3OmniTransformerTesterConfig, TrainingTesterMixin): + def test_gradient_checkpointing_is_applied(self): + super().test_gradient_checkpointing_is_applied(expected_set={"Cosmos3OmniTransformer"}) + + @pytest.mark.skip("The transformer returns one tensor list per generated modality.") + def test_training(self): + super().test_training() + + @pytest.mark.skip("The transformer returns one tensor list per generated modality.") + def test_training_with_ema(self): + super().test_training_with_ema() + @pytest.mark.skip("The transformer returns one tensor list per generated modality.") + def test_gradient_checkpointing_equivalence(self): + super().test_gradient_checkpointing_equivalence() -def test_cosmos3_nemotron_rms_norm_multiplies_in_float32(): - hidden_states = torch.randn(2, 3, 8, dtype=torch.bfloat16) - norm = Cosmos3NemotronRMSNorm(8, eps=1e-5).bfloat16() - norm.weight.data.copy_(torch.randn(8, dtype=torch.bfloat16)) + @pytest.mark.skip("The transformer returns one tensor list per generated modality.") + def test_mixed_precision_training(self): + super().test_mixed_precision_training() - expected = hidden_states.float() - expected = expected * torch.rsqrt(expected.pow(2).mean(-1, keepdim=True) + 1e-5) - expected = (norm.weight.float() * expected).to(hidden_states.dtype) - torch.testing.assert_close(norm(hidden_states), expected, rtol=0, atol=0) +class TestCosmos3OmniTransformerAttention(Cosmos3OmniTransformerTesterConfig, AttentionTesterMixin): + pass diff --git a/tests/pipelines/cosmos/test_cosmos3.py b/tests/pipelines/cosmos/test_cosmos3.py index 2c8e702c27bb..10a6275e9574 100644 --- a/tests/pipelines/cosmos/test_cosmos3.py +++ b/tests/pipelines/cosmos/test_cosmos3.py @@ -13,78 +13,166 @@ # limitations under the License. import json -from types import SimpleNamespace +import tempfile +import unittest +from unittest import mock import torch +from transformers import AutoTokenizer from diffusers import AutoencoderKLWan, Cosmos3OmniPipeline, Cosmos3OmniTransformer, UniPCMultistepScheduler +from ...testing_utils import enable_full_determinism, torch_device +from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS +from ..test_pipelines_common import PipelineTesterMixin -class DummyTokenizer: - eos_token_id = 11 - def __init__(self): - self.conversations = [] +enable_full_determinism() - def convert_tokens_to_ids(self, token): - return {"<|vision_start|>": 20}.get(token, 0) - def apply_chat_template(self, conversations, **kwargs): - self.conversations.append(conversations) - return SimpleNamespace(input_ids=[1, 2]) +class Cosmos3OmniPipelineWrapper(Cosmos3OmniPipeline): + @staticmethod + def from_pretrained(*args, **kwargs): + kwargs.setdefault("enable_safety_checker", False) + kwargs.setdefault("safety_checker", None) + return Cosmos3OmniPipeline.from_pretrained(*args, **kwargs) -def get_dummy_pipeline(**kwargs): - transformer = Cosmos3OmniTransformer( - hidden_size=16, - intermediate_size=32, - head_dim=4, - num_attention_heads=4, - num_key_value_heads=2, - num_hidden_layers=1, - latent_channel=1, - latent_patch_size=1, - patch_latent_dim=1, - vocab_size=32, +class Cosmos3OmniPipelineFastTests(PipelineTesterMixin, unittest.TestCase): + pipeline_class = Cosmos3OmniPipelineWrapper + params = TEXT_TO_IMAGE_PARAMS - {"cross_attention_kwargs", "negative_prompt_embeds", "prompt_embeds"} + batch_params = TEXT_TO_IMAGE_BATCH_PARAMS + required_optional_params = frozenset( + [ + "num_inference_steps", + "generator", + "latents", + "output_type", + "return_dict", + "callback_on_step_end", + "callback_on_step_end_tensor_inputs", + ] ) - - torch.manual_seed(0) - vae = AutoencoderKLWan( - base_dim=3, - z_dim=16, - dim_mult=[1, 1, 1, 1], - num_res_blocks=1, - temperal_downsample=[False, True, True], - ) - - return Cosmos3OmniPipeline( - transformer=transformer, - text_tokenizer=DummyTokenizer(), - vae=vae, - scheduler=UniPCMultistepScheduler(), - enable_safety_checker=False, - **kwargs, - ) - - -def test_cosmos3_pipeline_saves_edge_configuration(tmp_path): - pipeline = get_dummy_pipeline(default_use_system_prompt=False, use_native_flow_schedule=True) - - assert not pipeline.config.enable_safety_checker - assert not pipeline.config.default_use_system_prompt - assert pipeline.config.use_native_flow_schedule - assert pipeline.safety_checker is None - - pipeline.save_config(tmp_path) - model_index = json.loads((tmp_path / "model_index.json").read_text()) - assert model_index["enable_safety_checker"] is False - assert model_index["default_use_system_prompt"] is False - assert model_index["use_native_flow_schedule"] is True - - -def test_cosmos3_tokenize_prompt_uses_checkpoint_system_prompt_default(): - pipeline = get_dummy_pipeline(default_use_system_prompt=False) - - pipeline.tokenize_prompt("A prompt", num_frames=1, add_resolution_template=False) - - assert all(conversation[0]["role"] == "user" for conversation in pipeline.text_tokenizer.conversations) + supports_dduf = False + test_xformers_attention = False + test_layerwise_casting = False + test_group_offloading = True + + def get_dummy_components(self): + torch.manual_seed(0) + transformer = Cosmos3OmniTransformer( + head_dim=6, + hidden_act="relu2", + hidden_size=6, + intermediate_size=12, + latent_channel=16, + latent_patch_size=1, + num_attention_heads=1, + num_hidden_layers=1, + num_key_value_heads=1, + patch_latent_dim=16, + qk_norm_for_text=False, + rms_norm_eps=1e-5, + rope_axes_dim=[1, 1, 1], + vocab_size=151657, + ) + + torch.manual_seed(0) + vae = AutoencoderKLWan( + base_dim=3, + z_dim=16, + dim_mult=[1, 1, 1, 1], + num_res_blocks=1, + temperal_downsample=[False, True, True], + ) + + text_tokenizer = AutoTokenizer.from_pretrained( + "hf-internal-testing/tiny-cosmos3-modular-pipe", subfolder="text_tokenizer" + ) + + return { + "transformer": transformer, + "text_tokenizer": text_tokenizer, + "vae": vae, + "scheduler": UniPCMultistepScheduler(), + "sound_tokenizer": None, + "safety_checker": None, + "enable_safety_checker": False, + } + + def get_dummy_inputs(self, device, seed=0): + generator = torch.Generator(device="cpu").manual_seed(seed) + return { + "prompt": "a dog", + "negative_prompt": "bad quality", + "height": 16, + "width": 16, + "num_frames": 1, + "num_inference_steps": 2, + "guidance_scale": 1.0, + "generator": generator, + "output_type": "np", + "use_system_prompt": False, + "add_resolution_template": False, + "add_duration_template": False, + } + + def test_inference(self): + pipeline = self.pipeline_class(**self.get_dummy_components()).to(torch_device) + pipeline.set_progress_bar_config(disable=None) + + video = pipeline(**self.get_dummy_inputs(torch_device)).video + + self.assertEqual(video.shape, (1, 16, 16, 3)) + + def test_components_function(self): + init_components = self.get_dummy_components() + pipeline = self.pipeline_class(**init_components) + component_names = { + name for name, component in init_components.items() if not isinstance(component, (str, int, float)) + } + + self.assertTrue(hasattr(pipeline, "components")) + self.assertEqual(set(pipeline.components), component_names) + + def test_cosmos3_pipeline_saves_edge_configuration(self): + components = self.get_dummy_components() + components["enable_safety_checker"] = False + components["default_use_system_prompt"] = False + components["use_native_flow_schedule"] = True + pipeline = self.pipeline_class(**components) + + assert not pipeline.config.enable_safety_checker + assert not pipeline.config.default_use_system_prompt + assert pipeline.config.use_native_flow_schedule + assert pipeline.safety_checker is None + + with tempfile.TemporaryDirectory() as tmpdir: + pipeline.save_config(tmpdir) + with open(f"{tmpdir}/model_index.json") as model_index_file: + model_index = json.load(model_index_file) + assert model_index["enable_safety_checker"] is False + assert model_index["default_use_system_prompt"] is False + assert model_index["use_native_flow_schedule"] is True + + def test_cosmos3_tokenize_prompt_uses_checkpoint_system_prompt_default(self): + components = self.get_dummy_components() + components["default_use_system_prompt"] = False + pipeline = self.pipeline_class(**components) + + with mock.patch.object( + pipeline.text_tokenizer, + "apply_chat_template", + wraps=pipeline.text_tokenizer.apply_chat_template, + ) as apply_chat_template: + pipeline.tokenize_prompt("A prompt", num_frames=1, add_resolution_template=False) + + assert all(call.args[0][0]["role"] == "user" for call in apply_chat_template.call_args_list) + + @unittest.skip("Cosmos3 currently supports one prompt per pipeline call.") + def test_inference_batch_consistent(self): + pass + + @unittest.skip("Cosmos3 currently supports one prompt per pipeline call.") + def test_inference_batch_single_identical(self): + pass From 0f6a27d5d899fe3ffd90b02a04f231369d46ea99 Mon Sep 17 00:00:00 2001 From: Atharva Joshi Date: Wed, 15 Jul 2026 20:01:50 -0700 Subject: [PATCH 7/8] Fix Cosmos3 callbacks on Python 3.10 --- src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py index 059424b27e41..db1365210c93 100644 --- a/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py +++ b/src/diffusers/pipelines/cosmos/pipeline_cosmos3_omni.py @@ -1783,7 +1783,9 @@ def __call__( action_latents[:, raw_action_dim_resolved:] = 0 if callback_on_step_end is not None: - callback_kwargs = {k: locals()[k] for k in callback_on_step_end_tensor_inputs} + callback_kwargs = {} + for key in callback_on_step_end_tensor_inputs: + callback_kwargs[key] = locals()[key] callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) latents = callback_outputs.pop("latents", latents) From 0cfca10f5f4727883b9068632b9360e0174e0c8b Mon Sep 17 00:00:00 2001 From: Atharva Joshi Date: Thu, 16 Jul 2026 09:19:54 -0700 Subject: [PATCH 8/8] Address Cosmos3 modular test feedback --- .../modular_pipelines/cosmos/encoders.py | 6 +- .../cosmos/test_modular_pipeline_cosmos3.py | 352 ++++++++---------- 2 files changed, 159 insertions(+), 199 deletions(-) diff --git a/src/diffusers/modular_pipelines/cosmos/encoders.py b/src/diffusers/modular_pipelines/cosmos/encoders.py index 9358dc0158d0..d8c5412078b3 100644 --- a/src/diffusers/modular_pipelines/cosmos/encoders.py +++ b/src/diffusers/modular_pipelines/cosmos/encoders.py @@ -497,7 +497,7 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) block_state = self.get_block_state(state) device = components._execution_device - dtype = components.transformer.dtype + dtype = components.vae.dtype if block_state.image is None: raise ValueError("`Cosmos3ImageVaeEncoderStep` requires an `image` input.") @@ -601,7 +601,7 @@ def __call__(self, components: Cosmos3OmniModularPipeline, state: PipelineState) block_state = self.get_block_state(state) device = components._execution_device - dtype = components.transformer.dtype + dtype = components.vae.dtype if block_state.video is None: raise ValueError("`Cosmos3VideoVaeEncoderStep` requires a `video` input.") @@ -769,7 +769,7 @@ 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 - dtype = components.transformer.dtype + dtype = components.vae.dtype chunk_id = block_state.chunk_id chunk_frames = block_state.chunk_frames diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py index 4b5589199fed..4986c102ae66 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py @@ -13,14 +13,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -import json - import pytest import torch from PIL import Image -from diffusers import Cosmos3OmniTransformer, ModularPipeline, UniPCMultistepScheduler -from diffusers.modular_pipelines import Cosmos3OmniBlocks, Cosmos3OmniModularPipeline +from diffusers import ModularPipeline, UniPCMultistepScheduler +from diffusers.modular_pipelines import ( + Cosmos3OmniBlocks, + Cosmos3OmniModularPipeline, + SequentialPipelineBlocks, +) from diffusers.modular_pipelines.cosmos.before_denoise import ( Cosmos3ActionDenoiseInputStep, Cosmos3ActionPackSequenceStep, @@ -30,7 +32,6 @@ Cosmos3VisionPackSequenceStep, ) from diffusers.modular_pipelines.cosmos.encoders import Cosmos3TextEncoderStep -from diffusers.modular_pipelines.modular_pipeline import PipelineState from ...testing_utils import torch_device from ..test_modular_pipelines_common import ModularPipelineTesterMixin @@ -181,15 +182,18 @@ def test_vae_encoder_is_standalone_and_validates_conditioning_inputs(self): assert vae_encoder.select_block(action=None, image=None, video=object()) == "video_conditioning" assert vae_encoder.select_block(action=object(), image=None, video=None) == "action_conditioning" - state = PipelineState() - state.set("image", Image.new("RGB", (32, 32))) - state.set("num_frames", 5) - state.set("height", 32) - state.set("width", 32) - _, state = vae_encoder(pipe, state) + vae_pipe = vae_encoder.init_pipeline(self.pretrained_model_name_or_path) + vae_pipe.load_components(torch_dtype=torch.float32) + outputs = vae_pipe( + image=Image.new("RGB", (32, 32)), + num_frames=5, + height=32, + width=32, + output=["x0_tokens_vision", "vision_condition_frames"], + ) - assert state.get("x0_tokens_vision") is not None - assert state.get("vision_condition_frames") == [0] + assert outputs["x0_tokens_vision"] is not None + assert outputs["vision_condition_frames"] == [0] action_vae_encoder = vae_encoder.sub_blocks["action_conditioning"] assert [input_param.name for input_param in action_vae_encoder.inputs] == ["action"] @@ -220,196 +224,152 @@ def test_rejects_batched_prompts(self, prompt_name): with pytest.raises(ValueError, match="batched prompts are not supported"): pipe(**inputs, output=self.output_name) + def test_pack_steps_do_not_require_vae(self): + cond_text_segment = {"vision_start_temporal_offset": 0, "und_len": 2} + uncond_text_segment = {"vision_start_temporal_offset": 0, "und_len": 1} + + vision_pipe = Cosmos3VisionPackSequenceStep().init_pipeline(self.pretrained_model_name_or_path) + vision_pipe.load_components(torch_dtype=torch.float32) + vision_segments = vision_pipe( + cond_text_segment=cond_text_segment, + uncond_text_segment=uncond_text_segment, + latents=torch.zeros(1, vision_pipe.transformer.config.latent_channel, 2, 2, 2), + fps_vision=5.0, + vision_condition_indexes_for_pack=[], + output=["cond_vision_segment", "uncond_vision_segment"], + ) -def test_cosmos3_pack_steps_do_not_require_vae(): - transformer = Cosmos3OmniTransformer( - hidden_size=32, - intermediate_size=64, - num_hidden_layers=2, - num_attention_heads=4, - num_key_value_heads=2, - head_dim=8, - latent_channel=4, - latent_patch_size=2, - patch_latent_dim=16, - vocab_size=16, - rope_scaling={"mrope_section": [2, 1, 1]}, - action_gen=True, - action_dim=10, - ) - pipe = Cosmos3OmniModularPipeline(blocks=Cosmos3VisionPackSequenceStep()) - pipe.update_components(transformer=transformer) - - assert [component.name for component in Cosmos3VisionPackSequenceStep().expected_components] == ["transformer"] - assert [component.name for component in Cosmos3ActionPackSequenceStep().expected_components] == ["transformer"] - assert not hasattr(pipe, "vae") - - vision_segment = pipe._prepare_vision_segment( - input_vision_tokens=torch.zeros(1, transformer.config.latent_channel, 2, 2, 2), - has_image_condition=False, - mrope_offset=0, - vision_fps=5.0, - curr=1, - device="cpu", - ) - action_segment = pipe._prepare_action_segment( - input_action_tokens=torch.zeros(4, transformer.config.action_dim), - condition_frame_indexes=[], - mrope_offset=0, - action_fps=5.0, - curr=1, - device="cpu", - ) - - assert vision_segment["vision_mrope_ids"].shape == (3, 2) - assert action_segment["action_mrope_ids"].shape == (3, 4) - - -def test_cosmos3_denoise_input_steps_assemble_modality_segments(): - cond_text_segment = {"text_mrope_ids": torch.tensor([[1, 2], [3, 4], [5, 6]]), "und_len": 2} - uncond_text_segment = {"text_mrope_ids": torch.tensor([[7], [8], [9]]), "und_len": 1} - cond_vision_segment = {"vision_mrope_ids": torch.tensor([[10], [11], [12]]), "num_vision_tokens": 1} - uncond_vision_segment = {"vision_mrope_ids": torch.tensor([[13, 14], [15, 16], [17, 18]]), "num_vision_tokens": 2} - cond_sound_segment = {"sound_mrope_ids": torch.tensor([[19, 20], [21, 22], [23, 24]]), "sound_len": 2} - uncond_sound_segment = {"sound_mrope_ids": torch.tensor([[25], [26], [27]]), "sound_len": 1} - cond_action_segment = {"action_mrope_ids": torch.tensor([[28], [29], [30]]), "action_len": 1} - uncond_action_segment = {"action_mrope_ids": torch.tensor([[31, 32], [33, 34], [35, 36]]), "action_len": 2} - - state = PipelineState() - state.set("cond_text_segment", cond_text_segment) - state.set("uncond_text_segment", uncond_text_segment) - state.set("cond_vision_segment", cond_vision_segment) - state.set("uncond_vision_segment", uncond_vision_segment) - state.set("cond_sound_segment", cond_sound_segment) - state.set("uncond_sound_segment", uncond_sound_segment) - state.set("cond_action_segment", cond_action_segment) - state.set("uncond_action_segment", uncond_action_segment) - - _, state = Cosmos3VisionDenoiseInputStep()(None, state) - _, state = Cosmos3SoundDenoiseInputStep()(None, state) - _, state = Cosmos3ActionDenoiseInputStep()(None, state) - - torch.testing.assert_close( - state.get("cond_position_ids"), - torch.cat( - [ - cond_text_segment["text_mrope_ids"], - cond_vision_segment["vision_mrope_ids"], - cond_sound_segment["sound_mrope_ids"], - cond_action_segment["action_mrope_ids"], - ], - dim=1, - ), - ) - torch.testing.assert_close( - state.get("uncond_position_ids"), - torch.cat( - [ - uncond_text_segment["text_mrope_ids"], - uncond_vision_segment["vision_mrope_ids"], - uncond_sound_segment["sound_mrope_ids"], - uncond_action_segment["action_mrope_ids"], - ], - dim=1, - ), - ) - assert state.get("cond_sequence_length") == 6 - assert state.get("uncond_sequence_length") == 6 - assert state.get("cond_packed_static") is None - assert state.get("uncond_packed_static") is None - - -def test_cosmos3_text_step_uses_pipeline_system_prompt_and_safety_configs(monkeypatch): - text_encoder = Cosmos3TextEncoderStep() - pipe = Cosmos3OmniModularPipeline( - blocks=text_encoder, - config_dict={"default_use_system_prompt": False, "enable_safety_checker": False}, - ) - captured = {} - - def tokenize_prompt(*args, **kwargs): - captured["use_system_prompt"] = kwargs["use_system_prompt"] - return torch.zeros((1, 1), dtype=torch.long), torch.zeros((1, 1), dtype=torch.long) - - monkeypatch.setattr(pipe, "tokenize_prompt", tokenize_prompt) - - state = PipelineState() - state.set("prompt", "A small robot moves across a table.") - state.set("negative_prompt", "") - state.set("num_frames", 5) - state.set("height", 32) - state.set("width", 32) - state.set("fps", 24.0) - state.set("add_resolution_template", True) - state.set("add_duration_template", True) - _, state = text_encoder(pipe, state) - - assert captured["use_system_prompt"] is False - assert not pipe.requires_safety_checker - assert state.get("cond_input_ids").shape == (1, 1) - - -def test_cosmos3_native_flow_schedule_uses_edge_sigma_grid(monkeypatch): - set_timesteps = Cosmos3SetTimestepsStep() - pipe = Cosmos3OmniModularPipeline(blocks=set_timesteps, config_dict={"use_native_flow_schedule": True}) - scheduler = UniPCMultistepScheduler(num_train_timesteps=100, use_flow_sigmas=True) - pipe.update_components(scheduler=scheduler) - captured = {} - - def capture_set_timesteps(num_inference_steps, device=None, sigmas=None): - captured["num_inference_steps"] = num_inference_steps - captured["device"] = device - captured["sigmas"] = sigmas - scheduler.timesteps = torch.arange(num_inference_steps) - - monkeypatch.setattr(scheduler, "set_timesteps", capture_set_timesteps) - - state = PipelineState() - state.set("num_inference_steps", 4) - _, state = set_timesteps(pipe, state) - - assert captured["num_inference_steps"] == 4 - assert captured["sigmas"] == pytest.approx([0.99, 0.7425, 0.495, 0.2475]) - assert state.get("timesteps").tolist() == [0, 1, 2, 3] - - -def test_cosmos3_modular_model_index_takes_precedence(tmp_path): - (tmp_path / "model_index.json").write_text(json.dumps({"_class_name": "Cosmos3OmniDiffusersPipeline"})) - (tmp_path / "modular_model_index.json").write_text( - json.dumps( + assert set(vision_pipe.components) == {"transformer"} + assert vision_segments["cond_vision_segment"]["vision_mrope_ids"].shape == (3, 2) + assert vision_segments["uncond_vision_segment"]["vision_mrope_ids"].shape == (3, 2) + + action_pipe = Cosmos3ActionPackSequenceStep().init_pipeline(self.pretrained_model_name_or_path) + action_pipe.load_components(torch_dtype=torch.float32) + action_segments = action_pipe( + cond_text_segment=cond_text_segment, + uncond_text_segment=uncond_text_segment, + cond_sequence_length=2, + uncond_sequence_length=1, + action_latents=torch.zeros(4, action_pipe.transformer.config.action_dim), + action_condition_frame_indexes=[], + fps_vision=5.0, + output=["cond_action_segment", "uncond_action_segment"], + ) + + assert set(action_pipe.components) == {"transformer"} + assert action_segments["cond_action_segment"]["action_mrope_ids"].shape == (3, 4) + assert action_segments["uncond_action_segment"]["action_mrope_ids"].shape == (3, 4) + + def test_denoise_input_steps_assemble_modality_segments(self): + cond_text_segment = {"text_mrope_ids": torch.tensor([[1, 2], [3, 4], [5, 6]]), "und_len": 2} + uncond_text_segment = {"text_mrope_ids": torch.tensor([[7], [8], [9]]), "und_len": 1} + cond_vision_segment = {"vision_mrope_ids": torch.tensor([[10], [11], [12]]), "num_vision_tokens": 1} + uncond_vision_segment = { + "vision_mrope_ids": torch.tensor([[13, 14], [15, 16], [17, 18]]), + "num_vision_tokens": 2, + } + cond_sound_segment = { + "sound_mrope_ids": torch.tensor([[19, 20], [21, 22], [23, 24]]), + "sound_len": 2, + } + uncond_sound_segment = {"sound_mrope_ids": torch.tensor([[25], [26], [27]]), "sound_len": 1} + cond_action_segment = {"action_mrope_ids": torch.tensor([[28], [29], [30]]), "action_len": 1} + uncond_action_segment = { + "action_mrope_ids": torch.tensor([[31, 32], [33, 34], [35, 36]]), + "action_len": 2, + } + + blocks = SequentialPipelineBlocks.from_blocks_dict( { - "_class_name": "Cosmos3OmniModularPipeline", - "_blocks_class_name": "Cosmos3OmniBlocks", + "vision": Cosmos3VisionDenoiseInputStep(), + "sound": Cosmos3SoundDenoiseInputStep(), + "action": Cosmos3ActionDenoiseInputStep(), } ) - ) + pipe = blocks.init_pipeline() + state = pipe( + cond_text_segment=cond_text_segment, + uncond_text_segment=uncond_text_segment, + cond_vision_segment=cond_vision_segment, + uncond_vision_segment=uncond_vision_segment, + cond_sound_segment=cond_sound_segment, + uncond_sound_segment=uncond_sound_segment, + cond_action_segment=cond_action_segment, + uncond_action_segment=uncond_action_segment, + ) - modular_pipe = ModularPipeline.from_pretrained(str(tmp_path)) - explicit_modular_pipe = Cosmos3OmniModularPipeline.from_pretrained(str(tmp_path)) + torch.testing.assert_close( + state.get("cond_position_ids"), + torch.cat( + [ + cond_text_segment["text_mrope_ids"], + cond_vision_segment["vision_mrope_ids"], + cond_sound_segment["sound_mrope_ids"], + cond_action_segment["action_mrope_ids"], + ], + dim=1, + ), + ) + torch.testing.assert_close( + state.get("uncond_position_ids"), + torch.cat( + [ + uncond_text_segment["text_mrope_ids"], + uncond_vision_segment["vision_mrope_ids"], + uncond_sound_segment["sound_mrope_ids"], + uncond_action_segment["action_mrope_ids"], + ], + dim=1, + ), + ) + assert state.get("cond_sequence_length") == 6 + assert state.get("uncond_sequence_length") == 6 + assert state.get("cond_packed_static") is None + assert state.get("uncond_packed_static") is None + + def test_text_step_uses_pipeline_system_prompt_and_safety_configs(self): + text_pipe = Cosmos3TextEncoderStep().init_pipeline(self.pretrained_model_name_or_path) + text_pipe.load_components() + text_pipe.disable_safety_checker() + assert text_pipe.config.default_use_system_prompt + assert not text_pipe.requires_safety_checker + + inputs = { + "prompt": "A small robot moves across a table.", + "negative_prompt": "", + "num_frames": 5, + "height": 32, + "width": 32, + } + default_with_system_prompt = text_pipe(**inputs, output="cond_input_ids") + explicit_with_system_prompt = text_pipe(**inputs, use_system_prompt=True, output="cond_input_ids") + explicit_without_system_prompt = text_pipe(**inputs, use_system_prompt=False, output="cond_input_ids") - assert isinstance(modular_pipe, Cosmos3OmniModularPipeline) - assert isinstance(modular_pipe.blocks, Cosmos3OmniBlocks) - assert isinstance(explicit_modular_pipe, Cosmos3OmniModularPipeline) - assert isinstance(explicit_modular_pipe.blocks, Cosmos3OmniBlocks) + text_pipe.update_components(default_use_system_prompt=False) + assert not text_pipe.config.default_use_system_prompt + default_without_system_prompt = text_pipe(**inputs, output="cond_input_ids") + updated_with_system_prompt = text_pipe(**inputs, use_system_prompt=True, output="cond_input_ids") + updated_without_system_prompt = text_pipe(**inputs, use_system_prompt=False, output="cond_input_ids") -def test_cosmos3_model_index_fallback_resolves_modular_pipeline(tmp_path): - (tmp_path / "model_index.json").write_text( - json.dumps( - { - "_class_name": "Cosmos3OmniPipeline", - "default_use_system_prompt": False, - "enable_safety_checker": False, - "use_native_flow_schedule": True, - } - ) - ) + assert default_with_system_prompt == explicit_with_system_prompt == updated_with_system_prompt + assert explicit_without_system_prompt == default_without_system_prompt == updated_without_system_prompt + assert len(default_with_system_prompt) > len(default_without_system_prompt) - pipe = ModularPipeline.from_pretrained(str(tmp_path)) + def test_set_timesteps_native_flow_schedule(self): + timesteps_pipe = Cosmos3SetTimestepsStep().init_pipeline(self.pretrained_model_name_or_path) + timesteps_pipe.load_components() + assert not timesteps_pipe.config.use_native_flow_schedule + + default_timesteps = timesteps_pipe(num_inference_steps=4, output="timesteps") + + timesteps_pipe.update_components( + scheduler=UniPCMultistepScheduler(num_train_timesteps=100, use_flow_sigmas=True), + use_native_flow_schedule=True, + ) + native_timesteps = timesteps_pipe(num_inference_steps=4, output="timesteps") - assert isinstance(pipe, Cosmos3OmniModularPipeline) - assert isinstance(pipe.blocks, Cosmos3OmniBlocks) - assert not pipe.config.default_use_system_prompt - assert not pipe.config.enable_safety_checker - assert pipe.config.use_native_flow_schedule - assert not pipe.requires_safety_checker + expected_sigmas = torch.tensor([0.99, 0.7425, 0.495, 0.2475]) + torch.testing.assert_close(timesteps_pipe.scheduler.sigmas[:-1], expected_sigmas) + assert native_timesteps.tolist() == [99, 74, 49, 24] + assert not torch.equal(native_timesteps, default_timesteps)