Update model.py

model.py

@@ -418,6 +418,7 @@ class StreamMultiDiffusion(nn.Module):
             else:
                 raise AttributeError('[ERROR]    Could not access latents of provided encoder_output')
 
+        imgs = imgs.to(dtype=self.vae.dtype, device=self.vae.device)
         imgs = 2 * imgs - 1
         latents = self.vae.config.scaling_factor * _retrieve_latents(self.vae.encode(imgs), generator=generator)
         if add_noise: