RohitGandikota commited on
Commit
5a392b9
·
1 Parent(s): 2688a68

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -229,7 +229,7 @@ class Demo:
229
  n_steps=50,
230
  generator=generator
231
  )
232
-
233
  orig_image = images[0][0]
234
 
235
  torch.cuda.empty_cache()
@@ -243,7 +243,7 @@ class Demo:
243
  n_steps=50,
244
  generator=generator
245
  )
246
-
247
  edited_image = images[0][0]
248
 
249
  del finetuner
 
229
  n_steps=50,
230
  generator=generator
231
  )
232
+ images = self.diffuser.run_safety_checker(images=images, device='cuda')
233
  orig_image = images[0][0]
234
 
235
  torch.cuda.empty_cache()
 
243
  n_steps=50,
244
  generator=generator
245
  )
246
+ images = self.diffuser.run_safety_checker(images=images, device='cuda')
247
  edited_image = images[0][0]
248
 
249
  del finetuner