From cab7b285bb6c709c7e8c5e09718950c8575fb7d3 Mon Sep 17 00:00:00 2001 From: Nouamane Tazi Date: Tue, 13 Sep 2022 18:06:11 +0000 Subject: [PATCH] remove blocking from all .to() operations at beginning of pipeline --- .../pipelines/stable_diffusion/pipeline_stable_diffusion.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py index 451171aeee..37a2737c95 100644 --- a/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py +++ b/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py @@ -190,7 +190,7 @@ class StableDiffusionPipeline(DiffusionPipeline): uncond_input = self.tokenizer( [""] * batch_size, padding="max_length", max_length=max_length, return_tensors="pt" ) - uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device, non_blocking=True))[0] + uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0] # For classifier free guidance, we need to do two forward passes. # Here we concatenate the unconditional and text embeddings into a single batch @@ -208,7 +208,7 @@ class StableDiffusionPipeline(DiffusionPipeline): else: if latents.shape != latents_shape: raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}") - latents = latents.to(self.device, non_blocking=True) + latents = latents.to(self.device) # set timesteps accepts_offset = "offset" in set(inspect.signature(self.scheduler.set_timesteps).parameters.keys())