1 files changed, 3 insertions, 2 deletions

diff --git a/pipelines/stable_diffusion/vlpn_stable_diffusion.py b/pipelines/stable_diffusion/vlpn_stable_diffusion.py
index a198cf6..bfecd1c 100644
--- a/pipelines/stable_diffusion/vlpn_stable_diffusion.py
+++ b/pipelines/stable_diffusion/vlpn_stable_diffusion.py
@@ -234,7 +234,8 @@ class VlpnStableDiffusion(DiffusionPipeline):
                 latents = torch.randn(latents_shape, generator=generator, device=self.device, dtype=latents_dtype)
         elif isinstance(latents, PIL.Image.Image):
             latents = preprocess(latents, width, height)
-            latent_dist = self.vae.encode(latents.to(self.device)).latent_dist
+            latents = latents.to(device=self.device, dtype=latents_dtype)
+            latent_dist = self.vae.encode(latents).latent_dist
             latents = latent_dist.sample(generator=generator)
             latents = 0.18215 * latents
 
@@ -249,7 +250,7 @@ class VlpnStableDiffusion(DiffusionPipeline):
             timesteps = torch.tensor([timesteps] * batch_size, device=self.device)
 
             # add noise to latents using the timesteps
-            noise = torch.randn(latents.shape, generator=generator, device=self.device)
+            noise = torch.randn(latents.shape, generator=generator, device=self.device, dtype=latents_dtype)
             latents = self.scheduler.add_noise(latents, noise, timesteps)
         else:
             if latents.shape != latents_shape: