Add EDMEulerScheduler

[patil-suraj]

What does this PR do?

This PR is based on #4481 by @dg845 !

This PR adds Euler scheduler with EDM formulation. The difference between this and EulerDiscreteScheduler is that

• EulerDiscreteScheduler was essentially designed for DDPM models to use euler-style sampling algorithms with discrete timesteps.
• EDMEulerScheduler follows the EDM formulation as closely as possible and is solely intended for models that use EDM formulation, like SVD. It does not support epsilon scaling as that's already covered by EulerDiscreteScheduler. So models that still DDPM (the alphas and betas schedules) should keep using EulerDiscreteScheduler

Why add a new scheduler class ?
While it is possible to support this in EulerDiscreteScheduler by introducing an argument maybe called scaling_type, which could correspond to epsilon (ddpm) v-scaling or edm, the issue is that:

• the naming is inconsistent. With edm, the model is conditioned on continuous noise scales rather than discrete timesteps.
• There is some functionality in EulerDiscreteScheduler that is not required for EDM, such as computing betas, rescaling for zero terminal SNR, interpolating sigmas etc.

Hence with the current schedule API if we support full EDM formulation in EulerDiscreteScheduler, the code will be confusing to follow with lot's of if/else branches.

API

scheduler = EDMEulerScheduler(sigma_min=0.002, sigma_max=80.0, sigma_data=0.5)

scheduler.sigmas # computed directly as per the karras paper 
scheduler.timesteps # this is the `c_noise` which is fed to the model instead of discrete timestep

# to precondition input
sample = scheduler.precondition_inputs(sample, sigma)

# to get the c_noise
c_noise = scheduler.precondition_noise(sigma) 

# precondition output
output = scheduler.precondition_outputs(sample, model_output, sigma)

These pre-conditioning methods can be used during training to easily scale the input/output based on sigma values, which the user is free to sample in any way they want during training.

We could also consider adding a method to help with sigma sampling as per here

The rest of the API follows the existing scheduler API.

I would love to have your thoughts here @yiyixuxu @sayakpaul @dhruvrnaik @dg845

[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.

[sayakpaul]

Would be nice to have docstrings.

[sayakpaul]

👌

[sayakpaul]

I guess users will have to do this bit manually during training when experimenting with different prediction types?

[sayakpaul]

Design looks very nice to me!

[yiyixuxu]

looking good to me:)

question:
for v_prediction: does it work the same as in what's currently implemented in EulerDiscreteScheduler?

[patil-suraj]

for v_prediction: does it work the same as in what's currently implemented in EulerDiscreteScheduler?

@yiyixuxu Yes, that's right. EulerDiscreteScheduler would work the same way if sigma_min , sigma_max is provided and timestep_type==continuous

Actually not sure if we should put it here, it's used to SVD and follows EDM as well.

[sayakpaul]

I think okay to put it there as it helps the community to experiment with different prediction types and provide us feedback if needed. We can iterate on top of that.