Update model.py

model.py

@@ -382,7 +382,6 @@ class StreamMultiDiffusion(nn.Module):
         inputs = self.i2t_processor(image, question, return_tensors='pt')
         out = self.i2t_model.generate(**{k: v.to(self.i2t_model.device) for k, v in inputs.items()}, max_new_tokens=77)
         prompt = self.i2t_processor.decode(out[0], skip_special_tokens=True).strip()
-        print(prompt)
         return prompt
 
     @torch.no_grad()
@@ -474,7 +473,6 @@ class StreamMultiDiffusion(nn.Module):
                 if self.white is None:
                     self.white = self.encode_imgs(torch.ones(1, 3, self.height, self.width, dtype=self.dtype, device=self.device))
                 mix_ratio = self.bootstrap_mix_ratios[:, None, None, None]
-                print(mix_ratio, mix_ratio.dtype, self.white.dtype, self.white.device, self.state['background'].latent.dtype, self.state['background'].latent.device)
                 self.bootstrap_latent = mix_ratio * self.white + (1.0 - mix_ratio) * self.state['background'].latent
 
             self.ready_checklist['background_registered'] = True
@@ -1093,6 +1091,7 @@ class StreamMultiDiffusion(nn.Module):
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         p = self.num_layers
         x_t_latent = x_t_latent.repeat_interleave(p, dim=0)  # (T * p, 4, h, w)
+        print('111111111111111111111')
 
         if self.bootstrap_steps[0] > 0:
             # Background bootstrapping.
@@ -1101,6 +1100,8 @@ class StreamMultiDiffusion(nn.Module):
                 self.stock_noise,
                 torch.tensor(self.sub_timesteps_tensor, device=self.device),
             )
+            print('111111111111111111111', bootstrap_steps)
+
             x_t_latent = rearrange(x_t_latent, '(t p) c h w -> p t c h w', p=p)
             bootstrap_mask = (
                 self.masks * self.bootstrap_steps[None, :, None, None, None]