From 455cd83bed8adf5cd62a89130cb07a64ca4a2b2c Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sun, 7 Jun 2026 01:31:19 -0700 Subject: [PATCH] [bugfix] cosmos3: fix black-video UniPC config; wire E2E inference The first real-weights Cosmos3-Nano T2V run produced an all-black video. Root cause: the checkpoint's diffusers-style scheduler_config.json sets use_karras_sigmas=true, and FastVideo's vendored UniPC checks karras *before* use_flow_sigmas, so it built diffusion (beta) sigmas instead of flow-matching sigmas -> scheduler.step diverged to NaN latents -> black frames. The DiT/CFG velocity itself was clean. The framework samples with FlowUniPCMultistepScheduler (pure flow: shift + num_train_timesteps only). FastVideo's vendored UniPC is the same algorithm, so the fix is config, not a port: - Coerce the loaded scheduler to the flow setup in Cosmos3OmniDiffusersPipeline.initialize_pipeline (kill karras/exp/beta sigmas, force use_flow_sigmas + flow_prediction + final_sigmas_type=zero) and rebuild the runtime scheduler from it. - Use FastVideo's native UniPC (not diffusers) in the pipeline and tests, honoring the "no diffusers at runtime" constraint. Parity: - Add test_cosmos3_scheduler_parity.py: FastVideo vendored UniPC (flow config) vs the framework FlowUniPCMultistepScheduler. timesteps bit-exact, sigmas within fp32 epsilon (~1e-8), full multi-step trajectory < ~1e-6 across shift in {10 (video), 3 (t2i)} and steps in {4,10,35}. - Fix test_cosmos3_denoise_cfg_parity to use the framework scheduler as the oracle (it previously compared diffusers-vs-diffusers, so the scheduler was never actually checked against the framework). Integration fixes to run end-to-end on real weights: - registry: map checkpoint DiT class Cosmos3OmniTransformer -> Cosmos3VFMTransformer. - component_loader: route the "text_tokenizer" module to TokenizerLoader; filter scheduler config to the class's __init__ params (robust to diffusers schema drift like shift_terminal/sigma_min/sigma_max). - DiT: materialize_non_persistent_buffers (recompute rotary_emb.inv_freq after meta-device FSDP load) + cast vision latents / timestep embeds to the compute dtype (no-ops in the fp32 parity tests). - sequence_packing: move packed tensors to the model device in to_dit_kwargs. - config: drop the unused identity text-preprocess (Cosmos3 has no text encoder; the DiT tokenizes in the denoise stage). Adds examples/inference/basic/basic_cosmos3_new_api.py (new VideoGenerator API). Verified: 1280x704, 29 frames, 35 steps on B200 -> coherent video matching the prompt (no NaNs, real pixel variance). Full cosmos3 parity suite green. --- .../inference/basic/basic_cosmos3_new_api.py | 77 +++++++++++ fastvideo/configs/pipelines/cosmos3.py | 19 ++- fastvideo/models/dits/cosmos3.py | 31 ++++- fastvideo/models/loader/component_loader.py | 16 ++- fastvideo/models/registry.py | 3 + .../basic/cosmos3/cosmos3_pipeline.py | 48 ++++++- .../basic/cosmos3/sequence_packing.py | 29 ++-- tests/local_tests/cosmos3/conftest.py | 4 +- .../test_cosmos3_denoise_cfg_parity.py | 31 ++--- .../cosmos3/test_cosmos3_pipeline_smoke.py | 5 +- .../cosmos3/test_cosmos3_scheduler_parity.py | 124 ++++++++++++++++++ 11 files changed, 334 insertions(+), 53 deletions(-) create mode 100644 examples/inference/basic/basic_cosmos3_new_api.py create mode 100644 tests/local_tests/cosmos3/test_cosmos3_scheduler_parity.py diff --git a/examples/inference/basic/basic_cosmos3_new_api.py b/examples/inference/basic/basic_cosmos3_new_api.py new file mode 100644 index 000000000..08c1e02c3 --- /dev/null +++ b/examples/inference/basic/basic_cosmos3_new_api.py @@ -0,0 +1,77 @@ +import os +import time + +from fastvideo import VideoGenerator +from fastvideo.api import ( + EngineConfig, + GenerationRequest, + GeneratorConfig, + OffloadConfig, + OutputConfig, + SamplingConfig, +) + +# NVIDIA Cosmos3-Nano omni world model — this example exercises the +# text-to-video (T2V) path through FastVideo's native Cosmos3 pipeline. +# Point COSMOS3_MODEL_PATH at a local diffusers checkpoint (e.g. +# ``official_weights/cosmos3``) to skip the Hugging Face download. +OUTPUT_PATH = "video_samples_cosmos3" + + +def main(): + model_name = os.environ.get("COSMOS3_MODEL_PATH", "nvidia/Cosmos3-Nano") + + generator_config = GeneratorConfig( + model_path=model_name, + engine=EngineConfig( + num_gpus=1, + use_fsdp_inference=False, + offload=OffloadConfig( + text_encoder=True, + pin_cpu_memory=True, + dit=False, + vae=False, + ), + ), + ) + + load_start_time = time.perf_counter() + generator = VideoGenerator.from_config(generator_config) + load_time = time.perf_counter() - load_start_time + + prompt = ( + "A golden retriever puppy runs across a sunlit meadow toward the camera, " + "ears flopping and wildflowers swaying in the breeze. Shallow depth of " + "field, warm afternoon light, smooth cinematic tracking shot." + ) + request = GenerationRequest( + prompt=prompt, + sampling=SamplingConfig( + # cosmos3_nano native defaults are 704x1280, 189 frames, 35 steps; + # overridable via env for quick smoke runs. + num_frames=int(os.environ.get("COSMOS3_NUM_FRAMES", "189")), + height=int(os.environ.get("COSMOS3_HEIGHT", "704")), + width=int(os.environ.get("COSMOS3_WIDTH", "1280")), + num_inference_steps=int(os.environ.get("COSMOS3_STEPS", "35")), + guidance_scale=6.0, + fps=24, + seed=1024, + ), + output=OutputConfig( + output_path=OUTPUT_PATH, + save_video=True, + return_frames=False, + ), + ) + + start_time = time.perf_counter() + result = generator.generate(request) + gen_time = time.perf_counter() - start_time + + print(f"Time taken to load model: {load_time} seconds") + print(f"Time taken to generate video: {gen_time} seconds") + print(f"Output written to: {result.video_path}") + + +if __name__ == "__main__": + main() diff --git a/fastvideo/configs/pipelines/cosmos3.py b/fastvideo/configs/pipelines/cosmos3.py index f81c849bf..427ac9661 100644 --- a/fastvideo/configs/pipelines/cosmos3.py +++ b/fastvideo/configs/pipelines/cosmos3.py @@ -12,8 +12,11 @@ Cosmos 2.5: 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. +- Scheduler is FastVideo-native ``UniPCMultistepScheduler`` configured for + pure flow matching (flow_prediction, use_flow_sigmas), equivalent to the + framework's ``FlowUniPCMultistepScheduler``. The checkpoint's diffusers-style + scheduler config (karras/sigma_min/max) is coerced to the flow setup in + ``Cosmos3OmniDiffusersPipeline.initialize_pipeline``. - 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. """ @@ -29,10 +32,6 @@ from fastvideo.configs.models.vaes import Cosmos3VAEConfig # Wan2.2 Autoencoder from fastvideo.configs.pipelines.base import PipelineConfig -def _identity_preprocess_text(prompt: str) -> str: - return prompt - - @dataclass class Cosmos3Config(PipelineConfig): """Configuration for the Cosmos3 video generation pipeline (T2V/I2V/T2I). @@ -46,11 +45,11 @@ class Cosmos3Config(PipelineConfig): vae_config: VAEConfig = field(default_factory=Cosmos3VAEConfig) - # No separate text encoder: Cosmos3LanguageModel lives inside the DiT. + # No separate text encoder: the Qwen3-VL-text backbone lives inside the DiT + # and the pipeline tokenizes in Cosmos3DenoisingStage, so all three + # text-encoder lists are empty (the generic text-encode stage is not used). text_encoder_configs: tuple[EncoderConfig, ...] = field(default_factory=tuple) - - preprocess_text_funcs: tuple[Callable[[str], str], - ...] = field(default_factory=lambda: (_identity_preprocess_text, )) + preprocess_text_funcs: tuple[Callable[[str], str], ...] = field(default_factory=tuple) postprocess_text_funcs: tuple[Callable[[BaseEncoderOutput], torch.Tensor], ...] = field(default_factory=tuple) dit_precision: str = "bf16" diff --git a/fastvideo/models/dits/cosmos3.py b/fastvideo/models/dits/cosmos3.py index c4b5c9ac3..074f646de 100644 --- a/fastvideo/models/dits/cosmos3.py +++ b/fastvideo/models/dits/cosmos3.py @@ -161,11 +161,21 @@ class Cosmos3TextRotaryEmbedding(nn.Module): ) -> None: super().__init__() self.head_dim = head_dim + self.rope_theta = float(rope_theta) self.mrope_section = list(mrope_section) - inv_freq = 1.0 / (rope_theta**(torch.arange(0, head_dim, 2, dtype=torch.int64).float() / head_dim)) - self.register_buffer("inv_freq", inv_freq, persistent=False) + self.register_buffer("inv_freq", self._compute_inv_freq(), persistent=False) self.attention_scaling = 1.0 # default RoPE has unit attention scaling + def _compute_inv_freq(self, device: torch.device | None = None) -> torch.Tensor: + exponent = torch.arange(0, self.head_dim, 2, dtype=torch.int64, device=device).float() / self.head_dim + return 1.0 / (self.rope_theta**exponent) + + def reset_inv_freq(self, device: torch.device) -> None: + """Recompute the non-persistent ``inv_freq`` on ``device`` after a + meta-device weight load (it is derived from ``rope_theta`` and is not + part of the checkpoint).""" + self.register_buffer("inv_freq", self._compute_inv_freq(device), persistent=False) + def _apply_interleaved_mrope(self, freqs: torch.Tensor) -> torch.Tensor: """freqs: [3, N, head_dim//2] -> [N, head_dim//2] (interleaved T/H/W).""" freqs_t = freqs[0].clone() @@ -323,6 +333,9 @@ class Cosmos3TimestepEmbedder(nn.Module): def forward(self, t: torch.Tensor) -> torch.Tensor: t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + # Sinusoidal runs in fp32 (timestep_scale makes inputs tiny); cast to the + # MLP weight dtype (bf16 at inference; no-op in the fp32 parity tests). + t_freq = t_freq.to(self.linear_1.weight.dtype) h, _ = self.linear_1(t_freq) h = self.act(h) out, _ = self.linear_2(h) @@ -770,6 +783,16 @@ class Cosmos3VFMTransformer(BaseDiT): flat = flat.unsqueeze(-1).expand(-1, packed_tokens.shape[1]) return packed_tokens.scatter_add(0, flat, packed_timestep_embeds) + def materialize_non_persistent_buffers(self, device: torch.device, dtype: torch.dtype | None = None) -> None: + """Recompute non-persistent buffers lost by the meta-device FSDP load. + + Only ``rotary_emb.inv_freq`` is non-persistent (derived from + ``rope_theta``, absent from the checkpoint). The FastVideo loader calls + this after ``load_model_from_full_model_state_dict``. + """ + del dtype # inv_freq stays float32 regardless of compute dtype + self.rotary_emb.reset_inv_freq(device) + # ------------------------------------------------------------------ # forward # ------------------------------------------------------------------ @@ -809,7 +832,9 @@ class Cosmos3VFMTransformer(BaseDiT): original_shapes: list[tuple[int, int, int]] | None = None if vision_tokens: packed_vision, original_shapes = self._patchify_and_pack(vision_tokens, vision_token_shapes) - packed_vision, _ = self.proj_in(packed_vision) + # Vision latents arrive in fp32 (noise/VAE); cast to the model's + # compute dtype (bf16 at inference; no-op in the fp32 parity tests). + packed_vision, _ = self.proj_in(packed_vision.to(target_dtype)) if self.latent_pos_embed is not None: pos_emb = self.latent_pos_embed(vision_token_shapes, fps=fps_vision).to(target_dtype) diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 06adb841f..cd567d395 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -92,6 +92,8 @@ class ComponentLoader(ABC): "tokenizer": (TokenizerLoader, "transformers"), "tokenizer_2": (TokenizerLoader, "transformers"), "tokenizer_3": (TokenizerLoader, "transformers"), + # Cosmos3's model_index names its Qwen2 tokenizer "text_tokenizer". + "text_tokenizer": (TokenizerLoader, "transformers"), "image_processor": (ImageProcessorLoader, "transformers"), "feature_extractor": (ImageProcessorLoader, "transformers"), "image_encoder": (ImageEncoderLoader, "transformers"), @@ -1159,7 +1161,19 @@ class SchedulerLoader(ComponentLoader): scheduler_cls, _ = ModelRegistry.resolve_model_cls(class_name) - scheduler = scheduler_cls(**config) + # Diffusers checkpoints can carry newer scheduler config keys than the + # vendored scheduler accepts (e.g. shift_terminal / sigma_min / sigma_max + # from a newer diffusers release). Filter to the class's __init__ params, + # mirroring diffusers' ``from_config``, so loading is robust to schema + # drift instead of crashing on an unexpected kwarg. + import inspect + valid_params = set(inspect.signature(scheduler_cls.__init__).parameters) + filtered_config = {k: v for k, v in config.items() if k in valid_params} + dropped = sorted(set(config) - set(filtered_config)) + if dropped: + logger.warning("Scheduler %s: dropping unsupported config keys %s", class_name, dropped) + + scheduler = scheduler_cls(**filtered_config) if fastvideo_args.pipeline_config.flow_shift is not None: scheduler.set_shift(fastvideo_args.pipeline_config.flow_shift) return scheduler diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index eb3308c18..30743826c 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -37,6 +37,9 @@ _TEXT_TO_VIDEO_DIT_MODELS = { "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"), "CosmosTransformer3DModel": ("dits", "cosmos", "CosmosTransformer3DModel"), "Cosmos25Transformer3DModel": ("dits", "cosmos2_5", "Cosmos25Transformer3DModel"), + # Cosmos3-Nano's checkpoint model_index names the DiT "Cosmos3OmniTransformer"; + # map that HF class name to FastVideo's native Cosmos3VFMTransformer. + "Cosmos3OmniTransformer": ("dits", "cosmos3", "Cosmos3VFMTransformer"), "LongCatVideoTransformer3DModel": ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"), # Wrapper (Phase 1) "LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"), # Native (Phase 2) "LTX2Transformer3DModel": ("dits", "ltx2", "LTX2Transformer3DModel"), diff --git a/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py b/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py index 6dc4ab699..9d989734c 100644 --- a/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py +++ b/fastvideo/pipelines/basic/cosmos3/cosmos3_pipeline.py @@ -11,8 +11,10 @@ Cosmos3 components: encode normalizes ``(mu - mean) * inv_std`` and decode denormalizes + clamps, * sequence-packing: :func:`pack_cosmos3_video_sequence` (native, parity-tested), * DiT: ``Cosmos3VFMTransformer`` (native, bit-identical to the framework), -* scheduler: diffusers ``UniPCMultistepScheduler`` (the actual scheduler the - checkpoint ships, ``flow_prediction`` + ``use_flow_sigmas``). +* scheduler: FastVideo-native ``UniPCMultistepScheduler`` configured for pure + flow matching (``flow_prediction`` + ``use_flow_sigmas``), numerically + equivalent to the framework's ``FlowUniPCMultistepScheduler`` (parity-tested + in ``test_cosmos3_scheduler_parity``). The denoise/CFG glue is a faithful port of the framework's ``Cosmos3OmniDiffusersPipeline`` math (mirrored in the framework-equivalent @@ -40,10 +42,11 @@ from dataclasses import dataclass from typing import Any import torch -from diffusers import UniPCMultistepScheduler from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.logger import init_logger +from fastvideo.models.schedulers.scheduling_unipc_multistep import ( + UniPCMultistepScheduler, ) from fastvideo.pipelines.basic.cosmos3.sequence_packing import ( Cosmos3SampleInputs, Cosmos3VisionItem, @@ -265,7 +268,8 @@ def cosmos3_get_cfg_velocity( base_fps=base_fps, temporal_compression_factor=temporal_compression_factor, ) - out = transformer(**packed.to_dit_kwargs()) + device = next(transformer.parameters()).device + out = transformer(**packed.to_dit_kwargs(device=device)) preds = out.get("preds_vision") if preds is None: # Fully-conditioned clip: no noisy frames, zero velocity. @@ -424,6 +428,31 @@ class Cosmos3OmniDiffusersPipeline(ComposedPipelineBase): _base_scheduler_config: Any = None _current_flow_shift: float | None = None + @staticmethod + def _flow_scheduler_config(config: Any) -> dict[str, Any]: + """Coerce a loaded UniPC config to the framework's flow-matching setup. + + The checkpoint ``scheduler_config.json`` carries diffusers-style fields + (``use_karras_sigmas=True``, ``sigma_min``/``sigma_max``, beta schedule) + that do not describe the framework sampler. The framework uses + ``FlowUniPCMultistepScheduler`` (pure flow matching: ``shift`` + + ``num_train_timesteps`` only). FastVideo's vendored UniPC checks + ``use_karras_sigmas`` *before* ``use_flow_sigmas``, so leaving karras on + builds diffusion-style sigmas and the denoise diverges to NaN. Force the + flow config here (parity-verified in ``test_cosmos3_scheduler_parity``). + """ + cfg = dict(config) + cfg.update( + use_karras_sigmas=False, + use_exponential_sigmas=False, + use_beta_sigmas=False, + use_flow_sigmas=True, + prediction_type="flow_prediction", + predict_x0=True, + final_sigmas_type="zero", + ) + return cfg + def initialize_pipeline(self, fastvideo_args: FastVideoArgs) -> None: """Bind the loaded scheduler + snapshot its config so per-request flow_shift rebuilds are cheap and the engine-init shift is applied.""" @@ -433,9 +462,14 @@ class Cosmos3OmniDiffusersPipeline(ComposedPipelineBase): self._engine_init_flow_shift = float(engine_shift) scheduler = self.get_module("scheduler") if scheduler is not None: - self.scheduler = scheduler - self._base_scheduler_config = scheduler.config - self._current_flow_shift = float(getattr(scheduler.config, "flow_shift", 1.0)) + # Rebuild from a flow-coerced config so the runtime scheduler matches + # the framework sampler (the loaded checkpoint config is diffusers-style). + flow_config = self._flow_scheduler_config(scheduler.config) + self.scheduler = UniPCMultistepScheduler.from_config(flow_config) + if isinstance(self.modules, dict): + self.modules["scheduler"] = self.scheduler + self._base_scheduler_config = self.scheduler.config + self._current_flow_shift = float(getattr(self.scheduler.config, "flow_shift", 1.0)) def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: """Wire the Cosmos3 stages. diff --git a/fastvideo/pipelines/basic/cosmos3/sequence_packing.py b/fastvideo/pipelines/basic/cosmos3/sequence_packing.py index 16cd730f9..9be22f7c8 100644 --- a/fastvideo/pipelines/basic/cosmos3/sequence_packing.py +++ b/fastvideo/pipelines/basic/cosmos3/sequence_packing.py @@ -111,21 +111,30 @@ class Cosmos3PackedSequence: vision_condition_mask: list[torch.Tensor] fps_vision: torch.Tensor | None = None - def to_dit_kwargs(self) -> dict[str, Any]: - """Return the kwargs dict for ``Cosmos3VFMTransformer.forward``.""" + def to_dit_kwargs(self, device: torch.device | str | None = None) -> dict[str, Any]: + """Return the kwargs dict for ``Cosmos3VFMTransformer.forward``. + + Packing is device-agnostic (ids/indexes/position-ids are built on CPU). + When ``device`` is given, every tensor input is moved to it so the DiT + forward runs on a single device (e.g. the model's GPU at inference). + """ + + def _mv(x: Any) -> Any: + return x.to(device) if (device is not None and torch.is_tensor(x)) else x + return dict( - text_ids=self.text_ids, - text_indexes=self.text_indexes, - position_ids=self.position_ids, + text_ids=_mv(self.text_ids), + text_indexes=_mv(self.text_indexes), + position_ids=_mv(self.position_ids), sequence_length=int(self.sequence_length), split_lens=list(self.split_lens), attn_modes=list(self.attn_modes), - vision_tokens=list(self.vision_tokens), + vision_tokens=[_mv(t) for t in self.vision_tokens], vision_token_shapes=list(self.vision_token_shapes), - vision_sequence_indexes=self.vision_sequence_indexes, - vision_timesteps=self.vision_timesteps, - vision_mse_loss_indexes=self.vision_mse_loss_indexes, - vision_noisy_frame_indexes=list(self.vision_noisy_frame_indexes), + vision_sequence_indexes=_mv(self.vision_sequence_indexes), + vision_timesteps=_mv(self.vision_timesteps), + vision_mse_loss_indexes=_mv(self.vision_mse_loss_indexes), + vision_noisy_frame_indexes=[_mv(t) for t in self.vision_noisy_frame_indexes], fps_vision=self.fps_vision, ) diff --git a/tests/local_tests/cosmos3/conftest.py b/tests/local_tests/cosmos3/conftest.py index d21ff1c0a..ee2da0765 100644 --- a/tests/local_tests/cosmos3/conftest.py +++ b/tests/local_tests/cosmos3/conftest.py @@ -18,7 +18,9 @@ from typing import Any import pytest import torch -from diffusers import UniPCMultistepScheduler +from fastvideo.models.schedulers.scheduling_unipc_multistep import ( + UniPCMultistepScheduler, +) from torch import nn _LATENT_CHANNEL = 16 diff --git a/tests/local_tests/cosmos3/test_cosmos3_denoise_cfg_parity.py b/tests/local_tests/cosmos3/test_cosmos3_denoise_cfg_parity.py index cc080f172..d7f311d16 100644 --- a/tests/local_tests/cosmos3/test_cosmos3_denoise_cfg_parity.py +++ b/tests/local_tests/cosmos3/test_cosmos3_denoise_cfg_parity.py @@ -39,8 +39,6 @@ cosmos_framework = pytest.importorskip( 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, @@ -52,6 +50,10 @@ from .test_cosmos3_dit_parity_mrope import ( # noqa: E402 _build_tiny_fastvideo_dit_mrope, ) from .test_cosmos3_reference_forward import _apply_sdpa_patches # noqa: E402 +from .test_cosmos3_scheduler_parity import ( # noqa: E402 + _fastvideo_scheduler, + _framework_scheduler, +) pytestmark = [pytest.mark.local] @@ -64,16 +66,9 @@ _SPECIAL_TOKENS = { "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, - ) +# Cosmos3 video flow_shift; framework scheduler is the parity oracle, FastVideo's +# vendored UniPC (flow config) is the unit under test. +_FLOW_SHIFT = 10.0 # --------------------------------------------------------------------------- @@ -250,10 +245,8 @@ class TestCosmos3DenoiseCFGParity: 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")) + fw_sched = _framework_scheduler(4, _FLOW_SHIFT) + fv_sched = _fastvideo_scheduler(4, _FLOW_SHIFT) t = fw_sched.timesteps[0] fw_v = _framework_cfg_velocity( @@ -304,8 +297,7 @@ class TestCosmos3DenoiseCFGParity: num_steps = 3 # Manual framework loop (oracle). - fw_sched = _make_scheduler() - fw_sched.set_timesteps(num_steps, device=torch.device("cpu")) + fw_sched = _framework_scheduler(num_steps, _FLOW_SHIFT) fw_latent = flat_latent.clone() for t in fw_sched.timesteps: v = _framework_cfg_velocity( @@ -322,8 +314,7 @@ class TestCosmos3DenoiseCFGParity: return_dict=False)[0].squeeze(0) # FastVideo engine loop. - fv_sched = _make_scheduler() - fv_sched.set_timesteps(num_steps, device=torch.device("cpu")) + fv_sched = _fastvideo_scheduler(num_steps, _FLOW_SHIFT) engine = Cosmos3DenoiseEngine( transformer=dit, scheduler=fv_sched, diff --git a/tests/local_tests/cosmos3/test_cosmos3_pipeline_smoke.py b/tests/local_tests/cosmos3/test_cosmos3_pipeline_smoke.py index c75563107..d0b82d0a9 100644 --- a/tests/local_tests/cosmos3/test_cosmos3_pipeline_smoke.py +++ b/tests/local_tests/cosmos3/test_cosmos3_pipeline_smoke.py @@ -35,7 +35,10 @@ from typing import Any import pytest import torch -from diffusers import UniPCMultistepScheduler + +from fastvideo.models.schedulers.scheduling_unipc_multistep import ( + UniPCMultistepScheduler, +) pytestmark = [pytest.mark.local] diff --git a/tests/local_tests/cosmos3/test_cosmos3_scheduler_parity.py b/tests/local_tests/cosmos3/test_cosmos3_scheduler_parity.py new file mode 100644 index 000000000..c615aeb15 --- /dev/null +++ b/tests/local_tests/cosmos3/test_cosmos3_scheduler_parity.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Numerical-parity test: FastVideo's UniPC scheduler vs the framework's. + +The Cosmos3 video sampler is the framework's flow-matching UniPC +(``cosmos_framework.model.vfm.diffusion.samplers.fm_solvers_unipc.FlowUniPCMultistepScheduler``), +driven by ``UniPCSampler`` with config ``num_train_timesteps=1000``, +``use_dynamic_shifting=False`` and a per-mode ``shift`` (10.0 for T2V/I2V, +3.0 for T2I). FastVideo reuses its vendored +``UniPCMultistepScheduler`` configured for pure flow matching +(``use_flow_sigmas=True``, ``prediction_type="flow_prediction"``, +``predict_x0=True``, ``solver_type="bh2"``, ``solver_order=2``, +``final_sigmas_type="zero"``) with ``flow_shift`` set to the same shift. + +This pins the scheduler — the one Cosmos3 component whose earlier test +(``test_cosmos3_denoise_cfg_parity``) compared diffusers-vs-diffusers rather +than against the framework oracle — by: + + * asserting the discrete ``timesteps`` and ``sigmas`` match the framework; + * running a full multi-step UniPC trajectory with a fixed sequence of + pseudo-random "velocity" model outputs (identical on both sides, so the + DiT is factored out) and asserting every intermediate + final latent + matches the framework bit-for-bit. + +CPU / float32. The framework scheduler is the parity ORACLE. + +Run: + cd && -m pytest \ + tests/local_tests/cosmos3/test_cosmos3_scheduler_parity.py -q +""" +from __future__ import annotations + +import pytest +import torch + +# The official framework provides the parity oracle. +fm_unipc = pytest.importorskip( + "cosmos_framework.model.vfm.diffusion.samplers.fm_solvers_unipc", + reason="cosmos_framework not installed; run in fv-cosmos3 env.", +) +FlowUniPCMultistepScheduler = fm_unipc.FlowUniPCMultistepScheduler + +from fastvideo.models.schedulers.scheduling_unipc_multistep import ( # noqa: E402 + UniPCMultistepScheduler, +) + +pytestmark = [pytest.mark.local] + + +def _framework_scheduler(num_steps: int, shift: float) -> FlowUniPCMultistepScheduler: + """Exactly how ``UniPCSampler`` builds + primes its scheduler.""" + sched = FlowUniPCMultistepScheduler( + num_train_timesteps=1000, + shift=1.0, + use_dynamic_shifting=False, + ) + sched.set_timesteps(num_steps, device=torch.device("cpu"), shift=shift) + return sched + + +def _fastvideo_scheduler(num_steps: int, shift: float) -> UniPCMultistepScheduler: + """FastVideo's vendored UniPC configured for the framework's flow setup.""" + sched = UniPCMultistepScheduler( + num_train_timesteps=1000, + solver_order=2, + prediction_type="flow_prediction", + use_flow_sigmas=True, + predict_x0=True, + solver_type="bh2", + final_sigmas_type="zero", + flow_shift=shift, + ) + sched.set_timesteps(num_steps, device=torch.device("cpu")) + return sched + + +_SHIFTS = [pytest.param(10.0, id="shift10_video"), pytest.param(3.0, id="shift3_t2i")] +_STEPS = [pytest.param(4, id="4steps"), pytest.param(10, id="10steps"), pytest.param(35, id="35steps")] + + +class TestCosmos3SchedulerParity: + + @pytest.mark.parametrize("shift", _SHIFTS) + @pytest.mark.parametrize("num_steps", _STEPS) + def test_timesteps_and_sigmas_match_framework(self, num_steps, shift): + fw = _framework_scheduler(num_steps, shift) + fv = _fastvideo_scheduler(num_steps, shift) + + t_max = (fw.timesteps.float() - fv.timesteps.float()).abs().max().item() + s_max = (fw.sigmas.float() - fv.sigmas.float()).abs().max().item() + print(f"\n[sched n={num_steps} shift={shift}] timesteps max diff={t_max:.3e} " + f"sigmas max diff={s_max:.3e}") + assert fw.timesteps.shape == fv.timesteps.shape + assert fw.sigmas.shape == fv.sigmas.shape + torch.testing.assert_close(fv.timesteps, fw.timesteps) + torch.testing.assert_close(fv.sigmas, fw.sigmas) + + @pytest.mark.parametrize("shift", _SHIFTS) + @pytest.mark.parametrize("num_steps", _STEPS) + def test_full_trajectory_matches_framework(self, num_steps, shift): + # A small latent so order-2 einsum paths exercise; batch axis as the + # samplers expect ([B, C, T, H, W]). + shape = (1, 4, 2, 3, 3) + torch.manual_seed(123) + init = torch.randn(shape, dtype=torch.float32) + # One pseudo-random "velocity" per step, identical on both sides. + velocities = [torch.randn(shape, dtype=torch.float32) for _ in range(num_steps)] + + fw = _framework_scheduler(num_steps, shift) + fv = _fastvideo_scheduler(num_steps, shift) + + fw_lat = init.clone() + fv_lat = init.clone() + worst = 0.0 + for i, t in enumerate(fw.timesteps): + v = velocities[i] + fw_lat = fw.step(model_output=v, timestep=t, sample=fw_lat, return_dict=False)[0] + fv_lat = fv.step(model_output=v, timestep=t, sample=fv_lat, return_dict=False)[0] + step_max = (fw_lat - fv_lat).abs().max().item() + worst = max(worst, step_max) + assert not torch.isnan(fv_lat).any(), f"FastVideo latent NaN at step {i}" + mean_abs = (fw_lat - fv_lat).abs().mean().item() + print(f"\n[traj n={num_steps} shift={shift}] worst step max diff={worst:.3e} " + f"final mean diff={mean_abs:.3e}") + torch.testing.assert_close(fv_lat, fw_lat, atol=1e-5, rtol=1e-4)