[BugFix]: Fixed add_noise in LMSDiscreteScheduler

[nicolas-dufour]

The formulation of add noise is LMSDiscreteScheduler is wrong (see crowsonkb/k-diffusion#4). This PR fixes that

[patil-suraj]

Ah, good catch! Will let @anton-l take a final look, he knows the best here.

[anton-l]

Hi @nicolas-dufour! For this particular sampler, the model needs to be trained similarly to DDPM/DDIM (i.e. with the current add_noise). That's what works for Stable Diffusion, as far as I know (let me know if I misunderstand something!).
But if you find a way to convert the beta schedule closer to the original k-lms, feel free to update the PR :)

[nicolas-dufour]

Hi @anton-l, I have tried to use the current add_noise with stable diffusion doing some img2img but it did not work and the only way i managed to make it work was to change add_noise to the proposed change.

From what i can see from k-diffusion they also use this noise technique at train time (https://github.com/crowsonkb/k-diffusion/blob/3f93c28088890e4d6bc593072739e1d6e759b392/k_diffusion/layers.py#L28). It's also what is proposed by the original Karras et al paper on which k-diffusion scheduler is based (https://arxiv.org/pdf/2206.00364.pdf Algo 2).

[patrickvonplaten]

Hi @anton-l, I have tried to use the current add_noise with stable diffusion doing some img2img but it did not work and the only way i managed to make it work was to change add_noise to the proposed change.

From what i can see from k-diffusion they also use this noise technique at train time (https://github.com/crowsonkb/k-diffusion/blob/3f93c28088890e4d6bc593072739e1d6e759b392/k_diffusion/layers.py#L28). It's also what is proposed by the original Karras et al paper on which k-diffusion scheduler is based (https://arxiv.org/pdf/2206.00364.pdf Algo 2).

Thanks for the links here @nicolas-dufour - from reading the code a bit, I think the sampler can work with both cases, e.g. it can work for inference on a model that was trained on classic DDPM, but as @nicolas-dufour pointed out one should also be able to train the model with the original formula.

Could we maybe allow for both and add a config parameter? Something like: predict_ddpm_criterion=False/True ? What do you think @anton-l ?

[patil-suraj]

Could we maybe allow for both and add a config parameter? Something like: predict_ddpm_criterion=False/True

I'm in favor of this. As @nicolas-dufour said, the current add_noise doesn't work for imag2image, in-painting. So would be important to have this.

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[patrickvonplaten]

@anton-l let's merge this and then move onto your follow-up PR?

[anton-l]

@nicolas-dufour thank you for the fix! Your investigation spawned another PR for img2img and K-LMS compatibility in general, so feel free to join the conversation, if you'd like: #270

[patrickvonplaten]

Awesome addition @nicolas-dufour ❤️