[Python] How do I visualize the latent representation produced by the Stable Diffusion VAE?

I am trying to visualize the latent representation produced by the VAE inside a Stable Diffusion pipeline

from diffusers import StableDiffusionPipeline
import torch

# A CUDA ordinal is simply the integer ID of a GPU in a system that has one or more GPUs.
def get_device(cuda_ordinal=None):
if torch.cuda.is_available():
return torch.device("cuda", cuda_ordinal)
if torch.backends.mps.is_available():
return torch.device("mps")
return torch.device("cpu")

device=get_device()

pipe=StableDiffusionPipeline.from_pretrained(
"stable-diffusion-v1-5/stable-diffusion-v1-5",
torch_dtype=torch.float16,
variant="fp16"
).to(device)


I wrote load_image helper

from PIL import Image
from io import BytesIO
import requests

def load_image(url, size=None, return_tensor=False):
if url.startswith("http"):
response = requests.get(url)
img = Image.open(BytesIO(response.content))
else:
img = Image.open(url)
if size is not None:
img = img.resize(size)
if return_tensor:
return TF.to_tensor(img)
return img


I successfully loaded and plotted with matplotlib library

import matplotlib.pyplot as plt

im=load_image("https://media.gettyimages.com/id/2244914756/photo/a-small-island-shaped-like-a-straw-hat-by-the-blue-lake.jpg?s=2048x2048&w=gi&k=20&c=wzAXUuF06cnKTYrV7Xs8DG9jHPvrAf-1tW2Vs53VXOg=",size=(512,512))


plt.imshow(im)
plt.axis('off') # remove axes
plt.show()


The VAE compresses this image into a four-channel latent representation and i like to visualize each channel of VAE's latent representation:

from torchvision import transforms
with torch.inference_mode():
tensor_im=transforms.ToTensor()(im).unsqueeze(0).to(device)*2-1
tensor_im=tensor_im.half()
latent=pipe.vae.encode(tensor_im)
latents=latent.latent_dist.sample()
latents=latents*0.18215

latents.shape
#torch.Size([1, 4, 64, 64])


this is helper function to plot each representation

import matplotlib.pyplot as plt

def show_latent_channels(latents):
lat = latents[0] # [4, 64, 64]
num_ch = lat.shape[0]

fig, axes = plt.subplots(1, num_ch, figsize=(12,3))
for i in range(num_ch):
ch = lat
# normalize channel to [0,1] for visualization
ch = (ch - ch.min()) / (ch.max() - ch.min())
axes.imshow(ch.cpu(), cmap="viridis")
axes.set_title(f"Channel {i}")
axes.axis("off")

plt.show()


when I call show_latent_channels(latents) I get this


My goal is simply to “see” what the VAE is storing inside the latent space but matplotlib shows an empty plot

Continue reading...