nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[fix]: match official Cosmos3 video sampling defaults

Browse Source
This commit is contained in:
SolitaryThinker
2026-07-21 15:45:38 -07:00
parent cd52f90b34
commit c1992baa7a
6 changed files with 61 additions and 121 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py
+1 -1
View File
@@ -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
fastvideo/pipelines/stages/cosmos3_stages.py
+11 -22
View File
@@ -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
tests/local_tests/cosmos3/test_cosmos3_flow_shift_parity.py
+37 -68
View File
@@ -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 <worktree> && <fv-cosmos3 python> -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
tests/local_tests/cosmos3/test_cosmos3_pipeline_call_graph.py
+4 -8
View File
@@ -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()
tests/local_tests/cosmos3/test_cosmos3_scheduler_default_parity.py
+6 -21
View File
@@ -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
tests/local_tests/pipelines/test_cosmos3_pipeline_smoke.py
+2 -1
View File
@@ -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(