Skip to content

Commit aca5cea

Browse files
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/1
1 parent 709e4cd commit aca5cea

3 files changed

Lines changed: 742 additions & 0 deletions

File tree

max/python/max/pipelines/lib/diffusion_schedulers/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,3 +15,4 @@
1515
from .scheduling_flow_match_euler_discrete import (
1616
FlowMatchEulerDiscreteScheduler,
1717
)
18+
from .scheduling_unipc_multistep import UniPCMultistepScheduler

max/python/max/pipelines/lib/diffusion_schedulers/scheduler_factory.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,11 @@
1818
from .scheduling_flow_match_euler_discrete import (
1919
FlowMatchEulerDiscreteScheduler,
2020
)
21+
from .scheduling_unipc_multistep import UniPCMultistepScheduler
2122

2223
_SCHEDULER_REGISTRY: dict[str, type] = {
2324
"FlowMatchEulerDiscreteScheduler": FlowMatchEulerDiscreteScheduler,
25+
"UniPCMultistepScheduler": UniPCMultistepScheduler,
2426
}
2527

2628

0 commit comments

Comments
 (0)