Update model.py

model.py

@@ -77,11 +77,11 @@ class StreamMultiDiffusion(nn.Module):
         self.default_mask_std = default_mask_std
         self.default_mask_strength = default_mask_strength
         self.default_prompt_strength = default_prompt_strength
-        self.bootstrap_steps = (
-            bootstrap_steps > torch.arange(len(t_index_list))).to(dtype=self.dtype, device=self.device)
+        self.register_buffer('bootstrap_steps', (
+            bootstrap_steps > torch.arange(len(t_index_list))).to(dtype=self.dtype, device=self.device))
         self.bootstrap_mix_steps = bootstrap_mix_steps
-        self.bootstrap_mix_ratios = (
-            bootstrap_mix_steps - torch.arange(len(t_index_list), device=self.device)).clip_(0, 1).to(self.dtype)
+        self.register_buffer('bootstrap_mix_ratios', (
+            bootstrap_mix_steps - torch.arange(len(t_index_list), device=self.device)).clip_(0, 1).to(self.dtype))
         # self.bootstrap_leak_sensitivity = bootstrap_leak_sensitivity
         self.preprocess_mask_cover_alpha = preprocess_mask_cover_alpha
         self.mask_type = mask_type