From c2c3ee0d244e0c03030db81bb56c53044d13ac53 Mon Sep 17 00:00:00 2001 From: mvanhorn Date: Wed, 17 Jun 2026 07:45:12 -0700 Subject: [PATCH] fix: Use Karras sigma schedule for Cosmos Predict2.5 SFT inference --- fastgen/networks/cosmos_predict2/network.py | 62 +++++++++++++++- scripts/inference/video_model_inference.py | 12 ++-- tests/test_network.py | 78 +++++++++++++++++++++ 3 files changed, 147 insertions(+), 5 deletions(-) diff --git a/fastgen/networks/cosmos_predict2/network.py b/fastgen/networks/cosmos_predict2/network.py index ff85360..4ece874 100644 --- a/fastgen/networks/cosmos_predict2/network.py +++ b/fastgen/networks/cosmos_predict2/network.py @@ -51,6 +51,55 @@ from fastgen.utils.basic_utils import str2bool +_COSMOS_KARRAS_SIGMA_MAX = 200.0 +_COSMOS_KARRAS_SIGMA_MIN = 0.01 +_COSMOS_KARRAS_RHO = 7.0 + + +def _build_cosmos_predict2_karras_schedule( + num_steps: int, + num_train_timesteps: int, + device: Union[torch.device, str], +) -> Tuple[torch.Tensor, torch.Tensor]: + """Build the Karras sigma/timestep schedule used by official Cosmos Predict2.5 inference.""" + if num_steps <= 0: + raise ValueError(f"num_steps must be positive, got {num_steps}") + + ramp = torch.arange(num_steps + 1, dtype=torch.float64) / num_steps + min_inv_rho = _COSMOS_KARRAS_SIGMA_MIN ** (1 / _COSMOS_KARRAS_RHO) + max_inv_rho = _COSMOS_KARRAS_SIGMA_MAX ** (1 / _COSMOS_KARRAS_RHO) + sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** _COSMOS_KARRAS_RHO + sigmas = sigmas / (1 + sigmas) + + timesteps = (sigmas * num_train_timesteps).to(device=device, dtype=torch.int64) + sigmas = torch.cat([sigmas, sigmas.new_zeros(1)]).to(dtype=torch.float32) + return sigmas, timesteps + + +def _reset_unipc_scheduler_state(scheduler: UniPCMultistepScheduler) -> None: + scheduler.model_outputs = [None] * scheduler.config.solver_order + scheduler.lower_order_nums = 0 + scheduler.last_sample = None + scheduler._step_index = None + scheduler._begin_index = None + + +def _apply_cosmos_predict2_karras_schedule( + scheduler: UniPCMultistepScheduler, + num_steps: int, + device: Union[torch.device, str], +) -> None: + sigmas, timesteps = _build_cosmos_predict2_karras_schedule( + num_steps=num_steps, + num_train_timesteps=scheduler.config.num_train_timesteps, + device=device, + ) + scheduler.sigmas = sigmas + scheduler.timesteps = timesteps + scheduler.num_inference_steps = len(timesteps) + _reset_unipc_scheduler_state(scheduler) + + # ---------------------- DiT Network ----------------------- @@ -767,6 +816,7 @@ def __init__( use_wan_fp32_strategy: bool = True, # FPS for temporal position embeddings fps: float = 24.0, + use_karras_sigma_schedule: bool = True, # Video2world (image-to-video) mode settings is_video2world: bool = False, num_conditioning_frames: int = 1, @@ -824,6 +874,9 @@ def __init__( # Default FPS for temporal position embeddings self.fps = fps + # Match official Cosmos Predict2.5 SFT inference schedule by default. + self.use_karras_sigma_schedule = use_karras_sigma_schedule + # Initialize the sample scheduler for inference later (otherwise it causes issues with Meta init) self.sample_scheduler = None @@ -1105,6 +1158,7 @@ def sample( num_conditioning_frames: int = 1, conditional_frame_timestep: float = 0.0, denoise_replace_gt_frames: bool = True, + use_karras_sigma_schedule: Optional[bool] = None, **kwargs, ) -> torch.Tensor: """Multistep sampling using the FlowUniPC method. @@ -1138,11 +1192,15 @@ def sample( Use 0.0 to indicate clean frames with no noise. denoise_replace_gt_frames: Whether to replace velocity for conditioning frames with analytical velocity (noise - gt_frames). Default True. + use_karras_sigma_schedule: Whether to use the official Cosmos Predict2.5 + Karras sigma schedule for SFT inference. Defaults to self.use_karras_sigma_schedule. Returns: The denoised sample tensor. """ assert self.schedule_type == "rf", f"{self.schedule_type} is not supported" + if use_karras_sigma_schedule is None: + use_karras_sigma_schedule = self.use_karras_sigma_schedule # Set timesteps with shift parameter (matching official Cosmos inference) if self.sample_scheduler is None: @@ -1155,6 +1213,8 @@ def sample( else: self.sample_scheduler.config.flow_shift = shift self.sample_scheduler.set_timesteps(num_inference_steps=num_steps, device=noise.device) + if use_karras_sigma_schedule: + _apply_cosmos_predict2_karras_schedule(self.sample_scheduler, num_steps=num_steps, device=noise.device) timesteps = self.sample_scheduler.timesteps # Initialize latents with proper scaling based on the initial timestep @@ -1210,7 +1270,7 @@ def sample( # Store initial noise for velocity replacement initial_noise = latents.clone() - for idx, timestep in tqdm(enumerate(timesteps), total=num_steps, desc="Sampling"): + for idx, timestep in tqdm(enumerate(timesteps), total=len(timesteps), desc="Sampling"): # Normalize timestep to [0, 1] range t = (timestep / self.sample_scheduler.config.num_train_timesteps).expand(latents.shape[0]) t = self.noise_scheduler.safe_clamp(t, min=self.noise_scheduler.min_t, max=self.noise_scheduler.max_t).to( diff --git a/scripts/inference/video_model_inference.py b/scripts/inference/video_model_inference.py index 83b9856..2137520 100644 --- a/scripts/inference/video_model_inference.py +++ b/scripts/inference/video_model_inference.py @@ -14,30 +14,34 @@ # T2V: eval teacher only (Wan) PYTHONPATH=$(pwd) FASTGEN_OUTPUT_ROOT='FASTGEN_OUTPUT' torchrun --nproc_per_node=1 --standalone \\ scripts/inference/video_model_inference.py --do_student_sampling False \\ + --num_steps 50 --fps 16 --neg_prompt_file scripts/inference/prompts/negative_prompt.txt \\ --config fastgen/configs/experiments/WanT2V/config_dmd2.py \\ - trainer.seed=1 trainer.ddp=True model.guidance_scale=5.0 log_config.name=wan_t2v_inference # I2V: image-to-video (Wan I2V) PYTHONPATH=$(pwd) FASTGEN_OUTPUT_ROOT='FASTGEN_OUTPUT' torchrun --nproc_per_node=1 --standalone \\ scripts/inference/video_model_inference.py --do_student_sampling False \\ + --num_steps 50 --fps 16 --neg_prompt_file scripts/inference/prompts/negative_prompt.txt \\ --input_image_file scripts/inference/prompts/source_image_paths.txt \\ - --config fastgen/configs/experiments/WanI2V/config_dmd2_14b.py \\ + --config fastgen/configs/experiments/WanI2V/config_dmd2_wan22_5b.py \\ - trainer.seed=1 trainer.ddp=True model.guidance_scale=5.0 log_config.name=wan_i2v_inference # V2V: video-to-video with VACE PYTHONPATH=$(pwd) FASTGEN_OUTPUT_ROOT='FASTGEN_OUTPUT' torchrun --nproc_per_node=1 --standalone \\ scripts/inference/video_model_inference.py --do_student_sampling False \\ + --num_steps 50 --fps 16 --neg_prompt_file scripts/inference/prompts/negative_prompt.txt \\ --source_video_file scripts/inference/prompts/source_video_paths.txt \\ - --config fastgen/configs/experiments/WanV2V/config_sft_latent.py \\ + --config fastgen/configs/experiments/WanV2V/config_sft.py \\ - trainer.seed=1 trainer.ddp=True model.guidance_scale=5.0 log_config.name=vace_wan_inference # Video2World: Cosmos Predict2 PYTHONPATH=$(pwd) FASTGEN_OUTPUT_ROOT='FASTGEN_OUTPUT' torchrun --nproc_per_node=1 --standalone \\ - scripts/inference/video_model_inference.py --do_student_sampling False \\ + scripts/inference/video_model_inference.py --do_student_sampling False --num_steps 35 --fps 24 \\ + --neg_prompt_file scripts/inference/prompts/negative_prompt_cosmos.txt \\ --input_image_file scripts/inference/prompts/source_image_paths.txt --num_conditioning_frames 1 \\ --config fastgen/configs/experiments/CosmosPredict2/config_sft.py \\ - trainer.seed=1 trainer.ddp=True model.guidance_scale=5.0 model.net.is_video2world=True \\ - log_config.name=cosmos_v2w_inference + model.input_shape="[16, 24, 88, 160]" log_config.name=cosmos_v2w_inference # Eval with skip-layer guidance (SLG) PYTHONPATH=$(pwd) FASTGEN_OUTPUT_ROOT='FASTGEN_OUTPUT' torchrun --nproc_per_node=1 --standalone \\ diff --git a/tests/test_network.py b/tests/test_network.py index f14cb4b..df93ad9 100644 --- a/tests/test_network.py +++ b/tests/test_network.py @@ -2,6 +2,10 @@ # SPDX-License-Identifier: Apache-2.0 import os +from pathlib import Path +from types import SimpleNamespace + +import numpy as np import torch import pytest @@ -35,6 +39,10 @@ Discriminator_EDM_ImageNet64_Config, ) from fastgen.configs.config_utils import override_config_with_opts +from fastgen.networks.cosmos_predict2.network import ( + _apply_cosmos_predict2_karras_schedule, + _build_cosmos_predict2_karras_schedule, +) from fastgen.utils.basic_utils import clear_gpu_memory from fastgen.utils.test_utils import RunIf from fastgen.utils.io_utils import set_env_vars @@ -249,6 +257,76 @@ def validate_noise_scheduler_properties(teacher, device, expected_schedule_type= raise ValueError(f"Unrecognized schedule type: {expected_schedule_type}") +def _reference_cosmos_karras_schedule(num_steps, num_train_timesteps=1000): + sigma_max, sigma_min, rho = 200.0, 0.01, 7.0 + sigmas = np.arange(num_steps + 1, dtype=np.float64) / num_steps + min_inv_rho = sigma_min ** (1 / rho) + max_inv_rho = sigma_max ** (1 / rho) + sigmas = (max_inv_rho + sigmas * (min_inv_rho - max_inv_rho)) ** rho + sigmas = sigmas / (1 + sigmas) + timesteps = (sigmas * num_train_timesteps).astype(np.int64) + sigmas = np.concatenate([sigmas, [0.0]]).astype(np.float32) + return sigmas, timesteps + + +def test_cosmos_predict2_karras_schedule_matches_reference(): + num_steps = 35 + sigmas, timesteps = _build_cosmos_predict2_karras_schedule( + num_steps=num_steps, + num_train_timesteps=1000, + device=torch.device("cpu"), + ) + expected_sigmas, expected_timesteps = _reference_cosmos_karras_schedule(num_steps) + + assert sigmas.shape == (num_steps + 2,) + assert timesteps.shape == (num_steps + 1,) + assert sigmas.dtype == torch.float32 + assert timesteps.dtype == torch.int64 + assert sigmas[-1].item() == 0.0 + assert torch.all(sigmas[:-1][:-1] > sigmas[:-1][1:]) + np.testing.assert_allclose(sigmas.numpy(), expected_sigmas, rtol=0, atol=1e-6) + np.testing.assert_array_equal(timesteps.numpy(), expected_timesteps) + + +def test_cosmos_predict2_karras_schedule_small_steps_and_state_reset(): + for num_steps in [1, 2]: + scheduler = SimpleNamespace( + config=SimpleNamespace(num_train_timesteps=1000, solver_order=3), + sigmas=torch.ones(1), + timesteps=torch.ones(1, dtype=torch.int64), + num_inference_steps=1, + model_outputs=["stale"] * 3, + lower_order_nums=2, + last_sample=torch.ones(1), + _step_index=5, + _begin_index=4, + ) + + _apply_cosmos_predict2_karras_schedule(scheduler, num_steps=num_steps, device=torch.device("cpu")) + expected_sigmas, expected_timesteps = _reference_cosmos_karras_schedule(num_steps) + + assert scheduler.sigmas.shape == (num_steps + 2,) + assert scheduler.timesteps.shape == (num_steps + 1,) + assert scheduler.timesteps.dtype == torch.int64 + assert scheduler.num_inference_steps == num_steps + 1 + np.testing.assert_allclose(scheduler.sigmas.numpy(), expected_sigmas, rtol=0, atol=1e-6) + np.testing.assert_array_equal(scheduler.timesteps.numpy(), expected_timesteps) + assert scheduler.model_outputs == [None, None, None] + assert scheduler.lower_order_nums == 0 + assert scheduler.last_sample is None + assert scheduler._step_index is None + assert scheduler._begin_index is None + + +def test_video_model_inference_cosmos_example_uses_cosmos_negative_prompt(): + script = Path("scripts/inference/video_model_inference.py").read_text() + cosmos_example = script.split("# Video2World: Cosmos Predict2", 1)[1].split("# Eval with skip-layer guidance", 1)[0] + + assert "--neg_prompt_file scripts/inference/prompts/negative_prompt_cosmos.txt" in cosmos_example + assert '--config fastgen/configs/experiments/CosmosPredict2/config_sft.py \\' in cosmos_example + assert 'model.input_shape="[16, 24, 88, 160]"' in cosmos_example + + def test_network_edm_cifar10(): teacher_config = EDM_CIFAR10_Config