SD VAE编码卫星图像

我们使用稳定扩散 VAE 将卫星图像编码到潜在空间中。 然后我们使用 wandb.Table 将潜在变量可视化。 最后，我们将潜伏解码回图像空间，令人惊讶的是，我们得到了输入的几乎无损的副本。

这意味着我们可以使用这种技术训练潜在扩散模型，从而节省大量计算量。 此外，编码可以离线完成（在训练扩散管道之前）。

对图像进行编码和解码的代码如下：

from diffusers import AutoencoderKL
from PIL import Image


vae = AutoencoderKL.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="vae")

def encode_img(input_img):
    # Single image -> single latent in a batch (so size 1, 4, 64, 64)
    if len(input_img.shape)<4:
        input_img = input_img.unsqueeze(0)
    with torch.no_grad():
        latent = vae.encode(input_img*2 - 1) # Note scaling
    return 0.18215 * latent.latent_dist.sample()

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

执行这段代码来可视化：

import wandb

table = wandb.Table(columns=["input_image", "c1", "c2", "c3", "c4", "decoded_image"])

for img in gb_data[0:4]:
    row = [wandb.Image(img)]
    latent = encode_img(img)
    row += [wandb.Image(c_img) for c_img in latent.squeeze()]
    decoded_img = decode_img(latent)[0]
    row += [wandb.Image(decoded_img)]
    table.add_data(*row)

结果如下：

原文链接：Using Stable Diffusion VAE to encode satellite images

BimAnt翻译整理，转载请标明出处