[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.
This commit is contained in:
SolitaryThinker
@@ -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()
|
||||
@@ -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"
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"),
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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]
|
||||
|
||||
|
||||
@@ -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 <worktree> && <fv-cosmos3 python> -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)
|
||||
Reference in New Issue
Block a user