add train seed

StableDiffuser.py

@@ -4,6 +4,7 @@ import torch
 from baukit import TraceDict
 from diffusers import StableDiffusionPipeline
 from PIL import Image
+from torch.cuda.amp import autocast
 from tqdm.auto import tqdm
 from diffusers.schedulers.scheduling_ddim import DDIMScheduler
 from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
@@ -142,6 +143,7 @@ class StableDiffuser(torch.nn.Module):
                   pred_x0=False,
                   trace_args=None,                  
                   show_progress=True,
+                  use_amp=False,
                   **kwargs):
 
         latents_steps = []
@@ -153,11 +155,12 @@ class StableDiffuser(torch.nn.Module):
             if trace_args:
                 trace = TraceDict(self, **trace_args)
 
-            noise_pred = self.predict_noise(
-                iteration, 
-                latents, 
-                text_embeddings,
-                **kwargs)
+            with autocast(enabled=use_amp):
+                noise_pred = self.predict_noise(
+                    iteration,
+                    latents,
+                    text_embeddings,
+                    **kwargs)
 
             # compute the previous noisy sample x_t -> x_t-1
             output = self.scheduler.step(noise_pred, self.scheduler.timesteps[iteration], latents)

app.py

@@ -191,12 +191,20 @@ class Demo:
                             label="Learning Rate",
                             info='Learning rate used to train'
                         )
+                        self.train_seed_input = gr.Number(
+                            value=-1,
+                            label="Seed",
+                            info="Set to a fixed number for reproducible training results, or use -1 to pick randomly"
+                        )
 
-                        with gr.Row():
-                            self.train_use_adamw8bit_input = gr.Checkbox(label="8bit AdamW", value=False)
-                            self.train_use_xformers_input = gr.Checkbox(label="xformers", value=True)
-                            self.train_use_amp_input = gr.Checkbox(label="AMP", value=True)
-                            #self.train_use_gradient_checkpointing_input = gr.Checkbox(label="Gradient checkpointing", value=True)
+                        with gr.Column():
+                            self.train_memory_options = gr.Markdown(interactive=False,
+r                                value='Performance and VRAM usage optimizations, may not work on all devices.')
+                            with gr.Row():
+                                self.train_use_adamw8bit_input = gr.Checkbox(label="8bit AdamW", value=True)
+                                self.train_use_xformers_input = gr.Checkbox(label="xformers", value=True)
+                                self.train_use_amp_input = gr.Checkbox(label="AMP", value=True)
+                                self.train_use_gradient_checkpointing_input = gr.Checkbox(label="Gradient checkpointing", value=True)
 
                     with gr.Column(scale=1):
 
@@ -209,16 +217,13 @@ class Demo:
                         self.download = gr.Files()
 
             with gr.Tab("Export") as export_column:
-
                 with gr.Row():
-
                     self.explain_train= gr.Markdown(interactive=False,
-                                      value='Export a model to Diffusers format. Please enter the base model and select the editing weights.')
+                        value='Export a model to Diffusers format. Please enter the base model and select the editing weights.')
 
                 with gr.Row():
 
                     with gr.Column(scale=3):
-
                         self.base_repo_id_or_path_input_export = gr.Text(
                             label="Base model",
                             value="CompVis/stable-diffusion-v1-4",
@@ -272,7 +277,8 @@ class Demo:
             self.train_use_adamw8bit_input,
             self.train_use_xformers_input,
             self.train_use_amp_input,
-            #self.train_use_gradient_checkpointing_input
+            self.train_use_gradient_checkpointing_input,
+            self.train_seed_input,
         ],
         outputs=[self.train_button, self.train_status, self.download, self.model_dropdown]
         )
@@ -287,6 +293,7 @@ class Demo:
 
     def train(self, repo_id_or_path, img_size, prompt, train_method, neg_guidance, iterations, lr,
               use_adamw8bit=True, use_xformers=False, use_amp=False, use_gradient_checkpointing=False,
+              seed = -1,
               pbar = gr.Progress(track_tqdm=True)):
 
         if self.training:
@@ -311,19 +318,25 @@ class Demo:
             modules = ".*attn1$"
             frozen = []
 
-        randn = torch.randint(1, 10000000, (1,)).item()
+        # build a save path, ensure it isn't in use
+        while True:
+            randn = torch.randint(1, 10000000, (1,)).item()
+            options = f'{"a8" if use_adamw8bit else ""}{"AM" if use_amp else ""}{"xf" if use_xformers else ""}{"gc" if use_gradient_checkpointing else ""}'
+            save_path = f"models/{prompt.lower().replace(' ', '')}_{train_method}_ng{neg_guidance}_lr{lr}_iter{iterations}_seed{seed}_{options}__{randn}.pt"
+            if not os.path.exists(save_path):
+                break
+            # repeat until a not-in-use path is found
 
