fix: Use Karras sigma schedule for Cosmos Predict2.5 SFT inference#33
fix: Use Karras sigma schedule for Cosmos Predict2.5 SFT inference#33mvanhorn wants to merge 1 commit into
Conversation
Greptile SummaryThis PR fixes degraded Cosmos Predict2.5 SFT inference quality by overriding the default linear sigma schedule with the Karras schedule (σ_max=200, σ_min=0.01, ρ=7) actually used by the official Cosmos codebase, and aligns the inference script example commands with
Confidence Score: 3/5The core sigma-schedule logic is correct and the test coverage is solid, but the sigmas tensor is left on CPU after the Karras override, which will cause a device error the first time scheduler.step() is called on a GPU. The Karras schedule math and state-reset are correctly implemented and well tested. However, _build_cosmos_predict2_karras_schedule never moves sigmas to device: it calls .to(dtype=torch.float32) but omits device=device, while the timesteps tensor is correctly moved. After the Karras override, scheduler.sigmas is a CPU tensor and scheduler.timesteps is a CUDA tensor, undoing the device placement that set_timesteps(device=noise.device) had just established. Every GPU inference run will hit a RuntimeError in scheduler.step(). The tests all use torch.device('cpu') so this is not caught. fastgen/networks/cosmos_predict2/network.py — specifically _build_cosmos_predict2_karras_schedule (sigmas device placement) and the use_karras_sigma_schedule=True default which may affect distillation configs. Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant C as Caller
participant S as CosmosPredict2.sample()
participant K as _apply_cosmos_predict2_karras_schedule()
participant B as _build_cosmos_predict2_karras_schedule()
participant U as UniPCMultistepScheduler
C->>S: "sample(noise, condition, num_steps=35)"
S->>U: "set_timesteps(35, device=cuda)"
Note over U: sigmas on cuda, timesteps on cuda
S->>K: apply schedule(scheduler, 35, cuda)
K->>B: build schedule(35, 1000, cuda)
Note over B: sigmas computed on CPU
Note over B: timesteps moved to cuda
Note over B: sigmas NOT moved to cuda
B-->>K: sigmas (CPU), timesteps (cuda)
K->>U: "scheduler.sigmas = CPU tensor"
K->>U: "scheduler.timesteps = cuda tensor"
loop 36 denoising steps
S->>U: step(velocity_pred, timestep, sample)
Note over U: self.sigmas[i] is CPU
Note over U: RuntimeError on GPU
end
S-->>C: denoised latents
%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
participant C as Caller
participant S as CosmosPredict2.sample()
participant K as _apply_cosmos_predict2_karras_schedule()
participant B as _build_cosmos_predict2_karras_schedule()
participant U as UniPCMultistepScheduler
C->>S: "sample(noise, condition, num_steps=35)"
S->>U: "set_timesteps(35, device=cuda)"
Note over U: sigmas on cuda, timesteps on cuda
S->>K: apply schedule(scheduler, 35, cuda)
K->>B: build schedule(35, 1000, cuda)
Note over B: sigmas computed on CPU
Note over B: timesteps moved to cuda
Note over B: sigmas NOT moved to cuda
B-->>K: sigmas (CPU), timesteps (cuda)
K->>U: "scheduler.sigmas = CPU tensor"
K->>U: "scheduler.timesteps = cuda tensor"
loop 36 denoising steps
S->>U: step(velocity_pred, timestep, sample)
Note over U: self.sigmas[i] is CPU
Note over U: RuntimeError on GPU
end
S-->>C: denoised latents
Reviews (1): Last reviewed commit: "fix: Use Karras sigma schedule for Cosmo..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| 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) |
There was a problem hiding this comment.
sigmas never moved to device, breaking GPU inference
ramp is always created on CPU via torch.arange(...), so all derived sigmas computations stay on CPU. timesteps is correctly moved with .to(device=device, ...), but the final sigmas tensor that is returned (and ultimately stored in scheduler.sigmas) stays on CPU. When scheduler.step() is called during GPU inference it will index into the CPU sigmas tensor and use the resulting CPU 0-dim tensors in arithmetic with CUDA latent tensors, raising a RuntimeError: Expected all tensors to be on the same device. The prior set_timesteps(device=noise.device) call correctly places sigmas on device — this Karras override silently undoes that.
| sigmas = torch.cat([sigmas, sigmas.new_zeros(1)]).to(dtype=torch.float32) | |
| timesteps = (sigmas * num_train_timesteps).to(device=device, dtype=torch.int64) | |
| sigmas = torch.cat([sigmas, sigmas.new_zeros(1)]).to(dtype=torch.float32, device=device) | |
| return sigmas, timesteps |
| use_wan_fp32_strategy: bool = True, | ||
| # FPS for temporal position embeddings | ||
| fps: float = 24.0, | ||
| use_karras_sigma_schedule: bool = True, |
There was a problem hiding this comment.
True may silently change distillation inference
The PR description says to "gate this behind a sensible default or config flag so it only affects the SFT inference path, not distillation." Setting use_karras_sigma_schedule=True as the class-level default means every CosmosPredict2 instance — including distillation models — will use the Karras schedule unless the corresponding experiment config explicitly passes use_karras_sigma_schedule=False. If the distillation configs don't set that flag, evaluation of distilled checkpoints will silently start using a different sigma schedule, making before/after comparisons invalid.
| ) | ||
| scheduler.sigmas = sigmas | ||
| scheduler.timesteps = timesteps | ||
| scheduler.num_inference_steps = len(timesteps) |
There was a problem hiding this comment.
num_inference_steps set to num_steps + 1, denoising loop runs one extra iteration
_build_cosmos_predict2_karras_schedule uses torch.arange(num_steps + 1) which produces num_steps + 1 ramp points (including both endpoints 0 and 1). This yields a timesteps tensor of length num_steps + 1, so scheduler.num_inference_steps is set to num_steps + 1 and the denoising loop in sample() iterates num_steps + 1 times. When a user passes --num_steps 35, they will actually get 36 denoising steps. If this matches the official Cosmos reference implementation this is intentional, but it should be documented explicitly.
|
✅ Action performedReview finished.
|
|
Hi @mvanhorn, Thanks a lot for the fix! Is there a reason why we can't just use |
WalkthroughCosmos Predict2 sampling now supports an optional Karras sigma schedule, resets UniPC scheduler state before sampling, and uses the resolved timestep count for progress. The Cosmos inference examples and tests were updated to match the new schedule and CLI arguments. ChangesCosmos Predict2 Karras Sampling
Sequence Diagram(s)sequenceDiagram
participant Sample as "CosmosPredict2.sample"
participant Scheduler as "UniPCMultistepScheduler"
participant Apply as "_apply_cosmos_predict2_karras_schedule"
participant Progress as "tqdm"
Sample->>Scheduler: set_timesteps(num_steps, device)
Sample->>Apply: apply Karras sigma schedule
Apply->>Scheduler: assign sigmas, timesteps, and reset step state
Sample->>Progress: iterate with total=len(timesteps)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@fastgen/networks/cosmos_predict2/network.py`:
- Line 819: The shared CosmosPredict2 network default for
use_karras_sigma_schedule should not be enabled in the base sample path, because
CosmosPredict2_2B_Config and CosmosPredict2_14B_Config are reused by non-SFT
configs like DMD2 and video2world. Update CosmosPredict2.sample() and the
relevant network/config definitions so the shared default stays False, then set
use_karras_sigma_schedule to True only in the SFT-specific config path in
config_sft.py, keeping non-SFT callers on the existing behavior unless they
explicitly override it.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Enterprise
Run ID: 4e4c1b73-0452-4878-a8c4-382cb661fece
📒 Files selected for processing (3)
fastgen/networks/cosmos_predict2/network.pyscripts/inference/video_model_inference.pytests/test_network.py
| use_wan_fp32_strategy: bool = True, | ||
| # FPS for temporal position embeddings | ||
| fps: float = 24.0, | ||
| use_karras_sigma_schedule: bool = True, |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Inspect CosmosPredict2 call sites and config overrides for the new schedule flag.
rg -n -C2 'use_karras_sigma_schedule|CosmosPredict2|config_sft|distill|sample\(' .Repository: NVlabs/FastGen
Length of output: 50371
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the CosmosPredict2 constructor, sample method, and config usages.
for f in \
fastgen/networks/cosmos_predict2/network.py \
fastgen/configs/net.py \
fastgen/configs/experiments/CosmosPredict2/config_sft.py \
fastgen/configs/experiments/CosmosPredict2/config_sft_14b.py \
fastgen/configs/experiments/CosmosPredict2/config_sft_v2w.py \
fastgen/configs/experiments/CosmosPredict2/config_dmd2.py \
fastgen/configs/experiments/CosmosPredict2/config_dmd2_14b.py \
fastgen/configs/experiments/CosmosPredict2/config_dmd2_v2w.py \
scripts/inference/video_model_inference.py
do
if [ -f "$f" ]; then
echo "===== $f ====="
wc -l "$f"
fi
done
echo "===== network.py outline ====="
ast-grep outline fastgen/networks/cosmos_predict2/network.py --view expanded || true
echo "===== network.py relevant slices ====="
sed -n '740,930p' fastgen/networks/cosmos_predict2/network.py
echo "===== net.py cosmos slice ====="
sed -n '270,310p' fastgen/configs/net.py
echo "===== Cosmos SFT configs ====="
sed -n '1,120p' fastgen/configs/experiments/CosmosPredict2/config_sft.py
printf '\n'
sed -n '1,120p' fastgen/configs/experiments/CosmosPredict2/config_sft_14b.py
printf '\n'
sed -n '1,120p' fastgen/configs/experiments/CosmosPredict2/config_sft_v2w.py
printf '\n'
sed -n '1,160p' fastgen/configs/experiments/CosmosPredict2/config_dmd2.py
printf '\n'
sed -n '1,160p' fastgen/configs/experiments/CosmosPredict2/config_dmd2_14b.py
printf '\n'
sed -n '1,160p' fastgen/configs/experiments/CosmosPredict2/config_dmd2_v2w.py
echo "===== inference call sites around model.sample ====="
sed -n '660,720p' scripts/inference/video_model_inference.pyRepository: NVlabs/FastGen
Length of output: 24679
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Focused searches for explicit schedule overrides and direct CosmosPredict2 usage.
rg -n -C3 'use_karras_sigma_schedule\s*=|use_karras_sigma_schedule|CosmosPredict2\(' fastgen scripts testsRepository: NVlabs/FastGen
Length of output: 4873
Move the Karras default to the SFT config path. CosmosPredict2_2B_Config / CosmosPredict2_14B_Config are reused by DMD2 and video2world configs unchanged, so CosmosPredict2.sample() still sends non-SFT callers down the Karras path unless they override it explicitly. Keep the shared network default False and set it to True only in fastgen/configs/experiments/CosmosPredict2/config_sft.py.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@fastgen/networks/cosmos_predict2/network.py` at line 819, The shared
CosmosPredict2 network default for use_karras_sigma_schedule should not be
enabled in the base sample path, because CosmosPredict2_2B_Config and
CosmosPredict2_14B_Config are reused by non-SFT configs like DMD2 and
video2world. Update CosmosPredict2.sample() and the relevant network/config
definitions so the shared default stays False, then set
use_karras_sigma_schedule to True only in the SFT-specific config path in
config_sft.py, keeping non-SFT callers on the existing behavior unless they
explicitly override it.
2 participants
Summary
Cosmos-Predict2.5-2B video2world inference under FastGen produces noticeably degraded output compared to the official cosmos-predict2.5 codebase when using the same checkpoint, prompt, image, resolution, and step count. The reporter traced the primary cause to a wrong sigma/timestep schedule: the official Cosmos uses a Karras sigma schedule (sigma_max=200, sigma_min=0.01, rho=7) in its FlowUniPCMultistepScheduler.set_timesteps(use_kerras_sigma=True), while FastGen's CosmosPredict2DiT.sample() relies on the default diffusers flow-shift linear schedule from UniPCMultistepScheduler.
Changes
In fastgen/networks/cosmos_predict2/network.py, modify the sample() method (around lines 1148-1161) so that, for the Cosmos SFT sampling path, after constructing the UniPCMultistepScheduler and calling set_timesteps, it overrides sample_scheduler.sigmas and sample_scheduler.timesteps with a Karras sigma schedule (sigma_max=200, sigma_min=0.01, rho=7) exactly as the official Cosmos does, including the sigmas/(1+sigmas) reparameterization, int64 timestep truncation, and resetting the scheduler's internal solver state (model_outputs, lower_order_nums, last_sample, _step_index, _begin_index). Gate this behind a sensible default or config flag so it only affects the SFT inference path the maintainer called out, not distillation. Update the inline example command(s) in scripts/inference/video_model_inference.py to pass the Cosmos negative-prompt file and otherwise align them with scripts/README.md. Add a focused unit test in tests/test_network.py asserting the produced sigma/timestep schedule matches the reference Karras values.
Fixes #18
Summary by CodeRabbit