Functional Sampling for Modular Diffusers

[Beinsezii]

For some reason Diffusers has historically made all of the sampling methods inside-out, where the solver cannot actually access the ODE itself, and is instead fed a fixed stream of predictions one at a time. While this is fine for the Euler method and linear multi-step methods, it makes higher order single-step methods extremely challenging, and variable step methods might even be impossible.

Effectively, diffusers does

for t in timesteps:
  output = model(sample, t) # This is outside of the sampler. The sampler class is inside-out.
  sample = sampler.step(sample, output, t)

instead of

sample = sampler.sample_model(sample, model, timesteps) # sample_model will call model(x, t) as many times as it needs to

It's not terribly difficult to rewrite the denoise loop in a pipeline but doing so for all hundred-something pipes in the library is extremely impractical.

With the recent efforts to break up the monolithic pipelines, I've found it was actually fairly easy to convert the modular denoise wrapper block to a system where the sampler is responsible for the denoise loop itself. As part of my efforts to add single-step methods to my library, I've written a small demonstrator using Flux.1 and Runge-Kutta which seems a lot more reasonable.

However, even if the code is straightforward, doing this for every model family will still become very tedious once more architectures are made modular. Therefore I think it would be significantly better to move Modular Diffusers to a system where the sampler is a functional/stateless system like you would typically expect from a differential solver, especially since the project is still young. It is easy to make a one-size-fits-all adapter to map any current inside-out sampler/scheduler into a functional one for backwards-compatibility, so to me it makes sense to migrate.

Implementation details aside, I'm wondering foremost if this is something that's even being considered? I don't know why it was done this way in the first place, as I believe just about all other diffusion applications write sampling the traditional way.

cc @yiyixuxu probably as someone who works on Modular a lot

---

• Public interest
  • Implementation of DPMSolverSinglestepScheduler #7470
  • Flow Matching RK4 Discreet Scheduler #12054
  • Res samplers? Beinsezii/skrample#52
• My library https://github.com/Beinsezii/skrample, Fully Modular Schedulers #11158
  • Functional sampling Refactor sampling module with new functional submodule Beinsezii/skrample#53
  • SD1.5 RK4 https://github.com/Beinsezii/skrample/blob/ce9f845d4ba750a097edf044e76cac92fe68f2be/examples/functional.py
  • Modular Flux RK4 https://github.com/Beinsezii/skrample/blob/ce9f845d4ba750a097edf044e76cac92fe68f2be/examples/diffusers/functional.py
  • Inside-out sampler adapter https://github.com/Beinsezii/skrample/blob/ce9f845d4ba750a097edf044e76cac92fe68f2be/skrample/sampling/interface.py#L7-L50