-        save_path = f"models/{randn}_{prompt.lower().replace(' ', '')}_{train_method}_ng{neg_guidance}_lr{lr}_iter{iterations}.pt"
         try:
             self.training = True
             train(repo_id_or_path, img_size, prompt, modules, frozen, iterations, neg_guidance, lr, save_path,
-                use_adamw8bit, use_xformers, use_amp, use_gradient_checkpointing)
+                use_adamw8bit, use_xformers, use_amp, use_gradient_checkpointing, seed=seed)
         finally:
             self.training = False
 
         torch.cuda.empty_cache()
 
-        new_model_name = f'*new* {os.path.basename(save_path)}'
+        new_model_name = f'{os.path.basename(save_path)}'
         model_map[new_model_name] = save_path
 
         return [gr.update(interactive=True, value='Train'),

memory_efficiency.py

@@ -37,7 +37,7 @@ class MemoryEfficiencyWrapper:
                     print("failed to load xformers, using attention slicing instead")
                     self.diffuser.unet.set_attention_slice("auto")
                     pass
-            elif (not self.amp and self.is_sd1attn):
+            elif (not self.use_amp and self.is_sd1attn):
                 print("AMP is disabled but model is SD1.X, using attention slicing instead of xformers")
                 self.diffuser.unet.set_attention_slice("auto")
         else:

train.py

@@ -1,3 +1,6 @@
+from random import random
+
+from accelerate.utils import set_seed
 from torch.cuda.amp import autocast
 
 from StableDiffuser import StableDiffuser
@@ -9,7 +12,7 @@ from memory_efficiency import MemoryEfficiencyWrapper
 
 
 def train(repo_id_or_path, img_size, prompt, modules, freeze_modules, iterations, negative_guidance, lr, save_path,
-          use_adamw8bit=True, use_xformers=True, use_amp=True, use_gradient_checkpointing=False):
+          use_adamw8bit=True, use_xformers=True, use_amp=True, use_gradient_checkpointing=False, seed=-1):
 
     nsteps = 50
     diffuser = StableDiffuser(scheduler='DDIM', repo_id_or_path=repo_id_or_path).to('cuda')
@@ -47,6 +50,10 @@ def train(repo_id_or_path, img_size, prompt, modules, freeze_modules, iterations
 
         print(f"using img_size of {img_size}")
 
+        if seed == -1:
+            seed = random.randint(0, 2 ** 30)
+        set_seed(seed)
+
         for i in pbar:
             with torch.no_grad():
                 diffuser.set_scheduler_timesteps(nsteps)
@@ -55,14 +62,15 @@ def train(repo_id_or_path, img_size, prompt, modules, freeze_modules, iterations
                 iteration = torch.randint(1, nsteps - 1, (1,)).item()
                 latents = diffuser.get_initial_latents(1, width=img_size, height=img_size, n_prompts=1)
 
-                with autocast(enabled=use_amp), finetuner:
+                with finetuner:
                     latents_steps, _ = diffuser.diffusion(
                         latents,
                         positive_text_embeddings,
                         start_iteration=0,
                         end_iteration=iteration,
                         guidance_scale=3,
-                        show_progress=False
+                        show_progress=False,
+                        use_amp=use_amp
                     )
 
                 diffuser.set_scheduler_timesteps(1000)
@@ -82,7 +90,7 @@ def train(repo_id_or_path, img_size, prompt, modules, freeze_modules, iterations
             # loss = criteria(e_n, e_0) works the best try 5000 epochs
             loss = criteria(negative_latents, neutral_latents - (negative_guidance*(positive_latents - neutral_latents)))
             memory_efficiency_wrapper.step(optimizer, loss)
-            optimizer.step()
+            optimizer.zero_grad()
 
     torch.save(finetuner.state_dict(), save_path)
 
@@ -104,5 +112,11 @@ if __name__ == '__main__':
     parser.add_argument('--iterations', type=int, required=True)
     parser.add_argument('--lr', type=float, required=True)
     parser.add_argument('--negative_guidance', type=float, required=True)
+    parser.add_argument('--seed', type=int, required=False, default=-1,
+                        help='Training seed for reproducible results, or -1 to pick a random seed')
+    parser.add_argument('--use_adamw8bit', action='store_true')
+    parser.add_argument('--use_xformers', action='store_true')
+    parser.add_argument('--use_amp', action='store_true')
+    parser.add_argument('--use_gradient_checkpointing', action='store_true')
 
     train(**vars(parser.parse_args()))