ironjr commited on
Commit
4d4572e
·
verified ·
1 Parent(s): ace8b4d

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +1 -2
model.py CHANGED
@@ -1120,9 +1120,8 @@ class StreamMultiDiffusion(nn.Module):
1120
  else:
1121
  x_t_latent_plus_uc = x_t_latent # (T * p, 4, h, w)
1122
 
1123
- print(x_t_latent_plus_uc.dtype, t_list.dtype, self.prompt_embeds.dtype, self.masks.dtype)
1124
  model_pred = self.unet(
1125
- x_t_latent_plus_uc, # (B, 4, h, w)
1126
  t_list, # (B,)
1127
  encoder_hidden_states=self.prompt_embeds, # (B, 77, 768)
1128
  return_dict=False,
 
1120
  else:
1121
  x_t_latent_plus_uc = x_t_latent # (T * p, 4, h, w)
1122
 
 
1123
  model_pred = self.unet(
1124
+ x_t_latent_plus_uc.to(self.dtype), # (B, 4, h, w)
1125
  t_list, # (B,)
1126
  encoder_hidden_states=self.prompt_embeds, # (B, 77, 768)
1127
  return_dict=False,