update

app.py

@@ -7,11 +7,16 @@ import math
 #import spaces 
 import torch 
 from PIL import Image
+import gc
 
 if torch.backends.mps.is_available():
     DEVICE = "mps"
+    torch.mps.empty_cache()
+    gc.collect()
 elif torch.cuda.is_available():
     DEVICE = "cuda"
+    torch.cuda.empty_cache()
+    gc.collect()
 else:
     DEVICE = "cpu"
 
@@ -60,7 +65,14 @@ def run_edit(image, prompt, resolution, negative_prompt="", guidance_scale=7, pr
     #image.resize((resolution, resolution))
     #return pipe_edit(prompt=prompt,image=image,height=resolution,width=resolution,negative_prompt=negative_prompt, guidance_scale=guidance_scale,num_inference_steps=20).images[0]
     print(f"width={image.width}, height={image.height}")
-    return pipe_edit(prompt=prompt,image=image,height=image.height,width=image.width,negative_prompt=negative_prompt, guidance_scale=guidance_scale,num_inference_steps=20).images[0]
+    img = pipe_edit(prompt=prompt,image=image,height=image.height,width=image.width,negative_prompt=negative_prompt, guidance_scale=guidance_scale,num_inference_steps=20).images[0]
+    if DEVICE == "cuda":
+        torch.cuda.empty_cache()
+        gc.collect()
+    elif DEVICE == "mps":
+        torch.mps.empty_cache()
+        gc.collect()
+    return img
 css = '''
 .gradio-container{
 max-width: 768px !important;