[ModelOutputs] Replace dict outputs with Dict/Dataclass and allow to return tuples

[patrickvonplaten]

Model Outputs

This PR allows the user to switch between tuple and Dict/Dataclass outputs for schedulers, models and pipelines.
Tuples are now returned whenever the return_dict flag is set to False of the respective step, forward and __call__ functions.

New recommended API:

Having merged this PR, users are recommended to use schedulers as following:

sample = scheduler.step(...).prev_sample

instead of

sample = scheduler.step(...)["prev_sample"]  # which is still possible though!

,models as following:

sample = unet(...).sample

instead of

sample = unet(...)["sample"]  # which is still possible though!

, and pipelines:

image = pipe(...).images[0]

instead of

image = pipe(...)["sample"][0]  # which is still possible, **but now has a deprecating warning**

Tests

Extensive comparison tests are added for schedulers, models, and fast pipeline tests

🚨 Deprecation Warnings 🚨

In this PR we're deprecating the "sample" output keyword for pipelines. Pipelines are modality specific and therefore it is more intuitive to call pipe(...).images[0]

🚨🚨🚨 Breaking changes 🚨🚨🚨

This PR introduces breaking changes for the following public-facing methods:

• VQModel.encode -> we return a dict/dataclass instead of a single tensor. In the future it's very likely required to return more than just one tensor. Please make sure to change latents = model.encode(...) to latents = model.encode(...)[0] or latents = model.encode(...).latens
• VQModel.decode -> we return a dict/dataclass instead of a single tensor. In the future it's very likely required to return more than just one tensor. Please make sure to change sample = model.decode(...) to sample = model.decode(...)[0] or sample = model.decode(...).sample
• VQModel.forward -> we return a dict/dataclass instead of a single tensor. In the future it's very likely required to return more than just one tensor. Please make sure to change sample = model(...) to sample = model(...)[0] or sample = model(...).sample
• AutoencoderKL.encode -> we return a dict/dataclass instead of a single tensor. In the future it's very likely required to return more than just one tensor. Please make sure to change latent_dist = model.encode(...) to latent_dist = model.encode(...)[0] or latent_dist = model.encode(...).latent_dist
• AutoencoderKL.decode -> we return a dict/dataclass instead of a single tensor. In the future it's very likely required to return more than just one tensor. Please make sure to change sample = model.decode(...) to sample = model.decode(...)[0] or sample = model.decode(...).sample
• AutoencoderKL.forward -> we return a dict/dataclass instead of a single tensor. In the future it's very likely required to return more than just one tensor. Please make sure to change sample = model(...) to sample = model(...)[0] or sample = model(...).sample

TODOS once PR is merged

• Change all official google colabs to the forrmat as proposed here (cc @pcuenca @patil-suraj @anton-l) and change the HF Blog

[HuggingFaceDocBuilderDev]

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

[patrickvonplaten]

Checked that all slow tests run as expected ✔️ - some tests are failing but they are currently also failing on master

[pcuenca]

Really comprehensive work, very impressive! We should merge it soon and deal with the breaking changes as it permeates the whole codebase.

We also need to adapt the training examples when we merge.

[pcuenca]

This should be of type DiagonalGaussianDistribution as the comment says.

The problem is that DiagonalGaussianDistribution is defined later in the file. Solutions:

• Move this down. It would be out of place with respect to the other output dataclasses.
• Declare class DiagonalGaussianDistribution: pass before, and comment it's a "forward" declaration.
• Use object as the type, with a comment.

[patil-suraj]

+1

[patrickvonplaten]

Using the string version instead as suggested here: https://peps.python.org/pep-0484/#forward-references

[pcuenca]

Ohh, nice!

[patil-suraj]

Thanks a lot for working on this big PR! In general everything looks good, just left some nit.

My only big comment, same as Pedro's is, should we rename ModelOutput to something else, as this is used for pipelines, and schedulers as well.

[patil-suraj]

+1