ironjr commited on
Commit
190b413
1 Parent(s): f118b07

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +1 -0
model.py CHANGED
@@ -418,6 +418,7 @@ class StreamMultiDiffusion(nn.Module):
418
  else:
419
  raise AttributeError('[ERROR] Could not access latents of provided encoder_output')
420
 
 
421
  imgs = 2 * imgs - 1
422
  latents = self.vae.config.scaling_factor * _retrieve_latents(self.vae.encode(imgs), generator=generator)
423
  if add_noise:
 
418
  else:
419
  raise AttributeError('[ERROR] Could not access latents of provided encoder_output')
420
 
421
+ imgs = imgs.to(dtype=self.vae.dtype, device=self.vae.device)
422
  imgs = 2 * imgs - 1
423
  latents = self.vae.config.scaling_factor * _retrieve_latents(self.vae.encode(imgs), generator=generator)
424
  if add_noise: