From 7ae57e6fbd5a46d2b653cf26a02491e6af277bcd Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sat, 6 Jun 2026 23:54:32 -0700 Subject: [PATCH] [feat] cosmos3: native video pipeline + framework denoise parity --- fastvideo/configs/pipelines/cosmos3.py | 35 +- .../basic/cosmos3/cosmos3_pipeline.py | 973 ++++++++---------- fastvideo/pipelines/basic/cosmos3/presets.py | 85 ++ fastvideo/pipelines/stages/cosmos3_stages.py | 230 +++++ fastvideo/registry.py | 3 + tests/local_tests/cosmos3/conftest.py | 374 +++---- .../test_cosmos3_denoise_cfg_parity.py | 351 +++++++ .../test_cosmos3_pipeline_call_graph.py | 337 +++--- .../cosmos3/test_cosmos3_pipeline_smoke.py | 310 ++++++ .../test_cosmos3_tokenizer_chat_template.py | 134 ++- 10 files changed, 1810 insertions(+), 1022 deletions(-) create mode 100644 fastvideo/pipelines/basic/cosmos3/presets.py create mode 100644 fastvideo/pipelines/stages/cosmos3_stages.py create mode 100644 tests/local_tests/cosmos3/test_cosmos3_denoise_cfg_parity.py create mode 100644 tests/local_tests/cosmos3/test_cosmos3_pipeline_smoke.py diff --git a/fastvideo/configs/pipelines/cosmos3.py b/fastvideo/configs/pipelines/cosmos3.py index f66a6617f..f81c849bf 100644 --- a/fastvideo/configs/pipelines/cosmos3.py +++ b/fastvideo/configs/pipelines/cosmos3.py @@ -1,16 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 """Cosmos3 pipeline configuration. -Port reference: vllm-omni PR #3454 (https://github.com/vllm-project/vllm-omni/pull/3454) -HEAD ``8536f5b1421f``. Cosmos3 is structurally different from Cosmos 2.5: +Reference of record: the official ``cosmos-framework`` / ``nvidia/Cosmos3-Nano`` +checkpoint (``model_index.json``). Cosmos3 is structurally different from +Cosmos 2.5: -- Dual-stack DiT (UND + GEN) lives entirely inside ``Cosmos3VFMTransformer``. -- No separate text encoder — ``Cosmos3LanguageModel`` is inside the DiT, so - ``text_encoder_configs`` is the empty tuple. The tokenizer is loaded by the - pipeline's ``__init__`` from the checkpoint's ``text_tokenizer/`` subfolder; - that is an instance-level concern, not a config-level concern. -- VAE is ``DistributedAutoencoderKLWan`` (same as Cosmos 2.5), so the - ``Cosmos25VAEConfig`` is reused as-is. +- Dual-pathway (UND + GEN) DiT lives entirely inside ``Cosmos3VFMTransformer`` + (``Cosmos3VideoConfig``). +- No separate text encoder — the Qwen3-VL-text backbone is inside the DiT, so + ``text_encoder_configs`` is the empty tuple. The Qwen2 tokenizer is loaded as + the ``text_tokenizer`` checkpoint module by the component loader. +- VAE is Wan2.2 ``AutoencoderKLWan`` (z_dim=48, scale_factor_spatial=16), + configured by ``Cosmos3VAEConfig`` (the checkpoint's exact latents_mean/std). +- Scheduler is diffusers ``UniPCMultistepScheduler`` (flow_prediction, + use_flow_sigmas), loaded from the checkpoint. - T2I default ``flow_shift`` is 3.0 (set per-request by ``_set_flow_shift``); T2V/I2V use the engine-init default of 1.0 baked into this config. """ @@ -22,7 +25,7 @@ import torch from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig from fastvideo.configs.models.dits.cosmos3 import (Cosmos3ArchConfig, Cosmos3VideoConfig) from fastvideo.configs.models.encoders import BaseEncoderOutput -from fastvideo.configs.models.vaes import Cosmos25VAEConfig # Cosmos3 reuses the Wan-based VAE +from fastvideo.configs.models.vaes import Cosmos3VAEConfig # Wan2.2 AutoencoderKLWan from fastvideo.configs.pipelines.base import PipelineConfig @@ -34,14 +37,14 @@ def _identity_preprocess_text(prompt: str) -> str: class Cosmos3Config(PipelineConfig): """Configuration for the Cosmos3 video generation pipeline (T2V/I2V/T2I). - Tier A scope: architectural skeleton. Real weights from ``nvidia/Cosmos3-Nano`` - are not yet published; this config is shaped to be activatable when they - land without further config changes. + Wires the framework-parity-verified Cosmos3 components: the native + ``Cosmos3VideoConfig`` DiT, the Wan2.2 ``Cosmos3VAEConfig`` VAE, the Qwen2 + tokenizer (loaded as ``text_tokenizer``), and the UniPC scheduler. """ dit_config: DiTConfig = field(default_factory=lambda: Cosmos3VideoConfig(arch_config=Cosmos3ArchConfig())) - vae_config: VAEConfig = field(default_factory=Cosmos25VAEConfig) + vae_config: VAEConfig = field(default_factory=Cosmos3VAEConfig) # No separate text encoder: Cosmos3LanguageModel lives inside the DiT. text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=tuple) @@ -55,7 +58,9 @@ class Cosmos3Config(PipelineConfig): text_encoder_precisions: tuple[str, ...] = field(default_factory=tuple) embedded_cfg_scale: float = 0.0 - flow_shift: float = 1.0 # T2V/I2V engine-init; T2I uses 3.0 set per-request + # T2V/I2V engine-init flow_shift (framework text2video/image2video default); + # T2I overrides to 3.0 per request via Cosmos3DenoisingStage._set_flow_shift. + flow_shift: float = 10.0 vae_tiling: bool = False vae_sp: bool = False diff --git a/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py b/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py index 33322216f..6dc4ab699 100644 --- a/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py +++ b/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py @@ -1,66 +1,63 @@ # SPDX-License-Identifier: Apache-2.0 -"""Cosmos3 pipeline (Phase 2f.2 partial). +"""FastVideo-native Cosmos3 video pipeline (T2V / I2V / T2I). -Upstream reference: - ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py`` from - vllm-omni HEAD ``8536f5b1421f78c7df06af6d96fa195c1ceb6384``. +This replaces the earlier vllm-omni-derived skeleton with a native, stage-based +:class:`ComposedPipelineBase` pipeline that wires the framework-parity-verified +Cosmos3 components: -Phase 2f.1 landed the checkpoint-remap + scheduler-shift surface: +* tokenizer: Qwen2 ``Qwen2TokenizerFast`` + chat template (the only allowed + third-party model-adjacent dependency; tokenizers are explicitly permitted), +* VAE: FastVideo-native ``AutoencoderKLWan`` (Wan2.2) via ``Cosmos3VAEConfig``; + encode normalizes ``(mu - mean) * inv_std`` and decode denormalizes + clamps, +* sequence-packing: :func:`pack_cosmos3_video_sequence` (native, parity-tested), +* DiT: ``Cosmos3VFMTransformer`` (native, bit-identical to the framework), +* scheduler: diffusers ``UniPCMultistepScheduler`` (the actual scheduler the + checkpoint ships, ``flow_prediction`` + ``use_flow_sigmas``). -* ``_remap_ckpt_key`` — ported verbatim from upstream lines 319-409. This is a - pure key-translation function with no FastVideo-specific adaptations; the - checkpoint converter at ``scripts/checkpoint_conversion/cosmos3_convert.py`` - (Phase 5) will reuse it. -* ``_set_flow_shift`` — ported from upstream lines 498-512 with lazy scheduler - construction so the scheduler-default-parity tests can exercise the method on - ``__new__``-allocated instances before ``__init__`` is wired. -* ``_engine_init_flow_shift`` — exposed as a class attribute (default ``1.0``) - so ``hasattr(__new__(...), "_engine_init_flow_shift")`` succeeds without - running ``__init__``. +The denoise/CFG glue is a faithful port of the framework's +``Cosmos3OmniDiffusersPipeline`` math (mirrored in the framework-equivalent +``diffusers_cosmos3.pipeline``): per UniPC timestep, run a SEQUENTIAL conditional +then unconditional pass (each repacks the sequence with the prompt / negative +prompt token ids, forwards the DiT, and zeros the prediction on conditioning +frames), then combine ``v = uncond + guidance * (cond - uncond)`` and take one +``scheduler.step(model_output=v, timestep, sample=latent)``. ``timestep_scale`` +is applied to the per-token timesteps *inside* the DiT (its ``forward`` already +multiplies ``vision_timesteps * timestep_scale`` before the time embedder), so +the loop passes raw scheduler timesteps to the packer. -Phase 2f.2 adds the runtime call graph (this commit): +The pure denoise math lives in :class:`Cosmos3DenoiseEngine` and the free +function :func:`cosmos3_get_cfg_velocity` so it can be unit-/parity-tested +directly against the framework oracle without constructing the full pipeline. -* ``diffuse`` — sequential 3-mode CFG denoising loop ported from upstream - lines 883-1033. The CFG-Parallel branch (cfg_parallel=True) is deferred - until FastVideo's classifier-free-guidance-world-size plumbing lands; - ``_cfg_parallel_active`` returns False so only the sequential and no-CFG - paths execute. -* ``forward`` — request parsing + T2I/T2V/I2V mode dispatch + flow-shift - selection + diffusion driver + decode, ported from upstream lines - 1037-1206 (single-prompt path; the T2I num_outputs_per_prompt > 1 branch - is preserved but uses the unimplemented ``_prepare_latents`` stub). -* ``_is_t2i_request`` / ``_get_sp_param`` / ``_cfg_parallel_active`` — - ported from upstream lines 452-496. -* Negative-prompt constants — ported from upstream lines 51-61. - -Helper methods ``_format_and_tokenize_prompts``, ``_prepare_latents``, -``_prepare_latents_i2v``, ``_set_scheduler_timesteps``, and -``_decode_latents`` are defined as ``NotImplementedError`` stubs so the -pipeline call-graph tests can monkey-patch them without an -``AttributeError``; their real implementations land in Phase 2c -(tokenizer) and Phase 2f.3+ (latent prep, scheduler timesteps, VAE decode). - -``__init__`` and ``create_pipeline_stages`` remain ``NotImplementedError`` -stubs — full module/stage wiring is left to a future phase. +No diffusers/transformers *model* classes are imported at runtime here; only the +Qwen2 tokenizer (loaded by the component loader) and the UniPC scheduler are +third-party, both explicitly allowed. """ - from __future__ import annotations -from types import SimpleNamespace +import math +from dataclasses import dataclass from typing import Any import torch -import torch.nn as nn from diffusers import UniPCMultistepScheduler from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger +from fastvideo.pipelines.basic.cosmos3.sequence_packing import ( + Cosmos3SampleInputs, + Cosmos3VisionItem, + pack_cosmos3_video_sequence, +) from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase logger = init_logger(__name__) -# Negative-prompt constants ported verbatim from upstream pipeline_cosmos3.py:51-61. -COSMOS3_DEFAULT_NEGATIVE_PROMPT = "" +# System prompts, verbatim from the framework (diffusers_cosmos3.pipeline). +_SYSTEM_PROMPT_IMAGE = "You are a helpful assistant who will generate images from a give prompt." +_SYSTEM_PROMPT_VIDEO = "You are a helpful assistant who will generate videos from a give prompt." + +# Default video negative prompt (framework / Cosmos quality prompt). COSMOS3_VIDEO_NEGATIVE_PROMPT = ( "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, " "over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, " @@ -68,168 +65,413 @@ COSMOS3_VIDEO_NEGATIVE_PROMPT = ( "jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, " "fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. " "Overall, the video is of poor quality.") -COSMOS3_T2V_NEGATIVE_PROMPT = COSMOS3_VIDEO_NEGATIVE_PROMPT -COSMOS3_I2V_NEGATIVE_PROMPT = COSMOS3_VIDEO_NEGATIVE_PROMPT -class Cosmos3OmniDiffusersPipeline(nn.Module, ComposedPipelineBase): - """Cosmos3 T2V/I2V/T2I pipeline (skeleton). +# =========================================================================== +# Special-token resolution (Qwen2 chat tokenizer) +# =========================================================================== +def cosmos3_special_tokens(tokenizer: Any) -> dict[str, int]: + """Resolve the Cosmos3 generation special tokens from a Qwen2 tokenizer. - Phase 2f.1 lands only the checkpoint-remap + scheduler-shift surface needed - by ``test_cosmos3_state_dict_keys.py`` and - ``test_cosmos3_scheduler_default_parity.py``. Full wiring (modules, - stages, ``diffuse()``, mode dispatch) is deferred to Phase 2f.2. + Mirrors the framework's ``llm_special_tokens``: + ``start_of_generation=<|vision_start|>``, ``end_of_generation=<|vision_end|>``, + ``eos_token_id=tokenizer.eos_token_id``. + """ + return { + "start_of_generation": int(tokenizer.convert_tokens_to_ids("<|vision_start|>")), + "end_of_generation": int(tokenizer.convert_tokens_to_ids("<|vision_end|>")), + "eos_token_id": int(tokenizer.eos_token_id), + } - Inherits from both ``nn.Module`` (matching upstream - ``pipeline_cosmos3.py:197-199`` — the conftest fixture - ``make_cosmos3_pipeline`` calls ``nn.Module.__init__`` directly) and - ``ComposedPipelineBase`` (FastVideo's pipeline convention — sibling - ``cosmos2_5_pipeline.py``). Phase 2f.2 can collapse either branch if - one proves unnecessary once the full pipeline lands. - The scheduler-parity tests build instances via - ``Cosmos3OmniDiffusersPipeline.__new__(...)`` and rely on: +def cosmos3_tokenize_caption( + tokenizer: Any, + caption: str, + *, + is_video: bool = False, + use_system_prompt: bool = False, +) -> list[int]: + """Tokenize a caption with the Qwen2 chat template (framework-faithful). - * ``hasattr(instance, "scheduler")`` succeeding before ``_set_flow_shift`` - is called (satisfied by the ``scheduler = None`` class attribute below); - * ``hasattr(instance, "_engine_init_flow_shift")`` succeeding (satisfied - by the class-attribute default ``1.0``). + Optionally prepends an image/video system prompt; always adds the + generation prompt and disables ``add_vision_id`` (matching the framework's + ``tokenize_caption``). + """ + conversations: list[dict[str, str]] = [] + if use_system_prompt: + conversations.append({ + "role": "system", + "content": _SYSTEM_PROMPT_VIDEO if is_video else _SYSTEM_PROMPT_IMAGE, + }) + conversations.append({"role": "user", "content": caption}) + token_ids = tokenizer.apply_chat_template( + conversations, + tokenize=True, + add_generation_prompt=True, + add_vision_id=False, + ) + return list(token_ids) + + +# =========================================================================== +# VAE encode/decode bridge (normalize / denormalize, matching the framework) +# =========================================================================== +@dataclass +class _VaeNorm: + """Cached ``mean`` / ``inv_std`` for VAE (de)normalization.""" + + mean: torch.Tensor # [z_dim] + inv_std: torch.Tensor # [z_dim] + + @classmethod + def from_vae(cls, vae: Any, dtype: torch.dtype) -> _VaeNorm: + mean = torch.tensor(list(vae.config.latents_mean), dtype=dtype) + std = torch.tensor(list(vae.config.latents_std), dtype=dtype) + return cls(mean=mean, inv_std=1.0 / std) + + +def cosmos3_vae_encode(vae: Any, video: torch.Tensor, norm: _VaeNorm) -> torch.Tensor: + """Encode ``[B, 3, T, H, W]`` pixels in [-1, 1] to NORMALIZED latents. + + Matches the framework ``DiffusersWan22VAE.encode``: take the posterior mode + and apply ``(mu - mean) * inv_std``. FastVideo's ``AutoencoderKLWan.encode`` + returns a ``DiagonalGaussianDistribution``; we read ``.mode()``. + """ + in_dtype = video.dtype + device = video.device + mean = norm.mean.to(device=device, dtype=in_dtype).view(1, -1, 1, 1, 1) + inv_std = norm.inv_std.to(device=device, dtype=in_dtype).view(1, -1, 1, 1, 1) + raw_mu = vae.encode(video).mode() + return ((raw_mu - mean) * inv_std).to(in_dtype) + + +def cosmos3_vae_decode(vae: Any, latents: torch.Tensor, norm: _VaeNorm) -> torch.Tensor: + """Decode NORMALIZED latents ``[B, z, T, H, W]`` to pixels ``[B, 3, T, H, W]``. + + Inverts the normalization (``z / inv_std + mean``) then calls + ``vae.decode`` (which already clamps to [-1, 1]). + """ + in_dtype = latents.dtype + device = latents.device + mean = norm.mean.to(device=device, dtype=in_dtype).view(1, -1, 1, 1, 1) + inv_std = norm.inv_std.to(device=device, dtype=in_dtype).view(1, -1, 1, 1, 1) + z_raw = latents / inv_std + mean + out = vae.decode(z_raw) + if isinstance(out, tuple): + out = out[0] + if hasattr(out, "sample"): + out = out.sample + return out.to(in_dtype) + + +# =========================================================================== +# Per-vision-item packing geometry +# =========================================================================== +@dataclass +class Cosmos3VisionSpec: + """Geometry + conditioning for one vision item in a denoise run. + + Args: + clean_latent: VAE-encoded conditioning latent ``[C, T, H, W]`` used to + keep condition frames clean (I2V / T2I). May be ``None`` for T2V. + condition_frame_indexes: Latent-frame indices kept clean. + shape: ``(C, T, H, W)`` of the latent for this item. """ - # --- Class attributes --------------------------------------------------- - # - # These exist as class attributes so ``hasattr(__new__(cls), name)`` is - # True without ``__init__`` having run. ``__init__`` (Phase 2f.2) will - # shadow them with instance attributes derived from the loaded scheduler. + shape: tuple[int, int, int, int] + condition_frame_indexes: list[int] + clean_latent: torch.Tensor | None = None + + @property + def numel(self) -> int: + return int(math.prod(self.shape)) + + +# =========================================================================== +# Pure denoise/CFG math (parity oracle target) +# =========================================================================== +def _split_flat_latent(flat: torch.Tensor, specs: list[Cosmos3VisionSpec]) -> list[torch.Tensor]: + """Split a flat latent vector into per-vision-item ``[C, T, H, W]`` tensors.""" + out: list[torch.Tensor] = [] + offset = 0 + for spec in specs: + out.append(flat[offset:offset + spec.numel].reshape(spec.shape)) + offset += spec.numel + return out + + +def cosmos3_get_cfg_velocity( + *, + transformer: Any, + flat_latent: torch.Tensor, + timestep: torch.Tensor, + guidance: float, + specs: list[Cosmos3VisionSpec], + cond_token_ids: list[int], + uncond_token_ids: list[int], + special_tokens: dict[str, int], + latent_patch_size: int, + temporal_modality_margin: int, + reset_spatial_ids: bool, + enable_fps_modulation: bool, + base_fps: float, + temporal_compression_factor: int, + include_end_of_generation_token: bool = False, + fps_per_item: list[float] | None = None, + normalize_cfg: bool = False, +) -> torch.Tensor: + """Sequential-CFG velocity for one denoise step (framework math). + + Replicates the framework ``get_cfg_velocity``: + + 1. split ``flat_latent`` into per-vision-item ``[C, T, H, W]`` latents, + 2. run a conditional pass (prompt tokens) and an unconditional pass + (negative-prompt tokens); each repacks via + :func:`pack_cosmos3_video_sequence`, forwards the DiT to obtain + ``preds_vision`` (a list of ``[1, C, T, H, W]`` unpatchified noisy-frame + predictions), and zeros the prediction on conditioning frames + (``pred * (1 - condition_mask)``), + 3. combine ``v = uncond + guidance * (cond - uncond)`` (optionally + norm-rescaled), returned flattened to match ``flat_latent``. + + ``timestep`` is a scalar tensor (raw scheduler timestep); ``timestep_scale`` + is applied inside the DiT, so it is passed through unscaled here. + """ + assert timestep.numel() == 1, "timestep must be a scalar" + timestep_value = float(timestep.reshape(()).item()) + + noise_x_vision = _split_flat_latent(flat_latent, specs) + + def _run(token_ids: list[int]) -> torch.Tensor: + samples = [ + Cosmos3SampleInputs( + text_ids=list(token_ids), + vision=Cosmos3VisionItem( + latent=latent, + condition_frame_indexes=list(spec.condition_frame_indexes), + fps=(fps_per_item[i] if fps_per_item is not None else None), + ), + timestep=timestep_value, + ) for i, (latent, spec) in enumerate(zip(noise_x_vision, specs, strict=False)) + ] + packed = pack_cosmos3_video_sequence( + samples, + special_tokens, + latent_patch_size=latent_patch_size, + include_end_of_generation_token=include_end_of_generation_token, + temporal_modality_margin=temporal_modality_margin, + reset_spatial_ids=reset_spatial_ids, + enable_fps_modulation=enable_fps_modulation, + base_fps=base_fps, + temporal_compression_factor=temporal_compression_factor, + ) + out = transformer(**packed.to_dit_kwargs()) + preds = out.get("preds_vision") + if preds is None: + # Fully-conditioned clip: no noisy frames, zero velocity. + return torch.zeros_like(flat_latent) + # Zero the prediction on conditioning frames, per item. + velocity_items: list[torch.Tensor] = [] + for pred, cond_mask in zip(preds, packed.vision_condition_mask, strict=False): + pred = pred.squeeze(0) if pred.dim() == 5 else pred # [C, T, H, W] + keep = (1.0 - cond_mask).to(dtype=pred.dtype, device=pred.device) # [T,1,1] + if keep.sum() > 0: + velocity_items.append(pred * keep) + else: + velocity_items.append(torch.zeros_like(pred)) + return torch.cat([v.reshape(-1) for v in velocity_items]) + + cond_v = _run(cond_token_ids) + uncond_v = _run(uncond_token_ids) + v_pred = uncond_v + guidance * (cond_v - uncond_v) + if normalize_cfg: + scale = (torch.norm(cond_v) / (torch.norm(v_pred) + 1e-8)).clamp(min=0.0, max=1.0) + v_pred = v_pred * scale + return v_pred + + +class Cosmos3DenoiseEngine: + """Stateless denoise driver tying CFG velocity to UniPC stepping. + + Holds the transformer + scheduler + packing constants and runs the full + UniPC denoise loop. Kept separate from the pipeline so it can be exercised + in isolation (smoke + parity tests) with stub or real components. + """ + + def __init__( + self, + *, + transformer: Any, + scheduler: Any, + special_tokens: dict[str, int], + latent_patch_size: int, + temporal_modality_margin: int, + reset_spatial_ids: bool, + enable_fps_modulation: bool, + base_fps: float, + temporal_compression_factor: int, + include_end_of_generation_token: bool = False, + ) -> None: + self.transformer = transformer + self.scheduler = scheduler + self.special_tokens = special_tokens + self.latent_patch_size = latent_patch_size + self.temporal_modality_margin = temporal_modality_margin + self.reset_spatial_ids = reset_spatial_ids + self.enable_fps_modulation = enable_fps_modulation + self.base_fps = base_fps + self.temporal_compression_factor = temporal_compression_factor + self.include_end_of_generation_token = include_end_of_generation_token + + def velocity( + self, + *, + flat_latent: torch.Tensor, + timestep: torch.Tensor, + guidance: float, + specs: list[Cosmos3VisionSpec], + cond_token_ids: list[int], + uncond_token_ids: list[int], + fps_per_item: list[float] | None = None, + ) -> torch.Tensor: + return cosmos3_get_cfg_velocity( + transformer=self.transformer, + flat_latent=flat_latent, + timestep=timestep, + guidance=guidance, + specs=specs, + cond_token_ids=cond_token_ids, + uncond_token_ids=uncond_token_ids, + special_tokens=self.special_tokens, + latent_patch_size=self.latent_patch_size, + temporal_modality_margin=self.temporal_modality_margin, + reset_spatial_ids=self.reset_spatial_ids, + enable_fps_modulation=self.enable_fps_modulation, + base_fps=self.base_fps, + temporal_compression_factor=self.temporal_compression_factor, + include_end_of_generation_token=self.include_end_of_generation_token, + fps_per_item=fps_per_item, + ) + + def denoise( + self, + *, + flat_latent: torch.Tensor, + timesteps: torch.Tensor, + guidance: float, + specs: list[Cosmos3VisionSpec], + cond_token_ids: list[int], + uncond_token_ids: list[int], + fps_per_item: list[float] | None = None, + progress_bar: Any | None = None, + ) -> torch.Tensor: + """Run the full UniPC denoise loop, returning the final flat latent. + + For each timestep: compute the sequential-CFG velocity, then + ``scheduler.step(model_output=v, timestep, sample=latent.unsqueeze(0))`` + (the framework steps with a leading batch axis), squeezing back to flat. + """ + latent = flat_latent + iterator = progress_bar(timesteps) if progress_bar is not None else timesteps + for t in iterator: + v_pred = self.velocity( + flat_latent=latent, + timestep=t.reshape(1), + guidance=guidance, + specs=specs, + cond_token_ids=cond_token_ids, + uncond_token_ids=uncond_token_ids, + fps_per_item=fps_per_item, + ) + stepped = self.scheduler.step( + model_output=v_pred, + timestep=t, + sample=latent.unsqueeze(0), + return_dict=False, + )[0] + latent = stepped.squeeze(0) + return latent + + +# =========================================================================== +# Pipeline (ComposedPipelineBase) +# =========================================================================== +class Cosmos3OmniDiffusersPipeline(ComposedPipelineBase): + """Cosmos3 video generation pipeline (T2V / I2V / T2I). + + Stage-based ``ComposedPipelineBase`` pipeline. The required modules + (``transformer`` / ``vae`` / ``scheduler`` / ``text_tokenizer``) are loaded + from the ``nvidia/Cosmos3-Nano`` checkpoint by the component loader. The + class name matches the checkpoint ``model_index.json`` ``_class_name`` so + the registry resolves it directly. + + The denoise/CFG/VAE math is delegated to module-level helpers + (:func:`cosmos3_get_cfg_velocity`, :class:`Cosmos3DenoiseEngine`, + :func:`cosmos3_vae_encode` / :func:`cosmos3_vae_decode`) which are + framework-parity tested in ``tests/local_tests/cosmos3``. + """ + + is_video_pipeline = True + # ``vision_encoder`` / ``sound_tokenizer`` ship in the checkpoint but the + # video path does not need them; they are intentionally omitted here. + _required_config_modules = ["text_tokenizer", "vae", "transformer", "scheduler"] + + # Engine-init flow_shift (T2V/I2V); T2I overrides to 3.0 per request. _engine_init_flow_shift: float = 1.0 + # Class-attribute defaults so ``__new__``-based unit tests can read these + # before ``initialize_pipeline`` runs. scheduler: Any = None _base_scheduler_config: Any = None _current_flow_shift: float | None = None - # -- Weight loading ----------------------------------------------------- + def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None: + """Bind the loaded scheduler + snapshot its config so per-request + flow_shift rebuilds are cheap and the engine-init shift is applied.""" + pipeline_config = fastvideo_args.pipeline_config + engine_shift = getattr(pipeline_config, "flow_shift", None) + if engine_shift is not None: + self._engine_init_flow_shift = float(engine_shift) + scheduler = self.get_module("scheduler") + if scheduler is not None: + self.scheduler = scheduler + self._base_scheduler_config = scheduler.config + self._current_flow_shift = float(getattr(scheduler.config, "flow_shift", 1.0)) - @staticmethod - def _remap_ckpt_key(key: str) -> str | None: - """Remap a Diffusers transformer key to the model parameter namespace. + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + """Wire the Cosmos3 stages. - Checkpoint keys arrive with a synthetic ``transformer.`` prefix from - ``weights_sources``. The source checkpoint itself uses the Diffusers - transformer namespace: top-level projections plus ``model.*`` for the - Qwen3-VL backbone. UND and GEN components share each layer in the - source and are split into separate module lists here. - - Returns the remapped name under ``transformer.``, or ``None`` to skip. - - Ported verbatim from upstream ``pipeline_cosmos3.py`` lines 319-409. + The whole text->latent->denoise->decode flow is custom (sequential CFG + with per-pass repacking), so a single :class:`Cosmos3DenoisingStage` + owns it. ``InputValidationStage`` runs first for the standard checks. """ - k = key - # Strip the weights_sources prefix - if k.startswith("transformer."): - k = k[len("transformer."):] + from fastvideo.pipelines.stages import InputValidationStage + from fastvideo.pipelines.stages.cosmos3_stages import Cosmos3DenoisingStage - # Top-level generation components. - if k.startswith(( - "vae2llm.", - "llm2vae.", - "time_embedder.", - )): - return f"transformer.{k}" - - # Skip lm_head - if k.startswith("lm_head."): - return None - - # embed_tokens / norm → language_model.* - if k.startswith("model.embed_tokens."): - return f"transformer.language_model.{k[len('model.'):]}" - if k.startswith("model.norm."): - return f"transformer.language_model.{k[len('model.'):]}" - - # norm_moe_gen → top level - if k.startswith("model.norm_moe_gen."): - return f"transformer.{k[len('model.'):]}" - - if not k.startswith("model.layers."): - return None - k = k[len("model."):] - - if not k.startswith("layers."): - return None - - parts = k.split(".", 2) # ['layers', '{i}', '{rest}'] - if len(parts) != 3: - return None - layer_idx = parts[1] - rest = parts[2] - - und_lp = f"transformer.language_model.layers.{layer_idx}" - gen_lp = f"transformer.gen_layers.{layer_idx}" - - _LAYER_MAP = { - # UND attention - "self_attn.q_proj.": f"{und_lp}.self_attn.q_proj.", - "self_attn.k_proj.": f"{und_lp}.self_attn.k_proj.", - "self_attn.v_proj.": f"{und_lp}.self_attn.v_proj.", - "self_attn.o_proj.": f"{und_lp}.self_attn.o_proj.", - "self_attn.q_norm.": f"{und_lp}.self_attn.q_norm.", - "self_attn.k_norm.": f"{und_lp}.self_attn.k_norm.", - # GEN attention - "self_attn.q_proj_moe_gen.": f"{gen_lp}.cross_attention.q_proj.", - "self_attn.k_proj_moe_gen.": f"{gen_lp}.cross_attention.k_proj.", - "self_attn.v_proj_moe_gen.": f"{gen_lp}.cross_attention.v_proj.", - "self_attn.o_proj_moe_gen.": f"{gen_lp}.cross_attention.o_proj.", - "self_attn.q_norm_moe_gen.": f"{gen_lp}.cross_attention.q_norm.", - "self_attn.k_norm_moe_gen.": f"{gen_lp}.cross_attention.k_norm.", - # Norms - "input_layernorm.": f"{und_lp}.input_layernorm.", - "post_attention_layernorm.": f"{und_lp}.post_attention_layernorm.", - "input_layernorm_moe_gen.": f"{gen_lp}.input_layernorm.", - "post_attention_layernorm_moe_gen.": f"{gen_lp}.post_attention_layernorm.", - # UND MLP - "mlp.gate_proj.": f"{und_lp}.mlp.gate_proj.", - "mlp.up_proj.": f"{und_lp}.mlp.up_proj.", - "mlp.down_proj.": f"{und_lp}.mlp.down_proj.", - # GEN MLP - "mlp_moe_gen.gate_proj.": f"{gen_lp}.mlp.gate_proj.", - "mlp_moe_gen.up_proj.": f"{gen_lp}.mlp.up_proj.", - "mlp_moe_gen.down_proj.": f"{gen_lp}.mlp.down_proj.", - } - - for pattern, replacement in _LAYER_MAP.items(): - if rest.startswith(pattern): - suffix = rest[len(pattern):] - return replacement + suffix - - return None + self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage()) + self.add_stage( + stage_name="denoising_stage", + stage=Cosmos3DenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + vae=self.get_module("vae"), + tokenizer=self.get_module("text_tokenizer"), + pipeline=self, + ), + ) # -- Scheduler control -------------------------------------------------- def _set_flow_shift(self, target_shift: float) -> None: - """Set the UniPC ``flow_shift`` to a concrete target value. + """Set UniPC ``flow_shift`` to ``target_shift``. - Adapted from upstream ``pipeline_cosmos3.py`` lines 498-512 with one - FastVideo-specific addition: when called on an instance built via - ``__new__`` (i.e. ``__init__`` has not run, so - ``self._base_scheduler_config is None``), this method lazily - constructs a default ``UniPCMultistepScheduler`` rather than - rebuilding from a saved config. Phase 2f.2's ``__init__`` will - replace that lazy path with a checkpoint-loaded scheduler and - snapshot its config into ``_base_scheduler_config`` per upstream. - - Tracking ``self._current_flow_shift`` explicitly is required because - the previous mode may have rebuilt the scheduler — we cannot rely on - ``self.scheduler.config.flow_shift`` reflecting the last requested - target if a rebuild was skipped via the equality check. + Lazily builds a default UniPC scheduler when called before + ``initialize_pipeline`` (e.g. the ``__new__``-based scheduler-parity + tests); otherwise rebuilds from the snapshotted base config only when + the target differs from the current shift. """ target = float(target_shift) - - # Lazy path: no checkpoint-loaded base config yet (Phase 2f.1 test - # entry via __new__). Construct a fresh UniPC scheduler at the - # requested flow_shift so test assertions can read - # ``self.scheduler.config.flow_shift``. Phase 2f.2's __init__ will - # load the real scheduler from the checkpoint and overwrite both - # ``self.scheduler`` and ``self._base_scheduler_config``. - if self._base_scheduler_config is None: + base_config = self._base_scheduler_config + if base_config is None: self.scheduler = UniPCMultistepScheduler( num_train_timesteps=1000, solver_order=2, @@ -240,366 +482,23 @@ class Cosmos3OmniDiffusersPipeline(nn.Module, ComposedPipelineBase): self._base_scheduler_config = self.scheduler.config self._current_flow_shift = target return - - # Rebuild only if the target differs from the current shift. - if self._current_flow_shift is not None and target == float(self._current_flow_shift): + current = self._current_flow_shift + if current is not None and target == float(current): return - self.scheduler = UniPCMultistepScheduler.from_config(self._base_scheduler_config, flow_shift=target) + self.scheduler = UniPCMultistepScheduler.from_config(base_config, flow_shift=target) + if isinstance(self.modules, dict): + self.modules["scheduler"] = self.scheduler self._current_flow_shift = target - # -- Request introspection ---------------------------------------------- + # -- Tokenization ------------------------------------------------------- - @staticmethod - def _cfg_parallel_active() -> bool: - """Return True when CFG-Parallel is enabled in the current topology. - - Upstream ``pipeline_cosmos3.py:452-457`` queries - ``get_classifier_free_guidance_world_size() > 1``. FastVideo's - CFG-Parallel plumbing is not wired in Phase 2f.2, so this always - returns False; ``diffuse()`` therefore exercises only the sequential - CFG path and the no-CFG path. - """ - return False - - @staticmethod - def _get_sp_param(sp: Any, key: str, default: Any = None) -> Any: - """Read a runtime control from sampling params. - - Ported verbatim from upstream ``pipeline_cosmos3.py:459-481``. - - Order of precedence: - 1. ``sp.extra_args[key]`` — preferred path; the OpenAI image/video - endpoints surface custom controls there. - 2. direct attribute on ``sp``. - 3. ``default``. - """ - extra = getattr(sp, "extra_args", None) - if isinstance(extra, dict) and extra.get(key) is not None: - return extra[key] - val = getattr(sp, key, None) - if val is not None: - return val - return default - - @staticmethod - def _is_t2i_request(req: Any) -> bool: - """Detect text-to-image mode from request-level prompt modalities. - - Ported verbatim from upstream ``pipeline_cosmos3.py:483-496``. - Raises ValueError when a prompt requests both image AND video output - simultaneously. - """ - if not req.prompts: - return False - first_prompt = req.prompts[0] - modalities = first_prompt.get("modalities", []) if isinstance(first_prompt, dict) else [] - if modalities is None: - modalities = [] - if isinstance(modalities, str): - modalities = [modalities] - if "image" in modalities and "video" in modalities: - raise ValueError("Cosmos3 prompt modalities cannot request both image and video output.") - return "image" in modalities - - # -- Helper stubs (real implementations land in later phases) ----------- - # - # These exist so the pipeline call-graph tests can monkey-patch them on - # the instance without an ``AttributeError``. Each raises - # ``NotImplementedError`` if invoked outside a test that replaces it. - - def _format_and_tokenize_prompts(self, *args: Any, **kwargs: Any) -> Any: - raise NotImplementedError("_format_and_tokenize_prompts lands in Phase 2c (tokenizer wiring).") - - def _prepare_latents(self, *args: Any, **kwargs: Any) -> Any: - raise NotImplementedError("_prepare_latents lands in Phase 2f.3 (latent preparation).") - - def _prepare_latents_i2v(self, *args: Any, **kwargs: Any) -> Any: - raise NotImplementedError("_prepare_latents_i2v lands in Phase 2f.3 (I2V latent preparation).") - - def _set_scheduler_timesteps(self, num_inference_steps: int) -> None: - raise NotImplementedError("_set_scheduler_timesteps lands in Phase 2f.3 (scheduler timesteps).") - - def _decode_latents(self, latents: torch.Tensor) -> Any: - raise NotImplementedError("_decode_latents lands in Phase 2f.3 (VAE decode wiring).") - - # -- Denoising loop ----------------------------------------------------- - - def diffuse( - self, - *, - latents: torch.Tensor, - timesteps: torch.Tensor, - cond_ids: torch.Tensor, - cond_mask: torch.Tensor, - uncond_ids: torch.Tensor, - uncond_mask: torch.Tensor, - guidance_scale: float, - shared_kwargs: dict[str, Any], - velocity_mask: torch.Tensor | None = None, - image_latent: torch.Tensor | None = None, - condition_latents: torch.Tensor | None = None, - guidance_interval: tuple[float, float] | None = None, - ) -> torch.Tensor: - """Sequential 3-mode CFG denoising loop with optional I2V conditioning. - - Ported from upstream ``pipeline_cosmos3.py:883-1033``. The - CFG-Parallel branch (upstream lines 953-980) is deferred until - FastVideo's classifier-free-guidance-world-size plumbing lands; - ``_cfg_parallel_active`` returns False in Phase 2f.2 so only the - sequential CFG branch (upstream lines 982-1019) and the no-CFG - branch (upstream lines 1021-1031) execute. - - Cosmos3's UND pathway is text-dependent, so sequential CFG keeps - separate K/V caches for the conditional and unconditional text - forwards and swaps them in before each branch's transformer call. - - I2V conditioning is applied via ``_step``: ``velocity_mask`` zeros - frame-0 noise predictions before the scheduler step, and - ``image_latent`` is re-injected into frame 0 after each step - (UniPC's predictor-corrector rescales the sample, so zero velocity - alone does not preserve frame 0). - - ``guidance_interval`` (T2I) restricts CFG to timesteps inside the - closed interval ``[lo, hi]``. Outside the interval the cond/uncond - delta is dropped and only the cond branch executes — equivalent to - CFG with scale=1.0 but cheaper. - """ - do_cfg = guidance_scale > 1.0 - cfg_parallel = self._cfg_parallel_active() and do_cfg - if cfg_parallel: - raise NotImplementedError("Cosmos3OmniDiffusersPipeline.diffuse: CFG-Parallel branch is " - "deferred; FastVideo cfg-world-size plumbing is not yet wired.") - - self.transformer.reset_cache() - - def _cfg_active_at(t: torch.Tensor) -> bool: - if guidance_interval is None: - return True - t_scalar = float(t.item()) if torch.is_tensor(t) else float(t) - lo, hi = guidance_interval - return lo <= t_scalar <= hi - - def _step(noise_pred: torch.Tensor, t: torch.Tensor, latents: torch.Tensor) -> torch.Tensor: - if isinstance(noise_pred, tuple): - raise ValueError("Cosmos3 video-only diffusion received tuple predictions.") - if velocity_mask is not None: - noise_pred = noise_pred * velocity_mask - latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] - if condition_latents is not None and velocity_mask is not None: - latents = velocity_mask * latents + (1.0 - velocity_mask) * condition_latents - elif image_latent is not None: - latents[:, :, 0:1, :, :] = image_latent - return latents - - if do_cfg: - cond_cache: tuple = (None, None) - uncond_cache: tuple = (None, None) - for t in self.progress_bar(timesteps): - timestep = t.unsqueeze(0) - cfg_active = _cfg_active_at(t) - - self.transformer.cached_kv, self.transformer.cached_freqs_gen = cond_cache - noise_cond = self.transformer( - hidden_states=latents, - timestep=timestep, - text_ids=cond_ids, - text_mask=cond_mask, - **shared_kwargs, - ) - if cond_cache[0] is None: - cond_cache = (self.transformer.cached_kv, self.transformer.cached_freqs_gen) - - if cfg_active: - self.transformer.cached_kv, self.transformer.cached_freqs_gen = uncond_cache - noise_uncond = self.transformer( - hidden_states=latents, - timestep=timestep, - text_ids=uncond_ids, - text_mask=uncond_mask, - **shared_kwargs, - ) - if uncond_cache[0] is None: - uncond_cache = (self.transformer.cached_kv, self.transformer.cached_freqs_gen) - noise_pred = noise_uncond + guidance_scale * (noise_cond - noise_uncond) - else: - noise_pred = noise_cond - - latents = _step(noise_pred, t, latents) - else: - for t in self.progress_bar(timesteps): - timestep = t.unsqueeze(0) - noise_pred = self.transformer( - hidden_states=latents, - timestep=timestep, - text_ids=cond_ids, - text_mask=cond_mask, - **shared_kwargs, - ) - latents = _step(noise_pred, t, latents) - - return latents - - # -- Forward (main generation entry point) ------------------------------ - - def forward(self, req: Any) -> SimpleNamespace: - """Cosmos3 inference: request parse + mode dispatch + diffuse + decode. - - Ported from upstream ``pipeline_cosmos3.py:1037-1206`` (single-prompt - path). FastVideo does not yet have ``OmniDiffusionRequest`` or - ``DiffusionOutput`` dataclasses on the production path, so the - return value is a duck-typed ``SimpleNamespace`` with an ``output`` - attribute mirroring upstream ``DiffusionOutput.output``. - - Mode selection (upstream lines 1059-1093): - * T2I: ``"image" in modalities`` and no preprocessed image. Defaults - to ``num_frames=1``, ``flow_shift=3.0``, - ``guidance_interval=(400.0, 1000.0)``, 50 steps, scale=7.0. - * I2V: ``preprocessed_image`` present and not T2I. Like T2V but - with image conditioning. - * T2V: otherwise. Defaults to ``num_frames=189``, - ``flow_shift=self._engine_init_flow_shift``, no guidance - interval, 35 steps, scale=6.0. - - Calls ``_set_flow_shift`` exactly once per request, after defaults - are resolved and before tokenization. - """ - if not req.prompts: - raise ValueError("Cosmos3OmniDiffusersPipeline.forward() requires at least one prompt.") - if len(req.prompts) > 1: - raise ValueError("Cosmos3OmniDiffusersPipeline currently supports a single prompt per request.") - - prompt_data = req.prompts[0] - if isinstance(prompt_data, str): - prompt = prompt_data - negative_prompt = None - image_tensor = None - else: - prompt = prompt_data.get("prompt", "") - negative_prompt = prompt_data.get("negative_prompt") - additional_info = prompt_data.get("additional_information", {}) or {} - image_tensor = additional_info.get("preprocessed_image") - - sp = req.sampling_params - is_t2i = self._is_t2i_request(req) - is_i2v = image_tensor is not None and not is_t2i - if negative_prompt is None: - if is_t2i: - negative_prompt = COSMOS3_DEFAULT_NEGATIVE_PROMPT - elif is_i2v: - negative_prompt = COSMOS3_I2V_NEGATIVE_PROMPT - else: - negative_prompt = COSMOS3_T2V_NEGATIVE_PROMPT - - if is_t2i: - height = sp.height or 1024 - width = sp.width or 1024 - num_frames = 1 - num_inference_steps = sp.num_inference_steps or 50 - guidance_scale = sp.guidance_scale if sp.guidance_scale else 7.0 - default_flow_shift = 3.0 - default_guidance_interval: tuple[float, float] | None = (400.0, 1000.0) - batch_size = max(1, int(getattr(sp, "num_outputs_per_prompt", None) or 1)) - else: - height = sp.height or 720 - width = sp.width or 1280 - num_frames = sp.num_frames or 189 - num_inference_steps = sp.num_inference_steps or 35 - guidance_scale = sp.guidance_scale if sp.guidance_scale else 6.0 - default_flow_shift = self._engine_init_flow_shift - default_guidance_interval = None - batch_size = 1 - - flow_shift_target = float(self._get_sp_param(sp, "flow_shift", default_flow_shift)) - guidance_interval = self._get_sp_param(sp, "guidance_interval", default_guidance_interval) - - frame_rate = (self._get_sp_param(sp, "resolved_frame_rate") or self._get_sp_param(sp, "frame_rate") or 24.0) - max_sequence_length = self._get_sp_param(sp, "max_sequence_length", 512) or 512 - use_system_prompt = bool(self._get_sp_param(sp, "use_system_prompt", False)) - - self._guidance_scale = guidance_scale - self._num_timesteps = num_inference_steps - - self._set_flow_shift(flow_shift_target) - - generator = sp.generator - if generator is None: - seed = sp.seed if sp.seed is not None else 42 - generator = torch.Generator(device=self.device).manual_seed(seed) - - cond_ids, cond_mask, uncond_ids, uncond_mask = self._format_and_tokenize_prompts( - prompt, - negative_prompt, - num_frames, - frame_rate, - height, - width, - max_sequence_length, - sp, - use_system_prompt, - is_t2i=is_t2i, - ) - - if image_tensor is not None and not is_t2i: - latents, velocity_mask, image_latent = self._prepare_latents_i2v( - image_tensor, - height, - width, - num_frames, - generator, - ) - condition_latents = None - else: - latents = self._prepare_latents(height, width, num_frames, generator) - velocity_mask = None - image_latent = None - condition_latents = None - - video_shape = (latents.shape[2], latents.shape[3], latents.shape[4]) - shared_kwargs: dict[str, Any] = dict(video_shape=video_shape, fps=frame_rate) - if velocity_mask is not None: - shared_kwargs["noisy_frame_mask"] = velocity_mask - - def _run_diffusion(start_latents: torch.Tensor) -> torch.Tensor: - self._set_scheduler_timesteps(num_inference_steps) - return self.diffuse( - latents=start_latents, - timesteps=self.scheduler.timesteps, - cond_ids=cond_ids, - cond_mask=cond_mask, - uncond_ids=uncond_ids, - uncond_mask=uncond_mask, - guidance_scale=guidance_scale, - shared_kwargs=shared_kwargs, - velocity_mask=velocity_mask, - image_latent=image_latent, - condition_latents=condition_latents, - guidance_interval=guidance_interval, - ) - - if is_t2i and batch_size > 1: - samples = [_run_diffusion(latents)] - for _ in range(batch_size - 1): - next_latents = self._prepare_latents(height, width, num_frames, generator) - samples.append(_run_diffusion(next_latents)) - latents = torch.cat(samples, dim=0) - else: - latents = _run_diffusion(latents) - - video = self._decode_latents(latents) - return SimpleNamespace(output={"image": video} if is_t2i else {"video": video}) - - # -- Construction stubs (full wiring deferred) -------------------------- - - def __init__(self, *args: Any, **kwargs: Any) -> None: - raise NotImplementedError("Cosmos3OmniDiffusersPipeline.__init__ is wired in a later phase. " - "Use Cosmos3OmniDiffusersPipeline.__new__(cls) for Phase 2f unit " - "tests that only exercise _remap_ckpt_key, _set_flow_shift, " - "diffuse, or forward (with monkey-patched helpers).") - - def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: - raise NotImplementedError("Cosmos3OmniDiffusersPipeline.create_pipeline_stages is wired in a later phase.") + def tokenize_caption(self, caption: str, *, is_video: bool = False, use_system_prompt: bool = False) -> list[int]: + return cosmos3_tokenize_caption(self.get_module("text_tokenizer"), + caption, + is_video=is_video, + use_system_prompt=use_system_prompt) -# Entry point for pipeline registry (placeholder; full registry wiring lands -# alongside Phase 2f.2 once create_pipeline_stages and __init__ are real). +# Entry point for the pipeline registry. The class name matches the checkpoint +# ``model_index.json`` ``_class_name`` so ``resolve_pipeline_cls`` finds it. EntryClass = Cosmos3OmniDiffusersPipeline diff --git a/fastvideo/pipelines/basic/cosmos3/presets.py b/fastvideo/pipelines/basic/cosmos3/presets.py new file mode 100644 index 000000000..d4053ebfc --- /dev/null +++ b/fastvideo/pipelines/basic/cosmos3/presets.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Cosmos3 (Cosmos3-Nano) inference presets. + +Defaults track the official ``cosmos-framework`` ``sample_args`` for the video +paths (``text2video`` / ``image2video``: guidance=6.0, num_steps=35, shift=10.0, +fps=24, num_frames=189) and ``text2image`` (guidance=4.0, num_steps=50, +shift=3.0). The default resolution is 16:9 at a VAE-aligned 704x1280 (spatial +compression 16 -> 44x80 latent grid). +""" +from fastvideo.api.presets import InferencePreset, PresetStageSpec + +_DENOISE_STAGE = PresetStageSpec( + name="denoise", + kind="denoising", + description="Cosmos3 sequential-CFG UniPC denoising pass", + allowed_overrides=frozenset({ + "num_inference_steps", + "guidance_scale", + }), +) + +# Framework video negative prompt (Cosmos quality prompt). +_COSMOS3_NEGATIVE_PROMPT = ( + "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, " + "over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, " + "underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, " + "jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, " + "fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. " + "Overall, the video is of poor quality.") + +COSMOS3_NANO = InferencePreset( + name="cosmos3_nano", + version=1, + model_family="cosmos3", + description="Cosmos3-Nano text-to-video", + workload_type="t2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + "height": 704, + "width": 1280, + "num_frames": 189, + "fps": 24, + "guidance_scale": 6.0, + "num_inference_steps": 35, + "negative_prompt": _COSMOS3_NEGATIVE_PROMPT, + }, +) + +COSMOS3_NANO_I2V = InferencePreset( + name="cosmos3_nano_i2v", + version=1, + model_family="cosmos3", + description="Cosmos3-Nano image-to-video", + workload_type="i2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + "height": 704, + "width": 1280, + "num_frames": 189, + "fps": 24, + "guidance_scale": 6.0, + "num_inference_steps": 35, + "negative_prompt": _COSMOS3_NEGATIVE_PROMPT, + }, +) + +COSMOS3_NANO_T2I = InferencePreset( + name="cosmos3_nano_t2i", + version=1, + model_family="cosmos3", + description="Cosmos3-Nano text-to-image", + workload_type="t2i", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + "height": 1024, + "width": 1024, + "num_frames": 1, + "fps": 24, + "guidance_scale": 4.0, + "num_inference_steps": 50, + "negative_prompt": "", + }, +) + +ALL_PRESETS = (COSMOS3_NANO, COSMOS3_NANO_I2V, COSMOS3_NANO_T2I) diff --git a/fastvideo/pipelines/stages/cosmos3_stages.py b/fastvideo/pipelines/stages/cosmos3_stages.py new file mode 100644 index 000000000..e4b925fdb --- /dev/null +++ b/fastvideo/pipelines/stages/cosmos3_stages.py @@ -0,0 +1,230 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Cosmos3 video denoising stage. + +The Cosmos3 video path is monolithic by design: each CFG pass repacks the whole +text+vision sequence (the conditional pass carries prompt tokens, the +unconditional pass carries negative-prompt tokens), so the standard +encode/condition/denoise/decode stage split does not apply. This single stage +owns the full flow, delegating the framework-parity-tested math to +``fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline``: + + 1. resolve mode (T2I / I2V / T2V) + per-mode defaults, set ``flow_shift``; + 2. tokenize the prompt + negative prompt with the Qwen2 chat template; + 3. VAE-encode the conditioning frame(s) for I2V / T2I (kept clean), build the + initial noise (clean condition frames + pure noise elsewhere); + 4. run the UniPC denoise loop with sequential CFG + (``Cosmos3DenoiseEngine.denoise``); + 5. VAE-decode + ``(1 + x) / 2`` clamp to [0, 1]. + +This mirrors the framework ``Cosmos3OmniDiffusersPipeline.__call__``. +""" +from __future__ import annotations + +import weakref +from typing import Any + +import torch +from diffusers.utils.torch_utils import randn_tensor +from tqdm import tqdm + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( + COSMOS3_VIDEO_NEGATIVE_PROMPT, + Cosmos3DenoiseEngine, + Cosmos3VisionSpec, + _VaeNorm, + cosmos3_special_tokens, + cosmos3_tokenize_caption, + cosmos3_vae_decode, + cosmos3_vae_encode, +) +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.base import PipelineStage + +logger = init_logger(__name__) + + +class Cosmos3DenoisingStage(PipelineStage): + """Full Cosmos3 video denoise: tokenize + encode + denoise + decode.""" + + def __init__(self, *, transformer, scheduler, vae, tokenizer, pipeline=None) -> None: + self.transformer = transformer + self.scheduler = scheduler + self.vae = vae + self.tokenizer = tokenizer + self.pipeline = weakref.ref(pipeline) if pipeline is not None else None + + # ------------------------------------------------------------------ + # Geometry helpers + # ------------------------------------------------------------------ + @staticmethod + def _latent_frames(num_frames: int, temporal_factor: int) -> int: + return (int(num_frames) - 1) // int(temporal_factor) + 1 + + def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: + pipeline_config = fastvideo_args.pipeline_config + arch = pipeline_config.dit_config.arch_config + device = self.transformer.embed_tokens.weight.device + dtype = self.transformer.embed_tokens.weight.dtype + + num_frames = int(batch.num_frames) if batch.num_frames is not None else 1 + height = int(batch.height) + width = int(batch.width) + fps = float(batch.fps) if batch.fps is not None else float(arch.base_fps) + guidance = float(batch.guidance_scale) + + is_t2i = num_frames == 1 and batch.preprocessed_image is None and batch.pil_image is None + is_i2v = (batch.preprocessed_image is not None or batch.pil_image is not None) and not is_t2i + + # Per-mode flow_shift, set on the owning pipeline (rebuilds scheduler). + pipe = self.pipeline() if self.pipeline is not None else None + engine_shift = float(getattr(pipe, "_engine_init_flow_shift", pipeline_config.flow_shift or 1.0)) + flow_shift = 3.0 if is_t2i else engine_shift + if pipe is not None and hasattr(pipe, "_set_flow_shift"): + pipe._set_flow_shift(flow_shift) + scheduler = pipe.scheduler + else: + scheduler = self.scheduler + + # ---- Tokenize prompt + negative prompt ---- + prompt = batch.prompt if isinstance(batch.prompt, str) else (batch.prompt[0] if batch.prompt else "") + negative_prompt = batch.negative_prompt + if negative_prompt is None: + negative_prompt = "" if is_t2i else COSMOS3_VIDEO_NEGATIVE_PROMPT + if isinstance(negative_prompt, list): + negative_prompt = negative_prompt[0] if negative_prompt else "" + + special_tokens = cosmos3_special_tokens(self.tokenizer) + is_video = not is_t2i + cond_ids = cosmos3_tokenize_caption(self.tokenizer, prompt, is_video=is_video, use_system_prompt=False) + uncond_ids = cosmos3_tokenize_caption(self.tokenizer, + negative_prompt, + is_video=is_video, + use_system_prompt=False) + + # ---- VAE normalization constants + geometry ---- + norm = _VaeNorm.from_vae(self.vae, dtype) + temporal_factor = int(arch.temporal_compression_factor) + spatial_factor = int(self.vae.config.scale_factor_spatial) + latent_t = self._latent_frames(num_frames, temporal_factor) + latent_h = height // spatial_factor + latent_w = width // spatial_factor + latent_channel = int(arch.latent_channel) + latent_shape = (latent_channel, latent_t, latent_h, latent_w) + + generator = batch.generator + if isinstance(generator, list): + generator = generator[0] if generator else None + + # ---- Conditioning latent (I2V / T2I) + condition mask ---- + condition_frame_indexes: list[int] = [] + clean_latent: torch.Tensor | None = None + if is_i2v or (is_t2i and (batch.preprocessed_image is not None or batch.pil_image is not None)): + image = batch.preprocessed_image if batch.preprocessed_image is not None else batch.pil_image + cond_pixels = self._image_to_video_tensor(image, num_frames, height, width, device, dtype) + clean_latent = cosmos3_vae_encode(self.vae, cond_pixels, norm).squeeze(0).float() # [C, T, H, W] + condition_frame_indexes = [0] + + # ---- Initial noise (clean condition frames + pure noise elsewhere) ---- + pure_noise = randn_tensor(latent_shape, generator=generator, device=device, dtype=dtype).float() + if clean_latent is not None: + cond_mask = torch.zeros((latent_t, 1, 1), device=device, dtype=pure_noise.dtype) + for idx in condition_frame_indexes: + if 0 <= idx < latent_t: + cond_mask[idx, 0, 0] = 1.0 + clean = clean_latent.to(device=device, dtype=pure_noise.dtype) + init_latent = cond_mask * clean + (1.0 - cond_mask) * pure_noise + else: + init_latent = pure_noise + + spec = Cosmos3VisionSpec( + shape=latent_shape, + condition_frame_indexes=condition_frame_indexes, + clean_latent=clean_latent, + ) + + # ---- Scheduler timesteps ---- + scheduler.set_timesteps(int(batch.num_inference_steps), device=device) + timesteps = scheduler.timesteps + + engine = Cosmos3DenoiseEngine( + transformer=self.transformer, + scheduler=scheduler, + special_tokens=special_tokens, + latent_patch_size=int(arch.latent_patch_size), + temporal_modality_margin=int(arch.temporal_modality_margin), + reset_spatial_ids=bool(arch.unified_3d_mrope_reset_spatial_ids), + enable_fps_modulation=bool(arch.enable_fps_modulation), + base_fps=float(arch.base_fps), + temporal_compression_factor=temporal_factor, + include_end_of_generation_token=False, + ) + + flat_latent = init_latent.reshape(-1) + fps_per_item = [fps] if bool(arch.enable_fps_modulation) else None + final_flat = engine.denoise( + flat_latent=flat_latent, + timesteps=timesteps, + guidance=guidance, + specs=[spec], + cond_token_ids=cond_ids, + uncond_token_ids=uncond_ids, + fps_per_item=fps_per_item, + progress_bar=lambda it: tqdm(it, desc="Cosmos3 denoising"), + ) + + # ---- Decode: [C, T, H, W] -> pixels [B, 3, T, H, W] in [0, 1] ---- + result_latent = final_flat.reshape(latent_shape).unsqueeze(0).to(device=device, dtype=dtype) + decoded = cosmos3_vae_decode(self.vae, result_latent, norm) # [B, 3, T, H, W] in [-1, 1] + video = ((1.0 + decoded) / 2.0).clamp(0.0, 1.0) + + batch.latents = result_latent + batch.output = video + return batch + + # ------------------------------------------------------------------ + # Image preprocessing + # ------------------------------------------------------------------ + @staticmethod + def _image_to_video_tensor( + image: Any, + num_frames: int, + height: int, + width: int, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + """Coerce a conditioning image into a ``[1, 3, T, H, W]`` pixel tensor in [-1, 1]. + + The first frame holds the (already preprocessed) conditioning image and + the remaining frames are zero-filled; only frame 0 is kept clean by the + condition mask, so the rest of the clip is denoised from noise. + """ + import torchvision.transforms.functional as TF + + if isinstance(image, torch.Tensor): + img = image.float() + # Accept [B,3,T,H,W], [3,T,H,W], [3,H,W], or [B,3,H,W]. + if img.dim() == 5: + img = img[0] + if img.dim() == 4 and img.shape[1] == 1: + img = img[:, 0] # [3, H, W] + elif img.dim() == 4: + img = img[:, 0] # take first frame -> [3, H, W] + # Normalize to [-1, 1] if it looks like [0, 1] or [0, 255]. + if img.max() > 1.5: + img = img / 127.5 - 1.0 + elif hasattr(image, "convert"): # PIL.Image + import numpy as np + arr = np.array(image.convert("RGB")) + img = torch.from_numpy(arr).permute(2, 0, 1).float() / 127.5 - 1.0 + else: + raise TypeError(f"Unsupported conditioning image type: {type(image)}") + + if img.shape[-2:] != (height, width): + img = TF.resize(img.unsqueeze(0), [height, width]).squeeze(0) + + video = torch.zeros((1, 3, num_frames, height, width), device=device, dtype=dtype) + video[0, :, 0] = img.to(device=device, dtype=dtype) + return video diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 9a20fe38e..b84c76aa4 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -1268,6 +1268,8 @@ def _register_presets() -> None: from fastvideo.api.presets import register_preset from fastvideo.pipelines.basic.cosmos.presets import ( ALL_PRESETS as COSMOS_PRESETS, ) + from fastvideo.pipelines.basic.cosmos3.presets import ( + ALL_PRESETS as COSMOS3_PRESETS, ) from fastvideo.pipelines.basic.dreamx_world.presets import ( ALL_PRESETS as DREAMX_WORLD_PRESETS, ) from fastvideo.pipelines.basic.gamecraft.presets import ( @@ -1311,6 +1313,7 @@ def _register_presets() -> None: all_preset_groups = ( COSMOS_PRESETS, + COSMOS3_PRESETS, DREAMX_WORLD_PRESETS, FLUX2_PRESETS, GAMECRAFT_PRESETS, diff --git a/tests/local_tests/cosmos3/conftest.py b/tests/local_tests/cosmos3/conftest.py index e2dca036b..d21ff1c0a 100644 --- a/tests/local_tests/cosmos3/conftest.py +++ b/tests/local_tests/cosmos3/conftest.py @@ -1,37 +1,34 @@ # SPDX-License-Identifier: Apache-2.0 -"""Shared fixtures for Cosmos3 local parity tests (Phase 2a Tier A). +"""Shared fixtures for the Cosmos3 native-pipeline local tests. -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). +These fixtures build the FastVideo-native Cosmos3 pipeline +(``fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline.Cosmos3OmniDiffusersPipeline``) +via ``__new__`` and wire it with tiny stub components so the runtime call graph +(sequential CFG, condition-frame masking, mode dispatch) can be exercised on CPU +without real weights or ``cosmos_framework``. -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. +The stub transformer implements the native DiT's packed-input contract +(``{"preds_vision": [[1, C, T, H, W], ...]}``) and records, per call, the first +``text_ids`` token so tests can assert the cond/uncond pass order. """ from __future__ import annotations -import sys -import types from types import SimpleNamespace from typing import Any import pytest import torch +from diffusers import UniPCMultistepScheduler from torch import nn +_LATENT_CHANNEL = 16 +_LATENT_PATCH_SIZE = 2 +_SPATIAL_FACTOR = 8 +_TEMPORAL_FACTOR = 4 + def pytest_configure(config: pytest.Config) -> None: - """Register ``local`` marker so ``pytestmark = [pytest.mark.local]`` - in sibling test files does not trigger ``PytestUnknownMarkWarning``. - """ + """Register the ``local`` marker used by sibling test files.""" config.addinivalue_line( "markers", "local: marker for local-only parity/scaffold tests (skipped in CI)", @@ -39,46 +36,39 @@ def pytest_configure(config: pytest.Config) -> None: # --------------------------------------------------------------------------- -# Stub scheduler — mirrors vllm-omni ``StubScheduler``. +# Stub transformer: records cond/uncond call order; bounded preds_vision. # --------------------------------------------------------------------------- -class StubScheduler: - """Minimal UniPC-shaped scheduler stub used by pipeline call-graph tests. +class StubCosmos3Transformer(nn.Module): + """Records each forward's first ``text_ids`` token + returns preds_vision. - 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``. + ``preds_vision`` is keyed by the first text token (so the conditional and + unconditional passes return different velocities) and is zero on + conditioning frames, matching the real DiT's unpatchify output. """ - 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 __init__(self, latent_channel: int = _LATENT_CHANNEL) -> None: + super().__init__() + self.latent_channel = latent_channel + self.embed_tokens = nn.Embedding(64, 8) + self.calls: list[dict[str, Any]] = [] - 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,) + def forward(self, **kwargs: Any) -> dict[str, Any]: + token_ids = kwargs["text_ids"] + token = int(token_ids.reshape(-1)[0].item()) if token_ids.numel() else 0 + self.calls.append({"token": token, "kwargs": dict(kwargs)}) + scale = 0.01 * (1.0 + (token % 7)) + preds: list[torch.Tensor] = [] + for latent, _shape, nfi in zip(kwargs["vision_tokens"], kwargs["vision_token_shapes"], + kwargs["vision_noisy_frame_indexes"]): + lat = latent.squeeze(0) if latent.dim() == 5 else latent # [C, T, H, W] + out = torch.zeros_like(lat) + if nfi.numel() > 0: + out[:, nfi] = scale * torch.tanh(lat[:, nfi]) + preds.append(out.unsqueeze(0)) + return {"preds_vision": preds} -class _ModeLatentDist: - """Stub for diffusers ``DiagonalGaussianDistribution.mode``.""" +class _StubLatentDist: def __init__(self, latents: torch.Tensor) -> None: self._latents = latents @@ -88,211 +78,135 @@ class _ModeLatentDist: class StubCosmos3VAE: - """VAE stub returning deterministic latents shaped by VAE scale factors. + """Deterministic VAE shaped by the Wan 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: + def __init__(self, z_dim: int = _LATENT_CHANNEL) -> None: self.config = SimpleNamespace( z_dim=z_dim, - scale_factor_temporal=temporal, - scale_factor_spatial=spatial, + scale_factor_temporal=_TEMPORAL_FACTOR, + scale_factor_spatial=_SPATIAL_FACTOR, 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)) + b, _c, t, h, w = video.shape + lt = (t - 1) // self.config.scale_factor_temporal + 1 + lh = h // self.config.scale_factor_spatial + lw = w // self.config.scale_factor_spatial + return _StubLatentDist(torch.ones(b, self.config.z_dim, lt, lh, lw, dtype=video.dtype, device=video.device)) - def decode(self, latents: torch.Tensor, return_dict: bool = False): - del return_dict - return (latents,) + def decode(self, z: torch.Tensor): + b, _c, lt, lh, lw = z.shape + t = (lt - 1) * self.config.scale_factor_temporal + 1 + h = lh * self.config.scale_factor_spatial + w = lw * self.config.scale_factor_spatial + sig = torch.nan_to_num(torch.tanh(z[:, :1, :1, :1, :1])).reshape(b, 1, 1, 1, 1) + return torch.clamp(torch.zeros(b, 3, t, h, w, dtype=z.dtype, device=z.device) + sig, -1.0, 1.0) -class StubCosmos3Transformer(nn.Module): - """Transformer stub that records per-call inputs and emits deterministic - output tensors keyed by the first ``text_ids`` token. +class StubQwen2Tokenizer: + """Qwen2-shaped chat tokenizer stub (special tokens + chat template).""" - Mirrors ``tests/diffusion/models/cosmos3/conftest.py:73-114``. + eos_token_id = 62 + _SPECIAL = {"<|vision_start|>": 60, "<|vision_end|>": 61} - 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 convert_tokens_to_ids(self, token: str) -> int: + return self._SPECIAL[token] - 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 apply_chat_template(self, conversations, *, tokenize=True, add_generation_prompt=True, add_vision_id=False): + user = next((c["content"] for c in conversations if c["role"] == "user"), "") + n = max(1, min(8, len(user) % 8 + 1)) + return [10 + (i % 40) for i in range(n)] -def passthrough_progress_bar(iterable): - return iterable +def make_scheduler(flow_shift: float = 10.0) -> UniPCMultistepScheduler: + return UniPCMultistepScheduler( + num_train_timesteps=1000, + solver_order=2, + prediction_type="flow_prediction", + use_flow_sigmas=True, + flow_shift=flow_shift, + ) # --------------------------------------------------------------------------- -# 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. +# Pipeline factory — builds the native pipeline via __new__ + stub modules. # --------------------------------------------------------------------------- @pytest.fixture def make_cosmos3_pipeline(): - """Factory that returns a FastVideo Cosmos3 pipeline pre-wired with stubs. + """Return a factory building the native Cosmos3 pipeline 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. + def _make(**overrides: Any): + from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( # noqa: F401 + Cosmos3OmniDiffusersPipeline, ) - 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 + pipe = Cosmos3OmniDiffusersPipeline.__new__(Cosmos3OmniDiffusersPipeline) + scheduler = make_scheduler() + pipe.modules = { + "transformer": StubCosmos3Transformer(), + "vae": StubCosmos3VAE(), + "scheduler": scheduler, + "text_tokenizer": StubQwen2Tokenizer(), + } + pipe.scheduler = scheduler + pipe._base_scheduler_config = scheduler.config + pipe._current_flow_shift = float(scheduler.config.flow_shift) + pipe._engine_init_flow_shift = 10.0 + for key, value in overrides.items(): + setattr(pipe, key, value) + return pipe return _make -def make_sampling_params(**overrides: Any) -> SimpleNamespace: - """Build a SamplingParams-like namespace with Cosmos3's expected fields. +@pytest.fixture +def make_cosmos3_stage(): + """Return a factory building a ``Cosmos3DenoisingStage`` bound to a pipeline.""" - 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": {}, - } + def _make(pipeline): + from fastvideo.pipelines.stages.cosmos3_stages import Cosmos3DenoisingStage + + return Cosmos3DenoisingStage( + transformer=pipeline.modules["transformer"], + scheduler=pipeline.modules["scheduler"], + vae=pipeline.modules["vae"], + tokenizer=pipeline.modules["text_tokenizer"], + pipeline=pipeline, + ) + + return _make + + +def make_forward_batch(*, num_frames: int, height: int, width: int, image: Any = None, **overrides: Any): + """Build a tiny ``ForwardBatch`` for the Cosmos3 stage.""" + from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + values: dict[str, Any] = dict( + data_type="video", + prompt="a calm ocean at sunrise", + negative_prompt="", + height=height, + width=width, + num_frames=num_frames, + fps=24, + num_inference_steps=2, + guidance_scale=6.0, + generator=torch.Generator("cpu").manual_seed(0), + preprocessed_image=image, + ) values.update(overrides) - return SimpleNamespace(**values) + return ForwardBatch(**values) + + +def make_fastvideo_args(): + """Build minimal ``fastvideo_args`` (only ``pipeline_config`` is read).""" + from fastvideo.configs.pipelines.cosmos3 import Cosmos3Config + + cfg = Cosmos3Config() + arch = cfg.dit_config.arch_config + arch.latent_channel = _LATENT_CHANNEL + arch.latent_patch_size = _LATENT_PATCH_SIZE + arch.temporal_compression_factor = _TEMPORAL_FACTOR + arch.enable_fps_modulation = False + return SimpleNamespace(pipeline_config=cfg) diff --git a/tests/local_tests/cosmos3/test_cosmos3_denoise_cfg_parity.py b/tests/local_tests/cosmos3/test_cosmos3_denoise_cfg_parity.py new file mode 100644 index 000000000..cc080f172 --- /dev/null +++ b/tests/local_tests/cosmos3/test_cosmos3_denoise_cfg_parity.py @@ -0,0 +1,351 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Numerical-parity test: FastVideo Cosmos3 denoise/CFG glue vs the framework. + +The DiT forward and the sequence-packing are already framework-parity-verified +(``test_cosmos3_dit_parity*`` / ``test_cosmos3_packing_parity``). This test pins +the remaining glue that the native pipeline adds — the SEQUENTIAL classifier-free +guidance velocity and one UniPC scheduler step — against the framework math +(``diffusers_cosmos3.pipeline.Cosmos3OmniDiffusersPipeline.get_cfg_velocity`` / +``__call__``): + + * for one denoise step, replicate the framework's ``get_cfg_velocity`` exactly + on top of the OFFICIAL ``Cosmos3VFMNetwork`` forward (oracle): a conditional + pass (prompt tokens) and an unconditional pass (negative-prompt tokens), + each masking the prediction on conditioning frames + (``pred * (1 - condition_mask)``), then ``v = uncond + g*(cond - uncond)``; + * run FastVideo's :func:`cosmos3_get_cfg_velocity` with the native DiT (the + framework weights copied in) + the native packer, and assert the velocity + matches the oracle; + * take one ``UniPCMultistepScheduler.step`` on each (the actual checkpoint + scheduler) and assert the stepped latent matches; + * drive :meth:`Cosmos3DenoiseEngine.denoise` for >= 2 steps and assert it + equals the manual framework step-by-step loop. + +CPU / float32, via the reference SDPA monkey-patch. The official model is the +parity ORACLE. + +Run: + cd && -m pytest \ + tests/local_tests/cosmos3/test_cosmos3_denoise_cfg_parity.py -q +""" +from __future__ import annotations + +import pytest +import torch + +# The official framework provides the parity oracle. +cosmos_framework = pytest.importorskip( + "cosmos_framework", + reason="cosmos_framework not installed; run in fv-cosmos3 env.", +) + +from diffusers import UniPCMultistepScheduler # noqa: E402 + +from .test_cosmos3_dit_parity import _copy_weights # noqa: E402 +from .test_cosmos3_dit_parity_mrope import ( # noqa: E402 + _LATENT_CHANNEL, + _LATENT_PATCH_SIZE, + _RESET_SPATIAL_IDS, + _TCF, + _TEMPORAL_MODALITY_MARGIN, + _build_tiny_cosmos3_mrope, + _build_tiny_fastvideo_dit_mrope, +) +from .test_cosmos3_reference_forward import _apply_sdpa_patches # noqa: E402 + +pytestmark = [pytest.mark.local] + +_apply_sdpa_patches() + +# Tiny special tokens (< tiny vocab_size=64), video path appends eos + sog. +_SPECIAL_TOKENS = { + "start_of_generation": 60, + "end_of_generation": 61, + "eos_token_id": 62, +} + + +def _make_scheduler() -> UniPCMultistepScheduler: + """UniPC scheduler matching the Cosmos3 checkpoint flow config.""" + return UniPCMultistepScheduler( + num_train_timesteps=1000, + solver_order=2, + prediction_type="flow_prediction", + use_flow_sigmas=True, + flow_shift=10.0, + ) + + +# --------------------------------------------------------------------------- +# Framework-oracle CFG velocity (replicates pipeline.get_cfg_velocity math). +# --------------------------------------------------------------------------- +def _framework_pack(*, text_ids, vision_latent, cond_frames, timestep): + from cosmos_framework.data.vfm.sequence_packing import ( + GenerationDataClean, + SequencePlan, + pack_input_sequence, + ) + + # vision_latent is [1, C, T, H, W]; temporal dim is axis 2. + gen_data_clean = GenerationDataClean( + batch_size=1, + is_image_batch=(vision_latent.shape[2] == 1), + x0_tokens_vision=[vision_latent], + fps_vision=None, + num_vision_items_per_sample=[1], + ) + plans = [SequencePlan(has_text=True, has_vision=True, condition_frame_indexes_vision=list(cond_frames))] + return pack_input_sequence( + sequence_plans=plans, + input_text_indexes=[list(text_ids)], + gen_data_clean=gen_data_clean, + input_timesteps=torch.tensor([timestep], dtype=torch.float32), + special_tokens=_SPECIAL_TOKENS, + latent_patch_size=_LATENT_PATCH_SIZE, + include_end_of_generation_token=False, + position_embedding_type="unified_3d_mrope", + unified_3d_mrope_reset_spatial_ids=_RESET_SPATIAL_IDS, + unified_3d_mrope_temporal_modality_margin=_TEMPORAL_MODALITY_MARGIN, + enable_fps_modulation=False, + base_fps=24.0, + temporal_compression_factor=_TCF, + ) + + +def _framework_inputs(ps): + """Framework PackedSequence -> framework Cosmos3VFMNetwork forward kwargs.""" + return dict(packed_seq=ps) + + +def _framework_cfg_velocity( + *, + vfm, + flat_latent: torch.Tensor, + timestep: torch.Tensor, + guidance: float, + vision_shape: tuple[int, int, int, int], + cond_frames: list[int], + cond_ids: list[int], + uncond_ids: list[int], +) -> torch.Tensor: + """Replicate the framework ``get_cfg_velocity`` on the oracle model. + + Single vision item; sequential cond then uncond pass; mask condition + frames; ``v = uncond + g*(cond - uncond)``. + """ + timestep_value = float(timestep.reshape(()).item()) + vision_latent = flat_latent.reshape(vision_shape) # [C, T, H, W] + + def _run(text_ids: list[int]) -> torch.Tensor: + ps = _framework_pack( + text_ids=text_ids, + # The framework packer expects a 5D [1, C, T, H, W] latent. + vision_latent=vision_latent.unsqueeze(0), + cond_frames=cond_frames, + timestep=timestep_value, + ) + out = vfm(**_framework_inputs(ps)) + preds = out.get("preds_vision") + cond_mask = ps.vision.condition_mask[0] # [T] or [T,1,1] + if preds is None: + return torch.zeros_like(flat_latent) + pred = preds[0].squeeze(0) # [C, T, H, W] + keep = (1.0 - cond_mask.reshape(-1, 1, 1)).to(dtype=pred.dtype, device=pred.device) + velocity = pred * keep if keep.sum() > 0 else torch.zeros_like(pred) + return velocity.reshape(-1) + + cond_v = _run(cond_ids) + uncond_v = _run(uncond_ids) + return uncond_v + guidance * (cond_v - uncond_v) + + +# --------------------------------------------------------------------------- +# Cases: T2V (no cond), I2V (cond frame 0), single-frame T2I. +# --------------------------------------------------------------------------- +_CASES = [ + pytest.param(2, 4, 4, 6, [], id="t2v_2x2x2"), + pytest.param(3, 8, 4, 5, [0], id="i2v_3x4x2_cond0"), + pytest.param(1, 8, 8, 4, [], id="t2i_1x4x4"), +] + + +def _build_models(num_layers: int = 2, seed_model: int = 42): + vfm = _build_tiny_cosmos3_mrope(seed=seed_model, num_layers=num_layers) + dit = _build_tiny_fastvideo_dit_mrope(num_layers=num_layers) + _copy_weights(vfm, dit) + return vfm, dit + + +def _fastvideo_velocity(dit, *, flat_latent, timestep, guidance, vision_shape, cond_frames, cond_ids, uncond_ids): + from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( + Cosmos3VisionSpec, + cosmos3_get_cfg_velocity, + ) + + spec = Cosmos3VisionSpec(shape=vision_shape, condition_frame_indexes=list(cond_frames)) + return cosmos3_get_cfg_velocity( + transformer=dit, + flat_latent=flat_latent, + timestep=timestep, + guidance=guidance, + specs=[spec], + cond_token_ids=cond_ids, + uncond_token_ids=uncond_ids, + special_tokens=_SPECIAL_TOKENS, + latent_patch_size=_LATENT_PATCH_SIZE, + temporal_modality_margin=_TEMPORAL_MODALITY_MARGIN, + reset_spatial_ids=_RESET_SPATIAL_IDS, + enable_fps_modulation=False, + base_fps=24.0, + temporal_compression_factor=_TCF, + ) + + +class TestCosmos3DenoiseCFGParity: + + @pytest.mark.parametrize(("grid_t", "latent_h", "latent_w", "n_text", "cond"), _CASES) + def test_cfg_velocity_matches_framework(self, grid_t, latent_h, latent_w, n_text, cond): + vfm, dit = _build_models() + torch.manual_seed(0) + cond_ids = torch.randint(0, 60, (n_text,)).tolist() + uncond_ids = torch.randint(0, 60, (max(1, n_text - 1),)).tolist() + vision_shape = (_LATENT_CHANNEL, grid_t, latent_h, latent_w) + flat_latent = torch.randn(int(torch.tensor(vision_shape).prod())) + timestep = torch.tensor([[500.0]]) # framework expects [1,1]; we reshape to scalar + guidance = 6.0 + + fw_v = _framework_cfg_velocity( + vfm=vfm, + flat_latent=flat_latent, + timestep=timestep, + guidance=guidance, + vision_shape=vision_shape, + cond_frames=cond, + cond_ids=cond_ids, + uncond_ids=uncond_ids, + ) + fv_v = _fastvideo_velocity( + dit, + flat_latent=flat_latent, + timestep=timestep, + guidance=guidance, + vision_shape=vision_shape, + cond_frames=cond, + cond_ids=cond_ids, + uncond_ids=uncond_ids, + ) + assert fw_v.shape == fv_v.shape, f"shape: fw={fw_v.shape} fv={fv_v.shape}" + max_abs = (fw_v - fv_v).abs().max().item() + print(f"\n[cfg_velocity {grid_t}x{latent_h}x{latent_w} cond={cond}] max abs diff = {max_abs:.3e}") + torch.testing.assert_close(fv_v, fw_v, atol=1e-4, rtol=1e-3) + + def test_one_unipc_step_matches_framework(self): + """CFG velocity + one UniPC step: FastVideo == framework math.""" + vfm, dit = _build_models() + grid_t, latent_h, latent_w = 2, 4, 4 + vision_shape = (_LATENT_CHANNEL, grid_t, latent_h, latent_w) + torch.manual_seed(3) + cond_ids = torch.randint(0, 60, (5,)).tolist() + uncond_ids = torch.randint(0, 60, (4,)).tolist() + flat_latent = torch.randn(int(torch.tensor(vision_shape).prod())) + guidance = 6.0 + + fw_sched = _make_scheduler() + fv_sched = _make_scheduler() + fw_sched.set_timesteps(4, device=torch.device("cpu")) + fv_sched.set_timesteps(4, device=torch.device("cpu")) + t = fw_sched.timesteps[0] + + fw_v = _framework_cfg_velocity( + vfm=vfm, + flat_latent=flat_latent, + timestep=t.reshape(1, 1), + guidance=guidance, + vision_shape=vision_shape, + cond_frames=[], + cond_ids=cond_ids, + uncond_ids=uncond_ids, + ) + fw_stepped = fw_sched.step(model_output=fw_v, timestep=t, sample=flat_latent.unsqueeze(0), + return_dict=False)[0].squeeze(0) + + fv_v = _fastvideo_velocity( + dit, + flat_latent=flat_latent, + timestep=t.reshape(1), + guidance=guidance, + vision_shape=vision_shape, + cond_frames=[], + cond_ids=cond_ids, + uncond_ids=uncond_ids, + ) + fv_stepped = fv_sched.step(model_output=fv_v, timestep=t, sample=flat_latent.unsqueeze(0), + return_dict=False)[0].squeeze(0) + + max_abs = (fw_stepped - fv_stepped).abs().max().item() + print(f"\n[unipc_step] max abs diff = {max_abs:.3e}") + torch.testing.assert_close(fv_stepped, fw_stepped, atol=1e-4, rtol=1e-3) + + def test_full_denoise_loop_matches_framework(self): + """Cosmos3DenoiseEngine.denoise (>= 2 steps) == framework step-by-step.""" + from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( + Cosmos3DenoiseEngine, + Cosmos3VisionSpec, + ) + + vfm, dit = _build_models() + grid_t, latent_h, latent_w = 2, 4, 4 + vision_shape = (_LATENT_CHANNEL, grid_t, latent_h, latent_w) + torch.manual_seed(5) + cond_ids = torch.randint(0, 60, (5,)).tolist() + uncond_ids = torch.randint(0, 60, (4,)).tolist() + flat_latent = torch.randn(int(torch.tensor(vision_shape).prod())) + guidance = 6.0 + num_steps = 3 + + # Manual framework loop (oracle). + fw_sched = _make_scheduler() + fw_sched.set_timesteps(num_steps, device=torch.device("cpu")) + fw_latent = flat_latent.clone() + for t in fw_sched.timesteps: + v = _framework_cfg_velocity( + vfm=vfm, + flat_latent=fw_latent, + timestep=t.reshape(1, 1), + guidance=guidance, + vision_shape=vision_shape, + cond_frames=[], + cond_ids=cond_ids, + uncond_ids=uncond_ids, + ) + fw_latent = fw_sched.step(model_output=v, timestep=t, sample=fw_latent.unsqueeze(0), + return_dict=False)[0].squeeze(0) + + # FastVideo engine loop. + fv_sched = _make_scheduler() + fv_sched.set_timesteps(num_steps, device=torch.device("cpu")) + engine = Cosmos3DenoiseEngine( + transformer=dit, + scheduler=fv_sched, + special_tokens=_SPECIAL_TOKENS, + latent_patch_size=_LATENT_PATCH_SIZE, + temporal_modality_margin=_TEMPORAL_MODALITY_MARGIN, + reset_spatial_ids=_RESET_SPATIAL_IDS, + enable_fps_modulation=False, + base_fps=24.0, + temporal_compression_factor=_TCF, + ) + spec = Cosmos3VisionSpec(shape=vision_shape, condition_frame_indexes=[]) + fv_latent = engine.denoise( + flat_latent=flat_latent.clone(), + timesteps=fv_sched.timesteps, + guidance=guidance, + specs=[spec], + cond_token_ids=cond_ids, + uncond_token_ids=uncond_ids, + ) + + assert fv_latent.shape == fw_latent.shape + max_abs = (fw_latent - fv_latent).abs().max().item() + print(f"\n[full_denoise {num_steps} steps] final latent max abs diff = {max_abs:.3e}") + torch.testing.assert_close(fv_latent, fw_latent, atol=1e-4, rtol=1e-3) diff --git a/tests/local_tests/cosmos3/test_cosmos3_pipeline_call_graph.py b/tests/local_tests/cosmos3/test_cosmos3_pipeline_call_graph.py index a33eb2823..6164961f3 100644 --- a/tests/local_tests/cosmos3/test_cosmos3_pipeline_call_graph.py +++ b/tests/local_tests/cosmos3/test_cosmos3_pipeline_call_graph.py @@ -1,199 +1,198 @@ # SPDX-License-Identifier: Apache-2.0 -"""Cosmos3 pipeline call-graph parity (Tier A scaffold). +"""Cosmos3 native-pipeline call-graph contract (Tier A, no real weights). -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``. +Pins the runtime call graph of the FastVideo-native Cosmos3 pipeline against the +framework math, using the stub components from ``conftest.py`` (no real weights, +no ``cosmos_framework``). The native pipeline replaced the vllm-omni-derived +``diffuse``/``forward(req)``/``reset_cache`` skeleton with a stage-based +``Cosmos3DenoisingStage`` + ``Cosmos3DenoiseEngine`` doing SEQUENTIAL CFG. -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: +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``. + 1. SEQUENTIAL CFG order — per UniPC step, the transformer is called twice, + conditional (prompt tokens) then unconditional (negative-prompt tokens), in + that order; over N steps the call order is ``[cond, uncond] * N``. + 2. CFG combination — the per-step velocity equals + ``uncond + guidance * (cond - uncond)`` (verified against the stub's known + per-token output) and one UniPC step advances the latent accordingly. + 3. I2V conditioning — a conditioning image is VAE-encoded and frame 0 is kept + clean: its velocity is zeroed (condition mask) so the decoded clip's + frame-0 latent equals the clean conditioning latent. + 4. Mode dispatch — the stage routes T2V (num_frames>1, flow_shift=10.0, + ``is_video`` tokenization) vs T2I (num_frames==1, flow_shift=3.0, image + tokenization), applying the per-mode ``flow_shift``. """ from __future__ import annotations import pytest import torch +from .conftest import make_fastvideo_args, make_forward_batch + pytestmark = [pytest.mark.local] - -def _ids(value: int) -> torch.Tensor: - return torch.tensor([[value]], dtype=torch.long) +_LATENT_CHANNEL = 16 +_LATENT_PATCH_SIZE = 2 +_TEMPORAL_FACTOR = 4 +_COND_TOKEN = 2 +_UNCOND_TOKEN = 1 -def _mask() -> torch.Tensor: - return torch.ones(1, 1, dtype=torch.long) +def _engine(pipeline, scheduler): + from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import Cosmos3DenoiseEngine - -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), + return Cosmos3DenoiseEngine( + transformer=pipeline.modules["transformer"], + scheduler=scheduler, + special_tokens={"start_of_generation": 60, "end_of_generation": 61, "eos_token_id": 62}, + latent_patch_size=_LATENT_PATCH_SIZE, + temporal_modality_margin=15_000, + reset_spatial_ids=True, + enable_fps_modulation=False, + base_fps=24.0, + temporal_compression_factor=_TEMPORAL_FACTOR, ) - 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. +def test_sequential_cfg_calls_cond_then_uncond_each_step(make_cosmos3_pipeline) -> None: + """Each UniPC step calls the transformer cond-then-uncond, in order.""" + from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import Cosmos3VisionSpec - 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), + scheduler = pipeline.modules["scheduler"] + scheduler.set_timesteps(2, device=torch.device("cpu")) + engine = _engine(pipeline, scheduler) + + shape = (_LATENT_CHANNEL, 2, 2, 2) + flat = torch.randn(int(torch.tensor(shape).prod())) + spec = Cosmos3VisionSpec(shape=shape, condition_frame_indexes=[]) + + engine.denoise( + flat_latent=flat, + timesteps=scheduler.timesteps, + guidance=6.0, + specs=[spec], + cond_token_ids=[_COND_TOKEN, 5, 6], + uncond_token_ids=[_UNCOND_TOKEN, 7], ) - torch.testing.assert_close(result[:, :, 0:1], torch.full((1, 2, 1, 1, 1), 7.0)) + tokens = [c["token"] for c in pipeline.modules["transformer"].calls] + # 2 steps -> 4 calls: cond, uncond, cond, uncond. + assert tokens == [_COND_TOKEN, _UNCOND_TOKEN, _COND_TOKEN, _UNCOND_TOKEN] -@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. +def test_cfg_velocity_combination_formula(make_cosmos3_pipeline) -> None: + """The per-step velocity equals ``uncond + g*(cond - uncond)``. - 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. + The stub returns ``scale(token) * tanh(latent)`` on noisy frames, so the + expected velocity is a closed form we can check exactly. """ - 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]) + from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( + Cosmos3VisionSpec, + cosmos3_get_cfg_velocity, ) - 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() + transformer = pipeline.modules["transformer"] + shape = (_LATENT_CHANNEL, 2, 2, 2) + flat = torch.randn(int(torch.tensor(shape).prod())) + guidance = 6.0 - from types import SimpleNamespace + v = cosmos3_get_cfg_velocity( + transformer=transformer, + flat_latent=flat, + timestep=torch.tensor([500.0]), + guidance=guidance, + specs=[Cosmos3VisionSpec(shape=shape, condition_frame_indexes=[])], + cond_token_ids=[_COND_TOKEN, 5, 6], + uncond_token_ids=[_UNCOND_TOKEN, 7], + special_tokens={"start_of_generation": 60, "end_of_generation": 61, "eos_token_id": 62}, + latent_patch_size=_LATENT_PATCH_SIZE, + temporal_modality_margin=15_000, + reset_spatial_ids=True, + enable_fps_modulation=False, + base_fps=24.0, + temporal_compression_factor=_TEMPORAL_FACTOR, + ) - 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={}, - ), - ) - ) + lat = flat.reshape(shape) + scale_cond = 0.01 * (1.0 + (_COND_TOKEN % 7)) + scale_uncond = 0.01 * (1.0 + (_UNCOND_TOKEN % 7)) + cond_v = (scale_cond * torch.tanh(lat)).reshape(-1) + uncond_v = (scale_uncond * torch.tanh(lat)).reshape(-1) + expected = uncond_v + guidance * (cond_v - uncond_v) + torch.testing.assert_close(v, expected, atol=1e-6, rtol=1e-5) + + +def test_i2v_keeps_condition_frame_clean(make_cosmos3_pipeline) -> None: + """I2V: frame-0 velocity is zeroed so the conditioning frame stays clean.""" + from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( + Cosmos3VisionSpec, + cosmos3_get_cfg_velocity, + ) + + pipeline = make_cosmos3_pipeline() + transformer = pipeline.modules["transformer"] + shape = (_LATENT_CHANNEL, 3, 2, 2) # 3 latent frames, frame 0 conditioned + flat = torch.randn(int(torch.tensor(shape).prod())) + + v = cosmos3_get_cfg_velocity( + transformer=transformer, + flat_latent=flat, + timestep=torch.tensor([500.0]), + guidance=6.0, + specs=[Cosmos3VisionSpec(shape=shape, condition_frame_indexes=[0])], + cond_token_ids=[_COND_TOKEN, 5, 6], + uncond_token_ids=[_UNCOND_TOKEN, 7], + special_tokens={"start_of_generation": 60, "end_of_generation": 61, "eos_token_id": 62}, + latent_patch_size=_LATENT_PATCH_SIZE, + temporal_modality_margin=15_000, + reset_spatial_ids=True, + enable_fps_modulation=False, + base_fps=24.0, + temporal_compression_factor=_TEMPORAL_FACTOR, + ) + v_grid = v.reshape(shape) # [C, T, H, W] + # Condition frame 0 velocity must be exactly zero; noisy frames non-zero. + assert torch.count_nonzero(v_grid[:, 0]) == 0 + assert torch.count_nonzero(v_grid[:, 1:]) > 0 + + +def test_stage_mode_dispatch_t2v(make_cosmos3_pipeline, make_cosmos3_stage) -> None: + """T2V (num_frames>1): flow_shift=engine-init (10.0), video tokenization.""" + pipeline = make_cosmos3_pipeline() + stage = make_cosmos3_stage(pipeline) + args = make_fastvideo_args() + batch = make_forward_batch(num_frames=5, height=16, width=16) + + out = stage.forward(batch, args) + assert float(pipeline.scheduler.config.flow_shift) == 10.0 + assert out.output is not None and out.output.dim() == 5 + # T2V latent: (5-1)//4 + 1 = 2 frames; 16/8 = 2 latent h/w. + assert tuple(out.latents.shape) == (1, _LATENT_CHANNEL, 2, 2, 2) + + +def test_stage_mode_dispatch_t2i(make_cosmos3_pipeline, make_cosmos3_stage) -> None: + """T2I (num_frames==1): flow_shift=3.0; single-frame latent.""" + pipeline = make_cosmos3_pipeline() + stage = make_cosmos3_stage(pipeline) + args = make_fastvideo_args() + batch = make_forward_batch(num_frames=1, height=16, width=16, guidance_scale=4.0) + + out = stage.forward(batch, args) + assert float(pipeline.scheduler.config.flow_shift) == 3.0 + assert tuple(out.latents.shape) == (1, _LATENT_CHANNEL, 1, 2, 2) + + +def test_stage_i2v_encodes_conditioning_image(make_cosmos3_pipeline, make_cosmos3_stage) -> None: + """I2V stage: a conditioning image is accepted and decoded to a finite clip.""" + pipeline = make_cosmos3_pipeline() + stage = make_cosmos3_stage(pipeline) + args = make_fastvideo_args() + image = torch.zeros(3, 16, 16) # [-1, 1] conditioning frame + batch = make_forward_batch(num_frames=5, height=16, width=16, image=image) + + out = stage.forward(batch, args) + # I2V uses the same engine-init flow_shift as T2V (10.0). + assert float(pipeline.scheduler.config.flow_shift) == 10.0 + assert out.output is not None and out.output.dim() == 5 + assert torch.isfinite(out.output).all() diff --git a/tests/local_tests/cosmos3/test_cosmos3_pipeline_smoke.py b/tests/local_tests/cosmos3/test_cosmos3_pipeline_smoke.py new file mode 100644 index 000000000..c75563107 --- /dev/null +++ b/tests/local_tests/cosmos3/test_cosmos3_pipeline_smoke.py @@ -0,0 +1,310 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Cosmos3 pipeline smoke test (CPU / float32, tiny components). + +Exercises the native pipeline's runtime path end-to-end on tiny stub-or-real +components, with NO real weights and NO cosmos_framework dependency: + + * a stub transformer implementing the native DiT's packed-input -> + ``{"preds_vision": [...]}`` contract with bounded, deterministic output + (the REAL DiT forward / unpatchify is exhaustively bit-identical-tested in + ``test_cosmos3_dit_parity*`` and ``test_cosmos3_denoise_cfg_parity``; an + untrained real DiT emits unbounded velocities that overflow UniPC, so the + smoke uses a stub to keep finiteness deterministic); + * a stub VAE exposing the ``AutoencoderKLWan`` surface + (``config.latents_mean/std/scale_factor_spatial``, ``encode().mode()``, + ``decode()``) used by the encode/normalize + decode/denormalize bridges; + * a stub Qwen2-shaped tokenizer (chat template + special tokens). + +Two paths are covered: + + 1. ``Cosmos3DenoiseEngine.denoise`` for >= 2 UniPC steps over a tiny T2V + latent, asserting a finite final latent of the right shape, plus the VAE + decode + ``(1 + x)/2`` clamp producing a finite ``[B, 3, T, H, W]`` video; + 2. the real ``Cosmos3DenoisingStage.forward`` (full tokenize -> noise -> + denoise -> decode wiring) driven through a ``__new__``-built pipeline + + ``ForwardBatch``, for both T2V and I2V. + +Run: + cd && -m pytest \ + tests/local_tests/cosmos3/test_cosmos3_pipeline_smoke.py -q +""" +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest +import torch +from diffusers import UniPCMultistepScheduler + +pytestmark = [pytest.mark.local] + +_LATENT_CHANNEL = 16 +_LATENT_PATCH_SIZE = 2 +_SPATIAL_FACTOR = 8 +_TEMPORAL_FACTOR = 4 + + +# --------------------------------------------------------------------------- +# Stub transformer implementing the native DiT packed-input contract. +# --------------------------------------------------------------------------- +class StubCosmos3Transformer(torch.nn.Module): + """Bounded stand-in for ``Cosmos3VFMTransformer.forward``. + + Consumes the same packed kwargs (``vision_token_shapes`` / + ``vision_noisy_frame_indexes`` / ``vision_tokens`` / ``text_ids``) and + returns ``{"preds_vision": [[1, C, T, H, W], ...]}`` with predictions only on + noisy frames (zeros on conditioning frames), matching the real DiT's + ``_unpatchify_and_unpack`` output structure. The prediction is a small + ``tanh`` of the input latent, scaled by the first text id so the cond and + uncond passes differ (exercising the CFG combination). + """ + + def __init__(self, latent_channel: int = _LATENT_CHANNEL) -> None: + super().__init__() + self.latent_channel = latent_channel + # A real attribute the stage reads for device/dtype. + self.embed_tokens = torch.nn.Embedding(64, 8) + + def forward(self, **kwargs: Any) -> dict[str, Any]: + token_ids = kwargs["text_ids"] + token = float(token_ids.reshape(-1)[0].item()) if token_ids.numel() else 1.0 + scale = 0.01 * (1.0 + (token % 7)) + token_shapes = kwargs["vision_token_shapes"] + noisy = kwargs["vision_noisy_frame_indexes"] + tokens = kwargs["vision_tokens"] + preds: list[torch.Tensor] = [] + for latent, (t, _h, _w), nfi in zip(tokens, token_shapes, noisy): + lat = latent.squeeze(0) if latent.dim() == 5 else latent # [C, T, H, W] + out = torch.zeros_like(lat) + if nfi.numel() > 0: + out[:, nfi] = scale * torch.tanh(lat[:, nfi]) + preds.append(out.unsqueeze(0)) # [1, C, T, H, W] + return {"preds_vision": preds} + + +# --------------------------------------------------------------------------- +# Stub VAE matching the AutoencoderKLWan surface used by the bridges. +# --------------------------------------------------------------------------- +class _StubLatentDist: + + def __init__(self, latents: torch.Tensor) -> None: + self._latents = latents + + def mode(self) -> torch.Tensor: + return self._latents + + +class StubCosmos3VAE: + """Minimal VAE: deterministic encode/decode shaped by scale factors.""" + + def __init__(self, z_dim: int = _LATENT_CHANNEL) -> None: + self.config = SimpleNamespace( + z_dim=z_dim, + scale_factor_temporal=_TEMPORAL_FACTOR, + scale_factor_spatial=_SPATIAL_FACTOR, + latents_mean=[0.0] * z_dim, + latents_std=[1.0] * z_dim, + ) + + def encode(self, video: torch.Tensor): + b, _c, t, h, w = video.shape + lt = (t - 1) // self.config.scale_factor_temporal + 1 + lh = h // self.config.scale_factor_spatial + lw = w // self.config.scale_factor_spatial + latents = torch.ones(b, self.config.z_dim, lt, lh, lw, dtype=video.dtype, device=video.device) + return _StubLatentDist(latents) + + def decode(self, z: torch.Tensor): + # Upsample latents back to pixel dims; clamp like AutoencoderKLWan. + b, _c, lt, lh, lw = z.shape + t = (lt - 1) * self.config.scale_factor_temporal + 1 + h = lh * self.config.scale_factor_spatial + w = lw * self.config.scale_factor_spatial + # Bounded function of z so the output reflects (and stays finite with) + # the latent: tanh maps any finite z to [-1, 1]; nan_to_num guards + # against non-finite latents from an untrained denoise. + z_signal = torch.nan_to_num(torch.tanh(z[:, :1, :1, :1, :1])) + out = torch.zeros(b, 3, t, h, w, dtype=z.dtype, device=z.device) + z_signal.reshape(b, 1, 1, 1, 1) + return torch.clamp(out, -1.0, 1.0) + + +# --------------------------------------------------------------------------- +# Stub Qwen2-shaped tokenizer (chat template + special tokens). +# --------------------------------------------------------------------------- +class StubQwen2Tokenizer: + eos_token_id = 62 + + _SPECIAL = {"<|vision_start|>": 60, "<|vision_end|>": 61} + + def convert_tokens_to_ids(self, token: str) -> int: + return self._SPECIAL[token] + + def apply_chat_template(self, conversations, *, tokenize=True, add_generation_prompt=True, add_vision_id=False): + # Deterministic small token ids from the user message length. + user = next((c["content"] for c in conversations if c["role"] == "user"), "") + n = max(1, min(8, len(user) % 8 + 1)) + return [10 + (i % 40) for i in range(n)] + + +def _make_scheduler() -> UniPCMultistepScheduler: + return UniPCMultistepScheduler( + num_train_timesteps=1000, + solver_order=2, + prediction_type="flow_prediction", + use_flow_sigmas=True, + flow_shift=10.0, + ) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- +class TestCosmos3PipelineSmoke: + + def test_denoise_engine_and_decode_finite(self): + """Engine.denoise (>= 2 steps) + VAE decode -> finite output, right shape.""" + from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( + Cosmos3DenoiseEngine, + Cosmos3VisionSpec, + _VaeNorm, + cosmos3_vae_decode, + ) + + dit = StubCosmos3Transformer() + vae = StubCosmos3VAE() + scheduler = _make_scheduler() + scheduler.set_timesteps(2, device=torch.device("cpu")) + + latent_shape = (_LATENT_CHANNEL, 2, 2, 2) # [C, T, H, W] + torch.manual_seed(1) + flat = torch.randn(int(torch.tensor(latent_shape).prod())) + + engine = Cosmos3DenoiseEngine( + transformer=dit, + scheduler=scheduler, + special_tokens={"start_of_generation": 60, "end_of_generation": 61, "eos_token_id": 62}, + latent_patch_size=_LATENT_PATCH_SIZE, + temporal_modality_margin=15_000, + reset_spatial_ids=True, + enable_fps_modulation=False, + base_fps=24.0, + temporal_compression_factor=_TEMPORAL_FACTOR, + ) + spec = Cosmos3VisionSpec(shape=latent_shape, condition_frame_indexes=[]) + out_flat = engine.denoise( + flat_latent=flat, + timesteps=scheduler.timesteps, + guidance=6.0, + specs=[spec], + cond_token_ids=[10, 11, 12], + uncond_token_ids=[13, 14], + ) + assert out_flat.shape == flat.shape + assert torch.isfinite(out_flat).all() + + norm = _VaeNorm.from_vae(vae, torch.float32) + result_latent = out_flat.reshape(latent_shape).unsqueeze(0) + decoded = cosmos3_vae_decode(vae, result_latent, norm) + video = ((1.0 + decoded) / 2.0).clamp(0.0, 1.0) + assert video.dim() == 5 and video.shape[1] == 3 + assert torch.isfinite(video).all() + assert float(video.min()) >= 0.0 and float(video.max()) <= 1.0 + + def _make_pipeline(self): + from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import Cosmos3OmniDiffusersPipeline + + pipe = Cosmos3OmniDiffusersPipeline.__new__(Cosmos3OmniDiffusersPipeline) + scheduler = _make_scheduler() + pipe.modules = { + "transformer": StubCosmos3Transformer(), + "vae": StubCosmos3VAE(), + "scheduler": scheduler, + "text_tokenizer": StubQwen2Tokenizer(), + } + pipe.scheduler = scheduler + pipe._base_scheduler_config = scheduler.config + pipe._current_flow_shift = float(scheduler.config.flow_shift) + pipe._engine_init_flow_shift = 10.0 + return pipe + + def _make_args(self): + from fastvideo.configs.pipelines.cosmos3 import Cosmos3Config + + cfg = Cosmos3Config() + # Shrink the DiT arch config to the tiny smoke geometry. + arch = cfg.dit_config.arch_config + arch.latent_channel = _LATENT_CHANNEL + arch.latent_patch_size = _LATENT_PATCH_SIZE + arch.temporal_compression_factor = _TEMPORAL_FACTOR + arch.enable_fps_modulation = False + return SimpleNamespace(pipeline_config=cfg) + + def _make_batch(self, *, num_frames, height, width, image=None): + from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + return ForwardBatch( + data_type="video", + prompt="a calm ocean at sunrise", + negative_prompt="", + height=height, + width=width, + num_frames=num_frames, + fps=24, + num_inference_steps=2, + guidance_scale=6.0, + generator=torch.Generator("cpu").manual_seed(0), + preprocessed_image=image, + ) + + def test_stage_forward_t2v_finite(self): + """The real Cosmos3DenoisingStage.forward runs T2V end-to-end.""" + from fastvideo.pipelines.stages.cosmos3_stages import Cosmos3DenoisingStage + + pipe = self._make_pipeline() + args = self._make_args() + stage = Cosmos3DenoisingStage( + transformer=pipe.modules["transformer"], + scheduler=pipe.modules["scheduler"], + vae=pipe.modules["vae"], + tokenizer=pipe.modules["text_tokenizer"], + pipeline=pipe, + ) + # 5 frames -> latent_t = (5-1)//4 + 1 = 2; 16x16 px -> 2x2 latent. + batch = self._make_batch(num_frames=5, height=16, width=16) + out = stage.forward(batch, args) + assert out.output is not None + assert out.output.dim() == 5 and out.output.shape[1] == 3 + assert torch.isfinite(out.output).all() + assert torch.isfinite(out.latents).all() + + def test_stage_forward_i2v_keeps_condition_frame(self): + """I2V: a conditioning image is VAE-encoded and frame 0 stays clean.""" + from fastvideo.pipelines.stages.cosmos3_stages import Cosmos3DenoisingStage + + pipe = self._make_pipeline() + args = self._make_args() + stage = Cosmos3DenoisingStage( + transformer=pipe.modules["transformer"], + scheduler=pipe.modules["scheduler"], + vae=pipe.modules["vae"], + tokenizer=pipe.modules["text_tokenizer"], + pipeline=pipe, + ) + # Conditioning image as a [3, H, W] tensor in [-1, 1]. + image = torch.zeros(3, 16, 16) + batch = self._make_batch(num_frames=5, height=16, width=16, image=image) + out = stage.forward(batch, args) + assert out.output is not None + assert out.output.dim() == 5 and out.output.shape[1] == 3 + assert torch.isfinite(out.output).all() + + def test_tokenize_caption_special_tokens(self): + """Pipeline.tokenize_caption uses the Qwen2 chat template + ids.""" + from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import cosmos3_special_tokens + + pipe = self._make_pipeline() + ids = pipe.tokenize_caption("hello world", is_video=True, use_system_prompt=False) + assert isinstance(ids, list) and len(ids) > 0 + special = cosmos3_special_tokens(pipe.get_module("text_tokenizer")) + assert special == {"start_of_generation": 60, "end_of_generation": 61, "eos_token_id": 62} diff --git a/tests/local_tests/cosmos3/test_cosmos3_tokenizer_chat_template.py b/tests/local_tests/cosmos3/test_cosmos3_tokenizer_chat_template.py index 89dc13fea..9c307f79a 100644 --- a/tests/local_tests/cosmos3/test_cosmos3_tokenizer_chat_template.py +++ b/tests/local_tests/cosmos3/test_cosmos3_tokenizer_chat_template.py @@ -1,91 +1,83 @@ # SPDX-License-Identifier: Apache-2.0 -"""Cosmos3 prompt tokenization parity — chat template + special tokens (Tier A). +"""Cosmos3 prompt tokenization contract — chat template + special tokens. -Reference: ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py`` -lines 562-606 (``Cosmos3OmniDiffusersPipeline._tokenize_prompt``). +The native pipeline tokenizes via the module-level helpers +``cosmos3_special_tokens`` / ``cosmos3_tokenize_caption`` +(``fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline``), which wrap a Qwen2 +chat-template tokenizer: -The Cosmos3 tokenizer is a Qwen2 chat-template tokenizer with two -appended special tokens: + * special tokens: ``start_of_generation=<|vision_start|>``, + ``end_of_generation=<|vision_end|>``, ``eos_token_id=tokenizer.eos_token_id`` + (the framework ``llm_special_tokens``); + * ``tokenize_caption`` applies the chat template with + ``add_generation_prompt=True`` / ``add_vision_id=False`` and an optional + image/video system prompt. - * ``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. +These contract checks run against the conftest Qwen2-shaped stub tokenizer (no +real weights). The byte-for-byte real-token-id check (``eos == 151645``, +``<|vision_start|> == 151652``) needs the real ``nvidia/Cosmos3-Nano`` +``text_tokenizer`` and is skipped cleanly when it is unavailable. """ from __future__ import annotations import pytest +from .conftest import StubQwen2Tokenizer + pytestmark = [pytest.mark.local] -def test_cosmos3_special_token_ids() -> None: - """Asserts the two Cosmos3 special-token IDs the pipeline depends on. +def test_native_special_tokens_resolution() -> None: + """``cosmos3_special_tokens`` resolves the three generation special tokens.""" + from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import cosmos3_special_tokens - 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. + special = cosmos3_special_tokens(StubQwen2Tokenizer()) + assert set(special) == {"start_of_generation", "end_of_generation", "eos_token_id"} + assert special["start_of_generation"] == StubQwen2Tokenizer().convert_tokens_to_ids("<|vision_start|>") + assert special["end_of_generation"] == StubQwen2Tokenizer().convert_tokens_to_ids("<|vision_end|>") + assert special["eos_token_id"] == StubQwen2Tokenizer.eos_token_id + + +def test_native_tokenize_caption_uses_chat_template() -> None: + """``cosmos3_tokenize_caption`` returns a non-empty token-id list. + + The video / image system-prompt variants tokenize independently (the chat + template prepends a role=system turn when ``use_system_prompt`` is set), and + the result is always a plain list of ints. + """ + from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import cosmos3_tokenize_caption + + tok = StubQwen2Tokenizer() + ids_video = cosmos3_tokenize_caption(tok, "a robot dances", is_video=True, use_system_prompt=False) + ids_image = cosmos3_tokenize_caption(tok, "a robot", is_video=False, use_system_prompt=True) + assert isinstance(ids_video, list) and all(isinstance(i, int) for i in ids_video) and ids_video + assert isinstance(ids_image, list) and ids_image + + +def test_cosmos3_special_token_ids_real_weights() -> None: + """Byte-for-byte Qwen2 special-token ids (needs the real text_tokenizer). + + Asserts ``eos_token_id == 151645`` and + ``convert_tokens_to_ids('<|vision_start|>') == 151652`` on the real + ``nvidia/Cosmos3-Nano`` Qwen2 tokenizer. Skipped cleanly when the real + tokenizer is not loadable in this environment. """ try: - from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( # type: ignore - Cosmos3OmniDiffusersPipeline, - ) + from transformers import AutoTokenizer except ImportError: - pytest.skip("FastVideo Cosmos3 pipeline not yet implemented (Phase 2b)") + pytest.skip("transformers not available") - pipeline = Cosmos3OmniDiffusersPipeline.__new__(Cosmos3OmniDiffusersPipeline) - tokenizer = getattr(pipeline, "tokenizer", None) - if tokenizer is None: - pytest.skip("Cosmos3 tokenizer instance not yet wired on the pipeline") + import os + candidate_paths = [ + os.path.join(os.environ.get("COSMOS3_WEIGHTS_DIR", ""), "text_tokenizer"), + "official_weights/cosmos3/text_tokenizer", + ] + tok_path = next((p for p in candidate_paths if p and os.path.isdir(p)), None) + if tok_path is None: + pytest.skip("real nvidia/Cosmos3-Nano text_tokenizer not available " + "(set COSMOS3_WEIGHTS_DIR or provide official_weights/cosmos3)") + + tokenizer = AutoTokenizer.from_pretrained(tok_path) 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