"Common Diffusion Noise Schedules and Sample Steps are Flawed" integration

[vvvm23]

I was reading the paper Common Diffusion Noise Schedules and Sample Steps are Flawed and found it pretty interesting. It proposes a few simple changes that could be useful when integrated into existing schedulers in the diffusers library. Namely:

1. "rescale the noise schedule to enforce zero terminal SNR"
2. "change the sampler to always start from the last timestep"
3. "rescale classifier-free guidance to prevent over-exposure"

They demonstrate the proposed fixes by finetuning Stable Diffusion and showing it can generate new concepts it could not successfully do before, such as solid colour backgrounds. Attached a figure to demonstrate this.

Is this something that would be valuable to add to the diffusers library? Are these proposals already integrated? Thanks~

[Max-We]

FYI I played around with this here: https://github.com/Max-We/sf-zero-signal-to-noise

[vvvm23]

@Max-We nice write up! Do you have any plans to integrate the changes into this library?

I am wondering whether LoRA finetuning would be sufficient to adapt the model for the new scheduler, or whether we would need to do full scale finetuning.

I asked online about releasing the official finetuned model weights, but I don't think we will get them.

[rvorias]

I am wondering whether LoRA finetuning would be sufficient to adapt the model for the new scheduler, or whether we would need to do full scale finetuning.

Been also wondering about this, but Lora only acts on the (cross) attention layers which is nice for "what" and "where", but not sure if it can also recalibrate the total noise prediction. Keen to try out tho.

[vvvm23]

I think LoRA can be applied to arbitrary layers, but maybe the default is cross attention only. It is a good observation though that maybe more than that is needed.

[adammenges]

Huge vote up for this, looks like it'll solve many of the issues I've been experiencing in my work with images using a white background.

[Max-We]

@vvvm23 I will try to implement parts of it this Saturday

[patrickvonplaten]

Very cool! Let me know if you need help @Max-We

[patrickvonplaten]

BTW Stable Diffusion 2.x is already trained with v-prediction, so we should be able to see clear improvements there without having to retrain the model - did anybody check this by any chance?

[bghira]

i have checked that. it works, but there's even better results by fine-tuning 2.1 even just a little bit with captioned photos after enabling the fixes. and the deeper you go, the better it gets. but out of the box, it results in much more useable 2.1 outputs.

[eeyrw]

I have tried trick "rescale classifier-free guidance to prevent over-exposure" as follow implement but it seems just making image more photo realistic rather than fixing over exposure:

        # 8. Denoising loop
        for i, t in enumerate(self.progress_bar(timesteps)):
            # expand the latents if we are doing classifier free guidance
            latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)

            # predict the noise residual
            noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample

            # perform guidance
            if do_classifier_free_guidance:
                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
                noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)

                # Here we rescale
                sigma_cfg = torch.std(noise_pred)
                sigma_pos = torch.std(noise_pred_text)
                noise_pred_rescaled = noise_pred*(sigma_pos/sigma_cfg)
                phi = 0.7
                noise_pred_new = phi*noise_pred_rescaled + (1-phi)*noise_pred
                noise_pred = noise_pred_new

[Max-We]

FYI I suggested changes for

1. "rescale the noise schedule to enforce zero terminal SNR"
2. "rescale classifier-free guidance to prevent over-exposure"

in PR #3664