Update model.py

model.py

@@ -295,8 +295,12 @@ class StreamMultiDiffusion(nn.Module):
     def reset_latent(self) -> None:
         # initialize x_t_latent (it can be any random tensor)
         b = (self.denoising_steps_num - 1) * self.frame_bff_size
-        self.x_t_latent_buffer = torch.zeros(
-            (b, 4, self.latent_height, self.latent_width), dtype=self.dtype, device=self.device)
+        if not hasattr(self, 'x_t_latent_buffer'):
+            self.register_buffer('x_t_latent_buffer', torch.zeros(
+                (b, 4, self.latent_height, self.latent_width), dtype=self.dtype, device=self.device))
+        else:
+            self.x_t_latent_buffer = torch.zeros(
+                (b, 4, self.latent_height, self.latent_width), dtype=self.dtype, device=self.device))
 
     def reset_state(self) -> None:
         # TODO Reset states for context switch between multiple users.
@@ -305,24 +309,35 @@ class StreamMultiDiffusion(nn.Module):
     def prepare(self) -> None:
         # make sub timesteps list based on the indices in the t_list list and the values in the timesteps list
         self.timesteps = self.scheduler.timesteps.to(self.device)
-        self.sub_timesteps = []
+        sub_timesteps = []
         for t in self.t_list:
-            self.sub_timesteps.append(self.timesteps[t])
-        sub_timesteps_tensor = torch.tensor(self.sub_timesteps, dtype=torch.long, device=self.device)
-        self.sub_timesteps_tensor = sub_timesteps_tensor.repeat_interleave(self.frame_bff_size, dim=0)
+            sub_timesteps.append(self.timesteps[t])
+        sub_timesteps_tensor = torch.tensor(sub_timesteps, dtype=torch.long, device=self.device)
+        if not hasattr(self, 'sub_timesteps_tensor'):
+            self.register_buffer('sub_timesteps_tensor', sub_timesteps_tensor.repeat_interleave(self.frame_bff_size, dim=0)
+        else:
+            self.sub_timesteps_tensor = sub_timesteps_tensor.repeat_interleave(self.frame_bff_size, dim=0)
 
         c_skip_list = []
         c_out_list = []
-        for timestep in self.sub_timesteps:
+        for timestep in sub_timesteps:
             c_skip, c_out = self.scheduler.get_scalings_for_boundary_condition_discrete(timestep)
             c_skip_list.append(c_skip)
             c_out_list.append(c_out)
-        self.c_skip = torch.stack(c_skip_list).view(len(self.t_list), 1, 1, 1).to(dtype=self.dtype, device=self.device)
-        self.c_out = torch.stack(c_out_list).view(len(self.t_list), 1, 1, 1).to(dtype=self.dtype, device=self.device)
+        c_skip = torch.stack(c_skip_list).view(len(self.t_list), 1, 1, 1).to(dtype=self.dtype, device=self.device)
+        c_out = torch.stack(c_out_list).view(len(self.t_list), 1, 1, 1).to(dtype=self.dtype, device=self.device)
+        if not hasattr(self, 'c_skip'):
+            self.register_buffer('c_skip', c_skip)
+        else:
+            self.c_skip = c_skip
+        if not hasattr(self, 'c_out'):
+            self.register_buffer('c_out', c_out)
+        else:
+            self.c_out = c_out
 
         alpha_prod_t_sqrt_list = []
         beta_prod_t_sqrt_list = []
-        for timestep in self.sub_timesteps:
+        for timestep in sub_timesteps:
             alpha_prod_t_sqrt = self.scheduler.alphas_cumprod[timestep].sqrt()
             beta_prod_t_sqrt = (1 - self.scheduler.alphas_cumprod[timestep]).sqrt()
             alpha_prod_t_sqrt_list.append(alpha_prod_t_sqrt)
@@ -331,12 +346,24 @@ class StreamMultiDiffusion(nn.Module):
             .to(dtype=self.dtype, device=self.device))
         beta_prod_t_sqrt = (torch.stack(beta_prod_t_sqrt_list).view(len(self.t_list), 1, 1, 1)
             .to(dtype=self.dtype, device=self.device))
-        self.alpha_prod_t_sqrt = alpha_prod_t_sqrt.repeat_interleave(self.frame_bff_size, dim=0)
-        self.beta_prod_t_sqrt = beta_prod_t_sqrt.repeat_interleave(self.frame_bff_size, dim=0)
+        if not hasattr(self, 'alpha_prod_t_sqrt'):
+            self.register_buffer('alpha_prod_t_sqrt', alpha_prod_t_sqrt.repeat_interleave(self.frame_bff_size, dim=0))
+        else:
+            self.alpha_prod_t_sqrt = alpha_prod_t_sqrt.repeat_interleave(self.frame_bff_size, dim=0)
+        if not hasattr(self, 'beta_prod_t_sqrt'):
+            self.register_buffer('beta_prod_t_sqrt', beta_prod_t_sqrt.repeat_interleave(self.frame_bff_size, dim=0))
+        else:
+            self.beta_prod_t_sqrt = beta_prod_t_sqrt.repeat_interleave(self.frame_bff_size, dim=0)
 
         noise_lvs = ((1 - self.scheduler.alphas_cumprod.to(self.device)[self.sub_timesteps_tensor]) ** 0.5)
-        self.noise_lvs = noise_lvs[None, :, None, None, None]
-        self.next_noise_lvs = torch.cat([noise_lvs[1:], noise_lvs.new_zeros(1)])[None, :, None, None, None]
+        if not hasattr(self, 'noise_lvs'):
+            self.register_buffer('noise_lvs', noise_lvs[None, :, None, None, None])
+        else:
+            self.noise_lvs = noise_lvs[None, :, None, None, None]
+        if not hasattr(self, 'next_noise_lvs'):
+            self.register_buffer('next_noise_lvs', torch.cat([noise_lvs[1:], noise_lvs.new_zeros(1)])[None, :, None, None, None])
+        else:
+            self.next_noise_lvs = torch.cat([noise_lvs[1:], noise_lvs.new_zeros(1)])[None, :, None, None, None]
 
     @torch.no_grad()
     def get_text_prompts(self, image: Image.Image) -> str: