RohitGandikota commited on
Commit
75cb83b
1 Parent(s): 0a3fdbd

initialise and del models in inference

Browse files
Files changed (1) hide show
  1. app.py +6 -3
app.py CHANGED
@@ -18,7 +18,7 @@ class Demo:
18
 
19
  self.training = False
20
 
21
- self.diffuser = StableDiffuser(scheduler='DDIM', seed=42).to('cuda').eval().half()
22
 
23
  with gr.Blocks() as demo:
24
  self.layout()
@@ -152,7 +152,7 @@ class Demo:
152
  )
153
 
154
  def train(self, prompt, train_method, neg_guidance, iterations, lr, pbar = gr.Progress(track_tqdm=True)):
155
-
156
  if self.training:
157
  return [gr.update(interactive=True, value='Train'), gr.update(value='Someone else is training... Try again soon'), None, gr.update()]
158
 
@@ -185,11 +185,13 @@ class Demo:
185
 
186
  model_map['Custom'] = save_path
187
 
 
 
188
  return [gr.update(interactive=True, value='Train'), gr.update(value='Done Training'), save_path, gr.Dropdown.update(choices=list(model_map.keys()), value='Custom')]
189
 
190
 
191
  def inference(self, prompt, seed, model_name, pbar = gr.Progress(track_tqdm=True)):
192
-
193
  self.diffuser._seed = seed or 42
194
 
195
  model_path = model_map[model_name]
@@ -220,6 +222,7 @@ class Demo:
220
 
221
  edited_image = images[0][0]
222
 
 
223
  torch.cuda.empty_cache()
224
 
225
  return edited_image, orig_image
 
18
 
19
  self.training = False
20
 
21
+
22
 
23
  with gr.Blocks() as demo:
24
  self.layout()
 
152
  )
153
 
154
  def train(self, prompt, train_method, neg_guidance, iterations, lr, pbar = gr.Progress(track_tqdm=True)):
155
+ # self.diffuser = StableDiffuser(scheduler='DDIM', seed=42).to('cuda').eval().half()
156
  if self.training:
157
  return [gr.update(interactive=True, value='Train'), gr.update(value='Someone else is training... Try again soon'), None, gr.update()]
158
 
 
185
 
186
  model_map['Custom'] = save_path
187
 
188
+ # del self.diffuser
189
+ # torch.cuda.empty_cache()
190
  return [gr.update(interactive=True, value='Train'), gr.update(value='Done Training'), save_path, gr.Dropdown.update(choices=list(model_map.keys()), value='Custom')]
191
 
192
 
193
  def inference(self, prompt, seed, model_name, pbar = gr.Progress(track_tqdm=True)):
194
+ self.diffuser = StableDiffuser(scheduler='DDIM', seed=42).to('cuda').eval().half()
195
  self.diffuser._seed = seed or 42
196
 
197
  model_path = model_map[model_name]
 
222
 
223
  edited_image = images[0][0]
224
 
225
+ del self.finetuner, self.diffuser
226
  torch.cuda.empty_cache()
227
 
228
  return edited_image, orig_image