[feat] cosmos3: native video pipeline + framework denoise parity
This commit is contained in:
SolitaryThinker
@@ -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
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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)
|
||||
@@ -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
|
||||
@@ -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,
|
||||
|
||||
@@ -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__(<FastVideoCosmos3Pipeline>)`` + ``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)
|
||||
|
||||
@@ -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 <worktree> && <fv-cosmos3 python> -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)
|
||||
@@ -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()
|
||||
|
||||
@@ -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 <worktree> && <fv-cosmos3 python> -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}
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user