Update model.py

model.py

@@ -1121,9 +1121,8 @@ class StreamMultiDiffusion(nn.Module):
         else:
             x_t_latent_plus_uc = x_t_latent  # (T * p, 4, h, w)
 
-        print('1111111111111111111111', x_t_latent_plus_uc.dtype, self.unet.dtype, self.prompt_embeds.dtype)
-        model_pred = self.unet(
-            x_t_latent_plus_uc,  # (B, 4, h, w)
+        model_pred = self.pipe.unet(
+            x_t_latent_plus_uc.to(self.pipe.unet.dtype),  # (B, 4, h, w)
             t_list,  # (B,)
             encoder_hidden_states=self.prompt_embeds,  # (B, 77, 768)
             return_dict=False,