Commit 36b459f6 authored by anton-l's avatar anton-l
Browse files

Make tqdm calls notebook-compatible - follow-up

parent 18200240
...@@ -817,7 +817,7 @@ class GlidePipeline(DiffusionPipeline): ...@@ -817,7 +817,7 @@ class GlidePipeline(DiffusionPipeline):
num_trained_timesteps = self.upscale_scheduler.timesteps num_trained_timesteps = self.upscale_scheduler.timesteps
inference_step_times = range(0, num_trained_timesteps, num_trained_timesteps // num_inference_steps_upscale) inference_step_times = range(0, num_trained_timesteps, num_trained_timesteps // num_inference_steps_upscale)
for t in tqdm.tqdm(reversed(range(num_inference_steps_upscale)), total=num_inference_steps_upscale): for t in tqdm(reversed(range(num_inference_steps_upscale)), total=num_inference_steps_upscale):
# 1. predict noise residual # 1. predict noise residual
with torch.no_grad(): with torch.no_grad():
time_input = torch.tensor([inference_step_times[t]] * image.shape[0], device=torch_device) time_input = torch.tensor([inference_step_times[t]] * image.shape[0], device=torch_device)
......
...@@ -53,7 +53,7 @@ class PNDMPipeline(DiffusionPipeline): ...@@ -53,7 +53,7 @@ class PNDMPipeline(DiffusionPipeline):
image = self.scheduler.step_prk(model_output, t, image, num_inference_steps)["prev_sample"] image = self.scheduler.step_prk(model_output, t, image, num_inference_steps)["prev_sample"]
timesteps = self.scheduler.get_time_steps(num_inference_steps) timesteps = self.scheduler.get_time_steps(num_inference_steps)
for t in tqdm.tqdm(range(len(timesteps))): for t in tqdm(range(len(timesteps))):
t_orig = timesteps[t] t_orig = timesteps[t]
model_output = self.unet(image, t_orig) model_output = self.unet(image, t_orig)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment