Skip to content

Fix DDIMInverseScheduler inf at first inverse step with prediction_type=sample - #14822

Open
DawnofGenX wants to merge 1 commit into
huggingface:mainfrom
DawnofGenX:fix/ddim-inverse-zero-beta
Open

DawnofGenX wants to merge 1 commit into
huggingface:mainfrom
DawnofGenX:fix/ddim-inverse-zero-beta

Conversation

@DawnofGenX

Copy link
Copy Markdown

What does this PR do?

Fixes #10920

DDIMInverseScheduler.step() produces inf values in prev_sample at the first inverse step when prediction_type="sample" and set_alpha_to_one=True.

Root cause: with set_alpha_to_one=True, alpha_prod_t is exactly 1.0 at the first inverse step, so beta_prod_t = 1 - alpha_prod_t is 0.0. The sample-prediction branch computes

pred_epsilon = (sample - alpha_prod_t ** 0.5 * model_output) / beta_prod_t ** 0.5

which divides by zero → inf (or nan when the numerator is also 0). Every subsequent inverse step is then poisoned.

Fix: guard the division. At a zero noise level the sample carries no noise, so pred_epsilon is not identifiable from sample + the x0 prediction; we fall back to zero noise instead of dividing by zero. This only changes behavior in the previously-divergent degenerate case (beta_prod_t == 0); all non-degenerate configs compute exactly the same values as before. This mirrors how the schedulers already guard divide-by-zero (e.g. torch.clamp(variance, min=1e-20) in scheduling_ddpm.py).

Regression test

tests/schedulers/test_scheduler_ddim_inverse.py::test_sample_prediction_first_step_no_inf uses the issue's exact repro (tiny CPU tensors, seeded) and additionally checks a full prediction_type="sample" inverse loop stays finite.

Before the fix (issue repro):

out finite: False | has inf: True
AssertionError: FAILED: inf in prev_sample

After the fix:

out finite: True | has inf: False | has nan: False
PASSED

Test results (CPU):

  • pytest tests/schedulers/test_scheduler_ddim_inverse.py -q35 passed, 1 skipped (baseline before change: 34 passed, 1 skipped; +1 = the new regression test)
  • pytest tests/schedulers/ -q -k ddim → 120 passed, 1 skipped; the 2 failing DDIMParallelSchedulerTest::test_compatibles cases fail identically on unmodified main (pre-existing, unrelated to this change)

Who can review?

Anyone in the community is free to review the PR until the process has been found stable.

…tion

With prediction_type="sample" and set_alpha_to_one=True, alpha_prod_t is
exactly 1.0 at the first inverse step, so beta_prod_t is 0 and the
pred_epsilon division by beta_prod_t ** 0.5 overflows to inf. Guard the
division and treat pred_epsilon as unidentifiable (zero) at zero noise
level, matching the convention used elsewhere for degenerate alphas.

Fixes huggingface#10920
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.

[DDIMInverseScheduler] inf values at first iteration when set_alpha_to_one=True and prediction_type="sample"

1 participant