[BE] adds internal keys support with base_model key

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,

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)

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,
         }

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,
         }

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,
         }