fix: Use Karras sigma schedule for Cosmos Predict2.5 SFT inference

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(

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 \\

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