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

app.py

@@ -5,11 +5,19 @@ from StableDiffuser import StableDiffuser
 from tqdm import tqdm
 from train import train
 
-model_map = {'Car' : 'models/car.pt',
-             'Van Gogh' : 'models/vangogh.pt', 
+#import os
+#os.system("nvidia-smi topo -m")
+
+model_map = {'Van Gogh' : 'models/vangogh.pt', 
+             'Pablo Picasso': 'models/pablopicasso.pt',
+             'Car' : 'models/car.pt',
+             'Garbage Truck': 'models/garbagetruck.pt',
+             'French Horn': 'models/frenchhorn.pt',
              'Kilian Eng' : 'models/kilianeng.pt',
              'Thomas Kinkade' : 'models/thomaskinkade.pt',
-             'Tyler Edlin' : 'models/tyleredlin.pt'}
+             'Tyler Edlin' : 'models/tyleredlin.pt',
+             'Kelly McKernan': 'models/kellymckernan.pt',
+             'Rembrandt': 'models/rembrandt.pt' }
 
 
 class Demo:
@@ -185,7 +193,7 @@ class Demo:
 
         model_map['Custom'] = save_path
 
-        return [gr.update(interactive=True, value='Train'), gr.update(value='Done Training'), save_path, gr.Dropdown.update(choices=list(model_map.keys()), value='Custom')]
+        return [gr.update(interactive=True, value='Train'), gr.update(value='Done Training! \n Try your custom model in the "Test" tab'), 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)):

models/frenchhorn.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:48dc787885e54bbad818b57205ccb39796051b07b98a00b43776fb7d7e375fc0
+size 3438372469

models/garbagetruck.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9e8bc8d2d973e941a16a03cae81113b5f5da07245888dd29934ea77af9242aba
+size 3438373845

models/kellymckernan.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dee79125ec560c09fcd8820f2df9e3bee94b8a8e8ef3c8e9b330a80ec87cf45e
+size 175879857

models/pablopicasso.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2988ba3a00af8e7796fdd459d31c63e45470d5b00b88fa0a18f1722c8c55fd9a
+size 175879775

models/rembrandt.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b95b9bd8d13f2cd8fc389ec2ec6246a3da334cc0c09cba323a3743ae2453cf58
+size 175879529