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

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

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

fastgen/configs/experiments/CosmosPredict2/config_dmd2_v2w.py

@@ -3,13 +3,16 @@
 
 """DMD2 config for Cosmos-Predict2.5-2B video2world model."""
 
+import copy
+
 import fastgen.configs.experiments.CosmosPredict2.config_dmd2 as config_dmd2_base
 
 
 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
 

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
@@ -27,13 +27,13 @@ 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
+    # 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

fastgen/configs/experiments/CosmosPredict2/config_sft_v2w.py

@@ -3,15 +3,20 @@
 
 """Configs for SFT on Cosmos-Predict2.5-2B video2world model."""
 
+import copy
+
 import fastgen.configs.experiments.CosmosPredict2.config_sft as config_sft_base
 
 
 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"
 

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

fastgen/networks/cosmos_predict2/network.py

@@ -767,6 +767,7 @@ def __init__(
         use_wan_fp32_strategy: bool = True,
         # FPS for temporal position embeddings
         fps: float = 24.0,
+        rope_enable_fps_modulation: bool = True,
         # Video2world (image-to-video) mode settings
         is_video2world: bool = False,
         num_conditioning_frames: int = 1,
@@ -807,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,
         )
 
@@ -1097,13 +1099,12 @@ 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,
         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,
         **kwargs,
     ) -> torch.Tensor:
@@ -1125,18 +1126,18 @@ 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.
             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.
+            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.
 
@@ -1145,17 +1146,19 @@ def sample(
         """
         assert self.schedule_type == "rf", f"{self.schedule_type} is not supported"
 
-        # 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)
+        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
@@ -1211,7 +1214,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 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(
@@ -1264,14 +1267,8 @@ 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)
+            # 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
 
@@ -1316,9 +1313,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:

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

scripts/inference/video_model_inference.py

@@ -28,16 +28,15 @@
     PYTHONPATH=$(pwd) FASTGEN_OUTPUT_ROOT='FASTGEN_OUTPUT' torchrun --nproc_per_node=1 --standalone \\
         scripts/inference/video_model_inference.py --do_student_sampling False \\
         --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 \\
-        --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
+        scripts/inference/video_model_inference.py --do_student_sampling False --num_steps 35 \\
+        --neg_prompt_file scripts/inference/prompts/negative_prompt_cosmos.txt \\
+        --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 \\
@@ -721,7 +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, matches Wan base_fps)",
+        help="Frames per second for saved video and model temporal encoding (default: 16)",
     )
     parser.add_argument(
         "--save_high_quality",
@@ -801,7 +800,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)