supports image output folder instead of storing it in database

.gitignore

@@ -8,6 +8,7 @@ __pycache__/
 
 # sqlite3 db
 *.db
+data/
 
 # Distribution / packaging
 .Python

backend.py

@@ -68,7 +68,7 @@ def backend(model, is_debugging: bool):
         if is_debugging:
             pending_jobs = database.get_jobs()
         else:
-            pending_jobs = database.get_all_pending_jobs()
+            pending_jobs = database.get_one_pending_job()
         if len(pending_jobs) == 0:
             continue
 
@@ -97,9 +97,6 @@ def backend(model, is_debugging: bool):
                     logger.info(f"translated {negative_prompt} to {negative_prompt_en}")
                     negative_prompt = negative_prompt_en
 
-        prompt += "RAW photo, (high detailed skin:1.2), 8k uhd, dslr, high quality, film grain, Fujifilm XT3"
-        negative_prompt += "(deformed iris, deformed pupils:1.4), worst quality, low quality, jpeg artifacts, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck"
-
         config = Config().set_config(next_job)
 
         try:
@@ -135,6 +132,7 @@ def backend(model, is_debugging: bool):
 
 
 def main(args):
+    database.set_image_output_folder(args.image_output_folder)
     database.connect(args.db)
 
     model = load_model(logger, args.gpu)
@@ -157,6 +155,15 @@ if __name__ == "__main__":
     # Add an argument to set the path of the database file
     parser.add_argument("--gpu", action="store_true", help="Enable to use GPU device")
 
+    # Add an argument to set the path of the database file
+    parser.add_argument(
+        "--image-output-folder",
+        "-o",
+        type=str,
+        default="",
+        help="Path to output images",
+    )
+
     args = parser.parse_args()
 
     main(args)

frontend.py

@@ -25,7 +25,6 @@ from utilities.database import Database
 
 logger = Logger(name=LOGGER_NAME_FRONTEND)
 database = Database(logger)
-
 app = Flask(__name__)
 
 
@@ -50,7 +49,7 @@ def add_job():
 
     if req[KEY_JOB_TYPE] == VALUE_JOB_IMG2IMG and REFERENCE_IMG not in req:
         return jsonify({"msg": "missing reference image"}), 404
-    
+
     if KEY_LANGUAGE in req and req[KEY_LANGUAGE] not in SUPPORTED_LANGS:
         return jsonify({"msg": f"not suporting {req[KEY_LANGUAGE]}"}), 404
 
@@ -128,6 +127,7 @@ def index():
 
 
 def main(args):
+    database.set_image_output_folder(args.image_output_folder)
     database.connect(args.db)
 
     if args.debug:
@@ -149,6 +149,15 @@ if __name__ == "__main__":
         "--db", type=str, default="happysd.db", help="Path to SQLite database file"
     )
 
+    # Add an argument to set the path of the database file
+    parser.add_argument(
+        "--image-output-folder",
+        "-o",
+        type=str,
+        default="",
+        help="Path to output images",
+    )
+
     args = parser.parse_args()
 
     main(args)

utilities/database.py

@@ -23,11 +23,17 @@ from utilities.constants import REQUIRED_KEYS
 
 from utilities.constants import REFERENCE_IMG
 from utilities.constants import BASE64IMAGE
+from utilities.constants import IMAGE_NOT_FOUND_BASE64
 
 from utilities.constants import HISTORY_TABLE_NAME
 from utilities.constants import USERS_TABLE_NAME
 from utilities.logger import DummyLogger
 
+from utilities.times import get_epoch_now
+from utilities.times import epoch_to_string
+from utilities.images import save_image
+from utilities.images import load_image
+
 
 # Function to acquire a lock on the database file
 def acquire_lock():
@@ -45,13 +51,25 @@ def release_lock():
 class Database:
     """This class represents a SQLite database and assumes single-thread usage."""
 
-    def __init__(self, logger: DummyLogger = DummyLogger()):
+    def __init__(self, logger: DummyLogger = DummyLogger(), image_folderpath=""):
         """Initialize the class with a logger instance, but without a database connection or cursor."""
         self.__connect = None  # the database connection object
         self.is_connected = False
         self.__cursor = None  # the cursor object for executing SQL statements
         self.__logger = logger  # the logger object for logging messages
 
