Views
No views yet
1!pip install diffusers
2from diffusers import DiffusionPipeline
3
4model_id = "CompVis/ldm-celebahq-256"
5
6# load model and scheduler
7pipeline = DiffusionPipeline.from_pretrained(model_id)
8
9# run pipeline in inference (sample random noise and denoise)
10image = pipeline(num_inference_steps=200)["sample"]
11
12# save image
13image[0].save("ldm_generated_image.png")1!pip install diffusers
2from diffusers import UNet2DModel, DDIMScheduler, VQModel
3import torch
4import PIL.Image
5import numpy as np
6import tqdm
7
8seed = 3
9
10# load all models
11unet = UNet2DModel.from_pretrained("CompVis/ldm-celebahq-256", subfolder="unet")
12vqvae = VQModel.from_pretrained("CompVis/ldm-celebahq-256", subfolder="vqvae")
13scheduler = DDIMScheduler.from_config("CompVis/ldm-celebahq-256", subfolder="scheduler")
14
15# set to cuda
16torch_device = "cuda" if torch.cuda.is_available() else "cpu"
17
18unet.to(torch_device)
19vqvae.to(torch_device)
20
21# generate gaussian noise to be decoded
22generator = torch.manual_seed(seed)
23noise = torch.randn(
24 (1, unet.in_channels, unet.sample_size, unet.sample_size),
25 generator=generator,
26).to(torch_device)
27
28# set inference steps for DDIM
29scheduler.set_timesteps(num_inference_steps=200)
30
31image = noise
32for t in tqdm.tqdm(scheduler.timesteps):
33 # predict noise residual of previous image
34 with torch.no_grad():
35 residual = unet(image, t)["sample"]
36
37 # compute previous image x_t according to DDIM formula
38 prev_image = scheduler.step(residual, t, image, eta=0.0)["prev_sample"]
39
40 # x_t-1 -> x_t
41 image = prev_image
42
43# decode image with vae
44with torch.no_grad():
45 image = vqvae.decode(image)
46
47# process image
48image_processed = image.cpu().permute(0, 2, 3, 1)
49image_processed = (image_processed + 1.0) * 127.5
50image_processed = image_processed.clamp(0, 255).numpy().astype(np.uint8)
51image_pil = PIL.Image.fromarray(image_processed[0])
52
53image_pil.save(f"generated_image_{seed}.png")


