cleanup and try to get cancellation working

README.md

@@ -10,7 +10,21 @@ pinned: false
 license: mit
 ---
 
+# A GUI with custom model support, validation, and sample generation for "Erasing Concepts from Diffusion Models"
 
+Enables xformers, 8 bit AdamW via bitsandbytes, and AMP - editing SD1.5 models works with 16GB VRAM, and 2.5 models including the ESD-u training works with 24GB VRAM.
+
+## Quick start
+
+To run on vast.ai, use eg `pytorch/pytorch:2.0.1-cuda11.7-cudnn8-devel` - you need `-devel` for 8bit AdamW to work.
+
+On the dev machine:
+```
+pip install -r requirements.txt
+python app.py
+```
+
+then use the Gradio interface at port 7860.
 
 # Erasing Concepts from Diffusion Models
 

StableDiffuser.py

@@ -107,6 +107,7 @@ class StableDiffuser(torch.nn.Module):
         return latents
 
     def get_cond_and_uncond_embeddings(self, prompts, negative_prompts=None, n_imgs=1):
+        assert n_imgs == 1
         text_tokens = self.text_tokenize(prompts)
         text_embeddings = self.text_encode(text_tokens)
         if negative_prompts is None:
@@ -115,8 +116,9 @@ class StableDiffuser(torch.nn.Module):
             negative_prompts.append("")
         unconditional_tokens = self.text_tokenize(negative_prompts)
         unconditional_embeddings = self.text_encode(unconditional_tokens)
-        text_embeddings = torch.cat([unconditional_embeddings, text_embeddings]).repeat_interleave(n_imgs, dim=0)
-        return text_embeddings
+        combined_embeddings = [torch.cat([unconditional_embeddings[i:i+1], text_embeddings[i:i+1]]) for i in range(len(prompts))]
+        combined_embeddings = torch.cat(combined_embeddings)
+        return combined_embeddings
 
     def predict_noise(self,
              iteration,

train.py

@@ -34,7 +34,11 @@ def validate(diffuser: StableDiffuser, finetuner: FineTunedModel,
 
         nsteps=50
         num_validation_prompts = validation_embeddings.shape[0] // 2
-        for i in range(0, num_validation_prompts):
+
+        for i in tqdm(range(num_validation_prompts))
+            if training_should_cancel:
+                print("cancel requested, bailing")
+                return
             accumulated_loss = None
             this_validation_embeddings = validation_embeddings[i*2:i*2+2]
             for j in range(val_count):
@@ -51,10 +55,14 @@ 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 range(0, num_samples):
+        for i in tqdm(range(0, num_samples));
             print(f'making sample {i}...')
+            if training_should_cancel:
+                print("cancel requested, bailing")
+                return
             with finetuner:
                 pipeline = StableDiffusionPipeline(vae=diffuser.vae,
                                                text_encoder=diffuser.text_encoder,
@@ -93,6 +101,8 @@ 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')
@@ -137,13 +147,19 @@ def train(repo_id_or_path, img_size, prompt, modules, freeze_modules, iterations
             seed = random.randint(0, 2 ** 30)
         set_seed(int(seed))
 
+        validate(diffuser, finetuner,
+                 validation_embeddings=validation_embeddings,
+                 sample_embeddings=sample_embeddings,
+                 neutral_embeddings=neutral_text_embeddings,
+                 logger=logger, use_amp=False, global_step=0)
+
         prev_losses = []
         start_loss = None
         max_prev_loss_count = 10
         try:
             for i in pbar:
                 if training_should_cancel:
-                    print("received cancellation request")
+                    print("cancel requested, bailing")
                     return None
 
                 with torch.no_grad():