feat(model): Add ER SDE / DPM++ 2M Scheduler Support For Anima#9125
Open
kappacommit wants to merge 42 commits intoinvoke-ai:mainfrom
Open
feat(model): Add ER SDE / DPM++ 2M Scheduler Support For Anima#9125kappacommit wants to merge 42 commits intoinvoke-ai:mainfrom
kappacommit wants to merge 42 commits intoinvoke-ai:mainfrom
Conversation
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…ch code Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…rminal-step test Address code review feedback on order-1 ER-SDE stepper: - Add docstring preconditions to integral helpers noting the logarithmic singularity at lam=0 (callers must guard sigma_next>0). - Tighten terminal-step test from atol=1e-5 to torch.equal — the algebra is exact when sigma_next=0, not approximate.
…ssertion Address code review feedback on the 2nd-order Taylor extension tests: - Assert state.old_d_x0 equals the analytically-computed d_x0 = 0.2v / (1.5 - 4.0) rather than just checking it's non-None. - Document that x_t is intentionally re-used across calls (state threading test, not trajectory correctness). - Document the order-2 correction coefficient magnitude that justifies the atol=1e-3 threshold in the engagement test.
… margin Address code review feedback on Task 3: - Comment why have_two_back checks both old_d_x0 and sigma_prev_prev (the sigma~=1 boundary path can break the joint invariant). - Document the analytically verified ~0.0004 per-element correction magnitude that justifies the atol=1e-3 threshold in the order-3 test.
…test Address code review feedback on Task 4: - Add an in-file comment above ANIMA_SCHEDULER_MAP explaining the convention: schedulers with custom code paths (er_sde) live in the Literal+labels only, not the map. - Hoist `import typing` to module-level in test_anima_schedulers.py (was inline-imported in two test functions). - Pin the er_sde label value (== "ER-SDE"), not just key existence.
…aming Address code review feedback on Task 5: - Comment explaining why fresh_noise is float32 (matches er_sde_rf_step's dtype contract with x_t.to(float32)). - Bridging comment at the inpaint extension call clarifying that sigma_next here means the same thing as sigma_prev in the Euler branch and the AnimaInpaintExtension API.
Address code review on Task 6 — er_sde was registered in the OpenAPI
schema but missing from the frontend's own Zod enums and Redux
PayloadAction types, so:
- The combobox dropdown didn't list it.
- setAnimaScheduler('er_sde') would fail TypeScript at the call site.
- Metadata recall for er_sde-generated images would silently no-op
(the scheduler value couldn't pass zParameterScheduler validation).
Changes:
- Add er_sde to zAnimaSchedulerField (the per-Anima Zod enum).
- Add er_sde to the animaScheduler state-shape Zod enum.
- Widen setAnimaScheduler's PayloadAction union.
- Add ER-SDE option to the ParamAnimaScheduler combobox.
- Make the metadata Scheduler handler accept ParameterAnimaScheduler too,
with a fallback parse and a narrowing guard before the SD/SDXL dispatch.
When step 0 goes through the sigma_curr=1 closed-form limit branch it writes state.sigma_prev_curr=1.0. On step 1, have_one_back was True and the 2nd-order path called _lambda(1.0) = 1.0/(1.0-1.0), crashing with ZeroDivisionError in every real denoise run. Fix: extend the have_one_back guard to require that sigma_prev_curr is more than _SIGMA_ONE_TOLERANCE below 1.0. The finite-difference derivative across the limit step is not meaningful, so skipping the 2nd-order term on that transition is correct. Order-3 is already gated behind old_d_x0 being set, which this path never sets, so no additional guard is needed there. Adds a regression test that runs the full sigma=1.0->0.9->0.7->0.0 sequence and asserts no ZeroDivisionError and all-finite output. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
ER-SDE solver (Cui et al., arXiv:2309.06169) usable across SD/SDXL (VP-SDE) and rectified-flow models. Anima migration follows in subsequent commits. - ERSDEScheduler(SchedulerMixin, ConfigMixin) with prediction_type (epsilon | v_prediction | flow_prediction), use_flow_sigmas, solver_order (1/2/3 with auto-warmup), and stochastic toggle - set_timesteps(sigmas=) for pre-shifted Anima/FLUX/Z-Image schedules - Closed-form limit at sigma=1 in flow mode - Unit tests + VP smoke + 5/5 Anima parity vs er_sde_rf_step (worst delta 5.137e-07) - Frontend wiring: zSchedulerField, SCHEDULER_OPTIONS, OpenAPI regen - parsing.tsx cleanup: removes the AnyScheduler widening since er_sde is now a first-class general scheduler Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Adds er_sde to the standard scheduler dispatch map with rectified-flow kwargs (flow_prediction, use_flow_sigmas=True, flow_shift=3.0, solver_order=3, stochastic=True). Anima still routes through the legacy elif is_er_sde: branch — that's removed in a follow-up commit. This is the additive prerequisite that lets the cutover happen without a window where Anima can't use ER-SDE. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Verifies the ANIMA_SCHEDULER_MAP['er_sde'] entry instantiates correctly, accepts pre-shifted sigmas via set_timesteps(sigmas=...), and resets multistep state. Catches wiring regressions that the algorithm-level parity test cannot. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
…_SCHEDULER_MAP Anima ER-SDE now flows through the same standard scheduler path as dpmpp_2m_sde — pre-shifted sigmas via scheduler.set_timesteps(sigmas=...), inpaint extension via inpaint_extension.merge_intermediate_latents_with_init_latents, step_callback per-step. The custom code path was the only thing keeping ER-SDE off the universal pipeline. Bumps invocation version to 1.5.0. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
ERSDEScheduler is now the production code path. er_sde_rf_step is kept as the comparison oracle for the scheduler's parity test, and as a self-contained mathematical reference for the rectified-flow algebra. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
ERSDEScheduler is the production code path. The pure-function helper was retained as a parity oracle but YAGNI — keeping ~200 lines of code purely as a regression net for hypothetical future drift isn't worth the maintenance signal it generates. Removes: - invokeai/backend/rectified_flow/er_sde.py - tests/backend/rectified_flow/test_er_sde.py - tests/backend/rectified_flow/test_er_sde_scheduler_anima_parity.py ERSDEScheduler's own tests (test_er_sde_scheduler.py) remain — they exercise both VP-SDE and rectified-flow paths directly. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Windows-side typegen run flipped these to backslashes. Restore the canonical forward-slash form to match upstream. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
kappacommit
requested review from
JPPhoto,
Pfannkuchensack,
blessedcoolant,
dunkeroni and
lstein
as code owners
github-actions
Bot
added
python
Sort imports + format per project ruff config (line-length 120). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
The animaScheduler field inlined its enum (originally to add er_sde
when the shared schema didn't have it yet). Now that zAnimaSchedulerField
already includes er_sde, reference the shared zParameterAnimaScheduler
to match the pattern used by scheduler/fluxScheduler/zImageScheduler.
Drops the redundant .default('euler') — initial value comes from
getInitialParamsState.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
kappacommit
changed the title
5 tasks
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
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
1 participant
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.
Summary
ER SDE is the primary recommended scheduler for Anima. After testing, it produces the best and most coherent outputs I've seen for the model.
Structurally, the easiest way to add support for ER SDE for Anima, was to also add support for SD (1.5, 2, SDXL). So this PR adds support for Anima and SDXL.
ER SDE is not available in diffusers, so a custom implementation was required.
Reference implementation was here: https://github.com/QinpengCui/ER-SDE-Solver this is where all of the fancy math and logic comes from (MIT licensed, free to use)
I also add support for DPM++ 2M Scheduler
Related Issues / Discussions
Requested in discord https://discord.com/channels/1020123559063990373/1149506274971631688/1501016017009246379
QA Instructions
--
--
Merge Plan
This is built on top of #9123 and should only be merged after that.
Checklist
What's Newcopy (if doing a release after this PR)