Update model.py

model.py

@@ -173,7 +173,8 @@ class StreamMultiDiffusion(nn.Module):
         self.latent_width = int(width // self.pipe.vae_scale_factor)
 
         # For bootstrapping.
-        self.white = self.encode_imgs(torch.ones(1, 3, height, width, dtype=self.dtype, device=self.device))
+        # self.white = self.encode_imgs(torch.ones(1, 3, height, width, dtype=self.dtype, device=self.device))
+        self.white = None
 
         # StreamDiffusion setting.
         self.t_list = t_index_list
@@ -441,6 +442,8 @@ class StreamMultiDiffusion(nn.Module):
             self.state['background'].embed = embed
 
             if self.bootstrap_steps[0] > 0:
+                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]
                 self.bootstrap_latent = mix_ratio * self.white + (1.0 - mix_ratio) * self.state['background'].latent