Fully Modular Schedulers #11158

Beinsezii · 2025-03-26T09:20:38Z

Beinsezii
Mar 26, 2025

I've been working on a library for modular samplers as a side project over at skrample, which started as a simple way to use any sampler on Flux but has gained some extra features since then.

It is fully generic and extensible using traits/abc, with the only hard dependency being numpy.

Of particular interest here is the skrample.diffusers module, which contains a wrapper class that presents the underlying skrample components in a way that's 100% compatible with every diffusers pipeline I was able to run in quickdif; no more having FluxPipeline yeeting errors on anything that's not FlowMatchEulerDiscreteScheduler.

For example,

from skrample.diffusers import SkrampleWrapperScheduler
from skrample.pytorch.noise import Brownian
from skrample.sampling import DPM, FLOW
from skrample.scheduling import Beta, Flow

pipeline: FluxPipeline = ...

pipeline.scheduler = SkrampleWrapperScheduler(
    sampler=DPM(order=2, add_noise=True, predictor=FLOW),
    schedule=Beta(Flow(shift=3)),
)

Should set up scheduler such that it should be roughly equivalent to

DPMSolverMultistepScheduler(
    algorithm_type="sde-dpmsolver++",
    flow_shift=3,
    prediction_type="flow_prediction",
    solver_order=2,
    use_beta_sigmas=True,
)

And will produce clean images as expected

But maybe you'd prefer the Brownian stochastic noise from torchsde as found in diffusers' DPMSolverSDEScheduler

Let's also be lazy and just re-use the perfectly good pipeline.scheduler.config, because knowing whether a model is epsilon/vpred/zsnr/flow is annoying.

pipeline.scheduler = SkrampleWrapperScheduler.from_diffusers_config(
    sampler=DPM,
    sampler_props={"order": 2},  # Takes priority over the diffusers config
    schedule_modifier=Beta,
    noise_type=Brownian,
    **pipeline.scheduler.config  # Also reads shift value
)

And now we have a brownian kitten

Overall, I've been quite happy playing with it on quickdif, and since I just tagged the first 0.1.0, I thought it would be fun to share what I've made for the three people that read this.