batching sample generation and cancellation support

app.py

@@ -76,7 +76,7 @@ class Demo:
                             label="Seed",
                             value=42
                         )
-                        with gr.Row(scale=1):
+                        with gr.Row():
                             self.img_width_infr = gr.Slider(
                                 label="Image width",
                                 minimum=256,
@@ -92,7 +92,7 @@ class Demo:
                                 step=64
                             )
 
-                        with gr.Row(scale=1):
+                        with gr.Row():
                             self.model_dropdown = gr.Dropdown(
                                 label="ESD Model",
                                 choices= list(model_map.keys()),
@@ -152,6 +152,15 @@ class Demo:
                             info="Image size for training, should match the model's native image size"
                         )
 
+                        self.train_sample_batch_size_input = gr.Slider(
+                            value=1,
+                            step=1,
+                            minimum=1,
+                            maximum=32,
+                            label="Sample generation batch size",
+                            info="Batch size for sample generation, larger needs more VRAM"
+                        )
+
                         self.prompt_input = gr.Text(
                             placeholder="Enter prompt...",
                             label="Prompt to Erase",
@@ -313,6 +322,7 @@ class Demo:
             self.train_use_gradient_checkpointing_input,
             self.train_seed_input,
             self.train_save_every_input,
+            self.train_sample_batch_size_input,
             self.train_validation_prompts,
             self.train_sample_positive_prompts,
             self.train_sample_negative_prompts,
@@ -322,7 +332,8 @@ class Demo:
         )
         self.train_cancel_button.click(self.cancel_training,
                                        inputs=[],
-                                       outputs=[self.train_cancel_button])
+                                       outputs=[self.train_cancel_button],
+                                       cancels=[train_event])
 
         self.export_button.click(self.export, inputs = [
             self.model_dropdown_export,
@@ -340,12 +351,14 @@ class Demo:
         return [self.model_dropdown.update(choices=list(model_map.keys()), value=current_model_name)]
 
     def cancel_training(self):
-        train.training_should_cancel = True
-        return [gr.update(value="Cancelling...", interactive=False)]
+        if self.training:
+            training_should_cancel.release()
+            print("cancellation requested...")
+        return [gr.update(value="Cancelling...", interactive=True)]
 
     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, save_every=-1,
+              seed=-1, save_every=-1, sample_batch_size=1,
               validation_prompts: str=None, sample_positive_prompts: str=None, sample_negative_prompts: str=None, validate_every_n_steps=-1,
               pbar=gr.Progress(track_tqdm=True)):
         """
@@ -373,8 +386,6 @@ class Demo:
         if self.training:
             return [gr.update(interactive=True, value='Train'), gr.update(value='Someone else is training... Try again soon'), None, gr.update()]
 
-        train.training_should_cancel = False
-
         print(f"Training {repo_id_or_path} at {img_size} to remove '{prompt}'.")
         print(f"  {train_method}, negative guidance {neg_guidance}, lr {lr}, {iterations} iterations.")
         print(f" {'✅' if use_gradient_checkpointing else '❌'} gradient checkpointing")
@@ -403,8 +414,8 @@ class Demo:
                 break
             # repeat until a not-in-use path is found
 
-        validation_prompts = [] if validation_prompts is None else validation_prompts.split('\n')
-        sample_positive_prompts = [] if sample_positive_prompts is None else sample_positive_prompts.split('\n')
+        validation_prompts = [] if validation_prompts is None else [p for p in validation_prompts.split('\n') if len(p)>0]
+        sample_positive_prompts = [] if sample_positive_prompts is None else [p for p in sample_positive_prompts.split('\n') if len(p)>0]
         sample_negative_prompts = [] if sample_negative_prompts is None else sample_negative_prompts.split('\n')
         print(f"validation prompts: {validation_prompts}")
         print(f"sample positive prompts: {sample_positive_prompts}")
@@ -413,9 +424,11 @@ class Demo:
         try:
             self.training = True
             self.train_cancel_button.update(interactive=True)
+            batch_size = 1 # other batch sizes are non-functional
             save_path = 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,
                   seed=int(seed), save_every_n_steps=int(save_every),
+                              batch_size=int(batch_size), sample_batch_size=int(sample_batch_size),
                   validate_every_n_steps=validate_every_n_steps, validation_prompts=validation_prompts,
                   sample_positive_prompts=sample_positive_prompts, sample_negative_prompts=sample_negative_prompts)
             if save_path is None:

train.py

@@ -1,5 +1,6 @@
 import os.path
 import random
+import multiprocessing
 
 from accelerate.utils import set_seed
 from diffusers import StableDiffusionPipeline
@@ -15,7 +16,7 @@ from isolate_rng import isolate_rng
 from memory_efficiency import MemoryEfficiencyWrapper
 from torch.utils.tensorboard import SummaryWriter
 
-training_should_cancel = False
+training_should_cancel = multiprocessing.Semaphore(0)
 
 def validate(diffuser: StableDiffuser, finetuner: FineTunedModel,
              validation_embeddings: torch.FloatTensor,
@@ -24,8 +25,11 @@ def validate(diffuser: StableDiffuser, finetuner: FineTunedModel,
              logger: SummaryWriter, use_amp: bool,
              global_step: int,
              validation_seed: int = 555,
+             batch_size: int = 1,
+             sample_batch_size: int = 1 # might need to be smaller than batch_size
              ):
     print("validating...")
+    assert batch_size==1, "batch_size != 1 not implemented work"
     with isolate_rng(include_cuda=True), torch.no_grad():
         set_seed(validation_seed)
         criteria = torch.nn.MSELoss()
@@ -33,14 +37,14 @@ def validate(diffuser: StableDiffuser, finetuner: FineTunedModel,
         val_count = 5
 
         nsteps=50
-        num_validation_prompts = validation_embeddings.shape[0] // 2
+        num_validation_batches = validation_embeddings.shape[0] // (batch_size*2)
 
-        for i in tqdm(range(num_validation_prompts)):
-            if training_should_cancel:
+        for i in tqdm(range(num_validation_batches)):
+            if training_should_cancel.acquire(block=False):
                 print("cancel requested, bailing")
                 return
             accumulated_loss = None
-            this_validation_embeddings = validation_embeddings[i*2:i*2+2]
+            this_validation_embeddings = validation_embeddings[i*batch_size*2:(i+1)*batch_size*2]
             for j in range(val_count):
                 iteration = random.randint(1, nsteps)
                 diffused_latents = get_diffused_latents(diffuser, nsteps, this_validation_embeddings, iteration, use_amp)
@@ -55,12 +59,11 @@ def validate(diffuser: StableDiffuser, finetuner: FineTunedModel,
                 loss = criteria(negative_latents, neutral_latents - (negative_guidance*(positive_latents - neutral_latents)))
                 accumulated_loss = (accumulated_loss or 0) + loss.item()
             logger.add_scalar(f"loss/val_{i}", accumulated_loss/val_count, global_step=global_step)
-            pbar.step()
 
-        num_samples = sample_embeddings.shape[0] // 2
-        for i in tqdm(range(0, num_samples)):
-            print(f'making sample {i}...')
-            if training_should_cancel:
+        num_sample_batches = sample_embeddings.shape[0] // (sample_batch_size*2)
+        for i in tqdm(range(0, num_sample_batches)):
+            print(f'making sample batch {i}...')
+            if training_should_cancel.acquire(block=False):
                 print("cancel requested, bailing")
                 return
             with finetuner:
@@ -72,10 +75,16 @@ def validate(diffuser: StableDiffuser, finetuner: FineTunedModel,
                                                    safety_checker=None,
                                                    feature_extractor=None,
                                                requires_safety_checker=False)
-                images = pipeline(prompt_embeds=sample_embeddings[i*2+1:i*2+2], negative_prompt_embeds=sample_embeddings[i*2:i*2+1],
+                batch_start = (i * sample_batch_size)*2
+                next_batch_start = batch_start + sample_batch_size*2 + 1
+                batch_negative_prompt_embeds = torch.cat([sample_embeddings[i+0:i+1] for i in range(batch_start, next_batch_start, 2)])
+                batch_prompt_embeds = torch.cat([sample_embeddings[i+1:i+2] for i in range(batch_start, next_batch_start, 2)])
+                images = pipeline(prompt_embeds=batch_prompt_embeds, #sample_embeddings[i*2+1:i*2+2],
+                                  negative_prompt_embeds=batch_negative_prompt_embeds, # sample_embeddings[i*2:i*2+1],
                                   num_inference_steps=50)
-                image_tensor = transforms.ToTensor()(images.images[0])
-                logger.add_image(f"samples/{i}", img_tensor=image_tensor, global_step=global_step)
+                for j in range(sample_batch_size):
+                    image_tensor = transforms.ToTensor()(images.images[j])
+                    logger.add_image(f"samples/{i*sample_batch_size+j}", img_tensor=image_tensor, global_step=global_step)
 
             """
             with finetuner, torch.cuda.amp.autocast(enabled=use_amp):
@@ -90,6 +99,7 @@ def validate(diffuser: StableDiffuser, finetuner: FineTunedModel,
 
 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, seed=-1,
+          batch_size=1, sample_batch_size=1,
           save_every_n_steps=-1, validate_every_n_steps=-1,
           validation_prompts=[], sample_positive_prompts=[], sample_negative_prompts=[]):
 
@@ -101,8 +111,6 @@ def train(repo_id_or_path, img_size, prompt, modules, freeze_modules, iterations
     neutral_latents = None
     positive_latents = None
 
-    global training_should_cancel
-
     nsteps = 50
     print(f"using img_size of {img_size}")
     diffuser = StableDiffuser(scheduler='DDIM', repo_id_or_path=repo_id_or_path, native_img_size=img_size).to('cuda')
@@ -135,6 +143,13 @@ def train(repo_id_or_path, img_size, prompt, modules, freeze_modules, iterations
             validation_embeddings = diffuser.get_cond_and_uncond_embeddings(validation_prompts, n_imgs=1)
             sample_embeddings = diffuser.get_cond_and_uncond_embeddings(sample_positive_prompts, sample_negative_prompts, n_imgs=1)
 
+        for i, validation_prompt in enumerate(validation_prompts):
+            logger.add_text(f"val/{i}", f"validation prompt: \"{validation_prompt}\"")
+        for i in range(len(sample_positive_prompts)):
+            positive_prompt = sample_positive_prompts[i]
+            negative_prompt = "" if i >= len(sample_negative_prompts) else sample_negative_prompts[i]
+            logger.add_text(f"sample/{i}", f"sample prompt: \"{positive_prompt}\", negative: \"{negative_prompt}\"")
+
         #if use_amp:
         #    diffuser.vae = diffuser.vae.to(diffuser.vae.device, dtype=torch.float16)
 
@@ -151,14 +166,15 @@ def train(repo_id_or_path, img_size, prompt, modules, freeze_modules, iterations
                  validation_embeddings=validation_embeddings,
                  sample_embeddings=sample_embeddings,
                  neutral_embeddings=neutral_text_embeddings,
-                 logger=logger, use_amp=False, global_step=0)
+                 logger=logger, use_amp=False, global_step=0,
+                 batch_size=batch_size, sample_batch_size=sample_batch_size)
 
         prev_losses = []
         start_loss = None
         max_prev_loss_count = 10
         try:
             for i in pbar:
-                if training_should_cancel:
+                if training_should_cancel.acquire(block=False):
                     print("cancel requested, bailing")
                     return None
 
@@ -210,7 +226,8 @@ def train(repo_id_or_path, img_size, prompt, modules, freeze_modules, iterations
                              validation_embeddings=validation_embeddings,
                              sample_embeddings=sample_embeddings,
                              neutral_embeddings=neutral_text_embeddings,
-                             logger=logger, use_amp=False, global_step=i)
+                             logger=logger, use_amp=False, global_step=i,
+                             batch_size=batch_size, sample_batch_size=sample_batch_size)
             torch.save(finetuner.state_dict(), save_path)
             return save_path
         finally:
@@ -220,7 +237,7 @@ def train(repo_id_or_path, img_size, prompt, modules, freeze_modules, iterations
 
 def get_diffused_latents(diffuser, nsteps, text_embeddings, end_iteration, use_amp):
     diffuser.set_scheduler_timesteps(nsteps)
-    latents = diffuser.get_initial_latents(1, n_prompts=1)
+    latents = diffuser.get_initial_latents(len(text_embeddings)//2, n_prompts=1)
     latents_steps, _ = diffuser.diffusion(
         latents,
         text_embeddings,