RohitGandikota commited on
Commit
e3abd80
1 Parent(s): 87cd610

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -2
app.py CHANGED
@@ -7,9 +7,11 @@ from train import train
7
 
8
  model_map = {'Car' : 'models/car.pt',
9
  'Van Gogh' : 'models/vangogh.pt',
 
10
  'Kilian Eng' : 'models/kilianeng.pt',
11
  'Thomas Kinkade' : 'models/thomaskinkade.pt',
12
- 'Tyler Edlin' : 'models/tyleredlin.pt'}
 
13
 
14
 
15
  class Demo:
@@ -191,7 +193,7 @@ class Demo:
191
  # del self.diffuser
192
  torch.cuda.empty_cache()
193
  # self.diffuser = StableDiffuser(scheduler='DDIM', seed=42).to('cuda').eval().half()
194
- return [gr.update(interactive=True, value='Train'), gr.update(value='Done Training'), save_path, gr.Dropdown.update(choices=list(model_map.keys()), value='Custom')]
195
 
196
 
197
  def inference(self, prompt, seed, model_name, pbar = gr.Progress(track_tqdm=True)):
 
7
 
8
  model_map = {'Car' : 'models/car.pt',
9
  'Van Gogh' : 'models/vangogh.pt',
10
+ 'Pablo Picasso': 'models/pablopicasso.pt',
11
  'Kilian Eng' : 'models/kilianeng.pt',
12
  'Thomas Kinkade' : 'models/thomaskinkade.pt',
13
+ 'Tyler Edlin' : 'models/tyleredlin.pt',
14
+ 'Kelly McKernan': '/models/kellymckernan.pt'}
15
 
16
 
17
  class Demo:
 
193
  # del self.diffuser
194
  torch.cuda.empty_cache()
195
  # self.diffuser = StableDiffuser(scheduler='DDIM', seed=42).to('cuda').eval().half()
196
+ 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')]
197
 
198
 
199
  def inference(self, prompt, seed, model_name, pbar = gr.Progress(track_tqdm=True)):