+        self.__image_output_folder = ""
+        self.set_image_output_folder(image_folderpath)
+
+    def set_image_output_folder(self, image_folderpath):
+        self.__image_output_folder = image_folderpath  # output image to a folder instead of storing it in sqlite3
+        if image_folderpath:
+            try:
+                os.makedirs(image_folderpath, exist_ok=True)
+            except OSError as err:
+                self.__logger.warn(f"{image_folderpath} failed to create: {err}")
+                self.__image_output_folder = ""
+
     def connect(self, db_filepath) -> bool:
         """
         Connect to the SQLite database file specified by `db_filepath`.
@@ -90,8 +108,8 @@ class Database:
             return result[0]
         return ""
 
-    def get_all_pending_jobs(self, apikey: str = "") -> list:
-        return self.get_jobs(apikey=apikey, job_status=VALUE_JOB_PENDING)
+    def get_one_pending_job(self, apikey: str = "") -> list:
+        return self.get_jobs(apikey=apikey, job_status=VALUE_JOB_PENDING, limit_count=1)
 
     def count_all_pending_jobs(self, apikey: str) -> int:
         """
@@ -108,7 +126,7 @@ class Database:
         result = c.execute(query_string, query_args).fetchone()
         return result[0]
 
-    def get_jobs(self, job_uuid="", apikey="", job_status="") -> list:
+    def get_jobs(self, job_uuid="", apikey="", job_status="", limit_count=0) -> list:
         """
         Get a list of jobs from the HISTORY_TABLE_NAME table based on optional filters.
 
@@ -133,6 +151,8 @@ class Database:
         query = f"SELECT {', '.join(columns)} FROM {HISTORY_TABLE_NAME}"
         if query_filters:
             query += f" WHERE {' AND '.join(query_filters)}"
+        if limit_count:
+            query += f" LIMIT {limit_count}"
 
         # execute the query and return the results
         c = self.get_cursor()
@@ -143,6 +163,11 @@ class Database:
             job = {
                 columns[i]: row[i] for i in range(len(columns)) if row[i] is not None
             }
+            # load image to job if has one
+            for key in [BASE64IMAGE, REFERENCE_IMG]:
+                if key in job and "base64" not in job[key]:
+                    data = load_image(job[key], to_base64=True)
+                    job[key] = data if data else IMAGE_NOT_FOUND_BASE64
             jobs.append(job)
 
         return jobs
@@ -159,6 +184,17 @@ class Database:
             job_uuid = str(uuid.uuid4())
         self.__logger.info(f"inserting a new job with {job_uuid}")
 
+        # store image to job_dict if has one
+        if (
+            self.__image_output_folder
+            and REFERENCE_IMG in job_dict
+            and "base64" in job_dict[REFERENCE_IMG]
+        ):
+            ref_img_filepath = f"{self.__image_output_folder}/{get_epoch_now()}_ref.png"
+            self.__logger.info(f"saving reference image to {ref_img_filepath}")
+            if save_image(job_dict[REFERENCE_IMG], ref_img_filepath):
+                job_dict[REFERENCE_IMG] = ref_img_filepath
+
         values = [job_uuid, VALUE_JOB_PENDING, datetime.datetime.now()]
         columns = [UUID, KEY_JOB_STATUS, "created_at"] + REQUIRED_KEYS + OPTIONAL_KEYS
         for column in REQUIRED_KEYS + OPTIONAL_KEYS:
@@ -181,6 +217,17 @@ class Database:
 
         Returns True if the update was successful, otherwise False.
         """
+        # store image to job_dict if has one
+        if (
+            self.__image_output_folder
+            and BASE64IMAGE in job_dict
+            and "base64" in job_dict[BASE64IMAGE]
+        ):
+            out_img_filepath = f"{self.__image_output_folder}/{get_epoch_now()}_out.png"
+            self.__logger.info(f"saving output image to {out_img_filepath}")
+            if save_image(job_dict[BASE64IMAGE], out_img_filepath):
+                job_dict[BASE64IMAGE] = out_img_filepath
+
         values = []
         columns = []
         for column in OUTPUT_ONLY_KEYS + REQUIRED_KEYS + OPTIONAL_KEYS:

utilities/images.py

@@ -6,22 +6,26 @@ import numpy as np
 from PIL import Image
 
 
-def load_image(image: Union[str, bytes]) -> Union[Image.Image, None]:
+def load_image(image: Union[str, bytes], to_base64: bool=False) -> Union[Image.Image, str, None]:
     if isinstance(image, bytes):
         return Image.open(io.BytesIO(image))
     elif os.path.isfile(image):
+        if to_base64:
+            return image_to_base64(image)
         with Image.open(image) as im:
             return Image.fromarray(np.asarray(im))
     return None
 
 
 def save_image(
-    image: Union[bytes, Image.Image], filepath: str, override: bool = False
+    image: Union[bytes, Image.Image, str], filepath: str, override: bool = False
 ) -> bool:
     if os.path.isfile(filepath) and not override:
         return False
     try:
-        if isinstance(image, Image.Image):
+        if isinstance(image, str):
+            base64_to_image(image).save(filepath)
+        elif isinstance(image, Image.Image):
             # this is an Image
             image.save(filepath)
         else: