[test]: Cosmos3 local-test parity scaffold (Tier A, 15 tests)
Adds 8 CPU-only parity test files + conftest stubs (StubScheduler, StubCosmos3VAE, StubCosmos3Transformer) under tests/local_tests/cosmos3/ mirroring the vllm-omni reference at tests/diffusion/models/cosmos3/conftest.py. All 15 tests skip until the FastVideo modules land (subsequent commits). Reference: vllm-omni PR #3454 @ 8536f5b1421f.
This commit is contained in:
SolitaryThinker
@@ -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: <https://github.com/vllm-project/vllm-omni/pull/3454>
|
||||
- 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.
|
||||
@@ -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__(<FastVideoCosmos3Pipeline>)`` + ``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)
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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={},
|
||||
),
|
||||
)
|
||||
)
|
||||
@@ -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)
|
||||
@@ -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."
|
||||
)
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user