[Flax] Add Vae for Stable Diffusion

[patrickvonplaten]

No description provided.

[patrickvonplaten]

Will add tests and docstring later:

To verify that it works:

#!/usr/bin/env python3
from diffusers import FlaxAutoencoderKL, AutoencoderKL
import torch
import numpy as np
import jax.numpy as jnp

model_flax, params = FlaxAutoencoderKL.from_pretrained("fusing/sd-v1-4-flax", subfolder="vae", use_auth_token=True)
model_pt = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae", use_auth_token=True)

sample_shape = (1, model_flax.config.in_channels, model_flax.config.sample_size, model_flax.config.sample_size)
sample = jnp.ones(sample_shape, dtype=jnp.float32)

sample_out_flax = model_flax.apply({"params": params}, sample)

sample_pt = torch.from_numpy(np.asarray(sample))

with torch.no_grad():
    sample_out_pt = model_pt(sample_pt)

max_diff = np.max(np.abs(np.asarray(sample_out_flax.sample) - sample_out_pt.sample.numpy()))
print("Max diff", max_diff)

[HuggingFaceDocBuilderDev]

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

[pcuenca]

Looks great. I'd suggest to test it with pmap if you haven't done it already, to ensure it works in multiple devices in jitted mode as well as in a single device.

[patil-suraj]

Looks great! Just left some nits to align the names with PyTorch.

Two things to be addressed before merging:

• Use lists for downsamplers and unsamplers, same way we do in PT. This will make auto-conversion possible.
• transpose the output in the decode method as that method is used in the pipeline.

For follow-up PR:

• We should move the blocks in its own fileunet_blocks_flax.py
• Make sure all module names are identical to it's PT counterpart.