Update model.py

model.py

@@ -437,6 +437,7 @@ class StreamMultiDiffusion(nn.Module):
             encoder. Shape: (B, 3, H, W).
         """
         latents = 1 / self.vae.config.scaling_factor * latents
+        latents = latents.to(dtype=self.vae.dtype, device=self.vae.device)
         imgs = self.vae.decode(latents).sample
         imgs = (imgs / 2 + 0.5).clip_(0, 1)
         return imgs
@@ -1125,8 +1126,6 @@ class StreamMultiDiffusion(nn.Module):
             t_list,  # (B,)
             encoder_hidden_states=self.prompt_embeds,  # (B, 77, 768)
             return_dict=False,
-            # TODO: Add SDXL Support.
-            # added_cond_kwargs={'text_embeds': add_text_embeds, 'time_ids': add_time_ids},
         )[0]  # (B, 4, h, w)
 
         if self.bootstrap_steps[0] > 0:
@@ -1233,7 +1232,7 @@ class StreamMultiDiffusion(nn.Module):
         if no_decode:
             return latent
 
-        imgs = self.decode_latents(latent.half())  # (1, 3, H, W)
+        imgs = self.decode_latents(latent.half()).float()  # (1, 3, H, W)
         img = T.ToPILImage()(imgs[0].cpu())
         return img