Update README.md

This commit is contained in:
Patrick von Platen 2022-06-10 14:50:57 +02:00 committed by GitHub
parent 4569f75880
commit abbbc27e88
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 27 additions and 26 deletions

View File

@ -45,7 +45,7 @@ torch_device = "cuda" if torch.cuda.is_available() else "cpu"
# 1. Load models # 1. Load models
noise_scheduler = GaussianDDPMScheduler.from_config("fusing/ddpm-lsun-church") noise_scheduler = GaussianDDPMScheduler.from_config("fusing/ddpm-lsun-church")
model = UNetModel.from_pretrained("fusing/ddpm-lsun-church").to(torch_device) unet = UNetModel.from_pretrained("fusing/ddpm-lsun-church").to(torch_device)
# 2. Sample gaussian noise # 2. Sample gaussian noise
image = noise_scheduler.sample_noise((1, model.in_channels, model.resolution, model.resolution), device=torch_device, generator=generator) image = noise_scheduler.sample_noise((1, model.in_channels, model.resolution, model.resolution), device=torch_device, generator=generator)
@ -53,21 +53,21 @@ image = noise_scheduler.sample_noise((1, model.in_channels, model.resolution, mo
# 3. Denoise # 3. Denoise
num_prediction_steps = len(noise_scheduler) num_prediction_steps = len(noise_scheduler)
for t in tqdm.tqdm(reversed(range(num_prediction_steps)), total=num_prediction_steps): for t in tqdm.tqdm(reversed(range(num_prediction_steps)), total=num_prediction_steps):
# predict noise residual # predict noise residual
with torch.no_grad(): with torch.no_grad():
residual = unet(image, t) residual = unet(image, t)
# predict previous mean of image x_t-1 # predict previous mean of image x_t-1
pred_prev_image = noise_scheduler.compute_prev_image_step(residual, image, t) pred_prev_image = noise_scheduler.compute_prev_image_step(residual, image, t)
# optionally sample variance # optionally sample variance
variance = 0 variance = 0
if t > 0: if t > 0:
noise = noise_scheduler.sample_noise(image.shape, device=image.device, generator=generator) noise = noise_scheduler.sample_noise(image.shape, device=image.device, generator=generator)
variance = noise_scheduler.get_variance(t).sqrt() * noise variance = noise_scheduler.get_variance(t).sqrt() * noise
# set current image to prev_image: x_t -> x_t-1 # set current image to prev_image: x_t -> x_t-1
image = pred_prev_image + variance image = pred_prev_image + variance
# 5. process image to PIL # 5. process image to PIL
image_processed = image.cpu().permute(0, 2, 3, 1) image_processed = image.cpu().permute(0, 2, 3, 1)
@ -93,7 +93,7 @@ torch_device = "cuda" if torch.cuda.is_available() else "cpu"
# 1. Load models # 1. Load models
noise_scheduler = DDIMScheduler.from_config("fusing/ddpm-celeba-hq") noise_scheduler = DDIMScheduler.from_config("fusing/ddpm-celeba-hq")
model = UNetModel.from_pretrained("fusing/ddpm-celeba-hq").to(torch_device) unet = UNetModel.from_pretrained("fusing/ddpm-celeba-hq").to(torch_device)
# 2. Sample gaussian noise # 2. Sample gaussian noise
image = noise_scheduler.sample_noise((1, model.in_channels, model.resolution, model.resolution), device=torch_device, generator=generator) image = noise_scheduler.sample_noise((1, model.in_channels, model.resolution, model.resolution), device=torch_device, generator=generator)
@ -103,21 +103,22 @@ num_inference_steps = 50
eta = 0.0 # <- deterministic sampling eta = 0.0 # <- deterministic sampling
for t in tqdm.tqdm(reversed(range(num_inference_steps)), total=num_inference_steps): for t in tqdm.tqdm(reversed(range(num_inference_steps)), total=num_inference_steps):
# 1. predict noise residual # 1. predict noise residual
with torch.no_grad(): orig_t = noise_scheduler.get_orig_t(t, num_inference_steps)
residual = unet(image, inference_step_times[t]) with torch.no_grad():
residual = unet(image, orig_t)
# 2. predict previous mean of image x_t-1 # 2. predict previous mean of image x_t-1
pred_prev_image = noise_scheduler.compute_prev_image_step(residual, image, t, num_inference_steps, eta) pred_prev_image = noise_scheduler.compute_prev_image_step(residual, image, t, num_inference_steps, eta)
# 3. optionally sample variance # 3. optionally sample variance
variance = 0 variance = 0
if eta > 0: if eta > 0:
noise = noise_scheduler.sample_noise(image.shape, device=image.device, generator=generator) noise = noise_scheduler.sample_noise(image.shape, device=image.device, generator=generator)
variance = noise_scheduler.get_variance(t).sqrt() * eta * noise variance = noise_scheduler.get_variance(t).sqrt() * eta * noise
# 4. set current image to prev_image: x_t -> x_t-1 # 4. set current image to prev_image: x_t -> x_t-1
image = pred_prev_image + variance image = pred_prev_image + variance
# 5. process image to PIL # 5. process image to PIL
image_processed = image.cpu().permute(0, 2, 3, 1) image_processed = image.cpu().permute(0, 2, 3, 1)