From 004de62940c41fc575bde6507491750e9666e358 Mon Sep 17 00:00:00 2001 From: Andrew White Date: Mon, 27 Jul 2026 07:38:44 -0500 Subject: [PATCH 1/2] fix(network): sample ignores skip_layers and (+1 more) Signed-off-by: Julius Berner --- fastgen/networks/cosmos_predict2/network.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/fastgen/networks/cosmos_predict2/network.py b/fastgen/networks/cosmos_predict2/network.py index b6215f7..d27b661 100644 --- a/fastgen/networks/cosmos_predict2/network.py +++ b/fastgen/networks/cosmos_predict2/network.py @@ -1212,15 +1212,18 @@ def sample( ) # Store initial noise for velocity replacement - initial_noise = latents.clone() + initial_noise = noise - for timestep in tqdm(timesteps, total=len(timesteps), desc="Sampling"): + total_steps = len(timesteps) + for step_idx, timestep in enumerate(tqdm(timesteps, total=total_steps, 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( latents.dtype ) + active_skip_layers = skip_layers if (step_idx / total_steps) >= skip_layers_start_percent else None + if video2world_mode: # Replace conditioning frames with clean latents using preserve_conditioning v2w_condition = {"conditioning_latents": conditioning_latents, "condition_mask": condition_mask} @@ -1249,6 +1252,7 @@ def sample( cond_with_mask, fps=fps, conditional_frame_timestep=conditional_frame_timestep, + skip_layers=active_skip_layers, ) # Classifier-free guidance @@ -1259,6 +1263,7 @@ def sample( neg_cond_with_mask, fps=fps, conditional_frame_timestep=conditional_frame_timestep, + skip_layers=active_skip_layers, ) velocity_pred = velocity_uncond + guidance_scale * (velocity_pred - velocity_uncond) From 0c20d60ba989d5088c193cf532896271ec775019 Mon Sep 17 00:00:00 2001 From: Julius Berner Date: Mon, 3 Aug 2026 17:56:18 +0000 Subject: [PATCH 2/2] fix(cosmos): correct skip-layer guidance and align the sampler Skip-layer guidance must degrade only the unconditional branch, matching Wan, WanI2V and dmd2; applying it to the conditional pass as well turns SLG into plain CFG on a truncated network. Supersedes the previous commit's placement. Also in this change: - FlowKarrasUniPCScheduler builds the official Cosmos Karras ramp in flow-matching units, so the schedule is identical on diffusers < 0.37, where sigma_min/sigma_max and the EDM->flow conversion do not yet exist. Verified bitwise equal to diffusers 0.38.0. - Rename skip_layers_start_percent to skip_layers_start_fraction across cosmos_predict2, Wan and WanI2V; it was always compared as a fraction. - Use the initial latents over [0, t_init] for the conditioning-frame velocity instead of the sigma-scaled latents. - Let tqdm derive its total from timesteps rather than a hand-computed value that could drift out of sync with the loop length. Signed-off-by: Julius Berner --- fastgen/networks/VaceWan/network_causal.py | 2 +- fastgen/networks/Wan/network.py | 9 +-- fastgen/networks/Wan/network_causal.py | 2 +- fastgen/networks/WanI2V/network.py | 6 +- fastgen/networks/WanI2V/network_causal.py | 2 +- fastgen/networks/cosmos_predict2/network.py | 73 ++++++++++++++++----- 6 files changed, 66 insertions(+), 28 deletions(-) diff --git a/fastgen/networks/VaceWan/network_causal.py b/fastgen/networks/VaceWan/network_causal.py index 45c0ed1..729c117 100644 --- a/fastgen/networks/VaceWan/network_causal.py +++ b/fastgen/networks/VaceWan/network_causal.py @@ -924,7 +924,7 @@ def sample( self.unipc_scheduler.config.flow_shift = shift self.unipc_scheduler.set_timesteps(num_inference_steps=sample_steps, device=noise.device) timesteps = self.unipc_scheduler.timesteps - for timestep in tqdm(timesteps, total=sample_steps - 1): + for timestep in tqdm(timesteps): t = (timestep / time_rescale_factor).expand(batch_size) x_cur = x_next flow_pred = self( diff --git a/fastgen/networks/Wan/network.py b/fastgen/networks/Wan/network.py index f6c3668..5181764 100644 --- a/fastgen/networks/Wan/network.py +++ b/fastgen/networks/Wan/network.py @@ -920,7 +920,7 @@ def sample( num_steps: int = 50, shift: float = 5.0, skip_layers: Optional[List[int]] = None, - skip_layers_start_percent: float = 0.0, + skip_layers_start_fraction: float = 0.0, **kwargs, ) -> torch.Tensor: """Multistep sample using the UniPC method @@ -933,7 +933,8 @@ def sample( num_steps (int): The number of sampling steps. shift (float): Noise schedule shift parameter. Affects temporal dynamics. skip_layers (Optional[List[int]]): List of transformer layers to skip (used by SLG) during sampling. - skip_layers_start_percent (float): The percentage of the sampling steps to start skipping layers. + skip_layers_start_fraction (float): Fraction in [0, 1] of the sampling steps to complete + before skip-layer guidance becomes active. Returns: torch.Tensor: The sample output. @@ -949,7 +950,7 @@ def sample( latents = self.noise_scheduler.latents(noise=noise, t_init=t_init) # main sampling loop - for idx, timestep in tqdm(enumerate(timesteps), total=num_steps - 1): + for idx, timestep in enumerate(tqdm(timesteps)): t = (timestep / self.unipc_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( latents.dtype @@ -974,7 +975,7 @@ def sample( return_features_early=False, feature_indices={}, return_logvar=False, - skip_layers=skip_layers if idx >= skip_layers_start_percent * num_steps else None, + skip_layers=skip_layers if idx >= skip_layers_start_fraction * num_steps else None, ) flow_pred = flow_uncond + guidance_scale * (flow_pred - flow_uncond) diff --git a/fastgen/networks/Wan/network_causal.py b/fastgen/networks/Wan/network_causal.py index 469ff45..4a3313f 100644 --- a/fastgen/networks/Wan/network_causal.py +++ b/fastgen/networks/Wan/network_causal.py @@ -1231,7 +1231,7 @@ def sample( x_next = x[:, :, start:end] # Reset scheduler state (model_outputs, lower_order_nums, etc.) self.unipc_scheduler.set_timesteps(num_inference_steps=sample_steps, device=noise.device) - for timestep in tqdm(timesteps, total=sample_steps - 1): + for timestep in tqdm(timesteps): t = (timestep / time_rescale_factor).expand(batch_size) x_cur = x_next flow_pred = self( diff --git a/fastgen/networks/WanI2V/network.py b/fastgen/networks/WanI2V/network.py index 7fa3138..a2d3017 100644 --- a/fastgen/networks/WanI2V/network.py +++ b/fastgen/networks/WanI2V/network.py @@ -344,7 +344,7 @@ def sample( num_steps: int = 40, shift: float = 3.0, skip_layers: Optional[List[int]] = None, - skip_layers_start_percent: float = 0.0, + skip_layers_start_fraction: float = 0.0, **kwargs, ) -> torch.Tensor: """Sample from the WanI2V model with proper first-frame conditioning. @@ -368,7 +368,7 @@ def sample( latents = self.noise_scheduler.latents(noise=noise, t_init=t_init) # Main sampling loop - for idx, timestep in tqdm(enumerate(timesteps), total=num_steps - 1): + for idx, timestep in enumerate(tqdm(timesteps)): t = (timestep / self.unipc_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( latents.dtype @@ -393,7 +393,7 @@ def sample( return_features_early=False, feature_indices={}, return_logvar=False, - skip_layers=skip_layers if idx >= skip_layers_start_percent * num_steps else None, + skip_layers=skip_layers if idx >= skip_layers_start_fraction * num_steps else None, ) flow_pred = flow_uncond + guidance_scale * (flow_pred - flow_uncond) diff --git a/fastgen/networks/WanI2V/network_causal.py b/fastgen/networks/WanI2V/network_causal.py index 0f65147..908828d 100644 --- a/fastgen/networks/WanI2V/network_causal.py +++ b/fastgen/networks/WanI2V/network_causal.py @@ -454,7 +454,7 @@ def sample( x_next = x[:, :, start:end] # Reset scheduler state (model_outputs, lower_order_nums, etc.) self.unipc_scheduler.set_timesteps(num_inference_steps=sample_steps, device=noise.device) - for timestep in tqdm(timesteps, total=sample_steps - 1): + for timestep in tqdm(timesteps): t = (timestep / time_rescale_factor).expand(batch_size) x_cur = x_next flow_pred = self( diff --git a/fastgen/networks/cosmos_predict2/network.py b/fastgen/networks/cosmos_predict2/network.py index d27b661..7b3d026 100644 --- a/fastgen/networks/cosmos_predict2/network.py +++ b/fastgen/networks/cosmos_predict2/network.py @@ -13,10 +13,12 @@ """ from typing import Any, Dict, List, Optional, Set, Tuple, Union, Mapping +import inspect import os from tqdm.auto import tqdm from einops import rearrange +import numpy as np import torch import torch.nn as nn import torch.nn.functional as F @@ -54,6 +56,44 @@ # ---------------------- DiT Network ----------------------- +class FlowKarrasUniPCScheduler(UniPCMultistepScheduler): + """UniPC whose sigmas follow the official Cosmos Predict2.5 Karras ramp. + + A Karras schedule over ``[sigma_min, sigma_max]`` mapped into flow-matching + units by ``sigma / (sigma + 1)``. Built here rather than through diffusers' + ``use_karras_sigmas``, which only gained that conversion in 0.37.0. + """ + + # Hardcoded in diffusers' _convert_to_karras; left unconfigurable to match. + _rho: float = 7.0 + + def __init__(self, *args, sigma_min: float = 0.01, sigma_max: float = 200.0, **kwargs): + # diffusers >= 0.37 declares these itself. Forward them there so they are not + # recorded in the config's `_use_default_values`, which `from_config` discards; + # register them ourselves on older versions that lack the parameters. + base_params = inspect.signature(UniPCMultistepScheduler.__init__).parameters + forwarded = {k: v for k, v in (("sigma_min", sigma_min), ("sigma_max", sigma_max)) if k in base_params} + super().__init__(*args, **forwarded, **kwargs) + if not forwarded: + self.register_to_config(sigma_min=sigma_min, sigma_max=sigma_max) + + def set_timesteps( + self, num_inference_steps: Optional[int] = None, device: Union[str, torch.device] = None, **kwargs + ): + if kwargs.get("sigmas") is not None: + raise ValueError("FlowKarrasUniPCScheduler builds its own ramp; an explicit `sigmas` is unsupported.") + assert num_inference_steps is not None, "num_inference_steps is required" + super().set_timesteps(num_inference_steps=num_inference_steps, device=device, **kwargs) + ramp = np.linspace(0.0, 1.0, num_inference_steps) + min_inv_rho = self.config.sigma_min ** (1 / self._rho) + max_inv_rho = self.config.sigma_max ** (1 / self._rho) + sigmas = (max_inv_rho + ramp * (min_inv_rho - max_inv_rho)) ** self._rho + sigmas = sigmas / (sigmas + 1) + # sigmas stay on CPU to avoid per-step host/device traffic, as in diffusers. + self.sigmas = torch.from_numpy(np.concatenate([sigmas, [0.0]]).astype(np.float32)) + self.timesteps = torch.from_numpy(sigmas * self.config.num_train_timesteps).to(device=device, dtype=torch.int64) + + class CosmosPredict2DiT(nn.Module): """ Cosmos Predict2 DiT (Diffusion Transformer) for video generation. @@ -1100,7 +1140,7 @@ def sample( guidance_scale: Optional[float] = 5.0, num_steps: int = 50, skip_layers: Optional[List[int]] = None, - skip_layers_start_percent: float = 0.0, + skip_layers_start_fraction: float = 0.0, fps: Optional[torch.Tensor] = None, conditioning_latents: Optional[torch.Tensor] = None, num_conditioning_frames: int = 1, @@ -1131,7 +1171,8 @@ def sample( 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. + skip_layers_start_fraction: Fraction in [0, 1] of the sampling steps to complete + before skip-layer guidance becomes active. 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. @@ -1146,15 +1187,14 @@ def sample( """ assert self.schedule_type == "rf", f"{self.schedule_type} is not supported" - # 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. + # Match official Cosmos Predict2.5 inference: a Karras ramp over [0.01, 200] in + # flow-matching units, and `num_steps` denotes intervals, hence `num_steps + 1` + # ramp points. if self.sample_scheduler is None: - self.sample_scheduler = UniPCMultistepScheduler( + self.sample_scheduler = FlowKarrasUniPCScheduler( num_train_timesteps=1000, prediction_type="flow_prediction", use_flow_sigmas=True, - use_karras_sigmas=True, sigma_min=0.01, sigma_max=200.0, ) @@ -1189,7 +1229,7 @@ def sample( conditioning_latents_full = None condition_mask = None condition_mask_C = None - initial_noise = None + initial_latents = None if video2world_mode: B, C, T, H, W = latents.shape @@ -1211,19 +1251,16 @@ def sample( latents, v2w_condition ) - # Store initial noise for velocity replacement - initial_noise = noise + # Store the initial latents for velocity replacement + initial_latents = latents.clone() - total_steps = len(timesteps) - for step_idx, timestep in enumerate(tqdm(timesteps, total=total_steps, desc="Sampling")): + for step_idx, timestep in enumerate(tqdm(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( latents.dtype ) - active_skip_layers = skip_layers if (step_idx / total_steps) >= skip_layers_start_percent else None - if video2world_mode: # Replace conditioning frames with clean latents using preserve_conditioning v2w_condition = {"conditioning_latents": conditioning_latents, "condition_mask": condition_mask} @@ -1252,7 +1289,6 @@ def sample( cond_with_mask, fps=fps, conditional_frame_timestep=conditional_frame_timestep, - skip_layers=active_skip_layers, ) # Classifier-free guidance @@ -1263,13 +1299,14 @@ def sample( neg_cond_with_mask, fps=fps, conditional_frame_timestep=conditional_frame_timestep, - skip_layers=active_skip_layers, + skip_layers=skip_layers if step_idx >= skip_layers_start_fraction * len(timesteps) else None, ) velocity_pred = velocity_uncond + guidance_scale * (velocity_pred - velocity_uncond) - # Replace velocity for conditioning frames with analytical velocity: v = noise - x0 + # Replace velocity for conditioning frames with the constant velocity that carries + # the initial latents onto the conditioning frames over the interval [0, t_init]. if video2world_mode and denoise_replace_gt_frames: - gt_velocity = initial_noise - conditioning_latents_full + gt_velocity = (initial_latents - conditioning_latents_full) / t_init velocity_pred = gt_velocity * condition_mask_C + velocity_pred * (1 - condition_mask_C) # Keep clean frames in the DiT input while UniPC evolves raw latents.