mikeee/diffusers-test

1import torch2from diffusers import UNetUnconditionalModel, DDIMScheduler3import PIL.Image4import numpy as np5import tqdm6 7torch_device = "cuda" if torch.cuda.is_available() else "cpu"8 9# 1. Load models10scheduler = DDIMScheduler.from_config("fu sing/ddpm-celeba-hq", tensor_format="pt")11unet = UNetUnconditionalModel.from_pretrained("fu sing/ddpm-celeba-hq", ddpm=True).to(torch_device)12 13# 2. Sample gaussian noise14generator = torch.manual_seed(23)15generator = torch.manual_seed(43)16unet.image_size = unet.resolution17image = torch.randn(18   (1, unet.in_channels, unet.image_size, unet.image_size),19   generator=generator,20)21image = image.to(torch_device)22 23# 3. Denoise24num_inference_steps = 5025num_inference_steps = 526eta = 0.0  # <- deterministic sampling27scheduler.set_timesteps(num_inference_steps)28 29for t in tqdm.tqdm(scheduler.timesteps):30    # 1. predict noise residual31    with torch.no_grad():32        residual = unet(image, t)["sample"]33 34    prev_image = scheduler.step(residual, t, image, eta)["prev_sample"]35 36    # 3. set current image to prev_image: x_t -> x_t-137    image = prev_image38 39# 4. process image to PIL40image_processed = image.cpu().permute(0, 2, 3, 1)41image_processed = (image_processed + 1.0) * 127.542image_processed = image_processed.numpy().astype(np.uint8)43image_pil = PIL.Image.fromarray(image_processed[0])44 45# 5. save image46image_pil.save("generated_image.png")47