ironjr commited on
Commit
af07318
·
verified ·
1 Parent(s): 7ec6cab

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +2 -2
model.py CHANGED
@@ -1123,9 +1123,9 @@ class StreamMultiDiffusion(nn.Module):
1123
 
1124
  print('1111111111111111111111', x_t_latent_plus_uc.dtype, self.unet.dtype, self.prompt_embeds.dtype)
1125
  model_pred = self.unet(
1126
- x_t_latent_plus_uc.to(self.dtype), # (B, 4, h, w)
1127
  t_list, # (B,)
1128
- encoder_hidden_states=self.prompt_embeds.float(), # (B, 77, 768)
1129
  return_dict=False,
1130
  )[0] # (B, 4, h, w)
1131
  print('222222222222222', model_pred.dtype)
 
1123
 
1124
  print('1111111111111111111111', x_t_latent_plus_uc.dtype, self.unet.dtype, self.prompt_embeds.dtype)
1125
  model_pred = self.unet(
1126
+ x_t_latent_plus_uc, # (B, 4, h, w)
1127
  t_list, # (B,)
1128
+ encoder_hidden_states=self.prompt_embeds, # (B, 77, 768)
1129
  return_dict=False,
1130
  )[0] # (B, 4, h, w)
1131
  print('222222222222222', model_pred.dtype)