Updated diffuser to use seperate generator

StableDiffuser.py

@@ -34,14 +34,11 @@ def default_parser():
 class StableDiffuser(torch.nn.Module):
 
     def __init__(self,
-                scheduler='LMS',
-                seed=None
+                scheduler='LMS'
         ):
 
         super().__init__()
 
-        self._seed = seed
-
         # Load the autoencoder model which will be used to decode the latents into image space.
         self.vae = AutoencoderKL.from_pretrained(
             "CompVis/stable-diffusion-v1-4", subfolder="vae")
@@ -62,25 +59,16 @@ class StableDiffuser(torch.nn.Module):
             self.scheduler = DDIMScheduler.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="scheduler")
         elif scheduler == 'DDPM':
             self.scheduler = DDPMScheduler.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="scheduler")    
-        self.generator = torch.Generator()
-
-        if self._seed is not None:
-
-            self.seed(seed)
 
         self.eval()
 
-    def seed(self, seed):
-
-        self.generator = torch.manual_seed(seed)
-
-    def get_noise(self, batch_size, img_size):
+    def get_noise(self, batch_size, img_size, generator=None):
 
         param = list(self.parameters())[0]
 
         return torch.randn(
             (batch_size, self.unet.in_channels, img_size // 8, img_size // 8),
-            generator=self.generator).type(param.dtype).to(param.device)
+            generator=generator).type(param.dtype).to(param.device)
 
     def add_noise(self, latents, noise, step):
 
@@ -118,9 +106,9 @@ class StableDiffuser(torch.nn.Module):
     def set_scheduler_timesteps(self, n_steps):
         self.scheduler.set_timesteps(n_steps, device=self.unet.device)
 
-    def get_initial_latents(self, n_imgs, img_size, n_prompts):
+    def get_initial_latents(self, n_imgs, img_size, n_prompts, generator=None):
 
-        noise = self.get_noise(n_imgs, img_size).repeat(n_prompts, 1, 1, 1)
+        noise = self.get_noise(n_imgs, img_size, generator=generator).repeat(n_prompts, 1, 1, 1)
 
         latents = noise * self.scheduler.init_noise_sigma
 
@@ -221,7 +209,7 @@ class StableDiffuser(torch.nn.Module):
                  n_steps=50,
                  n_imgs=1,
                  end_iteration=None,
-                 reseed=False,
+                 generator=None,
                  **kwargs
                  ):
 
@@ -233,11 +221,7 @@ class StableDiffuser(torch.nn.Module):
 
         self.set_scheduler_timesteps(n_steps)
 
-        if reseed:
-
-            self.seed(self._seed)
-
-        latents = self.get_initial_latents(n_imgs, img_size, len(prompts))
+        latents = self.get_initial_latents(n_imgs, img_size, len(prompts), generator=generator)
 
         text_embeddings = self.get_text_embeddings(prompts,n_imgs=n_imgs)
 

app.py

@@ -17,18 +17,18 @@ class Demo:
     def __init__(self) -> None:
 
         self.training = False
+        self.generating = False
 
-        self.diffuser = StableDiffuser(scheduler='DDIM', seed=42).to('cuda').eval().half()
+        self.diffuser = StableDiffuser(scheduler='DDIM').to('cuda').eval().half()
 
         with gr.Blocks() as demo:
             self.layout()
-            demo.queue(concurrency_count=2).launch()
+            demo.queue(concurrency_count=5).launch()
 
 
     def layout(self):
 
         with gr.Row():
-            
                 
             with gr.Tab("Test") as inference_column:
 
@@ -152,13 +152,10 @@ class Demo:
         )
 
     def train(self, prompt, train_method, neg_guidance, iterations, lr, pbar = gr.Progress(track_tqdm=True)):
-        # self.diffuser = StableDiffuser(scheduler='DDIM', seed=42).to('cuda').eval().half()
+
         if self.training:
             return [gr.update(interactive=True, value='Train'), gr.update(value='Someone else is training... Try again soon'), None, gr.update()]
-        # clear the diffusers
-        # del self.diffuser
-        # torch.cuda.empty_cache()
-        
+
         if train_method == 'ESD-x':
 
             modules = ".*attn2$"
@@ -188,45 +185,46 @@ class Demo:
 
         model_map['Custom'] = save_path
 
-        # del self.diffuser
-        torch.cuda.empty_cache()
-        # self.diffuser = StableDiffuser(scheduler='DDIM', seed=42).to('cuda').eval().half()
         return [gr.update(interactive=True, value='Train'), gr.update(value='Done Training'), save_path, gr.Dropdown.update(choices=list(model_map.keys()), value='Custom')]
 
 
     def inference(self, prompt, seed, model_name, pbar = gr.Progress(track_tqdm=True)):
         
-        self.diffuser._seed = seed or 42
+        seed = seed or 42
+
+        generator = torch.manual_seed(seed)
 
         model_path = model_map[model_name]
         
         checkpoint = torch.load(model_path)
 
-        self.finetuner = FineTunedModel.from_checkpoint(self.diffuser, checkpoint).eval().half()
+        finetuner = FineTunedModel.from_checkpoint(self.diffuser, checkpoint).eval().half()
 
         torch.cuda.empty_cache()
 
         images = self.diffuser(
             prompt,
             n_steps=50,
-            reseed=True
+            generator=generator
         )
 
         orig_image = images[0][0]
 
         torch.cuda.empty_cache()
 
-        with self.finetuner:
+        generator = torch.manual_seed(seed)
+
+        with finetuner:
 
             images = self.diffuser(
                 prompt,
                 n_steps=50,
-                reseed=True
+                generator=generator
             )
 
         edited_image = images[0][0]
 
-        del self.finetuner
+        del finetuner
         torch.cuda.empty_cache()
 
         return edited_image, orig_image

train.py

@@ -10,9 +10,6 @@ def train(prompt, modules, freeze_modules, iterations, negative_guidance, lr, sa
     diffuser = StableDiffuser(scheduler='DDIM').to('cuda')
     diffuser.train()
 
-
-
-
     finetuner = FineTunedModel(diffuser, modules, frozen_modules=freeze_modules)
 
     optimizer = torch.optim.Adam(finetuner.parameters(), lr=lr)
@@ -69,14 +66,13 @@ def train(prompt, modules, freeze_modules, iterations, negative_guidance, lr, sa
 
         loss = criteria(negative_latents, neutral_latents - (negative_guidance*(positive_latents - neutral_latents))) #loss = criteria(e_n, e_0) works the best try 5000 epochs
         
-        del negative_latents, neutral_latents, positive_latents, latents_steps, latents
-        torch.cuda.empty_cache()
-
         loss.backward()
         optimizer.step()
 
     torch.save(finetuner.state_dict(), save_path)
-    del diffuser
+
+    del diffuser, loss, optimizer, finetuner, negative_latents, neutral_latents, positive_latents, latents_steps, latents
+
     torch.cuda.empty_cache()
 if __name__ == '__main__':