Clear torch resources wen done #5

h2oai · Oct 31, 2022 · 1387a24 · 1387a24
1 parent f719aec
commit 1387a24
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/img_styler/image_prompt/stable_diffusion.py b/img_styler/image_prompt/stable_diffusion.py
@@ -1,3 +1,4 @@
+import gc
 import torch
 from torch import autocast
 from PIL import Image
@@ -24,4 +25,6 @@ def generate_image_with_prompt(input_img_path: str, prompt_txt: str = "Face port
 
     file_name = output_path + '/result.jpg'
     images[0].save(file_name)
+    gc.collect()
+    torch.cuda.empty_cache()
     return file_name