diff --git a/tests/local_tests/cosmos3/README.md b/tests/local_tests/cosmos3/README.md new file mode 100644 index 000000000..95a38eecf --- /dev/null +++ b/tests/local_tests/cosmos3/README.md @@ -0,0 +1,39 @@ +# Cosmos3 local parity workspace + +## Overview + +Cosmos3 Phase-1 prep tracks a future FastVideo port from the vllm-omni Cosmos3 integration in PR #3454. The current reference supports a single `Cosmos3OmniDiffusersPipeline` for text-to-video (T2V), image-to-video (I2V), and text-to-image (T2I). Follow-up Cosmos3 capabilities mentioned in the PR body, such as sound generation and action-generation modes, are out of scope for the initial FastVideo port. + +## Reference code + +- Reference checkout: `/home/william5lin/cosmos3-reference` +- Source PR: +- Pinned HEAD: `8536f5b1421f78c7df06af6d96fa195c1ceb6384` +- Key files: + - `vllm_omni/deploy/cosmos3.yaml` + - `vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py` + - `vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py` + - `vllm_omni/diffusion/models/cosmos3/guardrails.py` + - `tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py` + - `tests/diffusion/models/cosmos3/test_cosmos3_transformer.py` + +## Weight status + +PENDING. Do not download weights during Phase 1. + +- Candidate serving/model id from the PR body: `nvidia/Cosmos3-Nano` +- Hugging Face API status on 2026-05-22: `401` for `https://huggingface.co/api/models/nvidia/Cosmos3-Nano`; API body reports `Invalid username or password.` +- Hugging Face author search `author=nvidia&search=Cosmos3`: empty list (`[]`). +- NGC API URL requested by handoff returned a Next.js 404 HTML page, not model metadata. + +## Parity-test placeholder + +Phase 2 should add local parity tests here after the official weights become accessible and FastVideo component prototypes exist. Suggested first targets: + +1. Transformer state-dict key/shape inventory versus `Cosmos3VFMTransformer`. +2. Scheduler/default-parameter parity for T2V, I2V, and T2I request modes. +3. Prompt metadata-template parity for duration/resolution and image-vs-video modalities. + +## SSIM placeholder + +No SSIM references seeded yet. Add SSIM coverage only after a FastVideo inference path can load resolved Cosmos3 weights and generate stable T2V/T2I/I2V outputs. diff --git a/tests/local_tests/cosmos3/__init__.py b/tests/local_tests/cosmos3/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/local_tests/cosmos3/conftest.py b/tests/local_tests/cosmos3/conftest.py new file mode 100644 index 000000000..e2dca036b --- /dev/null +++ b/tests/local_tests/cosmos3/conftest.py @@ -0,0 +1,298 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared fixtures for Cosmos3 local parity tests (Phase 2a Tier A). + +Ported from the vllm-omni reference suite so FastVideo's Tier A parity +coverage can land before NVIDIA publishes the real Cosmos3 weights. +The fixtures here are intentionally self-contained: they do NOT import +``vllm_omni`` and do NOT import any FastVideo Cosmos3 modules (those +do not exist yet; they are introduced in Phase 2b). + +Reference: +- ``tests/diffusion/models/cosmos3/conftest.py`` lines 1-176 from + ``vllm-omni`` HEAD ``8536f5b1421f78c7df06af6d96fa195c1ceb6384``. + +The stubs reproduce the same name + protocol as the upstream fixtures +so test code can be ported with minimal renaming once FastVideo's +Cosmos3 pipeline class lands and a ``make_cosmos3_pipeline`` fixture +can be re-pointed at the real class. +""" +from __future__ import annotations + +import sys +import types +from types import SimpleNamespace +from typing import Any + +import pytest +import torch +from torch import nn + + +def pytest_configure(config: pytest.Config) -> None: + """Register ``local`` marker so ``pytestmark = [pytest.mark.local]`` + in sibling test files does not trigger ``PytestUnknownMarkWarning``. + """ + config.addinivalue_line( + "markers", + "local: marker for local-only parity/scaffold tests (skipped in CI)", + ) + + +# --------------------------------------------------------------------------- +# Stub scheduler — mirrors vllm-omni ``StubScheduler``. +# --------------------------------------------------------------------------- +class StubScheduler: + """Minimal UniPC-shaped scheduler stub used by pipeline call-graph tests. + + Mirrors ``tests/diffusion/models/cosmos3/conftest.py:16-30`` from + the vllm-omni reference: records ``set_timesteps`` and ``step`` calls, + advances ``latents`` by ``+noise_pred``, and exposes a ``config`` + namespace carrying ``num_train_timesteps`` and ``flow_shift``. + """ + + def __init__( + self, + timesteps: list[int] | None = None, + *, + flow_shift: float = 1.0, + ) -> None: + self.timesteps = torch.tensor(timesteps or [9, 3], dtype=torch.int64) + self.config = SimpleNamespace(num_train_timesteps=1000, flow_shift=flow_shift) + self.set_timesteps_calls: list[tuple[int, torch.device]] = [] + self.step_calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] + + def set_timesteps(self, num_steps: int, device: torch.device) -> None: + self.set_timesteps_calls.append((num_steps, device)) + self.timesteps = torch.arange(num_steps, 0, -1, dtype=torch.int64, device=device) + + def step( + self, + noise_pred: torch.Tensor, + timestep: torch.Tensor, + latents: torch.Tensor, + **kwargs: Any, + ): + del kwargs + self.step_calls.append((noise_pred.clone(), timestep.clone(), latents.clone())) + return (latents + noise_pred,) + + +class _ModeLatentDist: + """Stub for diffusers ``DiagonalGaussianDistribution.mode``.""" + + def __init__(self, latents: torch.Tensor) -> None: + self._latents = latents + + def mode(self) -> torch.Tensor: + return self._latents + + +class StubCosmos3VAE: + """VAE stub returning deterministic latents shaped by VAE scale factors. + + Mirrors ``tests/diffusion/models/cosmos3/conftest.py:41-70``. + """ + + dtype = torch.float32 + + def __init__(self, z_dim: int = 2, *, temporal: int = 4, spatial: int = 8) -> None: + self.config = SimpleNamespace( + z_dim=z_dim, + scale_factor_temporal=temporal, + scale_factor_spatial=spatial, + latents_mean=[0.0] * z_dim, + latents_std=[1.0] * z_dim, + ) + + def encode(self, video: torch.Tensor): + latent_frames = (video.shape[2] - 1) // self.config.scale_factor_temporal + 1 + latent_height = video.shape[-2] // self.config.scale_factor_spatial + latent_width = video.shape[-1] // self.config.scale_factor_spatial + latents = torch.ones( + video.shape[0], + self.config.z_dim, + latent_frames, + latent_height, + latent_width, + dtype=video.dtype, + device=video.device, + ) + return SimpleNamespace(latent_dist=_ModeLatentDist(latents)) + + def decode(self, latents: torch.Tensor, return_dict: bool = False): + del return_dict + return (latents,) + + +class StubCosmos3Transformer(nn.Module): + """Transformer stub that records per-call inputs and emits deterministic + output tensors keyed by the first ``text_ids`` token. + + Mirrors ``tests/diffusion/models/cosmos3/conftest.py:73-114``. + + The stub also exposes ``cached_kv`` / ``cached_freqs_gen`` so pipeline + tests can assert that: + * ``reset_cache()`` is called before each diffusion loop; + * UND/cond and uncond branches each populate their own cache exactly + once and reuse it on subsequent timesteps. + """ + + def __init__(self, *, latent_channel_size: int = 2) -> None: + super().__init__() + self.latent_channel_size = latent_channel_size + self.cached_kv: Any | None = None + self.cached_freqs_gen: Any | None = None + self.calls: list[dict[str, Any]] = [] + self.reset_calls = 0 + + def reset_cache(self) -> None: + self.reset_calls += 1 + self.cached_kv = None + self.cached_freqs_gen = None + + def forward( + self, + *, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + text_ids: torch.Tensor, + text_mask: torch.Tensor, + **kwargs: Any, + ) -> torch.Tensor: + token = int(text_ids.reshape(-1)[0].item()) if text_ids.numel() else 0 + self.calls.append( + { + "token": token, + "timestep": timestep.clone(), + "text_mask": text_mask.clone(), + "cache_before": self.cached_kv, + "kwargs": dict(kwargs), + } + ) + if self.cached_kv is None: + marker = torch.tensor([token], dtype=torch.float32) + self.cached_kv = [(marker, marker + 100)] + self.cached_freqs_gen = (marker + 200, marker + 300) + return torch.full_like(hidden_states, float(token)) + + +def passthrough_progress_bar(iterable): + return iterable + + +# --------------------------------------------------------------------------- +# Tiny config — mirrors test_cosmos3_transformer.py:15-29 +# --------------------------------------------------------------------------- +def _tiny_cosmos3_config(**overrides: Any) -> dict: + """Minimal Cosmos3 transformer config sufficient for shape-only construction. + + Mirrors ``tests/diffusion/models/cosmos3/test_cosmos3_transformer.py:15-29``. + """ + config: dict = { + "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 + + +# --------------------------------------------------------------------------- +# Guardrail no-op stub — mirrors conftest.py:121-129 +# --------------------------------------------------------------------------- +@pytest.fixture(autouse=True) +def fake_cosmos3_guardrails(monkeypatch: pytest.MonkeyPatch): + """Install a no-op replacement for the Cosmos3 guardrails module. + + The vllm-omni reference imports guardrails eagerly; FastVideo's port + may either skip guardrails or use a different module path. Either + way, an autouse stub avoids accidental network-dependent imports + during scaffold-test collection. + """ + module = types.ModuleType("vllm_omni.diffusion.models.cosmos3.guardrails") + module.is_guardrails_enabled = lambda od_config, sampling_params=None: False # type: ignore[attr-defined] + module.ensure_initialized = lambda od_config: None # type: ignore[attr-defined] + module.check_text_safety = lambda text: None # type: ignore[attr-defined] + module.check_video_safety = lambda video: video # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, module.__name__, module) + return module + + +# --------------------------------------------------------------------------- +# Pipeline factory — kept as a scaffold; activate in Phase 2b when the +# FastVideo Cosmos3 pipeline class exists. +# --------------------------------------------------------------------------- +@pytest.fixture +def make_cosmos3_pipeline(): + """Factory that returns a FastVideo Cosmos3 pipeline pre-wired with stubs. + + In Phase 2a (Tier A, no real weights), tests that consume this factory + should ``pytest.skip`` if the FastVideo Cosmos3 pipeline class does not + yet exist. In Phase 2b we replace the placeholder construction with + ``object.__new__()`` + ``nn.Module.__init__`` + and re-point the stubs onto the real attribute names. + + Mirrors the upstream layout at + ``tests/diffusion/models/cosmos3/conftest.py:132-157``. + """ + + def _make(): + try: + from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( # type: ignore + Cosmos3OmniDiffusersPipeline, + ) + except ImportError: + pytest.skip( + "FastVideo Cosmos3 pipeline class not yet implemented " + "(Phase 2b will provide fastvideo.pipelines.basic.cosmos3)." + ) + + pipeline = object.__new__(Cosmos3OmniDiffusersPipeline) + nn.Module.__init__(pipeline) + pipeline.od_config = SimpleNamespace() + pipeline.device = torch.device("cpu") + pipeline.dtype = torch.float32 + pipeline.transformer = StubCosmos3Transformer(latent_channel_size=2) + pipeline.vae = StubCosmos3VAE(z_dim=2) + pipeline.vae_scale_factor_temporal = 4 + pipeline.vae_scale_factor_spatial = 8 + pipeline.scheduler = StubScheduler([9, 3], flow_shift=1.0) + pipeline._base_scheduler_config = pipeline.scheduler.config + pipeline._engine_init_flow_shift = 1.0 + pipeline._current_flow_shift = 1.0 + pipeline._guidance_scale = None + pipeline._num_timesteps = None + pipeline.progress_bar = passthrough_progress_bar + return pipeline + + return _make + + +def make_sampling_params(**overrides: Any) -> SimpleNamespace: + """Build a SamplingParams-like namespace with Cosmos3's expected fields. + + Mirrors ``tests/diffusion/models/cosmos3/conftest.py:160-176``. + """ + values = { + "height": None, + "width": None, + "num_frames": None, + "num_inference_steps": None, + "guidance_scale": None, + "generator": None, + "seed": 123, + "num_outputs_per_prompt": 1, + "frame_rate": None, + "resolved_frame_rate": None, + "max_sequence_length": None, + "extra_args": {}, + } + values.update(overrides) + return SimpleNamespace(**values) diff --git a/tests/local_tests/cosmos3/test_cosmos3_mrope_parity.py b/tests/local_tests/cosmos3/test_cosmos3_mrope_parity.py new file mode 100644 index 000000000..bb29a18e5 --- /dev/null +++ b/tests/local_tests/cosmos3/test_cosmos3_mrope_parity.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Cosmos3 unified 3D mRoPE position-ID parity (Tier A scaffold). + +Reference: ``vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py`` +lines 113-177 (``compute_mrope_position_ids_text`` / +``compute_mrope_position_ids_vision``). The reference test asserting +these invariants lives at +``tests/diffusion/models/cosmos3/test_cosmos3_transformer.py:32-57``. + +The three invariants under test: + + 1. Text tokens broadcast the same monotonically-increasing positions + across all three (t, h, w) axes. With ``num_tokens=3`` and + ``temporal_offset=5`` the result is ``[[5,6,7], [5,6,7], [5,6,7]]`` + and the next-offset is ``8``. + + 2. Vision tokens (no FPS modulation) flatten a ``(grid_t, grid_h, grid_w)`` + position grid in t-major order. With ``(2, 2, 3)`` and offset ``10`` + the resulting shape is ``(3, 12)`` and the temporal row begins + ``[10]*6 + [11]*6``; next-offset is ``12``. + + 3. FPS-modulated vision tokens scale the temporal axis by + ``base_fps / temporal_compression_factor / (fps / tcf)``. With + ``fps=12``, ``base_fps=24``, ``tcf=4``, ``grid_t=2`` the first row is + ``[10.0, 12.0]``. + +The FastVideo side currently does NOT exist; the test is wrapped in +``try/except ImportError`` and skips. Phase 2b replaces the skip with +the real import + assertion path. +""" +from __future__ import annotations + +import pytest +import torch + +pytestmark = [pytest.mark.local] + + +def test_compute_mrope_position_ids_text_and_vision() -> None: + """Asserts the 3 invariants of unified 3D mRoPE position-ID generation. + + Once FastVideo's ``fastvideo.models.dits.cosmos3`` exports + ``compute_mrope_position_ids_text`` and + ``compute_mrope_position_ids_vision``, this test verifies they produce + output tensors identical to the vllm-omni reference at + transformer_cosmos3.py:113-177. + """ + try: + from fastvideo.models.dits.cosmos3 import ( # type: ignore + compute_mrope_position_ids_text, + compute_mrope_position_ids_vision, + ) + except ImportError: + pytest.skip("FastVideo Cosmos3 not yet implemented (Phase 2b)") + + 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 tuple(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 diff --git a/tests/local_tests/cosmos3/test_cosmos3_patchify_unpatchify_parity.py b/tests/local_tests/cosmos3/test_cosmos3_patchify_unpatchify_parity.py new file mode 100644 index 000000000..7ee2978ba --- /dev/null +++ b/tests/local_tests/cosmos3/test_cosmos3_patchify_unpatchify_parity.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Cosmos3 ``[B,C,T,H,W] <-> [B, T*hp*wp, p*p*C]`` patchify roundtrip (Tier A). + +Reference: ``vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py`` +lines 1009-1036 (``Cosmos3VFMTransformer.patchify`` / +``Cosmos3VFMTransformer.unpatchify``) and the reference assertion at +``tests/diffusion/models/cosmos3/test_cosmos3_transformer.py:98-101``. + +Invariant: ``unpatchify(patchify(x)) == x`` for any ``x`` with shape +``(B, C, t, h, w)`` where ``h, w`` are divisible by ``latent_patch_size``. +Also exercises a non-trivial channel count (3) to ensure the +``permute([0, 2, 3, 5, 4, 6, 1])`` reordering is correct. +""" +from __future__ import annotations + +import pytest +import torch + +pytestmark = [pytest.mark.local] + + +def test_patchify_unpatchify_roundtrip() -> None: + """Asserts that the FastVideo Cosmos3 transformer's patchify/unpatchify + pair are exact inverses for ``latent_patch_size=2``, ``latent_channel=3``. + + Once FastVideo's ``fastvideo.models.dits.cosmos3.Cosmos3VFMTransformer`` + lands, replace the skip with the upstream-equivalent assertion path. + """ + try: + from fastvideo.models.dits.cosmos3 import Cosmos3VFMTransformer # type: ignore + except ImportError: + pytest.skip("FastVideo Cosmos3 not yet implemented (Phase 2b)") + + from torch import nn + + model = object.__new__(Cosmos3VFMTransformer) + nn.Module.__init__(model) + 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) + roundtrip = model.unpatchify(model.patchify(latents, t=1, h=3, w=5), t=1, h=3, w=5) + torch.testing.assert_close(roundtrip, latents) + + +def test_patchify_default_patch_size() -> None: + """Asserts shape contract for the default ``latent_patch_size=[1,2,2]`` + (i.e. spatial-only patching) with a representative video latent. + + With ``[B,C,T,H,W] = [1, 16, 2, 8, 8]`` and patch=2 on H/W, expected + flattened tokens = ``T * (H/2) * (W/2) = 2 * 4 * 4 = 32`` and each token + carries ``2*2*C = 4*16 = 64`` channels. + """ + try: + from fastvideo.models.dits.cosmos3 import Cosmos3VFMTransformer # type: ignore + except ImportError: + pytest.skip("FastVideo Cosmos3 not yet implemented (Phase 2b)") + + from torch import nn + + model = object.__new__(Cosmos3VFMTransformer) + nn.Module.__init__(model) + model.latent_patch_size = 2 + model.latent_channel_size = 16 + + latents = torch.zeros(1, 16, 2, 8, 8) + tokens = model.patchify(latents, t=2, h=8, w=8) + assert tuple(tokens.shape) == (1, 32, 64) + restored = model.unpatchify(tokens, t=2, h=8, w=8) + assert tuple(restored.shape) == (1, 16, 2, 8, 8) diff --git a/tests/local_tests/cosmos3/test_cosmos3_pipeline_call_graph.py b/tests/local_tests/cosmos3/test_cosmos3_pipeline_call_graph.py new file mode 100644 index 000000000..a33eb2823 --- /dev/null +++ b/tests/local_tests/cosmos3/test_cosmos3_pipeline_call_graph.py @@ -0,0 +1,199 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Cosmos3 pipeline call-graph parity (Tier A scaffold). + +Reference: + * ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py:883-1033`` — + ``Cosmos3OmniDiffusersPipeline.diffuse``: 3-mode CFG denoising loop + with UND-cache management. + * ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py:1037-1067`` — + ``forward``: parses the request, selects T2I/T2V/I2V mode, dispatches + defaults. + * Reference test invariants at + ``tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py:126-156,196-251``. + +The Tier A scaffold uses the stubs from ``conftest.py`` (StubScheduler, +StubCosmos3VAE, StubCosmos3Transformer) so the call-graph can be tested +without a real DiT or VAE. The 4 invariants under test: + + 1. ``diffuse(...)`` calls ``transformer.reset_cache()`` exactly once + before iterating timesteps; + 2. With ``do_cfg=True`` and ``guidance_interval=None``, each timestep + invokes the transformer twice — once with ``cond_ids`` and once with + ``uncond_ids`` — and reuses the cached UND K/V from step 1 onward + (asserted via ``StubCosmos3Transformer.calls`` order); + 3. I2V mode: ``velocity_mask`` zeros frame-0 noise predictions, and + ``image_latent`` is re-injected into frame 0 after each scheduler step; + 4. ``forward`` selects T2I vs T2V mode from ``prompt["modalities"]`` + and applies the per-mode default ``flow_shift`` / ``num_frames``. +""" +from __future__ import annotations + +import pytest +import torch + +pytestmark = [pytest.mark.local] + + +def _ids(value: int) -> torch.Tensor: + return torch.tensor([[value]], dtype=torch.long) + + +def _mask() -> torch.Tensor: + return torch.ones(1, 1, dtype=torch.long) + + +def test_diffuse_resets_cache_and_calls_cfg_in_order(make_cosmos3_pipeline) -> None: + """Asserts the 3-mode CFG call graph: ``reset_cache`` first, then + interleaved ``cond / uncond`` transformer calls with cache reuse. + + With timesteps ``[900, 100]`` and ``guidance_scale=3.0``, the upstream + reference (test_cosmos3_pipeline.py:126-142) asserts the call order is + ``[2, 1, 2]`` — i.e. step 0 calls cond (token=2) then uncond (token=1), + step 1 reuses the cached UND K/V and only calls cond again (token=2). + """ + pipeline = make_cosmos3_pipeline() + latents = torch.zeros(1, 2, 1, 1, 1) + + result = pipeline.diffuse( + latents=latents, + timesteps=torch.tensor([900, 100]), + cond_ids=_ids(2), + cond_mask=_mask(), + uncond_ids=_ids(1), + uncond_mask=_mask(), + guidance_scale=3.0, + shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0}, + guidance_interval=(500.0, 1000.0), + ) + assert pipeline.transformer.reset_calls == 1 + assert [call["token"] for call in pipeline.transformer.calls] == [2, 1, 2] + torch.testing.assert_close(result, torch.full_like(latents, 6.0)) + + +def test_diffuse_i2v_velocity_mask_zeros_frame_zero(make_cosmos3_pipeline) -> None: + """Asserts I2V velocity-mask + image-latent re-injection contract. + + Cross-check: pipeline_cosmos3.py:937-951. The velocity_mask zeroes + noise predictions on conditioning frames before stepping, and + ``image_latent`` is overwritten into frame 0 of the output latents + after each scheduler step. + """ + pipeline = make_cosmos3_pipeline() + result = pipeline.diffuse( + latents=torch.zeros(1, 2, 2, 1, 1), + timesteps=torch.tensor([7]), + cond_ids=_ids(2), + cond_mask=_mask(), + uncond_ids=_ids(1), + uncond_mask=_mask(), + guidance_scale=1.0, + shared_kwargs={"video_shape": (2, 1, 1), "fps": 24.0}, + velocity_mask=torch.tensor([[[[[0.0]], [[1.0]]]]]), + image_latent=torch.full((1, 2, 1, 1, 1), 7.0), + ) + torch.testing.assert_close(result[:, :, 0:1], torch.full((1, 2, 1, 1, 1), 7.0)) + + +@pytest.mark.parametrize( + ("modalities", "expected_is_t2i", "expected_default_flow_shift", "expected_default_frames"), + [ + (["image"], True, 3.0, 1), + (["video"], False, 1.0, 189), + ], +) +def test_forward_mode_dispatch_t2i_vs_t2v( + make_cosmos3_pipeline, + modalities: list[str], + expected_is_t2i: bool, + expected_default_flow_shift: float, + expected_default_frames: int, +) -> None: + """Asserts forward() routes to T2I vs T2V mode based on prompt modalities, + and applies per-mode defaults. + + Cross-check: pipeline_cosmos3.py:1069-1093. T2I defaults: + ``num_frames=1``, ``flow_shift=3.0``, ``num_inference_steps=50``, + ``guidance_interval=[400, 1000]``. T2V defaults: ``num_frames=189``, + ``flow_shift=engine_init`` (1.0 here), ``num_inference_steps=35``, + no guidance_interval. + """ + pipeline = make_cosmos3_pipeline() + + from types import SimpleNamespace + + captured: dict[str, object] = {"flow_shifts": [], "format_calls": []} + + def fake_format(prompt, negative_prompt, num_frames, frame_rate, height, width, *args, **kwargs): + captured["format_calls"].append( + { + "is_t2i": kwargs.get("is_t2i"), + "num_frames": num_frames, + } + ) + return _ids(2), _mask(), _ids(1), _mask() + + pipeline._format_and_tokenize_prompts = fake_format + pipeline._prepare_latents = lambda *a, **kw: torch.zeros(1, 2, 1, 1, 1) + pipeline._set_flow_shift = lambda target: captured["flow_shifts"].append(target) + pipeline._set_scheduler_timesteps = lambda steps: setattr( + pipeline.scheduler, "timesteps", torch.tensor([7]) + ) + pipeline.diffuse = lambda **kw: kw["latents"] + pipeline._decode_latents = lambda latents: latents + + output = pipeline.forward( + SimpleNamespace( + prompts=[{"prompt": "test", "modalities": modalities}], + sampling_params=SimpleNamespace( + height=None, + width=None, + num_frames=None, + num_inference_steps=None, + guidance_scale=None, + generator=None, + seed=123, + num_outputs_per_prompt=1, + frame_rate=None, + resolved_frame_rate=None, + max_sequence_length=None, + extra_args={}, + ), + ) + ) + assert captured["format_calls"][-1]["is_t2i"] is expected_is_t2i + assert captured["format_calls"][-1]["num_frames"] == expected_default_frames + assert captured["flow_shifts"] == [expected_default_flow_shift] + expected_output_key = "image" if expected_is_t2i else "video" + assert expected_output_key in output.output + + +def test_forward_rejects_both_image_and_video_modalities(make_cosmos3_pipeline) -> None: + """Asserts ``_is_t2i_request`` raises ValueError when a prompt + requests both image and video modalities simultaneously. + + Cross-check: pipeline_cosmos3.py:490-496. + """ + pipeline = make_cosmos3_pipeline() + + from types import SimpleNamespace + + with pytest.raises(ValueError, match="both image and video"): + pipeline.forward( + SimpleNamespace( + prompts=[{"prompt": "x", "modalities": ["image", "video"]}], + sampling_params=SimpleNamespace( + height=None, + width=None, + num_frames=None, + num_inference_steps=None, + guidance_scale=None, + generator=None, + seed=123, + num_outputs_per_prompt=1, + frame_rate=None, + resolved_frame_rate=None, + max_sequence_length=None, + extra_args={}, + ), + ) + ) diff --git a/tests/local_tests/cosmos3/test_cosmos3_scheduler_default_parity.py b/tests/local_tests/cosmos3/test_cosmos3_scheduler_default_parity.py new file mode 100644 index 000000000..aaf90f6c7 --- /dev/null +++ b/tests/local_tests/cosmos3/test_cosmos3_scheduler_default_parity.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Cosmos3 scheduler default + per-request override parity (Tier A scaffold). + +Reference: + * ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py:275-307`` — + initial UniPCMultistepScheduler load (preserves solver_order, + timestep_spacing, beta_schedule, sigma bounds, flow_shift) and + one-time override at engine-init if ``od_config.flow_shift`` is set. + * ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py:498-512`` — + ``_set_flow_shift(target_shift)``: rebuild the scheduler via + ``UniPCMultistepScheduler.from_config(base_config, flow_shift=target)`` + when the requested target differs from the current shift. + * ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py:1069-1110`` — + per-request mode defaults: T2I uses ``shift=3.0``; T2V/I2V use the + engine-init shift (typically 1.0); ``flow_shift`` may be overridden + per request via ``sampling_params.extra_args["flow_shift"]``. + +The invariant under test: for the same RNG seed and the same number of +inference steps, the scheduler's ``timesteps`` tensor must be identical +whenever the ``flow_shift`` is identical, and must change deterministically +when ``flow_shift`` is overridden via ``_set_flow_shift``. +""" +from __future__ import annotations + +import pytest +import torch + +pytestmark = [pytest.mark.local] + + +def test_t2i_default_flow_shift_is_3() -> None: + """Asserts that T2I requests rebuild the scheduler at ``flow_shift=3.0``. + + Cross-check: pipeline_cosmos3.py:1073-1080 sets + ``default_flow_shift = 3.0`` for T2I, and + pipeline_cosmos3.py:1110 calls ``self._set_flow_shift(flow_shift_target)`` + which rebuilds the scheduler via ``UniPCMultistepScheduler.from_config( + base_config, flow_shift=3.0)``. + """ + try: + from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( # type: ignore + Cosmos3OmniDiffusersPipeline, + ) + except ImportError: + pytest.skip("FastVideo Cosmos3 pipeline not yet implemented (Phase 2b)") + + pipeline = Cosmos3OmniDiffusersPipeline.__new__(Cosmos3OmniDiffusersPipeline) + if not hasattr(pipeline, "_set_flow_shift") or not hasattr(pipeline, "scheduler"): + pytest.skip("FastVideo Cosmos3 scheduler/_set_flow_shift not yet wired") + + pipeline._set_flow_shift(3.0) + assert float(pipeline.scheduler.config.flow_shift) == 3.0 + + +def test_t2v_default_flow_shift_is_engine_init() -> None: + """Asserts T2V/I2V use the engine-init shift (e.g. 1.0), NOT a fixed default. + + Cross-check: pipeline_cosmos3.py:1091 sets + ``default_flow_shift = self._engine_init_flow_shift`` for T2V/I2V + (NOT ``None`` — passing ``None`` would leak a prior T2I rebuild + forward). + """ + try: + from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( # type: ignore + Cosmos3OmniDiffusersPipeline, + ) + except ImportError: + pytest.skip("FastVideo Cosmos3 pipeline not yet implemented (Phase 2b)") + + pipeline = Cosmos3OmniDiffusersPipeline.__new__(Cosmos3OmniDiffusersPipeline) + if not hasattr(pipeline, "_engine_init_flow_shift") or not hasattr(pipeline, "_set_flow_shift"): + pytest.skip("FastVideo Cosmos3 _engine_init_flow_shift not yet wired") + + init_shift = float(pipeline._engine_init_flow_shift) + pipeline._set_flow_shift(init_shift) + assert float(pipeline.scheduler.config.flow_shift) == init_shift + + +def test_scheduler_timesteps_deterministic_under_seed() -> None: + """Asserts that ``scheduler.set_timesteps(N)`` is deterministic given the + same N and the same flow_shift. + + The UniPC scheduler's timestep sequence does not depend on a torch + seed (it's a closed-form function of N + scheduler config), so + invoking ``set_timesteps`` twice should produce identical tensors. + """ + try: + from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( # type: ignore + Cosmos3OmniDiffusersPipeline, + ) + except ImportError: + pytest.skip("FastVideo Cosmos3 pipeline not yet implemented (Phase 2b)") + + pipeline = Cosmos3OmniDiffusersPipeline.__new__(Cosmos3OmniDiffusersPipeline) + if not hasattr(pipeline, "scheduler") or not hasattr(pipeline, "_set_flow_shift"): + pytest.skip("FastVideo Cosmos3 scheduler not yet wired") + + pipeline._set_flow_shift(3.0) + pipeline.scheduler.set_timesteps(35, device=torch.device("cpu")) + seq_a = pipeline.scheduler.timesteps.clone() + + pipeline.scheduler.set_timesteps(35, device=torch.device("cpu")) + seq_b = pipeline.scheduler.timesteps.clone() + + torch.testing.assert_close(seq_a, seq_b) diff --git a/tests/local_tests/cosmos3/test_cosmos3_state_dict_keys.py b/tests/local_tests/cosmos3/test_cosmos3_state_dict_keys.py new file mode 100644 index 000000000..45b31e9da --- /dev/null +++ b/tests/local_tests/cosmos3/test_cosmos3_state_dict_keys.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Cosmos3 state-dict-key contract for checkpoint conversion (Tier A scaffold). + +Reference: ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py:319-409`` +(``Cosmos3OmniDiffusersPipeline._remap_ckpt_key``). The vllm-omni loader's +remap is the canonical source of truth: it defines which Diffusers +checkpoint keys become which target module-tree parameter names. + +The FastVideo conversion script at +``scripts/checkpoint_conversion/cosmos3_convert.py`` (to be authored in +Phase 5) must produce a state-dict that loads cleanly into the FastVideo +Cosmos3 DiT module tree. This test pins the target side of that contract: + + * ``embed_tokens``/``norm`` live under ``language_model.`` + * ``norm_moe_gen`` is a top-level (not per-layer) module + * Per-layer keys split into UND (``language_model.layers.{i}``) and + GEN (``gen_layers.{i}``) sub-trees: + - UND attention: ``self_attn.{q,k,v,o}_proj``, ``self_attn.{q,k}_norm`` + - GEN cross-attention: ``cross_attention.{q,k,v,o}_proj``, + ``cross_attention.{q,k}_norm`` (from the ``*_moe_gen`` source keys) + - UND/GEN norms: ``input_layernorm``, ``post_attention_layernorm`` + - UND/GEN MLPs: ``mlp.{gate,up,down}_proj`` + * Top-level adapters ``vae2llm``, ``llm2vae``, ``time_embedder`` live + under ``transformer.`` (no remapping inside the transformer namespace). + * ``lm_head.weight`` is skipped (mapped to ``None``). +""" +from __future__ import annotations + +import pytest + +pytestmark = [pytest.mark.local] + + +EXPECTED_REMAPS: dict[str, str | None] = { + "model.embed_tokens.weight": "transformer.language_model.embed_tokens.weight", + "model.norm.weight": "transformer.language_model.norm.weight", + "model.norm_moe_gen.weight": "transformer.norm_moe_gen.weight", + "model.layers.3.self_attn.q_proj.weight": "transformer.language_model.layers.3.self_attn.q_proj.weight", + "model.layers.3.self_attn.q_proj_moe_gen.weight": "transformer.gen_layers.3.cross_attention.q_proj.weight", + "model.layers.3.self_attn.k_norm_moe_gen.weight": "transformer.gen_layers.3.cross_attention.k_norm.weight", + "model.layers.3.input_layernorm.weight": "transformer.language_model.layers.3.input_layernorm.weight", + "model.layers.3.input_layernorm_moe_gen.weight": "transformer.gen_layers.3.input_layernorm.weight", + "model.layers.3.mlp.gate_proj.weight": "transformer.language_model.layers.3.mlp.gate_proj.weight", + "model.layers.3.mlp_moe_gen.up_proj.weight": "transformer.gen_layers.3.mlp.up_proj.weight", + "vae2llm.weight": "transformer.vae2llm.weight", + "llm2vae.weight": "transformer.llm2vae.weight", + "time_embedder.linear_1.weight": "transformer.time_embedder.linear_1.weight", + "lm_head.weight": None, +} + + +def test_fastvideo_cosmos3_remap_matches_reference() -> None: + """Asserts the FastVideo Cosmos3 checkpoint remap table matches the + canonical vllm-omni reference at pipeline_cosmos3.py:319-409. + + Once the FastVideo Cosmos3 pipeline (or a standalone remap utility + in ``scripts/checkpoint_conversion/cosmos3_convert.py``) exposes a + ``_remap_ckpt_key`` callable, this test verifies every key in + ``EXPECTED_REMAPS`` maps to the expected target name (or ``None`` + for skipped keys like ``lm_head.weight``). + """ + try: + from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( # type: ignore + Cosmos3OmniDiffusersPipeline, + ) + except ImportError: + pytest.skip("FastVideo Cosmos3 pipeline not yet implemented (Phase 2b)") + + remap = getattr(Cosmos3OmniDiffusersPipeline, "_remap_ckpt_key", None) + if remap is None: + pytest.skip("Cosmos3OmniDiffusersPipeline._remap_ckpt_key not yet defined") + + actual = {key: remap(key) for key in EXPECTED_REMAPS} + assert actual == EXPECTED_REMAPS + + +def test_fastvideo_cosmos3_dit_module_tree_param_names() -> None: + """Asserts the FastVideo Cosmos3 DiT module tree produces the param + names that the checkpoint converter is expected to write. + + Once ``fastvideo.models.dits.cosmos3.Cosmos3VFMTransformer`` is + constructible with a tiny config (see ``conftest._tiny_cosmos3_config``), + this test instantiates a 1-layer model and verifies that the + parameter-name set contains the expected UND/GEN/adapter prefixes: + * ``language_model.layers.0.self_attn.q_proj.weight`` + * ``gen_layers.0.cross_attention.q_proj.weight`` + * ``norm_moe_gen.weight`` + * ``vae2llm.*`` and ``llm2vae.*`` + """ + try: + from fastvideo.models.dits.cosmos3 import Cosmos3VFMTransformer # type: ignore + except ImportError: + pytest.skip("FastVideo Cosmos3 not yet implemented (Phase 2b)") + + from types import SimpleNamespace + import torch + + from .conftest import _tiny_cosmos3_config + + config = _tiny_cosmos3_config(num_hidden_layers=1) + model = Cosmos3VFMTransformer(SimpleNamespace(tf_model_config=config, dtype=torch.float32)) + names = set(name for name, _ in model.named_parameters()) + + required_prefixes = [ + "language_model.layers.0.self_attn.q_proj", + "gen_layers.0.cross_attention.q_proj", + "norm_moe_gen", + "vae2llm", + "llm2vae", + ] + for prefix in required_prefixes: + assert any(name.startswith(prefix) for name in names), ( + f"Cosmos3 DiT module tree missing expected prefix {prefix!r}; " + f"converter at scripts/checkpoint_conversion/cosmos3_convert.py " + f"will fail to load weights for this subtree." + ) diff --git a/tests/local_tests/cosmos3/test_cosmos3_tokenizer_chat_template.py b/tests/local_tests/cosmos3/test_cosmos3_tokenizer_chat_template.py new file mode 100644 index 000000000..89dc13fea --- /dev/null +++ b/tests/local_tests/cosmos3/test_cosmos3_tokenizer_chat_template.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Cosmos3 prompt tokenization parity — chat template + special tokens (Tier A). + +Reference: ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py`` +lines 562-606 (``Cosmos3OmniDiffusersPipeline._tokenize_prompt``). + +The Cosmos3 tokenizer is a Qwen2 chat-template tokenizer with two +appended special tokens: + + * ``eos_token_id`` == 151645 + * ``<|vision_start|>`` == 151652 + +The reference pipeline: + 1. Wraps the prompt in a ``role=user`` conversation (optionally with a + ``role=system`` prefix carrying ``COSMOS3_SYSTEM_PROMPT``); + 2. Applies the chat template with ``add_generation_prompt=True``; + 3. Truncates to ``max_sequence_length``; + 4. Appends ``eos`` and ``<|vision_start|>``; + 5. Right-pads with ``pad_token_id`` and produces a matching + attention mask of ``[1]*seq_len + [0]*pad_len``. + +This Tier A scaffold validates the two special-token IDs and the +right-pad shape contract. Full byte-for-byte parity against an +official Cosmos3 tokenizer requires Phase 2b weights. +""" +from __future__ import annotations + +import pytest + +pytestmark = [pytest.mark.local] + + +def test_cosmos3_special_token_ids() -> None: + """Asserts the two Cosmos3 special-token IDs the pipeline depends on. + + Specifically ``eos_token_id == 151645`` and + ``convert_tokens_to_ids('<|vision_start|>') == 151652``. These IDs + are pinned by the Qwen2 base tokenizer and are appended in + ``Cosmos3OmniDiffusersPipeline._tokenize_prompt`` at + pipeline_cosmos3.py:596-597. + """ + try: + from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( # type: ignore + Cosmos3OmniDiffusersPipeline, + ) + except ImportError: + pytest.skip("FastVideo Cosmos3 pipeline not yet implemented (Phase 2b)") + + pipeline = Cosmos3OmniDiffusersPipeline.__new__(Cosmos3OmniDiffusersPipeline) + tokenizer = getattr(pipeline, "tokenizer", None) + if tokenizer is None: + pytest.skip("Cosmos3 tokenizer instance not yet wired on the pipeline") + + assert tokenizer.eos_token_id == 151645 + assert tokenizer.convert_tokens_to_ids("<|vision_start|>") == 151652 + + +def test_cosmos3_chat_template_shape_contract() -> None: + """Asserts the right-pad shape contract of ``_tokenize_prompt``. + + For ``max_sequence_length = N``, the returned ``input_ids`` and + ``attention_mask`` must each be ``[1, N]``, with the first + ``seq_len`` mask entries equal to 1 and the remainder equal to 0. + + Once Phase 2b lands, this test should: + * call ``_tokenize_prompt`` directly with a small max_seq_length; + * assert ``input_ids.shape == (1, max_seq_length)``; + * assert ``attention_mask.sum().item() == seq_len`` and that the + last two non-pad tokens are ``[eos, vision_start]``. + """ + try: + from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( # type: ignore + Cosmos3OmniDiffusersPipeline, + ) + except ImportError: + pytest.skip("FastVideo Cosmos3 pipeline not yet implemented (Phase 2b)") + + pipeline = Cosmos3OmniDiffusersPipeline.__new__(Cosmos3OmniDiffusersPipeline) + if not hasattr(pipeline, "_tokenize_prompt"): + pytest.skip("_tokenize_prompt not yet wired on the FastVideo pipeline") + + input_ids, attention_mask = pipeline._tokenize_prompt( + "A robot.", max_sequence_length=32, use_system_prompt=False + ) + assert tuple(input_ids.shape) == (1, 32) + assert tuple(attention_mask.shape) == (1, 32) + seq_len = int(attention_mask.sum().item()) + assert attention_mask[0, :seq_len].sum().item() == seq_len + assert attention_mask[0, seq_len:].sum().item() == 0 + assert int(input_ids[0, seq_len - 2].item()) == 151645 + assert int(input_ids[0, seq_len - 1].item()) == 151652