Merge branch 'main' of https://huggingface.co./spaces/baulab/Erasing-Concepts-In-Diffusion into main

app.py

@@ -5,10 +5,11 @@ from StableDiffuser import StableDiffuser
 from tqdm import tqdm
 from train import train
 
-model_map = {
-    'Car' : 'models/car.pt',
-    'Van Gogh' : 'models/vangogh.pt',
-}
+model_map = {'Car' : 'models/car.pt',
+             'Van Gogh' : 'models/vangogh.pt', 
+             'Kilian Eng' : 'models/kilianeng.pt',
+             'Thomas Kinkade' : 'models/thomaskinkade.pt',
+             'Tyler Edlin' : 'models/tyleredlin.pt'}
 
 
 class Demo:
@@ -50,7 +51,7 @@ class Demo:
 
                             self.model_dropdown = gr.Dropdown(
                                 label="ESD Model",
-                                choices=['Van Gogh', 'Car'],
+                                choices= list(model_map.keys()),
                                 value='Van Gogh',
                                 interactive=True
                             )
@@ -151,10 +152,13 @@ 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$"
@@ -184,11 +188,14 @@ 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
 
         model_path = model_map[model_name]
@@ -219,6 +226,7 @@ class Demo:
 
         edited_image = images[0][0]
 
+        del self.finetuner
         torch.cuda.empty_cache()
 
         return edited_image, orig_image

train.py

@@ -76,7 +76,8 @@ def train(prompt, modules, freeze_modules, iterations, negative_guidance, lr, sa
         optimizer.step()
 
     torch.save(finetuner.state_dict(), save_path)
-
+    del diffuser
+    torch.cuda.empty_cache()
 if __name__ == '__main__':
 
     import argparse