Fix DDIMInverseScheduler inf at first inverse step with prediction_type=sample - #14822
Open
DawnofGenX wants to merge 1 commit into
Open
DawnofGenX wants to merge 1 commit into
DawnofGenX wants to merge 1 commit into
Conversation
…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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Fixes #10920
DDIMInverseScheduler.step()producesinfvalues inprev_sampleat the first inverse step whenprediction_type="sample"andset_alpha_to_one=True.Root cause: with
set_alpha_to_one=True,alpha_prod_tis exactly1.0at the first inverse step, sobeta_prod_t = 1 - alpha_prod_tis0.0. Thesample-prediction branch computeswhich divides by zero →
inf(ornanwhen the numerator is also0). Every subsequent inverse step is then poisoned.Fix: guard the division. At a zero noise level the sample carries no noise, so
pred_epsilonis not identifiable fromsample+ 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)inscheduling_ddpm.py).Regression test
tests/schedulers/test_scheduler_ddim_inverse.py::test_sample_prediction_first_step_no_infuses the issue's exact repro (tiny CPU tensors, seeded) and additionally checks a fullprediction_type="sample"inverse loop stays finite.Before the fix (issue repro):
After the fix:
Test results (CPU):
pytest tests/schedulers/test_scheduler_ddim_inverse.py -q→ 35 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 failingDDIMParallelSchedulerTest::test_compatiblescases fail identically on unmodifiedmain(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.