Update app.py

app.py

@@ -7,9 +7,11 @@ from train import train
 
 model_map = {'Car' : 'models/car.pt',
              'Van Gogh' : 'models/vangogh.pt', 
+             'Pablo Picasso': 'models/pablopicasso.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'}
 
 
 class Demo:
@@ -191,7 +193,7 @@ class Demo:
         # 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')]
+        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)):