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

add torch.cuda.empty_cache() after model forward

Browse Source
This commit is contained in:
Sanster 2022-03-12 20:43:47 +08:00
parent 4e5903129c
commit 61fcffda94
1 changed files with 11 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
main.py
View File

@ -27,7 +27,8 @@ from lama_cleaner.helper import (
load_img, load_img,
norm_img, norm_img,
numpy_to_bytes, numpy_to_bytes,
resize_max_size, ) resize_max_size,
)
NUM_THREADS = str(multiprocessing.cpu_count()) NUM_THREADS = str(multiprocessing.cpu_count())
@ -74,6 +75,8 @@ def process():
res_np_img = model(image, mask) res_np_img = model(image, mask)
torch.cuda.empty_cache()
return send_file( return send_file(
io.BytesIO(numpy_to_bytes(res_np_img)), io.BytesIO(numpy_to_bytes(res_np_img)),
mimetype="image/jpeg", mimetype="image/jpeg",
@ -91,8 +94,13 @@ def get_args_parser():
parser = argparse.ArgumentParser() parser = argparse.ArgumentParser()
parser.add_argument("--port", default=8080, type=int) parser.add_argument("--port", default=8080, type=int)
parser.add_argument("--model", default="lama", choices=["lama", "ldm"]) parser.add_argument("--model", default="lama", choices=["lama", "ldm"])
parser.add_argument("--ldm-steps", default=50, type=int, help="Steps for DDIM sampling process." parser.add_argument(
"The larger the value, the better the result, but it will be more time-consuming") "--ldm-steps",
default=50,
type=int,
help="Steps for DDIM sampling process."
"The larger the value, the better the result, but it will be more time-consuming",
)
parser.add_argument("--device", default="cuda", type=str) parser.add_argument("--device", default="cuda", type=str)
parser.add_argument("--debug", action="store_true") parser.add_argument("--debug", action="store_true")
return parser.parse_args() return parser.parse_args()