模型地址：https://huggingface.co/CompVis/stable-diffusion-v1-4/tree/main/vae
主要参考:Using-Stable-Diffusion-VAE-to-encode-satellite-images
sd1.4 vae
下载到本地
from diffusers import AutoencoderKL
from PIL import Image
import torch
import torchvision.transforms as T
# ./huggingface/stable-diffusion-v1-4/vae 切换为任意本地路径
vae = AutoencoderKL.from_pretrained("./huggingface/stable-diffusion-v1-4/vae",variant='fp16')
# c:\Users\zeng\Downloads\vae_config.json
def encode_img(input_img):
# Single image -> single latent in a batch (so size 1, 4, 64, 64)
# Transform the image to a tensor and normalize it
transform = T.Compose([
# T.Resize((256, 256)),
T.ToTensor()
])
input_img = transform(input_img)
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().cpu()
# image = T.Resize(original_size)(image.squeeze())
return T.ToPILImage()(image.squeeze())
if __name__ == '__main__':
# Load an example image
input_img = Image.open("huge.jpg")
original_size = input_img.size
print('original_size',original_size)
# Encode and decode the image
latents = encode_img(input_img)
reconstructed_img = decode_img(latents)
# Save the reconstructed image
reconstructed_img.save("reconstructed_example2.jpg")
# Concatenate the original and reconstructed images
concatenated_img = Image.new('RGB', (original_size[0] * 2, original_size[1]))
concatenated_img.paste(input_img, (0, 0))
concatenated_img.paste(reconstructed_img, (original_size[0], 0))
# Save the concatenated image
concatenated_img.save("concatenated_example2.jpg")
总结
### 文章总结：使用Stable Diffusion VAE模型对卫星图像进行编码与解码本文详细介绍了如何使用Hugging Face平台的Stable Diffusion v1.4 VAE (Variational Autoencoder) 模型对卫星图像进行编码与解码的过程。Stable Diffusion VAE是一种强大的生成模型，特别适用于图像数据的压缩、传输或进一步分析等场景。
#### 关键步骤：
1. **模型下载与加载**：
- 首先，从Hugging Face的[模型库](https://huggingface.co/CompVis/stable-diffusion-v1-4/tree/main/vae)下载Stable Diffusion VAE模型，并通过`AutoencoderKL.from_pretrained`方法将其加载到本地Python环境中。
2. **图像编码函数** (`encode_img`)：
- 定义一个函数将输入图像转换成潜在表示（latent representation）。此函数首先对图像进行预处理（如转换为Tensor，归一化等），然后通过VAE模型进行编码。注意，编码时对输入Tensor进行了缩放（乘以2再减1），以符合VAE模型的输入要求。
- 编码结果是潜在分布的一个样本，通过乘以特定的系数（0.18215）进行调整，便于后续处理。
3. **图像解码函数** (`decode_img`)：
- 与编码相对，解码函数通过VAE模型将潜在表示转换回图像。首先，对潜在表示进行逆缩放处理，然后通过模型解码，最后将解码结果转换为PIL图像格式，对像素值进行必要的调整并转换为RGB格式。
4. **示例与测试**：
- 在主函数中，首先加载一张示例图像，记录其原始大小。
- 使用定义的编码函数对图像进行编码，再用解码函数进行解码，重建图像。
- 最后，将原始图像与重建的图像并排保存为一张图片，以便于对比观察。
#### 关键点回顾：
- 利用了Stable Diffusion VAE模型进行图像的压缩（编码）与解压（解码）。
- 图像预处理和后处理步骤确保符合VAE模型的输入和输出要求。
- 通过调整缩放系数，优化编码和解码过程中的潜在表示准确性。
此过程不仅适用于卫星图像，还可广泛应用于其他类型的图像数据，包括但不限于摄影作品、艺术画作等，展示了Stable Diffusion VAE作为通用图像处理工具的巨大潜力。