Commit aca5cea
committed
[MAX] Add UniPC multistep scheduler for Wan diffusion
## Summary
Add a numpy-only UniPC multistep scheduler for Wan diffusion pipelines.
## Description
- Implements the UniPC-BH2 algorithm with corrector and predictor steps
- Supports flow-matching sigma schedules (used by Wan 2.1/2.2)
- Provides `build_step_coefficients()` to precompute per-step coefficient matrices on the host, enabling on-device scheduler steps without Python-side numpy calls during denoising
- Registers `UniPCMultistepScheduler` in the diffusion scheduler factory
This is a numpy-only port of the diffusers `UniPCMultistepScheduler`, specialized for the Wan pipeline configuration.
## Dependencies
None — can be merged independently.
## Checklist
- [x] PR is small and focused
- [x] I ran `./bazelw run format` to format my changes
Assisted-by: Claude Code
Assisted-by: Claude Code
stack-info: PR: #13, branch: jglee-sqbits/stack/11 parent 709e4cd commit aca5cea
3 files changed
Lines changed: 742 additions & 0 deletions
File tree
- max/python/max/pipelines/lib/diffusion_schedulers
Lines changed: 1 addition & 0 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
15 | 15 | | |
16 | 16 | | |
17 | 17 | | |
| 18 | + | |
Lines changed: 2 additions & 0 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
18 | 18 | | |
19 | 19 | | |
20 | 20 | | |
| 21 | + | |
21 | 22 | | |
22 | 23 | | |
23 | 24 | | |
| 25 | + | |
24 | 26 | | |
25 | 27 | | |
26 | 28 | | |
| |||
0 commit comments