cocktailpeanut 6Morpheus6 commited on
Commit
f4d41e7
1 Parent(s): 5c2594c

empty cache (#2)

Browse files

- empty cache (cc5890be247ea3a953de00e5c45591e10496867b)


Co-authored-by: Morpheus <[email protected]>

Files changed (1) hide show
  1. app.py +7 -0
app.py CHANGED
@@ -6,6 +6,7 @@ import cv2
6
  import gradio as gr
7
  import numpy as np
8
  import PIL
 
9
  #import spaces
10
  import torch
11
  from diffusers import LCMScheduler
@@ -310,6 +311,12 @@ def generate_image(
310
  generator=generator,
311
  ).images
312
 
 
 
 
 
 
 
313
  return images[0], gr.update(visible=True)
314
 
315
 
 
6
  import gradio as gr
7
  import numpy as np
8
  import PIL
9
+ import gc
10
  #import spaces
11
  import torch
12
  from diffusers import LCMScheduler
 
311
  generator=generator,
312
  ).images
313
 
314
+ gc.collect()
315
+ if device == 'mps':
316
+ torch.mps.empty_cache()
317
+ elif device == 'cuda':
318
+ torch.cuda.empty_cache()
319
+
320
  return images[0], gr.update(visible=True)
321
 
322