Allow SD pipeline to use newer schedulers, eg: FlowMatch

[ppbrown]

Allow SD pipeline to use newer schedulers, eg: FlowMatch,
by skipping attribute that doesnt exist there
(scale_model_input)

I am currently experimenting with SD + FlowMatchEuler, and had to do
an ugly hand hack in my training code to get the pipeline to continue:

print("HAND HACKING FLOWMATCH MODULE")
from diffusers import FlowMatchEulerDiscreteScheduler
def scale_model_input(self, sample, timestep):
    return sample

FlowMatchEulerDiscreteScheduler.scale_model_input = scale_model_input

Would be nice to have core diffusers handle this so I can share training results with people cleanly.

Who can review?

• Pipelines and pipeline callbacks: @yiyixuxu and @asomoza

[yiyixuxu]

hi @ppbrown
Thanks for the issue, I think it's good idea to have consistent interface between flow match and regular schedulers, even though we may not yet have a strong use case for some of the features

feel free to end a PR :)

[ppbrown]

Hello,
i feel like you are suggesting I create a new PR... but I'm not sure what specifically you are suggesting?
It almost sounds like you are suggesting I open a PR about "the issue I just created"... But this IS a PR ?

[yiyixuxu]

thanks
sorry I got confused, and thought it was an issue

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[damian0815]

curious what you're doing exactly @ppbrown ! i have a project where i'm currently converting SD2 to flow matching (since v-pred is a good chunk of the way to flow-matching anyway)..

[ppbrown]

It... aint easy.

Initially I just tried a brute force "plug it in and retrain unet" strategy. Since that worked for prior experiments like swapping vae to sdxl VAE.

But the noise schedule is way more disruptive.So the naiive approach ended up doing wierd stuff. The initial steps threw off the unet knowledge so that it couldnt render things properly.... my guess is i would have to approach it like a "retrain from noise" project. No thanks!

So then I tried freeezing everythning except the "time" layers.
But that didnt make enough progress.

Next I tried unfreezing time layers, and a couple more.
Some progress.... but also some mishaps.

Now I'm attempting a hybrid approach, where I do the partial training at least until the loss is within some more reasonable range.
It STARTS at 2.x !!!
So far, it plateaus at around 0.6
I'm experimenting with various LR schedules to see if I can get it to 0.5

Then I'll do the switchover to full unet training from that, and see if that gets me where I want.

[damian0815]

Interesting @ppbrown ! in my case it's been really straightforward, but I'm resuming from SD2, which uses v-pred - which is already 90% of the way to flow-matching.

I've had to add QK-norm (#12051), I'm using an effective batch size of 1000, and I've been using an SD3-inspired loss-scale-per-timestep. all the code is over here → https://github.com/damian0815/EveryDream2trainer/tree/flow_matching (flow_matching branch).

WIP model on huggingface damian0815/sdsf-flowmatch-largeres-68k-m2x-qk_rms_1e6-16k-1e5-11k-lores-16k-flattimesteps-49k. I hang out sometimes in the EveryDream2 discord if you wanna discuss/collaborate!

[damian0815]

current samples: