Update model.py

model.py

@@ -78,7 +78,7 @@ class StreamMultiDiffusion(nn.Module):
         self.default_mask_strength = default_mask_strength
         self.default_prompt_strength = default_prompt_strength
         self.register_buffer('bootstrap_steps', (
-            bootstrap_steps > torch.arange(len(t_index_list))).to(dtype=self.dtype, device=self.device))
+            bootstrap_steps > torch.arange(len(t_index_list))).float().to(dtype=self.dtype, device=self.device))
         self.bootstrap_mix_steps = bootstrap_mix_steps
         self.register_buffer('bootstrap_mix_ratios', (
             bootstrap_mix_steps - torch.arange(len(t_index_list), device=self.device)).clip_(0, 1).to(self.dtype))
@@ -1091,8 +1091,6 @@ 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.
             bootstrap_latent = self.scheduler.add_noise(
@@ -1100,7 +1098,6 @@ class StreamMultiDiffusion(nn.Module):
                 self.stock_noise,
                 torch.tensor(self.sub_timesteps_tensor, device=self.device),
             )
-            print('111111111111111111111', self.bootstrap_steps)
 
             x_t_latent = rearrange(x_t_latent, '(t p) c h w -> p t c h w', p=p)
             bootstrap_mask = (
@@ -1109,11 +1106,9 @@ class StreamMultiDiffusion(nn.Module):
             ) # (p, t, c, h, w)
             x_t_latent = (1.0 - bootstrap_mask) * bootstrap_latent[None] + bootstrap_mask * x_t_latent
             x_t_latent = rearrange(x_t_latent, 'p t c h w -> (t p) c h w')
-            print('222222222222222222222')
 
             # Centering.
             x_t_latent = shift_to_mask_bbox_center(x_t_latent, rearrange(self.masks, 'p t c h w -> (t p) c h w'), reverse=True)
-        print('333333333333333333333')
 
         t_list = self.sub_timesteps_tensor_  # (T * p,)
         if self.guidance_scale > 1.0 and self.cfg_type == 'initialize':