# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from __future__ import annotations from types import SimpleNamespace from typing import Any import pytest import torch from torch import nn from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion] @pytest.fixture(autouse=True) def _single_rank_tensor_parallel(monkeypatch: pytest.MonkeyPatch) -> None: """Provide the TP metadata required by vLLM parallel linear layers.""" from vllm.model_executor import parameter from vllm.model_executor.layers import linear from vllm_omni.diffusion.models.cosmos3 import ( transformer_cosmos3, transformer_cosmos3_edge, ) monkeypatch.setattr(linear, "get_tensor_model_parallel_rank", lambda: 0) monkeypatch.setattr(linear, "get_tensor_model_parallel_world_size", lambda: 1) monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0) monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1) monkeypatch.setattr(transformer_cosmos3, "get_tensor_model_parallel_world_size", lambda: 1) monkeypatch.setattr(transformer_cosmos3_edge, "get_tensor_model_parallel_world_size", lambda: 1) @pytest.fixture def accelerator_device() -> torch.device: """Provide an accelerator device, skipping when none is available.""" if current_omni_platform.get_device_count() == 0: pytest.skip("Accelerator required for this test") return current_omni_platform.get_torch_device(0) def _tiny_cosmos3_config(**overrides): config = { "hidden_size": 8, "num_hidden_layers": 0, "num_attention_heads": 2, "num_key_value_heads": 2, "head_dim": 4, "intermediate_size": 16, "vocab_size": 32, "latent_patch_size": 1, "latent_channel": 2, "rope_scaling": {"mrope_section": [1, 1, 0]}, } config.update(overrides) return config def _tiny_cosmos3_edge_config(**overrides): from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( COSMOS3_EDGE_BACKBONE_TYPE, ) config = _tiny_cosmos3_config( latent_patch_size=2, latent_channel=48, temporal_compression_factor=4, backbone_type=COSMOS3_EDGE_BACKBONE_TYPE, qk_norm_for_text=False, ) config.update(overrides) return config def test_mrope_position_ids_cover_text_video_sound_and_action() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import ( compute_mrope_position_ids_action, compute_mrope_position_ids_sound, compute_mrope_position_ids_text, compute_mrope_position_ids_vision, ) text_ids, text_offset = compute_mrope_position_ids_text(num_tokens=3, temporal_offset=5) assert text_ids.tolist() == [[5, 6, 7], [5, 6, 7], [5, 6, 7]] assert text_offset == 8 vision_ids, vision_offset = compute_mrope_position_ids_vision(2, 2, 3, temporal_offset=10, fps=None) assert vision_ids.shape == (3, 12) assert vision_ids[0].tolist() == [10] * 6 + [11] * 6 assert vision_offset == 12 modulated_ids, modulated_offset = compute_mrope_position_ids_vision( 2, 1, 1, temporal_offset=10, fps=12.0, base_fps=24.0, temporal_compression_factor=4, ) torch.testing.assert_close(modulated_ids[0], torch.tensor([10.0, 12.0])) assert modulated_offset == 13 sound_ids, sound_offset = compute_mrope_position_ids_sound(3, temporal_offset=10, sound_latent_fps=25.0) torch.testing.assert_close(sound_ids[0], torch.tensor([10.0, 10.96, 11.92])) assert sound_offset == 12 action_ids, action_offset = compute_mrope_position_ids_action(3, temporal_offset=10, action_fps=None) assert action_ids.tolist() == [[11, 12, 13], [0, 0, 0], [0, 0, 0]] assert action_offset == 14 def test_timestep_embedder_stores_frequencies_in_fp32() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import ( TimestepEmbedder, ) embedder = TimestepEmbedder(hidden_size=8, frequency_embedding_size=16) assert embedder.freqs.dtype == torch.float32 @pytest.mark.parametrize( ("key", "value"), [ ("qk_norm_for_diffusion", False), ("qk_norm_for_text", False), ("position_embedding_type", "rotary"), ("unified_3d_mrope_reset_spatial_ids", False), ("joint_attn_implementation", "one_way"), ], ) def test_validate_supported_config_rejects_unsupported_flags(key: str, value) -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer with pytest.raises(ValueError, match=f"{key}="): Cosmos3VFMTransformer._validate_supported_config({key: value}) Cosmos3VFMTransformer._validate_supported_config({}) Cosmos3VFMTransformer._validate_supported_config(None) def test_edge_config_resolves_nemotron_defaults() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( Cosmos3EdgeVFMTransformer, ) model = Cosmos3EdgeVFMTransformer( SimpleNamespace( tf_model_config=_tiny_cosmos3_edge_config( layer_norm_epsilon=1e-5, rope_scaling={}, rope_parameters={"mrope_section": [24, 20, 20]}, ), dtype=torch.float32, ) ) assert model.rms_norm_eps == 1e-5 assert model.rope_theta == 100_000_000 assert model.mrope_section == [24, 20, 20] assert model.latent_channel_size == 48 assert model.latent_patch_size == 2 assert model.temporal_compression_factor == 4 def test_edge_config_requires_backbone_type() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( Cosmos3EdgeVFMTransformer, ) config = _tiny_cosmos3_edge_config() config.pop("backbone_type") with pytest.raises(ValueError, match="must declare backbone_type"): Cosmos3EdgeVFMTransformer( SimpleNamespace( tf_model_config=config, dtype=torch.float32, ) ) def test_edge_und_and_gen_use_relu2_and_no_und_qk_norm() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( Cosmos3EdgeVFMTransformer, Cosmos3Relu2MLP, ) model = Cosmos3EdgeVFMTransformer( SimpleNamespace( tf_model_config=_tiny_cosmos3_edge_config( num_hidden_layers=1, use_und_k_norm_for_gen=False, ), dtype=torch.float32, ) ) layer = model.language_model.layers[0] gen_layer = model.gen_layers[0] assert isinstance(layer.mlp, Cosmos3Relu2MLP) assert not hasattr(layer.mlp, "gate_proj") assert isinstance(gen_layer.mlp, Cosmos3Relu2MLP) assert not hasattr(gen_layer.mlp, "gate_proj") assert not hasattr(layer.self_attn, "norm_q") assert not hasattr(layer.self_attn, "norm_k") assert layer.self_attn.k_norm_und_for_gen is None def test_edge_creates_gen_k_norm_when_enabled() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import RMSNorm from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( Cosmos3EdgeVFMTransformer, ) model = Cosmos3EdgeVFMTransformer( SimpleNamespace( tf_model_config=_tiny_cosmos3_edge_config( num_hidden_layers=1, qk_norm_for_diffusion=True, use_und_k_norm_for_gen=True, ), dtype=torch.float32, ) ) k_norm = model.language_model.layers[0].self_attn.k_norm_und_for_gen assert model.use_und_k_norm_for_gen is True assert isinstance(k_norm, RMSNorm) assert isinstance(model.gen_layers[0].cross_attention.norm_q, RMSNorm) assert isinstance(model.gen_layers[0].cross_attention.norm_k, RMSNorm) def test_edge_supports_no_gen_qk_norm_variant() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( Cosmos3EdgeVFMTransformer, ) model = Cosmos3EdgeVFMTransformer( SimpleNamespace( tf_model_config=_tiny_cosmos3_edge_config( num_hidden_layers=1, qk_norm_for_diffusion=False, use_und_k_norm_for_gen=False, ), dtype=torch.float32, ) ) assert model.qk_norm_for_diffusion is False assert model.language_model.layers[0].self_attn.k_norm_und_for_gen is None assert model.gen_layers[0].cross_attention.qk_norm is False assert not hasattr(model.gen_layers[0].cross_attention, "norm_q") assert not hasattr(model.gen_layers[0].cross_attention, "norm_k") def test_edge_validates_required_relu2_weights() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( Cosmos3EdgeVFMTransformer, ) model = object.__new__(Cosmos3EdgeVFMTransformer) nn.Module.__init__(model) model.num_hidden_layers = 1 model.qk_norm_for_diffusion = True model.use_und_k_norm_for_gen = True complete = { "transformer.language_model.layers.0.mlp.up_proj.weight", "transformer.language_model.layers.0.mlp.down_proj.weight", "transformer.language_model.layers.0.self_attn.k_norm_und_for_gen.weight", "transformer.gen_layers.0.mlp.up_proj.weight", "transformer.gen_layers.0.mlp.down_proj.weight", } model.validate_loaded_weights(complete) mlp_weights = complete - { "transformer.language_model.layers.0.self_attn.k_norm_und_for_gen.weight", } for missing_weight in mlp_weights: with pytest.raises(ValueError, match="missing required weights"): model.validate_loaded_weights(complete - {missing_weight}) missing_k_norm = complete - {"transformer.language_model.layers.0.self_attn.k_norm_und_for_gen.weight"} with pytest.raises(ValueError, match=r"self_attn\.k_norm_und_for_gen"): model.validate_loaded_weights(missing_k_norm) model.use_und_k_norm_for_gen = False model.validate_loaded_weights(missing_k_norm) def test_edge_gen_cached_k_is_normalized_but_reasoner_uses_raw_k() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import RMSNorm from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( Cosmos3EdgeCausalAttention, ) class ScaleProjection(nn.Module): def __init__(self, scale: float) -> None: super().__init__() self.scale = scale def forward(self, x: torch.Tensor) -> torch.Tensor: return x * self.scale class IdentityOutput(nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: return x class RawKAttention(nn.Module): def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: del v return q + k attn = object.__new__(Cosmos3EdgeCausalAttention) nn.Module.__init__(attn) attn.hidden_size = 4 attn.num_heads = 2 attn.num_kv_heads = 2 attn.head_dim = 2 attn.num_heads_local = 2 attn.num_kv_heads_local = 2 attn.to_q = ScaleProjection(1.0) attn.to_k = ScaleProjection(1.0) attn.to_v = ScaleProjection(1.0) attn.to_out = IdentityOutput() attn.k_norm_und_for_gen = RMSNorm(2, eps=1e-5) attn.attn = RawKAttention() hidden_states = torch.tensor( [[[1.0, 2.0, 3.0, 4.0], [2.0, 1.0, 4.0, 3.0]]], dtype=torch.float32, ) freqs_cos = torch.ones(1, 2, 1, 2) freqs_sin = torch.zeros(1, 2, 1, 2) reasoner_out, gen_k, gen_v = attn(hidden_states, freqs_cos, freqs_sin) attn.to_k.scale = 3.0 scaled_reasoner_out, scaled_gen_k, scaled_gen_v = attn(hidden_states, freqs_cos, freqs_sin) torch.testing.assert_close(gen_k, scaled_gen_k) torch.testing.assert_close(gen_v, scaled_gen_v) assert not torch.allclose(reasoner_out, scaled_reasoner_out) def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3LanguageModel, Cosmos3VFMTransformer from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3_edge import ( Cosmos3EdgeLanguageModel, Cosmos3EdgeVFMTransformer, ) for transformer_cls, language_model_cls in ( (Cosmos3VFMTransformer, Cosmos3LanguageModel), (Cosmos3EdgeVFMTransformer, Cosmos3EdgeLanguageModel), ): model = object.__new__(transformer_cls) nn.Module.__init__(model) model.language_model = nn.Module() model.language_model.layers = nn.ModuleList([nn.Linear(2, 2) for _ in range(2)]) model.gen_layers = nn.ModuleList([nn.Linear(2, 2)]) model.norm_moe_gen = nn.LayerNorm(2) matched = [ name for name, module in model.named_modules() if any(condition(name, module) for condition in model._hsdp_shard_conditions) ] assert matched == ["language_model.layers.0", "language_model.layers.1", "gen_layers.0"] assert transformer_cls._layerwise_offload_blocks_attrs == ["gen_layers"] assert language_model_cls._layerwise_offload_blocks_attrs == ["layers"] assert transformer_cls._repeated_blocks == ["Cosmos3GenDecoderLayer"] model.latent_patch_size = 2 model.latent_channel_size = 3 latents = torch.arange(1 * 3 * 1 * 3 * 5, dtype=torch.float32).reshape(1, 3, 1, 3, 5) torch.testing.assert_close( model.unpatchify(model.patchify(latents, t=1, h=3, w=5), t=1, h=3, w=5), latents, ) def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> None: from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None)) output = transformer_cosmos3.Cosmos3VFMTransformer( SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32) )( hidden_states=torch.zeros(1, 2, 1, 2, 2), timestep=torch.tensor([1.0]), text_ids=torch.tensor([[1, 2]], dtype=torch.long), text_mask=torch.ones(1, 2, dtype=torch.long), video_shape=(1, 2, 2), fps=24.0, ) assert tuple(output.shape) == (1, 2, 1, 2, 2) def test_qwen_gen_mlp_remains_gated() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GatedMLP, Cosmos3VFMTransformer model = Cosmos3VFMTransformer( SimpleNamespace(tf_model_config=_tiny_cosmos3_config(num_hidden_layers=1), dtype=torch.float32) ) assert isinstance(model.gen_layers[0].mlp, Cosmos3GatedMLP) assert hasattr(model.gen_layers[0].mlp, "gate_proj") def test_model_cpu_offload_swaps_back_to_generator(monkeypatch: pytest.MonkeyPatch) -> None: from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None)) class ToyReasonerLayer(nn.Module): def __init__(self) -> None: super().__init__() self.weight = nn.Parameter(torch.ones(1)) def forward(self, hidden: torch.Tensor, freqs): del freqs kv = hidden.new_zeros(hidden.shape[0], hidden.shape[1], 1, 1) return hidden + self.weight.to(hidden.dtype), kv, kv class ToyGeneratorLayer(nn.Module): def __init__(self) -> None: super().__init__() self.weight = nn.Parameter(torch.ones(1)) def forward(self, hidden: torch.Tensor, **kwargs): del kwargs return hidden + self.weight.to(hidden.dtype) model = transformer_cosmos3.Cosmos3VFMTransformer( SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32) ) model.language_model.layers = nn.ModuleList([ToyReasonerLayer()]) model.gen_layers = nn.ModuleList([ToyGeneratorLayer()]) model.enable_model_cpu_offload(device=torch.device("cpu"), pin_memory=False) assert model._model_cpu_offload_enabled assert set(model._model_cpu_offload_components()) == {"reasoner", "generator"} assert model.device == torch.device("cpu") output = model( hidden_states=torch.zeros(1, 2, 1, 2, 2), timestep=torch.tensor([1.0]), text_ids=torch.tensor([[1, 2]], dtype=torch.long), text_mask=torch.ones(1, 2, dtype=torch.long), video_shape=(1, 2, 2), fps=24.0, ) assert tuple(output.shape) == (1, 2, 1, 2, 2) # The reasoner runs once then the generator component stays resident for GEN. assert model._active_model_cpu_offload_component == "generator" model.disable_model_cpu_offload() assert not model._model_cpu_offload_enabled def test_model_cpu_offload_moves_reasoner_and_generator_between_cpu_and_device( monkeypatch: pytest.MonkeyPatch, accelerator_device: torch.device ) -> None: from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None)) model = transformer_cosmos3.Cosmos3VFMTransformer( SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32) ) model.language_model.layers = nn.ModuleList([nn.Linear(2, 2)]) model.gen_layers = nn.ModuleList([nn.Linear(2, 2)]) model.to(accelerator_device) reasoner_param = model.language_model.layers[0].weight generator_param = model.gen_layers[0].weight model.enable_model_cpu_offload(device=accelerator_device, pin_memory=False) assert model._model_cpu_offload_enabled # On enable, every group is parked on CPU until a phase activates it. assert reasoner_param.device.type == "cpu" assert generator_param.device.type == "cpu" model._activate_model_cpu_offload_component("reasoner") assert reasoner_param.device == accelerator_device assert generator_param.device.type == "cpu" model._activate_model_cpu_offload_component("generator") assert reasoner_param.device.type == "cpu" assert generator_param.device == accelerator_device model.disable_model_cpu_offload() assert not model._model_cpu_offload_enabled # Disable restores both components onto the device. assert reasoner_param.device == accelerator_device assert generator_param.device == accelerator_device def test_multi_control_attention_weights_target_outputs() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3CrossAttention class ControlValueAttention: def __call__(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: del k # Each per-control K/V layout is [text, control_i, target]. return v[:, 1:2].expand(q.shape[0], q.shape[1], -1, -1) attention = SimpleNamespace(multi_control_attn=ControlValueAttention()) q = torch.zeros(1, 3, 1, 1) k = torch.zeros_like(q) v = torch.tensor([[[[10.0]], [[20.0]], [[30.0]]]]) text_kv = torch.zeros(1, 1, 1, 1) output = Cosmos3CrossAttention._forward_multi_control( attention, q, k, v, text_kv, text_kv, control_token_sizes=(1, 1), control_weights=(0.25, 0.75), ) torch.testing.assert_close(output, torch.tensor([[[10.0], [20.0], [17.5]]])) def test_forward_accepts_transfer_control_latents(monkeypatch: pytest.MonkeyPatch) -> None: from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None)) model = transformer_cosmos3.Cosmos3VFMTransformer( SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32) ) hidden_states = torch.zeros(1, 2, 1, 2, 2) output = model( hidden_states=hidden_states, timestep=torch.tensor([1.0]), text_ids=torch.tensor([[1, 2]], dtype=torch.long), text_mask=torch.ones(1, 2, dtype=torch.long), video_shape=(1, 2, 2), fps=24.0, control_latents=[torch.ones_like(hidden_states), torch.full_like(hidden_states, 2.0)], control_weights=[1.0, 3.0], ) assert tuple(output.shape) == tuple(hidden_states.shape) with pytest.raises(ValueError, match="control_weights length"): model( hidden_states=hidden_states, timestep=torch.tensor([1.0]), text_ids=torch.tensor([[1, 2]], dtype=torch.long), text_mask=torch.ones(1, 2, dtype=torch.long), video_shape=(1, 2, 2), fps=24.0, control_latents=[torch.ones_like(hidden_states), torch.full_like(hidden_states, 2.0)], control_weights=[1.0], ) with pytest.raises(ValueError, match="control latent shape"): model( hidden_states=hidden_states, timestep=torch.tensor([1.0]), text_ids=torch.tensor([[1, 2]], dtype=torch.long), text_mask=torch.ones(1, 2, dtype=torch.long), video_shape=(1, 2, 2), fps=24.0, control_latents=[torch.zeros(1, 2, 2, 2, 2)], ) def test_forward_gathers_gen_tokens_before_unpatchify(monkeypatch: pytest.MonkeyPatch) -> None: from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 class RecordingGather(nn.Module): def __init__(self) -> None: super().__init__() self.calls = 0 def forward(self, hidden_gen: torch.Tensor) -> torch.Tensor: self.calls += 1 return hidden_gen monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None)) model = transformer_cosmos3.Cosmos3VFMTransformer( SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32) ) gather = RecordingGather() model.gen_sp_gather = gather output = model( hidden_states=torch.zeros(1, 2, 1, 2, 2), timestep=torch.tensor([1.0]), text_ids=torch.tensor([[1, 2]], dtype=torch.long), text_mask=torch.ones(1, 2, dtype=torch.long), video_shape=(1, 2, 2), fps=24.0, ) assert gather.calls == 1 assert tuple(output.shape) == (1, 2, 1, 2, 2) def test_sound_and_action_modules_follow_config() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer tiny = _tiny_cosmos3_config() no_modal = Cosmos3VFMTransformer(SimpleNamespace(tf_model_config=tiny, dtype=torch.float32)) with_sound = Cosmos3VFMTransformer( SimpleNamespace(tf_model_config=tiny, dtype=torch.float32), sound_gen=True, sound_dim=5, sound_latent_fps=40.0, ) with_action = Cosmos3VFMTransformer( SimpleNamespace( tf_model_config={**tiny, "action_gen": True, "max_action_dim": 6, "num_embodiment_domains": 9}, dtype=torch.float32, ) ) assert no_modal.sound_gen is False assert no_modal.action_gen is False assert not hasattr(no_modal, "audio_proj_in") assert not hasattr(no_modal, "action_proj_in") assert with_sound.sound_dim == 5 assert with_sound.sound_latent_fps == 40.0 assert with_sound.audio_proj_in.in_features == 5 assert with_action.action_dim == 6 assert with_action.action_proj_in.num_domains == 9 @pytest.mark.parametrize( "kwargs", [ {"sound_gen": True}, {"sound_gen": True, "sound_dim": 5}, {"sound_gen": True, "sound_latent_fps": 40.0}, ], ) def test_transformer_requires_sound_dim_and_fps_when_sound_gen_true(kwargs: dict[str, Any]) -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer with pytest.raises(ValueError, match=r"requires an explicit sound_dim and sound_latent_fps"): Cosmos3VFMTransformer( SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32), **kwargs, ) def test_sound_and_action_pack_unpack_validate_shapes() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer model = object.__new__(Cosmos3VFMTransformer) nn.Module.__init__(model) model.sound_dim = 3 model.action_dim = 3 sound = torch.arange(2 * 3 * 4, dtype=torch.float32).reshape(2, 3, 4) action = torch.arange(2 * 4 * 3, dtype=torch.float32).reshape(2, 4, 3) torch.testing.assert_close(model.unpack_sound(model.pack_sound(sound)), sound) torch.testing.assert_close(model.unpack_action(model.pack_action(action)), action) with pytest.raises(ValueError, match="channel mismatch"): model.pack_sound(torch.zeros(1, 4, 2)) with pytest.raises(ValueError, match="dimension mismatch"): model.pack_action(torch.zeros(1, 2, 4)) @pytest.mark.parametrize( ("config", "transformer_kwargs", "extra_kwargs", "expected_shapes"), [ ( _tiny_cosmos3_config(), {"sound_gen": True, "sound_dim": 3, "sound_latent_fps": 24.0}, {"sound_latents": torch.zeros(1, 3, 4)}, [(1, 2, 1, 2, 2), (1, 3, 4)], ), ( _tiny_cosmos3_config(action_gen=True, max_action_dim=3, num_embodiment_domains=4), {}, {"action_latents": torch.zeros(1, 5, 3), "action_domain_ids": torch.tensor([2])}, [(1, 2, 1, 2, 2), (1, 5, 3)], ), ], ) def test_forward_returns_video_plus_optional_modality_predictions( monkeypatch: pytest.MonkeyPatch, config, transformer_kwargs, extra_kwargs, expected_shapes, ) -> None: from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None)) output = transformer_cosmos3.Cosmos3VFMTransformer( SimpleNamespace(tf_model_config=config, dtype=torch.float32), **transformer_kwargs, )( hidden_states=torch.zeros(1, 2, 1, 2, 2), timestep=torch.tensor([1.0]), text_ids=torch.tensor([[1, 2]], dtype=torch.long), text_mask=torch.ones(1, 2, dtype=torch.long), video_shape=(1, 2, 2), fps=24.0, action_noisy_mask=torch.ones(1, 5, 1), **extra_kwargs, ) assert isinstance(output, tuple) assert [tuple(tensor.shape) for tensor in output] == expected_shapes def test_forward_with_sound_ulysses_error_mentions_combined_sequence(monkeypatch: pytest.MonkeyPatch) -> None: import vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 as cosmos3_module model = cosmos3_module.Cosmos3VFMTransformer( SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32), sound_gen=True, sound_dim=3, sound_latent_fps=40.0, ) monkeypatch.setattr(cosmos3_module, "_get_ulysses_state", lambda: (2, 0, None)) with pytest.raises(ValueError, match=r"GEN sequence length \(3 = video tokens 2 \+ sound tokens 1\)"): model( hidden_states=torch.zeros(1, 2, 1, 1, 2), timestep=torch.tensor([1.0]), text_ids=torch.tensor([[1, 2]], dtype=torch.long), text_mask=torch.ones(1, 2, dtype=torch.long), video_shape=(1, 1, 2), fps=24.0, sound_latents=torch.zeros(1, 3, 1), ) def test_sound_latent_frames_padded_for_sequence_parallel(monkeypatch: pytest.MonkeyPatch) -> None: from vllm_omni.diffusion.distributed import parallel_state from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3 model = object.__new__(transformer_cosmos3.Cosmos3VFMTransformer) model.latent_patch_size = 2 vs = (3, 16, 16) monkeypatch.setattr(parallel_state, "get_ulysses_parallel_world_size", lambda: 1) assert model.sound_latent_frames_for_sequence_parallel(video_shape=vs, sound_frames=97) == 97 monkeypatch.setattr(parallel_state, "get_ulysses_parallel_world_size", lambda: 2) assert model.sound_latent_frames_for_sequence_parallel(video_shape=vs, sound_frames=97) == 98 assert model.sound_latent_frames_for_sequence_parallel(video_shape=vs, sound_frames=98) == 98 # ulysses=4, with a transfer control folded into the vision base. monkeypatch.setattr(parallel_state, "get_ulysses_parallel_world_size", lambda: 4) padded = model.sound_latent_frames_for_sequence_parallel(video_shape=vs, sound_frames=97, num_vision_items=2) assert (2 * 192 + padded) % 4 == 0 def test_compute_rope_freqs_places_text_video_action_and_sound_positions() -> None: from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer class FakeRotary: def __init__(self) -> None: self.position_ids: list[torch.Tensor] = [] def __call__(self, x, position_ids): del x self.position_ids.append(position_ids.detach().cpu()) batch, seq = position_ids.shape[1], position_ids.shape[2] return torch.zeros(batch, seq, 4), torch.ones(batch, seq, 4) rotary = FakeRotary() model = object.__new__(Cosmos3VFMTransformer) nn.Module.__init__(model) model.language_model = SimpleNamespace(rotary_emb=rotary) model.temporal_modality_margin = 100 model.base_fps = 24.0 model.temporal_compression_factor = 4 model.temporal_compression_factor_sound = 1 model.sound_latent_fps = 25.0 model.enable_fps_modulation = False freqs_und, freqs_gen = model._compute_rope_freqs( text_mask=torch.tensor([[1, 1, 0], [1, 0, 0]], dtype=torch.long), t=2, hp=1, wp=1, fps=24.0, device=torch.device("cpu"), dtype=torch.float32, ) text_pos, vision_pos = rotary.position_ids assert text_pos[:, 0, :].tolist() == [[0, 1, 0], [0, 1, 0], [0, 1, 0]] assert vision_pos[0, 0].tolist() == [102, 103] assert freqs_und[0].shape == (2, 3, 1, 4) assert freqs_gen[0].shape == (2, 2, 1, 4) rotary.position_ids.clear() model._compute_rope_freqs( text_mask=torch.tensor([[1, 1]], dtype=torch.long), t=2, hp=1, wp=1, fps=24.0, device=torch.device("cpu"), dtype=torch.float32, t_action=2, action_start_frame_offset=1, t_sound=1, ) _, gen_pos = rotary.position_ids assert gen_pos.shape == (3, 1, 5) assert gen_pos[0, 0].tolist() == [102, 103, 103, 104, 102] rotary.position_ids.clear() model._compute_rope_freqs( text_mask=torch.tensor([[1, 1]], dtype=torch.long), t=2, hp=1, wp=1, fps=24.0, device=torch.device("cpu"), dtype=torch.float32, num_vision_items=3, share_vision_temporal_positions=True, ) _, shared_gen_pos = rotary.position_ids assert shared_gen_pos[0, 0].tolist() == [102, 103, 102, 103, 102, 103] rotary.position_ids.clear() model._compute_rope_freqs( text_mask=torch.tensor([[1, 1]], dtype=torch.long), t=2, hp=1, wp=1, fps=24.0, device=torch.device("cpu"), dtype=torch.float32, num_vision_items=3, share_vision_temporal_positions=False, ) _, offset_gen_pos = rotary.position_ids assert offset_gen_pos[0, 0].tolist() == [102, 103, 104, 105, 106, 107]