Update model.py

model.py

@@ -28,6 +28,7 @@ import torch.nn.functional as F
 import torchvision.transforms as T
 from einops import rearrange
 
+import random
 from collections import deque
 from typing import Tuple, List, Literal, Optional, Union
 from PIL import Image
@@ -54,6 +55,7 @@ class StreamMultiDiffusion(nn.Module):
         delta: float = 1.0,
         cfg_type: Literal['none', 'full', 'self', 'initialize'] = 'none',
         seed: int = 2024,
+        seedfix: bool = False,
         autoflush: bool = True,
         default_mask_std: float = 8.0,
         default_mask_strength: float = 1.0,
@@ -71,6 +73,7 @@ class StreamMultiDiffusion(nn.Module):
         self.device = device
         self.dtype = dtype
         self.seed = seed
+        self.seedfix = seedfix
         self.sd_version = sd_version
 
         self.autoflush = autoflush
@@ -1255,12 +1258,16 @@ class StreamMultiDiffusion(nn.Module):
         if not ignore_check_ready and not self.check_ready():
             return
         if not ignore_check_ready and self.is_dirty:
-            print("I'm so dirty now!")
+            print("I'm dirty!")
             self.commit()
             self.flush()
 
-        latent = torch.randn((1, self.unet.config.in_channels, self.latent_height, self.latent_width),
-            dtype=self.dtype, device=self.device)  # (1, 4, h, w)
+        if self.seedfix:
+            self.reset_seed(self.generator, self.seed)
+            latent = self.init_noise[:1]
+        else:
+            latent = torch.randn((1, self.unet.config.in_channels, self.latent_height, self.latent_width),
+                dtype=self.dtype, device=self.device)  # (1, 4, h, w)
         latent = torch.cat((latent, self.x_t_latent_buffer), dim=0)  # (t, 4, h, w)
         self.stock_noise = torch.cat((self.init_noise[:1], self.stock_noise[:-1]), dim=0)  # (t, 4, h, w)
         if self.cfg_type in ('self', 'initialize'):