fix(network): sample ignores skip_layers and (+1 more)

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(

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)
 

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(

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)
 

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(

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,10 +1251,10 @@ def sample(
                 latents, v2w_condition
             )
 
-            # Store initial noise for velocity replacement
-            initial_noise = latents.clone()
+            # Store the initial latents for velocity replacement
+            initial_latents = latents.clone()
 
-        for timestep in tqdm(timesteps, total=len(timesteps), 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(
@@ -1259,12 +1299,14 @@ def sample(
                     neg_cond_with_mask,
                     fps=fps,
                     conditional_frame_timestep=conditional_frame_timestep,
+                    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.