elixir-nx / bumblebee

Pre-trained Neural Network models in Axon (+ 🤗 Models integration)
Apache License 2.0
1.27k stars 90 forks source link

Add LCM scheduler #320

Closed wtedw closed 6 months ago

wtedw commented 6 months ago

This is a simple version of LCMScheduler that supports txt2img.
Here's a livebook that makes use of the scheduler:

To support LCM-LoRA natively in Axon + Bumblebee, we'd need these PRs

also, seed + strength should be forwarded to the scheduler_init function, but I'll create a separate PR for this.

josevalim commented 6 months ago

I believe @seanmor5 wants LoRA support in Axon. :)

jonatanklosko commented 6 months ago

@wtedw thanks a lot for the PR! I'm doing minor some adjustments and docs, will push and merge tomorrow :)

seed + strength should be forwarded to the scheduler_init function

Good call, we should pass the prng key when calling any of the schedulers, I will change it afterwards. The strength is only relevant for img2img, so we can wait until we support it before we expose it.

wtedw commented 6 months ago

@jonatanklosko , Awesome thanks!

jonatanklosko commented 6 months ago

Btw. we will need a separate serving for LCM Diffusion, corresponding to LatentConsistencyModelPipeline in hf/diffusers. There are some differences from Stable Diffusion, in particular LCM doesn't really support negative prompt and the guidance scale is incorporated into the U-Net forward pass as additional timestamp embeddings (which is likely why using different guiding scales didn't work very well for you). But we will get to this once we support loading LoRA!

wtedw commented 6 months ago

@jonatanklosko Ah that's some good info. Thanks for looking into it!