From 6bf79f1e854f38b12baa72cd5969b6e0a3f3af89 Mon Sep 17 00:00:00 2001 From: Chenlei Hu Date: Wed, 6 Mar 2024 03:21:30 +0000 Subject: [PATCH] Fix decode output device/dtype (#29) --- lib_layerdiffusion/models.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/lib_layerdiffusion/models.py b/lib_layerdiffusion/models.py index fbd8656..9eb35c8 100644 --- a/lib_layerdiffusion/models.py +++ b/lib_layerdiffusion/models.py @@ -288,14 +288,19 @@ class TransparentVAEDecoder: return median @torch.no_grad() - def decode_pixel(self, pixel, latent): + def decode_pixel( + self, pixel: torch.TensorType, latent: torch.TensorType + ) -> torch.TensorType: # pixel.shape = [B, C=3, H, W] assert pixel.shape[1] == 3 + pixel_device = pixel.device + pixel_dtype = pixel.dtype + pixel = pixel.to(device=self.load_device, dtype=self.dtype) latent = latent.to(device=self.load_device, dtype=self.dtype) # y.shape = [B, C=4, H, W] y = self.estimate_augmented(pixel, latent) y = y.clip(0, 1) assert y.shape[1] == 4 - return y - + # Restore image to original device of input image. + return y.to(pixel_device, dtype=pixel_dtype)