From c1992baa7ab2ca4c64bb83fff2d34882bf72357f Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Mon, 20 Jul 2026 13:41:37 -0700 Subject: [PATCH] [fix]: match official Cosmos3 video sampling defaults --- .../basic/cosmos3/cosmos3_pipeline.py | 2 +- fastvideo/pipelines/stages/cosmos3_stages.py | 33 ++---- .../cosmos3/test_cosmos3_flow_shift_parity.py | 105 ++++++------------ .../test_cosmos3_pipeline_call_graph.py | 12 +- .../test_cosmos3_scheduler_default_parity.py | 27 +---- .../pipelines/test_cosmos3_pipeline_smoke.py | 3 +- 6 files changed, 61 insertions(+), 121 deletions(-) diff --git a/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py b/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py index e7ebde51d..b28b90586 100644 --- a/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py +++ b/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py @@ -591,7 +591,7 @@ class Cosmos3OmniDiffusersPipeline(ComposedPipelineBase): _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 + _engine_init_flow_shift: float = 10.0 # Class-attribute defaults so ``__new__``-based unit tests can read these # before ``initialize_pipeline`` runs. scheduler: Any = None diff --git a/fastvideo/pipelines/stages/cosmos3_stages.py b/fastvideo/pipelines/stages/cosmos3_stages.py index 1630eea11..54428cbd9 100644 --- a/fastvideo/pipelines/stages/cosmos3_stages.py +++ b/fastvideo/pipelines/stages/cosmos3_stages.py @@ -162,25 +162,15 @@ class Cosmos3DenoisingStage(PipelineStage): return (int(num_frames) - 1) // int(temporal_factor) + 1 @staticmethod - def _flow_shift_for_resolution(height: int, width: int) -> float: - """UniPC ``flow_shift`` for a given pixel resolution. + def _flow_shift_for_mode(*, is_video: bool) -> float: + """Return the official per-modality UniPC shift. - Mirrors the framework's ``_RESOLUTION_SHIFT_DEFAULTS`` (8B VLM backbone, - which Cosmos3-Nano uses): the shift is keyed by the named resolution - bucket the (H, W) belongs to, regardless of task (T2V/I2V/T2I): - - "256" -> 3.0, "480" -> 5.0, "704"/"720"/"768" -> 10.0 - - We invert the framework's ``{IMAGE,VIDEO}_RES_SIZE_INFO`` tables by the - longest side: <=320 is the 256 bucket, 640-832 the 480 bucket, and - 960-1360 the 704/720/768 buckets. + The current ``text2video`` and ``image2video`` defaults explicitly set + shift 10.0, while ``text2image`` explicitly sets shift 3.0. Because the + value is already configured, the framework's resolution-based fallback + is not consulted for these modes. """ - long_side = max(int(height), int(width)) - if long_side <= 480: - return 3.0 - if long_side <= 896: - return 5.0 - return 10.0 + return 10.0 if is_video else 3.0 def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch: pipeline_config = fastvideo_args.pipeline_config @@ -198,12 +188,11 @@ class Cosmos3DenoisingStage(PipelineStage): is_i2v = (batch.preprocessed_image is not None or batch.pil_image is not None) and not is_t2i is_video = not is_t2i - # Resolution-based flow_shift, set on the owning pipeline (rebuilds the - # scheduler). The framework picks the UniPC shift purely from the named - # resolution bucket (``_RESOLUTION_SHIFT_DEFAULTS``), NOT from the task, - # so T2V/I2V/T2I at the same resolution share a shift. + # Per-modality flow_shift, set on the owning pipeline (rebuilds the + # scheduler). The official defaults explicitly use 10.0 for T2V/I2V + # and 3.0 for T2I. pipe = self.pipeline() if self.pipeline is not None else None - flow_shift = self._flow_shift_for_resolution(height, width) + flow_shift = self._flow_shift_for_mode(is_video=is_video) if pipe is not None and hasattr(pipe, "_set_flow_shift"): pipe._set_flow_shift(flow_shift) scheduler = pipe.scheduler diff --git a/tests/local_tests/cosmos3/test_cosmos3_flow_shift_parity.py b/tests/local_tests/cosmos3/test_cosmos3_flow_shift_parity.py index fa78c75d2..90d9609fe 100644 --- a/tests/local_tests/cosmos3/test_cosmos3_flow_shift_parity.py +++ b/tests/local_tests/cosmos3/test_cosmos3_flow_shift_parity.py @@ -1,80 +1,49 @@ # SPDX-License-Identifier: Apache-2.0 -"""Numerical-parity test: FastVideo Cosmos3 UniPC flow_shift vs the framework. - -The framework selects the UniPC ``shift`` purely from the named resolution -bucket the (H, W) belongs to, via ``OmniSampleArgs._RESOLUTION_SHIFT_DEFAULTS`` -(keyed by the VLM model size — Cosmos3-Nano uses the 8B backbone — and the -resolution string), NOT from the task (T2V/I2V/T2I share a shift at a given -resolution). FastVideo gets raw pixel ``height``/``width`` and must map back to -the same shift. - -This pins ``Cosmos3DenoisingStage._flow_shift_for_resolution`` against the -framework's own tables: for every (resolution, aspect) entry in -``VIDEO_RES_SIZE_INFO`` whose resolution has an 8B shift default, the FastVideo -shift for that exact pixel size must equal the framework default. - -The framework tables are the parity ORACLE. - -Run: - cd && -m pytest \ - tests/local_tests/cosmos3/test_cosmos3_flow_shift_parity.py -q -""" +"""Cosmos3 per-modality UniPC shift parity with the pinned framework.""" from __future__ import annotations +import json +import os +import subprocess +from pathlib import Path + import pytest -# The official framework provides the parity oracle for the resolution->pixel -# tables. (``cosmos_framework.inference.args`` — which holds the shift constant — -# can't be imported here: it transitively requires ``multistorageclient``. The -# small shift table is mirrored verbatim below with its source location.) -_utils = pytest.importorskip( - "cosmos_framework.data.generator.utils", - reason="cosmos_framework not installed; run in fv-cosmos3 env.", -) - -from fastvideo.pipelines.stages.cosmos3_stages import ( # noqa: E402 - Cosmos3DenoisingStage, -) +from fastvideo.pipelines.stages.cosmos3_stages import Cosmos3DenoisingStage pytestmark = [pytest.mark.local] -# Cosmos3-Nano's VLM backbone is Qwen3-VL-8B (checkpoint config.json). -_MODEL_SIZE = "8B" -# Verbatim from cosmos_framework.inference.args.OmniSampleArgs -# ._RESOLUTION_SHIFT_DEFAULTS (args.py:770), restricted to the 8B rows. -_SHIFT_DEFAULTS = { - ("8B", "256"): 3.0, - ("8B", "480"): 5.0, - ("8B", "720"): 10.0, - ("8B", "768"): 10.0, - ("32B", "256"): 5.0, - ("32B", "480"): 5.0, - ("32B", "720"): 5.0, - ("32B", "768"): 5.0, -} -_VIDEO_RES = _utils.VIDEO_RES_SIZE_INFO -_IMAGE_RES = _utils.IMAGE_RES_SIZE_INFO +REPO_ROOT = Path(__file__).resolve().parents[3] +REFERENCE_REVISION = "ed8287fd7477113f8ac4f6b84290514d55cf0cdc" +REFERENCE_REPO = Path(os.getenv("COSMOS3_OFFICIAL_REPO", REPO_ROOT / "cosmos-framework")) -def _cases(): - seen = set() - for resolution, by_aspect in {**_VIDEO_RES, **_IMAGE_RES}.items(): - key = (_MODEL_SIZE, resolution) - if key not in _SHIFT_DEFAULTS: - continue - expected = _SHIFT_DEFAULTS[key] - for aspect, (a, b) in by_aspect.items(): - cid = f"{resolution}_{aspect.replace(',', '-')}_{a}x{b}" - if cid in seen: - continue - seen.add(cid) - yield pytest.param(a, b, expected, id=cid) +def _official_shift(mode: str) -> float: + if not REFERENCE_REPO.is_dir(): + pytest.skip(f"Pinned NVIDIA cosmos-framework clone not found: {REFERENCE_REPO}") + result = subprocess.run( + ["git", "-C", str(REFERENCE_REPO), "rev-parse", "HEAD"], + check=True, + capture_output=True, + text=True, + ) + actual = result.stdout.strip() + assert actual == REFERENCE_REVISION, f"Cosmos3 oracle must be pinned at {REFERENCE_REVISION}, got {actual}" + + defaults_path = REFERENCE_REPO / "cosmos_framework" / "inference" / "defaults" / mode / "sample_args.json" + defaults = json.loads(defaults_path.read_text(encoding="utf-8")) + return float(defaults["shift"]) -class TestCosmos3FlowShiftParity: - - @pytest.mark.parametrize(("dim_a", "dim_b", "expected_shift"), list(_cases())) - def test_flow_shift_matches_framework(self, dim_a, dim_b, expected_shift): - got = Cosmos3DenoisingStage._flow_shift_for_resolution(dim_a, dim_b) - assert got == expected_shift, ( - f"shift for {dim_a}x{dim_b}: got {got}, framework default {expected_shift}") +@pytest.mark.parametrize( + ("mode", "is_video"), + [ + pytest.param("text2video", True, id="t2v"), + pytest.param("image2video", True, id="i2v"), + pytest.param("text2image", False, id="t2i"), + ], +) +def test_flow_shift_matches_official_modality_default(mode: str, is_video: bool) -> None: + expected = _official_shift(mode) + actual = Cosmos3DenoisingStage._flow_shift_for_mode(is_video=is_video) + assert actual == expected 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 433f0476b..4a8990c97 100644 --- a/tests/local_tests/cosmos3/test_cosmos3_pipeline_call_graph.py +++ b/tests/local_tests/cosmos3/test_cosmos3_pipeline_call_graph.py @@ -158,30 +158,27 @@ def test_i2v_keeps_condition_frame_clean(make_cosmos3_pipeline) -> None: def test_stage_mode_dispatch_t2v(make_cosmos3_pipeline, make_cosmos3_stage) -> None: - """T2V (num_frames>1): resolution-based flow_shift, video tokenization.""" + """T2V (num_frames>1): shift 10, 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) - # flow_shift is resolution-based (not task-based): 16x16 -> "256" bucket -> 3.0. - # (full resolution->shift parity in test_cosmos3_flow_shift_parity.) - assert float(pipeline.scheduler.config.flow_shift) == 3.0 + 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): single-frame latent; resolution-based flow_shift.""" + """T2I (num_frames==1): single-frame latent and shift 3.""" 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) - # 16x16 -> "256" bucket -> 3.0 (resolution-based, not task-based). assert float(pipeline.scheduler.config.flow_shift) == 3.0 assert tuple(out.latents.shape) == (1, _LATENT_CHANNEL, 1, 2, 2) @@ -195,7 +192,6 @@ def test_stage_i2v_encodes_conditioning_image(make_cosmos3_pipeline, make_cosmos batch = make_forward_batch(num_frames=5, height=16, width=16, image=image) out = stage.forward(batch, args) - # flow_shift is resolution-based (not task-based): 16x16 -> 3.0. - assert float(pipeline.scheduler.config.flow_shift) == 3.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_scheduler_default_parity.py b/tests/local_tests/cosmos3/test_cosmos3_scheduler_default_parity.py index aaf90f6c7..3fe08f31f 100644 --- a/tests/local_tests/cosmos3/test_cosmos3_scheduler_default_parity.py +++ b/tests/local_tests/cosmos3/test_cosmos3_scheduler_default_parity.py @@ -1,19 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 """Cosmos3 scheduler default + per-request override parity (Tier A scaffold). -Reference: - * ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py:275-307`` — - initial UniPCMultistepScheduler load (preserves solver_order, - timestep_spacing, beta_schedule, sigma bounds, flow_shift) and - one-time override at engine-init if ``od_config.flow_shift`` is set. - * ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py:498-512`` — - ``_set_flow_shift(target_shift)``: rebuild the scheduler via - ``UniPCMultistepScheduler.from_config(base_config, flow_shift=target)`` - when the requested target differs from the current shift. - * ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py:1069-1110`` — - per-request mode defaults: T2I uses ``shift=3.0``; T2V/I2V use the - engine-init shift (typically 1.0); ``flow_shift`` may be overridden - per request via ``sampling_params.extra_args["flow_shift"]``. +The pinned ``cosmos-framework`` modality defaults use shift 10.0 for T2V/I2V +and shift 3.0 for T2I. ``_set_flow_shift`` rebuilds the FastVideo scheduler +when a request changes modes. The invariant under test: for the same RNG seed and the same number of inference steps, the scheduler's ``timesteps`` tensor must be identical @@ -52,14 +42,8 @@ def test_t2i_default_flow_shift_is_3() -> None: assert float(pipeline.scheduler.config.flow_shift) == 3.0 -def test_t2v_default_flow_shift_is_engine_init() -> None: - """Asserts T2V/I2V use the engine-init shift (e.g. 1.0), NOT a fixed default. - - Cross-check: pipeline_cosmos3.py:1091 sets - ``default_flow_shift = self._engine_init_flow_shift`` for T2V/I2V - (NOT ``None`` — passing ``None`` would leak a prior T2I rebuild - forward). - """ +def test_video_engine_init_flow_shift_is_10() -> None: + """T2V/I2V start from the official video shift of 10.0.""" try: from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( # type: ignore Cosmos3OmniDiffusersPipeline, @@ -72,6 +56,7 @@ def test_t2v_default_flow_shift_is_engine_init() -> None: pytest.skip("FastVideo Cosmos3 _engine_init_flow_shift not yet wired") init_shift = float(pipeline._engine_init_flow_shift) + assert init_shift == 10.0 pipeline._set_flow_shift(init_shift) assert float(pipeline.scheduler.config.flow_shift) == init_shift diff --git a/tests/local_tests/pipelines/test_cosmos3_pipeline_smoke.py b/tests/local_tests/pipelines/test_cosmos3_pipeline_smoke.py index 6cd55b7bb..3f361f1c3 100644 --- a/tests/local_tests/pipelines/test_cosmos3_pipeline_smoke.py +++ b/tests/local_tests/pipelines/test_cosmos3_pipeline_smoke.py @@ -45,7 +45,8 @@ def test_cosmos3_production_loader_video_latent(tmp_path: Path, workload: str) - output_type="latent", ) try: - height, width, num_frames = 256, 448, 25 + # Official ``resolution=256, aspect_ratio=16,9`` geometry. + height, width, num_frames = 192, 320, 25 initial_noise = torch.zeros(1, 48, 7, height // 16, width // 16) conditioning = Image.new("RGB", (width, height), color=(64, 128, 192)) if workload == "i2v" else None result = generator.generate_video(