Update README.md

README.md

@@ -45,7 +45,7 @@ torch_device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # 1. Load models
 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
 image = noise_scheduler.sample_noise((1, model.in_channels, model.resolution, model.resolution), device=torch_device, generator=generator)
@@ -93,7 +93,7 @@ torch_device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # 1. Load models
 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
 image = noise_scheduler.sample_noise((1, model.in_channels, model.resolution, model.resolution), device=torch_device, generator=generator)
@@ -104,8 +104,9 @@ eta = 0.0  # <- deterministic sampling
 
 for t in tqdm.tqdm(reversed(range(num_inference_steps)), total=num_inference_steps):
 	# 1. predict noise residual
+	orig_t = noise_scheduler.get_orig_t(t, num_inference_steps)
 	with torch.no_grad():
-				residual = unet(image, inference_step_times[t])
+	    residual = unet(image, orig_t)
 
 	# 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)