happyz
/
stable-diffusion-for-fun
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[BE] cleans up memory cache when job runs failed to avoid memory leak

This commit is contained in:
HappyZ 2023-05-24 15:19:49 -07:00
parent 99c5261c14
commit ca3fa187f4
3 changed files with 12 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
BUILD
View File

@ -28,6 +28,7 @@ par_binary(
deps=[ deps=[
"//utilities:constants", "//utilities:constants",
"//utilities:database", "//utilities:database",
"//utilities:memory",
"//utilities:logger", "//utilities:logger",
"//utilities:model", "//utilities:model",
"//utilities:config", "//utilities:config",

10
backend.py
View File

@ -1,4 +1,5 @@
import argparse import argparse
import torch
from utilities.constants import LOGGER_NAME_BACKEND from utilities.constants import LOGGER_NAME_BACKEND
from utilities.constants import LOGGER_NAME_TXT2IMG from utilities.constants import LOGGER_NAME_TXT2IMG
@ -30,6 +31,7 @@ from utilities.text2img import Text2Img
from utilities.img2img import Img2Img from utilities.img2img import Img2Img
from utilities.inpainting import Inpainting from utilities.inpainting import Inpainting
from utilities.times import wait_for_seconds from utilities.times import wait_for_seconds
from utilities.memory import empty_memory_cache
logger = Logger(name=LOGGER_NAME_BACKEND) logger = Logger(name=LOGGER_NAME_BACKEND)
@ -132,10 +134,10 @@ def backend(model, is_debugging: bool):
break break
except BaseException as e: except BaseException as e:
logger.error(e) logger.error(e)
if not is_debugging:
database.update_job( database.update_job(
{KEY_JOB_STATUS: VALUE_JOB_FAILED}, job_uuid=next_job[UUID] {KEY_JOB_STATUS: VALUE_JOB_FAILED}, job_uuid=next_job[UUID]
) )
empty_memory_cache()
continue continue
database.update_job(result_dict, job_uuid=next_job[UUID]) database.update_job(result_dict, job_uuid=next_job[UUID])
@ -172,7 +174,11 @@ if __name__ == "__main__":
parser.add_argument("--gpu", action="store_true", help="Enable to use GPU device") parser.add_argument("--gpu", action="store_true", help="Enable to use GPU device")
# Add an argument to reduce memory usage # Add an argument to reduce memory usage
parser.add_argument("--reduce-memory-usage", action="store_true", help="Reduce memory usage when using GPU") parser.add_argument(
"--reduce-memory-usage",
action="store_true",
help="Reduce memory usage when using GPU",
)
# Add an argument to set the path of the database file # Add an argument to set the path of the database file
parser.add_argument( parser.add_argument(

1
utilities/database.py
View File

@ -13,7 +13,6 @@ from utilities.constants import VALUE_JOB_IMG2IMG
from utilities.constants import VALUE_JOB_INPAINTING from utilities.constants import VALUE_JOB_INPAINTING
from utilities.constants import KEY_JOB_STATUS from utilities.constants import KEY_JOB_STATUS
from utilities.constants import VALUE_JOB_PENDING from utilities.constants import VALUE_JOB_PENDING
from utilities.constants import VALUE_JOB_RUNNING
from utilities.constants import VALUE_JOB_DONE from utilities.constants import VALUE_JOB_DONE
from utilities.constants import LOCK_FILEPATH from utilities.constants import LOCK_FILEPATH