Fully Modular Schedulers

[Beinsezii]

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.

[MunazaAshraf10]

insane!! excited to try it

[yiyixuxu]

hi @Beinsezii
thanks for sharing! super cool

[yiyixuxu]

cc @stevhliu
does it make sense to add a mention and link to this somewhere in the doc?

[Beinsezii]

There's no pypi package yet but it's usable with the git syntax like
pip install skrample[all]@git+https://github.com/Beinsezii/[email protected]

the master branch is mostly covered under both highest and lowest resolution CI but I still recommend pinning anything from git

[stevhliu]

Nice project, thanks for sharing!

I think you can add this to the https://huggingface.co/docs/diffusers/main/en/community_projects doc

[Beinsezii]

Ok I put it on pypi now for v0.2 https://pypi.org/project/skrample/

[Beinsezii]

Did you want me to add it to the doc or were one of you going to?