From 1436ca2bb9e7d8fb6b16e9d61149d574f5cbb672 Mon Sep 17 00:00:00 2001 From: mvanhorn Date: Wed, 17 Jun 2026 07:45:12 -0700 Subject: [PATCH 1/5] 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 0fd14cc..fcb0755 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. @@ -1139,11 +1193,15 @@ def sample( to disable the special conditioning-frame timestep. 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: @@ -1156,6 +1214,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 @@ -1211,7 +1271,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 13cc85b..a8793ef 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 From e23e4be6eb31193db4be8a349136bf60e28b4311 Mon Sep 17 00:00:00 2001 From: Julius Berner Date: Fri, 24 Jul 2026 01:05:08 +0000 Subject: [PATCH 2/5] fix: align Cosmos Predict2.5 inference with official settings --- .../experiments/CosmosPredict2/config_dmd2.py | 7 +- .../CosmosPredict2/config_dmd2_14b.py | 4 +- .../CosmosPredict2/config_dmd2_v2w.py | 3 + .../experiments/CosmosPredict2/config_sft.py | 7 +- .../CosmosPredict2/config_sft_v2w.py | 5 + fastgen/configs/net.py | 5 +- fastgen/networks/cosmos_predict2/network.py | 110 +++++------------- scripts/README.md | 9 +- scripts/inference/video_model_inference.py | 18 ++- tests/test_network.py | 78 ------------- 10 files changed, 58 insertions(+), 188 deletions(-) diff --git a/fastgen/configs/experiments/CosmosPredict2/config_dmd2.py b/fastgen/configs/experiments/CosmosPredict2/config_dmd2.py index 408e391..2ea477a 100644 --- a/fastgen/configs/experiments/CosmosPredict2/config_dmd2.py +++ b/fastgen/configs/experiments/CosmosPredict2/config_dmd2.py @@ -28,8 +28,8 @@ def create_config(): # Small resolution for testing: 320x176 video -> 40x22 latent, 93 frames -> 24 latent # Full 720p: [16, 24, 88, 160] (1280x704 @ 93 frames) # config.model.input_shape = [16, 24, 22, 40] # cthw - 256p (320x176 video) - config.model.input_shape = [16, 24, 60, 104] # cthw - 480p, 93 frames - # config.model.input_shape = [16, 24, 88, 160] # cthw - full 720p, 93 frames + # config.model.input_shape = [16, 24, 60, 104] # cthw - 480p, 93 frames + config.model.input_shape = [16, 24, 88, 160] # cthw - full 720p, 93 frames # Network and discriminator config.model.net = CosmosPredict2_2B_Config @@ -44,7 +44,8 @@ def create_config(): config.model.gan_use_same_t_noise = True config.model.fake_score_pred_type = "x0" config.model.student_sample_type = "ode" - config.model.guidance_scale = 3.0 + # Official 2B DMD2 teacher guidance scale + config.model.guidance_scale = 4.0 config.model.pretrained_model_path = f"{CKPT_ROOT_DIR}/cosmos_predict2/Cosmos-Predict2.5-2B/base/post-trained/81edfebe-bd6a-4039-8c1d-737df1a790bf_ema_bf16.pt" # Timestep sampling diff --git a/fastgen/configs/experiments/CosmosPredict2/config_dmd2_14b.py b/fastgen/configs/experiments/CosmosPredict2/config_dmd2_14b.py index bea7329..6a41eb6 100644 --- a/fastgen/configs/experiments/CosmosPredict2/config_dmd2_14b.py +++ b/fastgen/configs/experiments/CosmosPredict2/config_dmd2_14b.py @@ -17,8 +17,8 @@ def create_config(): # Small resolution for testing: 320x176 video -> 40x22 latent, 93 frames -> 24 latent # Full 720p: [16, 24, 88, 160] (1280x704 @ 93 frames) # config.model.input_shape = [16, 24, 22, 40] # cthw - 256p (320x176 video) - config.model.input_shape = [16, 24, 60, 104] # cthw - 480p, 93 frames - # config.model.input_shape = [16, 24, 88, 160] # cthw - full 720p, 93 frames + # config.model.input_shape = [16, 24, 60, 104] # cthw - 480p, 93 frames + config.model.input_shape = [16, 24, 88, 160] # cthw - full 720p, 93 frames # Network and discriminator for 14B config.model.net = CosmosPredict2_14B_Config diff --git a/fastgen/configs/experiments/CosmosPredict2/config_dmd2_v2w.py b/fastgen/configs/experiments/CosmosPredict2/config_dmd2_v2w.py index c428953..d084e8b 100644 --- a/fastgen/configs/experiments/CosmosPredict2/config_dmd2_v2w.py +++ b/fastgen/configs/experiments/CosmosPredict2/config_dmd2_v2w.py @@ -3,6 +3,8 @@ """DMD2 config for Cosmos-Predict2.5-2B video2world model.""" +import copy + import fastgen.configs.experiments.CosmosPredict2.config_dmd2 as config_dmd2_base @@ -10,6 +12,7 @@ def create_config(): config = config_dmd2_base.create_config() # Network for v2w + config.model.net = copy.deepcopy(config.model.net) config.model.net.is_video2world = True config.model.net.num_conditioning_frames = 1 diff --git a/fastgen/configs/experiments/CosmosPredict2/config_sft.py b/fastgen/configs/experiments/CosmosPredict2/config_sft.py index 6377bd6..b88b0b7 100644 --- a/fastgen/configs/experiments/CosmosPredict2/config_sft.py +++ b/fastgen/configs/experiments/CosmosPredict2/config_sft.py @@ -28,12 +28,13 @@ def create_config(): # Full 720p: [16, 24, 88, 160] (1280x704 @ 93 frames) # config.model.input_shape = [16, 24, 24, 40] # cthw - 256p # config.model.input_shape = [16, 21, 60, 104] # cthw - 480p, 81 frames - config.model.input_shape = [16, 24, 60, 104] # cthw - 480p - # config.model.input_shape = [16, 24, 88, 160] # cthw - full 720p + # config.model.input_shape = [16, 24, 60, 104] # cthw - 480p + config.model.input_shape = [16, 24, 88, 160] # cthw - full 720p config.model.net = CosmosPredict2_2B_Config config.model.pretrained_model_path = f"{CKPT_ROOT_DIR}/cosmos_predict2/Cosmos-Predict2.5-2B/base/post-trained/81edfebe-bd6a-4039-8c1d-737df1a790bf_ema_bf16.pt" - config.model.guidance_scale = 3.0 + # Official post-trained 2B/14B inference uses guidance scale 7 + config.model.guidance_scale = 5.0 config.model.student_sample_steps = 35 config.dataloader_train = VideoLoaderConfig diff --git a/fastgen/configs/experiments/CosmosPredict2/config_sft_v2w.py b/fastgen/configs/experiments/CosmosPredict2/config_sft_v2w.py index 5c38adc..211b3a3 100644 --- a/fastgen/configs/experiments/CosmosPredict2/config_sft_v2w.py +++ b/fastgen/configs/experiments/CosmosPredict2/config_sft_v2w.py @@ -3,6 +3,8 @@ """Configs for SFT on Cosmos-Predict2.5-2B video2world model.""" +import copy + import fastgen.configs.experiments.CosmosPredict2.config_sft as config_sft_base @@ -10,8 +12,11 @@ def create_config(): config = config_sft_base.create_config() # Network for v2w + config.model.net = copy.deepcopy(config.model.net) config.model.net.is_video2world = True config.model.net.num_conditioning_frames = 1 + # Official post-trained Video2World inference uses guidance scale 8 + config.model.guidance_scale = 5.0 config.log_config.group = "cosmos_predict2_sft_v2w" diff --git a/fastgen/configs/net.py b/fastgen/configs/net.py index 7833c0b..f1aa667 100644 --- a/fastgen/configs/net.py +++ b/fastgen/configs/net.py @@ -277,7 +277,10 @@ num_blocks=28, num_heads=16, sac_config=L(SACConfig)(mode=CheckpointMode.BLOCK_WISE), - fps=24, + # Checkpoint-specific override: the released post-trained 2B model uses + # 16 FPS with temporal RoPE FPS modulation disabled. + fps=16, + rope_enable_fps_modulation=False, is_video2world=False, num_conditioning_frames=1, enable_logvar_linear=False, # Enable logvar_linear for sCM-like models diff --git a/fastgen/networks/cosmos_predict2/network.py b/fastgen/networks/cosmos_predict2/network.py index fcb0755..dbb8290 100644 --- a/fastgen/networks/cosmos_predict2/network.py +++ b/fastgen/networks/cosmos_predict2/network.py @@ -51,55 +51,6 @@ 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 ----------------------- @@ -816,7 +767,7 @@ def __init__( use_wan_fp32_strategy: bool = True, # FPS for temporal position embeddings fps: float = 24.0, - use_karras_sigma_schedule: bool = True, + rope_enable_fps_modulation: bool = True, # Video2world (image-to-video) mode settings is_video2world: bool = False, num_conditioning_frames: int = 1, @@ -857,6 +808,7 @@ def __init__( rope_h_extrapolation_ratio=rope_h_extrapolation_ratio, rope_w_extrapolation_ratio=rope_w_extrapolation_ratio, rope_t_extrapolation_ratio=rope_t_extrapolation_ratio, + rope_enable_fps_modulation=rope_enable_fps_modulation, use_wan_fp32_strategy=use_wan_fp32_strategy, ) @@ -874,9 +826,6 @@ 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 @@ -1156,9 +1105,8 @@ def sample( fps: Optional[torch.Tensor] = None, conditioning_latents: Optional[torch.Tensor] = None, num_conditioning_frames: int = 1, - conditional_frame_timestep: float = 0.0, + conditional_frame_timestep: float = 1e-4, denoise_replace_gt_frames: bool = True, - use_karras_sigma_schedule: Optional[bool] = None, **kwargs, ) -> torch.Tensor: """Multistep sampling using the FlowUniPC method. @@ -1179,46 +1127,43 @@ def sample( noise: The noisy latents to start from, shape (B, C, T, H, W). condition: The conditioning embeddings. neg_condition: The negative/unconditional embeddings for CFG. - guidance_scale: Classifier-free guidance scale. Values > 1.0 enable CFG. - num_steps: The number of sampling steps. - shift: Noise schedule shift parameter. + guidance_scale: Conventional classifier-free guidance scale. + Values greater than 1.0 enable CFG. + num_steps: Number of intervals in the official Karras schedule. + The sampler evaluates ``num_steps + 1`` timesteps. + shift: Retained for scheduler API compatibility. The official + Karras path uses fixed sigma bounds and does not apply flow shift. skip_layers: List of transformer layers to skip (for skip-layer guidance). skip_layers_start_percent: Percentage of steps before starting to skip layers. fps: Frames per second tensor for temporal conditioning. conditioning_latents: Latent frames to condition on for video2world mode, shape (B, C, T, H, W). If provided, enables video2world mode. num_conditioning_frames: Number of frames to use for conditioning (default 1). - conditional_frame_timestep: Timestep value for conditioning frames (default 0.0). - Use 0.0 to indicate clean frames with no noise, or a negative value - to disable the special conditioning-frame timestep. 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) + # Match official Cosmos Predict2.5 inference. Diffusers uses the + # configured sigma bounds for its Karras conversion; the official + # `num_steps` denotes intervals, hence `num_steps + 1` ramp points. if self.sample_scheduler is None: self.sample_scheduler = UniPCMultistepScheduler( num_train_timesteps=1000, prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=shift, + use_karras_sigmas=True, + sigma_min=0.01, + sigma_max=200.0, ) - 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) + self.sample_scheduler.set_timesteps(num_inference_steps=num_steps + 1, device=noise.device) timesteps = self.sample_scheduler.timesteps - # Initialize latents with proper scaling based on the initial timestep + # Initialize latents using the FastGen noise-schedule convention. t_init = timesteps[0] / self.sample_scheduler.config.num_train_timesteps latents = self.noise_scheduler.latents(noise=noise, t_init=t_init) @@ -1271,7 +1216,7 @@ def sample( # Store initial noise for velocity replacement initial_noise = latents.clone() - for idx, timestep in tqdm(enumerate(timesteps), total=len(timesteps), desc="Sampling"): + for timestep in tqdm(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( @@ -1324,14 +1269,10 @@ def sample( gt_velocity = initial_noise - conditioning_latents_full velocity_pred = gt_velocity * condition_mask_C + velocity_pred * (1 - condition_mask_C) - # Scheduler step (use model_input for video2world since that's what velocity was computed on) - sample_input = model_input if video2world_mode else latents - latents = self.sample_scheduler.step(velocity_pred, timestep, sample_input, return_dict=False)[0] - - # Preserve conditioning frames after each step to prevent numerical drift - if video2world_mode: - v2w_condition = {"conditioning_latents": conditioning_latents, "condition_mask": condition_mask} - latents = self.preserve_conditioning(latents, v2w_condition) + # Only the DiT input receives clean conditioning frames. UniPC must + # evolve the original latent state; the analytical conditioning-frame + # velocity above drives those frames from their initial noise to x0. + latents = self.sample_scheduler.step(velocity_pred, timestep, latents, return_dict=False)[0] return latents @@ -1376,9 +1317,10 @@ def forward( fps: Frames per second tensor padding_mask: Padding mask tensor skip_layers: List of block indices to skip during forward pass - conditional_frame_timestep: Timestep value for conditioning frames (default 0.0). - Use 0.0 to indicate clean frames with no noise, or a negative value - to disable the special conditioning-frame timestep. + conditional_frame_timestep: Normalized timestep for conditioning frames. + ``forward()`` defaults to 0.0; ``sample()`` passes 1e-4, equivalent + to the official raw timestep 0.1 after its 0.001 network scale. + A negative value disables the special conditioning-frame timestep. Returns: Depending on the arguments: diff --git a/scripts/README.md b/scripts/README.md index a8d62a2..4c49159 100644 --- a/scripts/README.md +++ b/scripts/README.md @@ -82,12 +82,9 @@ python scripts/inference/video_model_inference.py \ ### Video2World (Cosmos) ```bash python scripts/inference/video_model_inference.py \ - --config fastgen/configs/experiments/CosmosPredict2/config_sft.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 \ - - model.guidance_scale=5.0 model.net.is_video2world=True model.input_shape="[16, 24, 88, 160]" + --config fastgen/configs/experiments/CosmosPredict2/config_sft_v2w.py \ + --do_student_sampling False --num_steps 35 \ + --neg_prompt_file scripts/inference/prompts/negative_prompt_cosmos.txt ``` ### Causal / Autoregressive diff --git a/scripts/inference/video_model_inference.py b/scripts/inference/video_model_inference.py index a8793ef..bd1fe65 100644 --- a/scripts/inference/video_model_inference.py +++ b/scripts/inference/video_model_inference.py @@ -14,34 +14,29 @@ # 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_wan22_5b.py \\ + --config fastgen/configs/experiments/WanI2V/config_dmd2_14b.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.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 --num_steps 35 --fps 24 \\ + scripts/inference/video_model_inference.py --do_student_sampling False --num_steps 35 \\ --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 \\ - model.input_shape="[16, 24, 88, 160]" log_config.name=cosmos_v2w_inference + --config fastgen/configs/experiments/CosmosPredict2/config_sft_v2w.py \\ + - trainer.ddp=True 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 \\ @@ -725,7 +720,8 @@ def main(args, config: BaseConfig): "--fps", default=16, type=int, - help="Frames per second for saved video and model temporal encoding (default: 16, matches Wan base_fps)", + help="Frames per second for saved video and model temporal encoding " + "(default: 16, matching Wan and Cosmos-Predict2.5 metadata)", ) parser.add_argument( "--save_high_quality", @@ -805,7 +801,7 @@ def main(args, config: BaseConfig): "--num_conditioning_frames", type=int, default=1, - help="Number of latent frames to condition on for I2V mode (default: 1).", + help="Number of latent frames to condition on for I2V/Video2World mode (default: 1).", ) args = parse_args(parser) diff --git a/tests/test_network.py b/tests/test_network.py index df93ad9..f14cb4b 100644 --- a/tests/test_network.py +++ b/tests/test_network.py @@ -2,10 +2,6 @@ # SPDX-License-Identifier: Apache-2.0 import os -from pathlib import Path -from types import SimpleNamespace - -import numpy as np import torch import pytest @@ -39,10 +35,6 @@ 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 @@ -257,76 +249,6 @@ 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 From a234d14ac2792b7550281570624aa45dfb11f604 Mon Sep 17 00:00:00 2001 From: Julius Berner Date: Fri, 24 Jul 2026 02:40:27 +0000 Subject: [PATCH 3/5] Adapt doc strings --- fastgen/networks/cosmos_predict2/network.py | 6 +----- scripts/inference/video_model_inference.py | 3 +-- 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/fastgen/networks/cosmos_predict2/network.py b/fastgen/networks/cosmos_predict2/network.py index dbb8290..a50d574 100644 --- a/fastgen/networks/cosmos_predict2/network.py +++ b/fastgen/networks/cosmos_predict2/network.py @@ -1099,7 +1099,6 @@ def sample( neg_condition: Optional[Any] = None, guidance_scale: Optional[float] = 5.0, num_steps: int = 50, - shift: float = 5.0, # Cosmos default skip_layers: Optional[List[int]] = None, skip_layers_start_percent: float = 0.0, fps: Optional[torch.Tensor] = None, @@ -1131,8 +1130,6 @@ def sample( Values greater than 1.0 enable CFG. num_steps: Number of intervals in the official Karras schedule. The sampler evaluates ``num_steps + 1`` timesteps. - shift: Retained for scheduler API compatibility. The official - Karras path uses fixed sigma bounds and does not apply flow shift. skip_layers: List of transformer layers to skip (for skip-layer guidance). skip_layers_start_percent: Percentage of steps before starting to skip layers. fps: Frames per second tensor for temporal conditioning. @@ -1155,7 +1152,6 @@ def sample( num_train_timesteps=1000, prediction_type="flow_prediction", use_flow_sigmas=True, - flow_shift=shift, use_karras_sigmas=True, sigma_min=0.01, sigma_max=200.0, @@ -1163,7 +1159,7 @@ def sample( self.sample_scheduler.set_timesteps(num_inference_steps=num_steps + 1, device=noise.device) timesteps = self.sample_scheduler.timesteps - # Initialize latents using the FastGen noise-schedule convention. + # Initialize latents with proper scaling based on the initial timestep t_init = timesteps[0] / self.sample_scheduler.config.num_train_timesteps latents = self.noise_scheduler.latents(noise=noise, t_init=t_init) diff --git a/scripts/inference/video_model_inference.py b/scripts/inference/video_model_inference.py index bd1fe65..f18ee4a 100644 --- a/scripts/inference/video_model_inference.py +++ b/scripts/inference/video_model_inference.py @@ -720,8 +720,7 @@ def main(args, config: BaseConfig): "--fps", default=16, type=int, - help="Frames per second for saved video and model temporal encoding " - "(default: 16, matching Wan and Cosmos-Predict2.5 metadata)", + help="Frames per second for saved video and model temporal encoding (default: 16)", ) parser.add_argument( "--save_high_quality", From beb6d4ebde7e3f66b515105dad12f6fd02b7b7d9 Mon Sep 17 00:00:00 2001 From: Julius Berner Date: Fri, 24 Jul 2026 05:23:40 +0000 Subject: [PATCH 4/5] Remove unused resolution and adapt docstrings --- fastgen/configs/experiments/CosmosPredict2/config_sft.py | 1 - fastgen/networks/cosmos_predict2/network.py | 6 +++--- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/fastgen/configs/experiments/CosmosPredict2/config_sft.py b/fastgen/configs/experiments/CosmosPredict2/config_sft.py index b88b0b7..00b7650 100644 --- a/fastgen/configs/experiments/CosmosPredict2/config_sft.py +++ b/fastgen/configs/experiments/CosmosPredict2/config_sft.py @@ -27,7 +27,6 @@ def create_config(): # Small resolution for testing: 320x176 video -> 40x22 latent, 21 frames -> 6 latent # Full 720p: [16, 24, 88, 160] (1280x704 @ 93 frames) # config.model.input_shape = [16, 24, 24, 40] # cthw - 256p - # config.model.input_shape = [16, 21, 60, 104] # cthw - 480p, 81 frames # config.model.input_shape = [16, 24, 60, 104] # cthw - 480p config.model.input_shape = [16, 24, 88, 160] # cthw - full 720p diff --git a/fastgen/networks/cosmos_predict2/network.py b/fastgen/networks/cosmos_predict2/network.py index a50d574..b6215f7 100644 --- a/fastgen/networks/cosmos_predict2/network.py +++ b/fastgen/networks/cosmos_predict2/network.py @@ -1136,6 +1136,8 @@ def sample( conditioning_latents: Latent frames to condition on for video2world mode, shape (B, C, T, H, W). If provided, enables video2world mode. num_conditioning_frames: Number of frames to use for conditioning (default 1). + conditional_frame_timestep: Normalized conditioning-frame timestep + forwarded to ``forward()`` (default 1e-4). denoise_replace_gt_frames: Whether to replace velocity for conditioning frames with analytical velocity (noise - gt_frames). Default True. @@ -1265,9 +1267,7 @@ def sample( gt_velocity = initial_noise - conditioning_latents_full velocity_pred = gt_velocity * condition_mask_C + velocity_pred * (1 - condition_mask_C) - # Only the DiT input receives clean conditioning frames. UniPC must - # evolve the original latent state; the analytical conditioning-frame - # velocity above drives those frames from their initial noise to x0. + # Keep clean frames in the DiT input while UniPC evolves raw latents. latents = self.sample_scheduler.step(velocity_pred, timestep, latents, return_dict=False)[0] return latents From 89d07e615834def4bd2c59bbb7486de9e01bf48c Mon Sep 17 00:00:00 2001 From: Julius Berner Date: Fri, 24 Jul 2026 18:34:07 +0000 Subject: [PATCH 5/5] Adapt defaults --- .../configs/experiments/CosmosPredict2/config_dmd2_14b.py | 7 +++++-- fastgen/configs/experiments/CosmosPredict2/config_sft.py | 2 +- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/fastgen/configs/experiments/CosmosPredict2/config_dmd2_14b.py b/fastgen/configs/experiments/CosmosPredict2/config_dmd2_14b.py index 6a41eb6..a1d403d 100644 --- a/fastgen/configs/experiments/CosmosPredict2/config_dmd2_14b.py +++ b/fastgen/configs/experiments/CosmosPredict2/config_dmd2_14b.py @@ -3,6 +3,8 @@ """DMD2 config for Cosmos-Predict2.5-14B model.""" +import copy + import fastgen.configs.experiments.CosmosPredict2.config_dmd2 as config_dmd2_base from fastgen.configs.discriminator import Discriminator_CosmosPredict2_14B_Config from fastgen.configs.net import CosmosPredict2_14B_Config, CosmosPredict2_14B_Aggressive_Config, CKPT_ROOT_DIR @@ -17,8 +19,8 @@ def create_config(): # Small resolution for testing: 320x176 video -> 40x22 latent, 93 frames -> 24 latent # Full 720p: [16, 24, 88, 160] (1280x704 @ 93 frames) # config.model.input_shape = [16, 24, 22, 40] # cthw - 256p (320x176 video) - # config.model.input_shape = [16, 24, 60, 104] # cthw - 480p, 93 frames - config.model.input_shape = [16, 24, 88, 160] # cthw - full 720p, 93 frames + config.model.input_shape = [16, 24, 60, 104] # cthw - 480p, 93 frames + # config.model.input_shape = [16, 24, 88, 160] # cthw - full 720p, 93 frames # Network and discriminator for 14B config.model.net = CosmosPredict2_14B_Config @@ -29,6 +31,7 @@ def create_config(): config.model.pretrained_model_path = f"{CKPT_ROOT_DIR}/cosmos_predict2/Cosmos-Predict2.5-14B/base/post-trained/e21d2a49-4747-44c8-ba44-9f6f9243715f_ema_bf16.pt" # Dataloader img_size = (W, H) = (latent_W * 8, latent_H * 8) + config.dataloader_train = copy.deepcopy(config.dataloader_train) config.dataloader_train.img_size = (config.model.input_shape[-1] * 8, config.model.input_shape[-2] * 8) config.dataloader_train.sequence_length = (config.model.input_shape[1] - 1) * 4 + 1 diff --git a/fastgen/configs/experiments/CosmosPredict2/config_sft.py b/fastgen/configs/experiments/CosmosPredict2/config_sft.py index 00b7650..2a62787 100644 --- a/fastgen/configs/experiments/CosmosPredict2/config_sft.py +++ b/fastgen/configs/experiments/CosmosPredict2/config_sft.py @@ -14,7 +14,7 @@ def create_config(): config.trainer.logging_iter = 100 config.trainer.save_ckpt_iter = 500 - config.model.net_optimizer.lr = 1e-5 + config.model.net_optimizer.lr = 5e-6 config.model.sample_t_cfg.time_dist_type = "uniform" config.model.sample_t_cfg.min_t = 0.001