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

[andrewwhitecdw]

Small fixes in fastgen/networks/cosmos_predict2/network.py:

fix: sample ignores skip_layers and skip_layers_start_percent args

Fix: Replace:

        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(
                latents.dtype
            )

            if video2world_mode:
                # Replace conditioning frames with clean latents using preserve_conditioning
                v2w_condition = {"conditioning_latents": conditioning_latents, "condition_mask": condition_mask}
                model_input = self.preserve_conditioning(latents, v2w_condition)

                # Wrap condition with mask for forward() to use
                cond_with_mask = {
                    "text_embeds": condition,
                    "conditioning_latents": conditioning_latents,
                    "condition_mask": condition_mask,
                }
                neg_cond_with_mask = {
                    "text_embeds": neg_condition,
                    "conditioning_latents": conditioning_latents,
                    "condition_mask": condition_mask,
                }
            else:
                model_input = latents
                cond_with_mask = condition
                neg_cond_with_mask = neg_condition

            # Forward pass
            velocity_pred = self(
                model_input,
                t,
                cond_with_mask,
                fps=fps,
                conditional_frame_timestep=conditional_frame_timestep,
            )

            # Classifier-free guidance
            if guidance_scale > 1.0:
                velocity_uncond = self(
                    model_input,
                    t,
                    neg_cond_with_mask,
                    fps=fps,
                    conditional_frame_timestep=conditional_frame_timestep,
                )
                velocity_pred = velocity_uncond + guidance_scale * (velocity_pred - velocity_uncond)

with:

        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}
                model_input = self.preserve_conditioning(latents, v2w_condition)

                # Wrap condition with mask for forward() to use
                cond_with_mask = {
                    "text_embeds": condition,
                    "conditioning_latents": conditioning_latents,
                    "condition_mask": condition_mask,
                }
                neg_cond_with_mask = {
                    "text_embeds": neg_condition,
                    "conditioning_latents": conditioning_latents,
                    "condition_mask": condition_mask,
                }
            else:
                model_input = latents
                cond_with_mask = condition
                neg_cond_with_mask = neg_condition

            # Forward pass
            velocity_pred = self(
                model_input,
                t,
                cond_with_mask,
                fps=fps,
                conditional_frame_timestep=conditional_frame_timestep,
                skip_layers=active_skip_layers,
            )

            # Classifier-free guidance
            if guidance_scale > 1.0:
                velocity_uncond = self(
                    model_input,
                    t,
                    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)

fix: sample stores scaled latents as initial noise for v2w velocity

Fix: Replace:

            # Store initial noise for velocity replacement
            initial_noise = latents.clone()

with:

            # Store initial noise for velocity replacement
            initial_noise = noise

Files changed

• fastgen/networks/cosmos_predict2/network.py

[greptile-apps]

Greptile Summary

The PR adjusts sampling-loop progress reporting, renames the skip-layer activation parameter, and updates Cosmos Predict2 scheduling and conditioned-frame velocity handling.

• Uses scheduler-derived timestep lengths in progress bars.
• Adds a Cosmos-specific flow-matching Karras scheduler.
• Applies skip-layer guidance according to a fractional sampling threshold.
• Reworks conditioned-frame velocity replacement around the initial latent state.

Confidence Score: 4/5

The PR is not yet safe to merge because the previously reported legacy keyword remains silently ignored.

The reply says commit 3ae2c86 added a backward-compatibility alias, but current HEAD still declares only skip_layers_start_fraction in all three samplers and leaves **kwargs to absorb skip_layers_start_percent; an existing caller using the old keyword therefore retains the default zero threshold and activates guidance from the first step.

Files Needing Attention: fastgen/networks/Wan/network.py, fastgen/networks/WanI2V/network.py, fastgen/networks/cosmos_predict2/network.py

Important Files Changed

Filename	Overview
fastgen/networks/Wan/network.py	Renames the skip-layer activation argument and updates timestep iteration, but the claimed compatibility mapping is absent.
fastgen/networks/WanI2V/network.py	Mirrors the Wan sampling-loop and skip-layer parameter changes, including the absent legacy-keyword mapping.
fastgen/networks/cosmos_predict2/network.py	Adds the custom Karras scheduler and revises skip-layer and conditioned-frame sampling behavior, while retaining the legacy-keyword compatibility problem.
fastgen/networks/VaceWan/network_causal.py	Removes an inaccurate explicit progress total and lets tqdm derive it from the timestep sequence.
fastgen/networks/Wan/network_causal.py	Lets tqdm derive sampling progress directly from the timestep sequence.
fastgen/networks/WanI2V/network_causal.py	Lets tqdm derive sampling progress directly from the timestep sequence.

_{Reviews (5): Last reviewed commit: "fix(cosmos): correct skip-layer guidance..." | Re-trigger Greptile}

[greptile-apps]

Old guidance keyword is ignored

When an existing caller passes skip_layers_start_percent, **kwargs silently consumes it while skip_layers_start_fraction remains 0.0, causing skip-layer guidance to activate from the first sampling step instead of the requested point. Preserve the old keyword as an alias or reject it explicitly; the same compatibility break exists in the WanI2V and Cosmos Predict2 samplers.

[andrewwhitecdw]

Addressed in 3ae2c86. Added a backward-compatibility alias so skip_layers_start_percent is honored in all three samplers (Wan, WanI2V, Cosmos Predict2) by mapping it to skip_layers_start_fraction and popping it from **kwargs before sampling. The docstrings now also note the deprecated alias.

[andrewwhitecdw]

Squashed to a single commit (dabb2be) with DCO Signed-off-by: Andrew White <andrewh@cdw.com> and a concise bulleted commit body.