diff --git a/BUILD b/BUILD index 0521958..d46061b 100644 --- a/BUILD +++ b/BUILD @@ -28,6 +28,7 @@ par_binary( deps=[ "//utilities:constants", "//utilities:database", + "//utilities:memory", "//utilities:logger", "//utilities:model", "//utilities:config", diff --git a/backend.py b/backend.py index e7b7cc3..a3235e4 100644 --- a/backend.py +++ b/backend.py @@ -1,4 +1,5 @@ import argparse +import torch from utilities.constants import LOGGER_NAME_BACKEND from utilities.constants import LOGGER_NAME_TXT2IMG @@ -30,6 +31,7 @@ from utilities.text2img import Text2Img from utilities.img2img import Img2Img from utilities.inpainting import Inpainting from utilities.times import wait_for_seconds +from utilities.memory import empty_memory_cache logger = Logger(name=LOGGER_NAME_BACKEND) @@ -132,10 +134,10 @@ def backend(model, is_debugging: bool): break except BaseException as e: logger.error(e) - if not is_debugging: - database.update_job( - {KEY_JOB_STATUS: VALUE_JOB_FAILED}, job_uuid=next_job[UUID] - ) + database.update_job( + {KEY_JOB_STATUS: VALUE_JOB_FAILED}, job_uuid=next_job[UUID] + ) + empty_memory_cache() continue 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") # 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 parser.add_argument( diff --git a/utilities/database.py b/utilities/database.py index 2c58f0f..19586b8 100644 --- a/utilities/database.py +++ b/utilities/database.py @@ -13,7 +13,6 @@ from utilities.constants import VALUE_JOB_IMG2IMG from utilities.constants import VALUE_JOB_INPAINTING from utilities.constants import KEY_JOB_STATUS 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 LOCK_FILEPATH