AutoencoderKL: embedding space distribution and image generation

[AlejandroBaron]

Hello I'm trying to get my feet wet with the VAE side of diffusers. From what I understand, VAEs should be able to generate samples from the prior (normal distribution). See this example

This is my attempt at doing so:

# %%
import torch
import matplotlib.pyplot as plt
from diffusers import AutoencoderKL
from diffusers.image_processor import VaeImageProcessor
from PIL import Image

# Instantiate AutoencoderKL object
vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae")
vae.eval()
vae_scale_factor = 2 ** (len(vae.config.block_out_channels) - 1)
pp = VaeImageProcessor(vae_scale_factor=vae_scale_factor, resample="bicubic")

def decode_img(latents):
    # bath of latents -> list of images
    latents = (1 / vae.config.scaling_factor) * latents
    with torch.no_grad():
        image = vae.decode(latents).sample
    image = (image / 2 + 0.5).clamp(0, 1)
    image = image.detach()
    return image

latent = torch.randn(1,4,58,103)
output = decode_img(latent)
plt.imshow(output.permute(0,2,3,1).numpy()[0])
plt.ylabel('KL Divergence')
plt.show()

But i just get random noise

Aren't these autoencoders actually VAEs in the sense that the KL is just used to organize the embedding space a bit?