Skip to content

Commit

Permalink
Fix image upcasting (#7858)
Browse files Browse the repository at this point in the history
Fix image's upcasting before `vae.encode()` when using `fp16`

Co-authored-by: YiYi Xu <[email protected]>
  • Loading branch information
tolgacangoz and yiyixuxu authored May 8, 2024
1 parent c221714 commit d50baf0
Show file tree
Hide file tree
Showing 2 changed files with 1 addition and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -1419,7 +1419,6 @@ def encode_image(self, image, dtype=None, height=None, width=None, resize_mode="
if needs_upcasting:
image = image.float()
self.upcast_vae()
image = image.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)

x0 = self.vae.encode(image).latent_dist.mode()
x0 = x0.to(dtype)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -525,8 +525,8 @@ def prepare_image_latents(
# make sure the VAE is in float32 mode, as it overflows in float16
needs_upcasting = self.vae.dtype == torch.float16 and self.vae.config.force_upcast
if needs_upcasting:
image = image.float()
self.upcast_vae()
image = image.to(next(iter(self.vae.post_quant_conv.parameters())).dtype)

image_latents = retrieve_latents(self.vae.encode(image), sample_mode="argmax")

Expand Down

0 comments on commit d50baf0

Please sign in to comment.