kamp/IOPaint
1
0
Fork 0
Code Issues Pull Requests Releases 1 Activity

use torch_gc

Browse Source
This commit is contained in:
Qing 2023-05-20 12:35:36 +08:00
parent 870376e4bf
commit 6674104742
1 changed files with 1 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
lama_cleaner/server.py
View File

@ -281,7 +281,6 @@ def process():
try: try:
res_np_img = model(image, mask, config) res_np_img = model(image, mask, config)
except RuntimeError as e: except RuntimeError as e:
torch.cuda.empty_cache()
if "CUDA out of memory. " in str(e): if "CUDA out of memory. " in str(e):
# NOTE: the string may change? # NOTE: the string may change?
return "CUDA out of memory", 500 return "CUDA out of memory", 500
@ -290,7 +289,7 @@ def process():
return f"{str(e)}", 500 return f"{str(e)}", 500
finally: finally:
logger.info(f"process time: {(time.time() - start) * 1000}ms") logger.info(f"process time: {(time.time() - start) * 1000}ms")
torch.cuda.empty_cache() torch_gc()
res_np_img = cv2.cvtColor(res_np_img.astype(np.uint8), cv2.COLOR_BGR2RGB) res_np_img = cv2.cvtColor(res_np_img.astype(np.uint8), cv2.COLOR_BGR2RGB)
if alpha_channel is not None: if alpha_channel is not None: