Make tqdm calls notebook-compatible - follow-up

aravind-h-v · Jul 18, 2022 · 36b459f · 36b459f
1 parent 1820024
commit 36b459f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/src/diffusers/pipelines/glide/pipeline_glide.py b/src/diffusers/pipelines/glide/pipeline_glide.py
@@ -817,7 +817,7 @@ def text_model_fn(x_t, timesteps, transformer_out, **kwargs):
         num_trained_timesteps = self.upscale_scheduler.timesteps
         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
             with torch.no_grad():
                 time_input = torch.tensor([inference_step_times[t]] * image.shape[0], device=torch_device)

diff --git a/src/diffusers/pipelines/pndm/pipeline_pndm.py b/src/diffusers/pipelines/pndm/pipeline_pndm.py
@@ -53,7 +53,7 @@ def __call__(self, batch_size=1, generator=None, torch_device=None, num_inferenc
             image = self.scheduler.step_prk(model_output, t, image, num_inference_steps)["prev_sample"]
 
         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]
             model_output = self.unet(image, t_orig)