File size: 2,236 Bytes
c9c864b 9379e34 c9c864b 3ee4870 a57626b c9c864b 9379e34 c9c864b 3ee4870 c9c864b ba75ac7 c9c864b 31e2197 2720c36 c9c864b a57626b 3ee4870 ba75ac7 a57626b 31e2197 2720c36 a57626b 3ee4870 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
#!/usr/bin/env python3
from diffusers import UNet2DModel, DDIMScheduler, VQModel
import torch
import PIL.Image
import numpy as np
import tqdm
seed = 3
# 1. Unroll the full loop
# ==================================================================
# load all models
unet = UNet2DModel.from_pretrained("./", subfolder="unet")
vqvae = VQModel.from_pretrained("./", subfolder="vqvae")
scheduler = DDIMScheduler.from_config("./", subfolder="scheduler")
# set to cuda
torch_device = "cuda" if torch.cuda.is_available() else "cpu"
unet.to(torch_device)
vqvae.to(torch_device)
# generate gaussian noise to be decoded
generator = torch.manual_seed(seed)
noise = torch.randn(
(1, unet.in_channels, unet.image_size, unet.image_size),
generator=generator,
).to(torch_device)
# set inference steps for DDIM
scheduler.set_timesteps(num_inference_steps=200)
image = noise
for t in tqdm.tqdm(scheduler.timesteps):
# predict noise residual of previous image
with torch.no_grad():
residual = unet(image, t)["sample"]
# compute previous image x_t according to DDIM formula
prev_image = scheduler.step(residual, t, image, eta=0.0)["prev_sample"]
# x_t-1 -> x_t
image = prev_image
# decode image with vae
with torch.no_grad():
image = vqvae.decode(image)
# process image
image_processed = image.cpu().permute(0, 2, 3, 1)
image_processed = (image_processed + 1.0) * 127.5
image_processed = image_processed.clamp(0, 255).numpy().astype(np.uint8)
image_pil = PIL.Image.fromarray(image_processed[0])
# 2. Use pipeline
# ==================================================================
from diffusers import LatentDiffusionUncondPipeline
import torch
import PIL.Image
import numpy as np
import tqdm
pipeline = LatentDiffusionUncondPipeline.from_pretrained("./")
# generatae image by calling the pipeline
generator = torch.manual_seed(seed)
image = pipeline(generator=generator, num_inference_steps=200)["sample"]
# process image
image_processed = image.cpu().permute(0, 2, 3, 1)
image_processed = (image_processed + 1.0) * 127.5
image_processed = image_processed.clamp(0, 255).numpy().astype(np.uint8)
image_pil = PIL.Image.fromarray(image_processed[0])
image_pil.save(f"generated_image_{seed}.png")
|