diff --git a/utilities/constants.py b/utilities/constants.py index 3a4c16d..587b014 100644 --- a/utilities/constants.py +++ b/utilities/constants.py @@ -94,6 +94,16 @@ OUTPUT_ONLY_KEYS = [ KEY_JOB_STATUS, # str ] +# -- internal +KEY_BASE_MODEL = "base_model" +INTERNAL_KEYS = [ + KEY_BASE_MODEL, +] + + +# +# language +# SUPPORTED_LANGS = [ VALUE_LANGUAGE_ZH_CN, VALUE_LANGUAGE_EN, diff --git a/utilities/database.py b/utilities/database.py index f4f0659..2c58f0f 100644 --- a/utilities/database.py +++ b/utilities/database.py @@ -20,6 +20,7 @@ from utilities.constants import LOCK_FILEPATH from utilities.constants import OUTPUT_ONLY_KEYS from utilities.constants import OPTIONAL_KEYS from utilities.constants import REQUIRED_KEYS +from utilities.constants import INTERNAL_KEYS from utilities.constants import REFERENCE_IMG from utilities.constants import MASK_IMG @@ -241,7 +242,7 @@ class Database: values = [] columns = [] - for column in OUTPUT_ONLY_KEYS + REQUIRED_KEYS + OPTIONAL_KEYS: + for column in INTERNAL_KEYS + OUTPUT_ONLY_KEYS + REQUIRED_KEYS + OPTIONAL_KEYS: value = job_dict.get(column, None) if value is not None: columns.append(column) diff --git a/utilities/img2img.py b/utilities/img2img.py index 19ee0e2..033ac92 100644 --- a/utilities/img2img.py +++ b/utilities/img2img.py @@ -8,6 +8,7 @@ from utilities.constants import KEY_SEED from utilities.constants import KEY_WIDTH from utilities.constants import KEY_HEIGHT from utilities.constants import KEY_STEPS +from utilities.constants import KEY_BASE_MODEL from utilities.config import Config from utilities.logger import DummyLogger from utilities.memory import empty_memory_cache @@ -152,4 +153,5 @@ class Img2Img: KEY_WIDTH: config.get_width(), KEY_HEIGHT: config.get_height(), KEY_STEPS: config.get_steps(), + KEY_BASE_MODEL: self.model.model_name, } diff --git a/utilities/inpainting.py b/utilities/inpainting.py index 632c87a..072186e 100644 --- a/utilities/inpainting.py +++ b/utilities/inpainting.py @@ -8,6 +8,7 @@ from utilities.constants import KEY_SEED from utilities.constants import KEY_WIDTH from utilities.constants import KEY_HEIGHT from utilities.constants import KEY_STEPS +from utilities.constants import KEY_BASE_MODEL from utilities.config import Config from utilities.logger import DummyLogger from utilities.memory import empty_memory_cache @@ -174,4 +175,5 @@ class Inpainting: KEY_WIDTH: config.get_width(), KEY_HEIGHT: config.get_height(), KEY_STEPS: config.get_steps(), + KEY_BASE_MODEL: self.model.inpainting_model_name, } diff --git a/utilities/text2img.py b/utilities/text2img.py index d058756..5008cc0 100644 --- a/utilities/text2img.py +++ b/utilities/text2img.py @@ -7,6 +7,7 @@ from utilities.constants import KEY_SEED from utilities.constants import KEY_WIDTH from utilities.constants import KEY_HEIGHT from utilities.constants import KEY_STEPS +from utilities.constants import KEY_BASE_MODEL from utilities.config import Config from utilities.logger import DummyLogger from utilities.memory import empty_memory_cache @@ -140,4 +141,5 @@ class Text2Img: KEY_WIDTH: config.get_width(), KEY_HEIGHT: config.get_height(), KEY_STEPS: config.get_steps(), + KEY_BASE_MODEL: self.model.model_name, }