Skip to content

Fix NaN samples from DDPM fixed_large_log variance - #14571

Open
lllakshit wants to merge 1 commit into
huggingface:mainfrom
lllakshit:fix/ddpm-fixed-large-log-nan
Open

Fix NaN samples from DDPM fixed_large_log variance#14571
lllakshit wants to merge 1 commit into
huggingface:mainfrom
lllakshit:fix/ddpm-fixed-large-log-nan

Conversation

@lllakshit

Copy link
Copy Markdown

What this PR does

DDPMScheduler._get_variance(..., variance_type="fixed_large_log") (Glide max_log) returns log(beta). step() treated that value like a linear variance and took sqrt(), which is NaN for beta < 1. The same path exists in DDPMParallelScheduler.

Sampling now converts log variance with exp(0.5 * log), the same transform already used for learned_range. That is equivalent to sqrt(beta) / fixed_large, so the intended Improved-DDPM / Glide scale is restored and samples stay finite.

Fixes #14569

Tests

  • tests/schedulers/test_scheduler_ddpm.py::DDPMSchedulerTest::test_fixed_large_log_sampling_is_finite_and_matches_fixed_large
  • tests/schedulers/test_scheduler_ddpm_parallel.py::DDPMParallelSchedulerTest::test_fixed_large_log_sampling_is_finite_and_matches_fixed_large

Both assert fixed_large_log samples are finite and match fixed_large at the same seed.

Self-review

  • Confirmed on current main with the issue reproduction: both schedulers produced NaNs at t=500 before the change.
  • _get_variance still returns log-space values for fixed_large_log (Glide max_log). Only the step() noise scale is corrected, so callers of _get_variance keep the old contract.
  • batch_step_no_noise does not add variance, so it did not need a change.
  • Existing fixed_small / fixed_large / fixed_small_log / learned_range branches are unchanged.

_get_variance(fixed_large_log) returns log(beta). step() treated that
value like a variance and took sqrt(), which is NaN for beta < 1.

Convert log variance with exp(0.5 * log) for both DDPMScheduler and
DDPMParallelScheduler, matching learned_range and the fixed_large scale.

Fixes huggingface#14569
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

schedulers size/S PR with diff < 50 LOC tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fixed_large_log variance sampling returns NaNs

1 participant