Flow Matching RK4 Discreet Scheduler

[rydeveraumn]

Has anyone seen an implementation of an flow matching RK4 discreet scheduler?

[a-r-r-o-w]

I do remember seeing RK4 implementation in NVIDIA Cosmos but don't see it any more. They do have third order solver though: https://github.com/nvidia-cosmos/cosmos-predict2/blob/main/cosmos_predict2/functional/runge_kutta.py

[Beinsezii]

The problem is diffusers' pipeline sampling code is all inside-out so single step methods are extremely difficult to implement, and adaptive methods are impossible.

I'm working on my own RK4 for common diffusion models
https://github.com/Beinsezii/skrample/blob/d0b7bdb9b5e2f04630d8b2b4235245924ad4ee5e/skrample/sampling/functional.py#L85-L152
but translating that to diffusers is basically impossible because 100% of the pipelines assume timesteps is a linear sequence and that each step only needs one model call.

Is there plans to fix this eventually? It's trivial to adapt already inside-out samplers to the functional method
https://github.com/Beinsezii/skrample/blob/d0b7bdb9b5e2f04630d8b2b4235245924ad4ee5e/skrample/sampling/interface.py#L11-L50
but not the other way around, and maintaining patches for 50+ pipelines out-of-tree would be unfeasible for an individual.