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

add try catch

Browse Source
This commit is contained in:
Qing 2023-03-30 21:02:18 +08:00
parent c52f733214
commit 6911d3ce16
1 changed files with 14 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
lama_cleaner/server.py
View File

@ -27,7 +27,7 @@ from lama_cleaner.plugins import (
RealESRGANUpscaler, RealESRGANUpscaler,
MakeGIF, MakeGIF,
GFPGANPlugin, GFPGANPlugin,
RestoreFormerPlugin RestoreFormerPlugin,
) )
from lama_cleaner.schema import Config from lama_cleaner.schema import Config
@ -314,7 +314,17 @@ def run_plugin():
rgb_np_img, alpha_channel, exif = load_img(origin_image_bytes, return_exif=True) rgb_np_img, alpha_channel, exif = load_img(origin_image_bytes, return_exif=True)
start = time.time() start = time.time()
bgr_res = plugins[name](rgb_np_img, files, form) try:
bgr_res = plugins[name](rgb_np_img, files, form)
except RuntimeError as e:
torch.cuda.empty_cache()
if "CUDA out of memory. " in str(e):
# NOTE: the string may change?
return "CUDA out of memory", 500
else:
logger.exception(e)
return "Internal Server Error", 500
logger.info(f"{name} process time: {(time.time() - start) * 1000}ms") logger.info(f"{name} process time: {(time.time() - start) * 1000}ms")
torch_gc() torch_gc()
@ -446,7 +456,8 @@ def build_plugins(args):
if args.enable_restoreformer: if args.enable_restoreformer:
logger.info(f"Initialize {RestoreFormerPlugin.name} plugin") logger.info(f"Initialize {RestoreFormerPlugin.name} plugin")
plugins[RestoreFormerPlugin.name] = RestoreFormerPlugin( plugins[RestoreFormerPlugin.name] = RestoreFormerPlugin(
args.restoreformer_device, upscaler=plugins.get(RealESRGANUpscaler.name, None) args.restoreformer_device,
upscaler=plugins.get(RealESRGANUpscaler.name, None),
) )
if args.enable_gif: if args.enable_gif:
logger.info(f"Initialize GIF plugin") logger.info(f"Initialize GIF plugin")