Update model.py

model.py

@@ -1122,25 +1122,12 @@ class StreamMultiDiffusion(nn.Module):
         else:
             x_t_latent_plus_uc = x_t_latent  # (T * p, 4, h, w)
 
-        ns = []
-        c1, c2, c3 = 0, 0, 0
-        for n, p in self.unet.named_parameters():
-            if p.data.dtype == torch.float:
-                c1 += 1
-                ns.append(n)
-            elif p.data.dtype == torch.half:
-                c2 += 1
-            else:
-                c3 += 1
-        print(c1, c2, c3)
-        print(ns)
         model_pred = self.unet(
             x_t_latent_plus_uc.to(self.unet.dtype),  # (B, 4, h, w)
             t_list,  # (B,)
             encoder_hidden_states=self.prompt_embeds,  # (B, 77, 768)
             return_dict=False,
         )[0]  # (B, 4, h, w)
-        print('222222222222222', model_pred.dtype)
 
         if self.bootstrap_steps[0] > 0:
             # Uncentering.
@@ -1151,6 +1138,7 @@ class StreamMultiDiffusion(nn.Module):
                 bootstrap_mask_ = torch.concat([bootstrap_mask, bootstrap_mask], dim=0)
             else:
                 bootstrap_mask_ = bootstrap_mask
+            print('2222222222222222222222222222222222222', model_pred.shape, bootstrap_mask_)
             model_pred = shift_to_mask_bbox_center(model_pred, bootstrap_mask_)
             x_t_latent = shift_to_mask_bbox_center(x_t_latent, bootstrap_mask)
 
@@ -1235,7 +1223,7 @@ class StreamMultiDiffusion(nn.Module):
             self.stock_noise_ = self.stock_noise.repeat_interleave(self.num_layers, dim=0)  # (T * p, 77, 768)
 
         x_0_pred_batch = self.unet_step(latent)
-
+        print('111111111111111111111111111111111')
         latent = x_0_pred_batch[-1:]
         self.x_t_latent_buffer = (
             self.alpha_prod_t_sqrt[1:] * x_0_pred_batch[:-1]