From 1a58958ab4f024dbc4c90a6404c2e66210db6d00 Mon Sep 17 00:00:00 2001 From: Daniel Gu Date: Thu, 11 May 2023 07:44:15 -0700 Subject: [PATCH] Fix mixed precision issue by wrapping the offending code with the torch.autocast context manager. --- src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py b/src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py index d5f598c311..e1184964ad 100644 --- a/src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py +++ b/src/diffusers/pipelines/unidiffuser/pipeline_unidiffuser.py @@ -1246,7 +1246,8 @@ class UniDiffuserPipeline(DiffusionPipeline): ) if reduce_text_emb_dim: - prompt_embeds = self.text_decoder.encode_prefix(prompt_embeds) + with torch.autocast(device_type=device.type): + prompt_embeds = self.text_decoder.encode_prefix(prompt_embeds) # 4. Encode image, if available; otherwise prepare image latents if mode in ["img2text"]: