playrth commited on
Commit
d2feaf6
1 Parent(s): 6f5909a

Update README.md

Browse files

Example about how to encode and decode image using the VAE.

Files changed (1) hide show
  1. README.md +53 -0
README.md CHANGED
@@ -17,6 +17,59 @@ model = "stabilityai/your-stable-diffusion-model"
17
  vae = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae")
18
  pipe = StableDiffusionPipeline.from_pretrained(model, vae=vae)
19
  ```
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
20
 
21
  ## Model
22
  [SDXL](https://huggingface.co/stabilityai/stable-diffusion-xl-base-0.9) is a [latent diffusion model](https://arxiv.org/abs/2112.10752), where the diffusion operates in a pretrained,
 
17
  vae = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae")
18
  pipe = StableDiffusionPipeline.from_pretrained(model, vae=vae)
19
  ```
20
+ #### How to encode and decode Image example
21
+ ```py
22
+ import torch
23
+ from PIL import Image
24
+ from diffusers import AutoencoderKL
25
+ from diffusers.image_processor import VaeImageProcessor
26
+ import matplotlib.pyplot as plt
27
+
28
+ device=torch.device("cuda" if torch.cuda.is_available else "cpu")
29
+ # Load the pre-trained VAE model
30
+ vae = AutoencoderKL.from_pretrained("stabilityai/sdxl-vae")
31
+ vae.to(device)
32
+ vae.eval()
33
+
34
+ # Load Image processor
35
+ image_processor = VaeImageProcessor()
36
+
37
+ # Load an image
38
+ image = Image.open("Paste Image here")
39
+
40
+ # Preprocess the image
41
+ image_tensor =image_processor.preprocess(image,height=256,width=256,resize_mode="fill").to(device)
42
+
43
+ # Encode the image
44
+ with torch.no_grad():
45
+ latent_representation = vae.encode(image_tensor).latent_dist.sample()
46
+
47
+ # Decode the latent representation back to image
48
+ with torch.no_grad():
49
+ reconstructed_image = vae.decode(latent_representation).sample
50
+
51
+ # Convert the decoded tensor to a displayable image
52
+ reconstructed_image = reconstructed_image.cpu()
53
+ reconstructed_image=image_processor.postprocess(reconstructed_image,output_type='pil')
54
+ reconstructed_image=reconstructed_image[0]
55
+
56
+ # Plot the original and reconstructed images side by side
57
+ plt.figure(figsize=(10, 5))
58
+
59
+ # Original image
60
+ plt.subplot(1, 2, 1)
61
+ plt.imshow(image)
62
+ plt.title("Original Image")
63
+ plt.axis("off")
64
+
65
+ # Reconstructed image
66
+ plt.subplot(1, 2, 2)
67
+ plt.imshow(reconstructed_image)
68
+ plt.title("Reconstructed Image")
69
+ plt.axis("off")
70
+
71
+ plt.show()
72
+ ```
73
 
74
  ## Model
75
  [SDXL](https://huggingface.co/stabilityai/stable-diffusion-xl-base-0.9) is a [latent diffusion model](https://arxiv.org/abs/2112.10752), where the diffusion operates in a pretrained,