temp save

scripts/iib/api.py

@@ -5,14 +5,9 @@ import shutil
 import sqlite3
 
 from scripts.iib.tool import (
-    comfyui_exif_data_to_str,
-    get_comfyui_exif_data,
     human_readable_size,
-    is_img_created_by_comfyui,
-    is_img_created_by_comfyui_with_webui_gen_info,
-    is_valid_media_path,
+    is_valid_image_path,
     temp_path,
-    read_sd_webui_gen_info_from_image,
     get_formatted_date,
     is_win,
     cwd,
@@ -58,6 +53,7 @@ from scripts.iib.logger import logger
 from scripts.iib.seq import seq
 import urllib.parse
 from scripts.iib.fastapi_video import range_requests_response
+from scripts.iib.vendor.index import parse_image_info
 
 index_html_path = os.path.join(cwd, "vue/dist/index.html")  # 在app.py也被使用
 
@@ -578,20 +574,7 @@ def infinite_image_browsing_api(app: FastAPI, **kwargs):
 
     @app.get(api_base + "/image_geninfo", dependencies=[Depends(verify_secret)])
     async def image_geninfo(path: str):
-        with Image.open(path) as img:
-            if is_img_created_by_comfyui(img):
-                if is_img_created_by_comfyui_with_webui_gen_info(img):
-                    return read_sd_webui_gen_info_from_image(img, path)
-                else:
-                    try:                    
-                        params = get_comfyui_exif_data(img)
-                        return comfyui_exif_data_to_str(params)
-                    except:
-                        logger.error('parse comfyui image failed. prompt:')
-                        logger.error(img.info.get('prompt'))
-                        return ''
-            else:
-                return read_sd_webui_gen_info_from_image(img, path)
+        return parse_image_info(path).raw_info
 
     class GeninfoBatchReq(BaseModel):
         paths: List[str]

scripts/iib/db/update_image_data.py

@@ -2,7 +2,6 @@ from contextlib import closing
 from typing import Dict, List
 from scripts.iib.db.datamodel import Image as DbImg, Tag, ImageTag, DataBase, Folder
 import os
-from PIL import Image
 from scripts.iib.tool import (
     read_sd_webui_gen_info_from_image,
     parse_generation_parameters,
@@ -14,33 +13,24 @@ from scripts.iib.tool import (
     comfyui_exif_data_to_str,
     is_img_created_by_comfyui,
     is_img_created_by_comfyui_with_webui_gen_info
+    is_valid_media_path,
+    get_modified_date,
+    is_dev
 )
-
+from scripts.iib.vendor.model import ImageGenerationInfo, ImageGenerationParams
 from scripts.iib.logger import logger
-
+from scripts.iib.vendor.index import parse_image_info
 
 # 定义一个函数来获取图片文件的EXIF数据
 def get_exif_data(file_path):
-    info = ''
-    params = None
     if get_video_type(file_path):
-        return params, info
+        return ImageGenerationInfo()
     try:
-        with Image.open(file_path) as img:
-            if is_img_created_by_comfyui(img):
-                if is_img_created_by_comfyui_with_webui_gen_info(img):
-                    info = read_sd_webui_gen_info_from_image(img, file_path)
-                    params = parse_generation_parameters(info)
-                else:
-                    params = get_comfyui_exif_data(img)
-                    info = comfyui_exif_data_to_str(params)
-            else:
-                info = read_sd_webui_gen_info_from_image(img, file_path)
-                params = parse_generation_parameters(info)
+        return parse_image_info(file_path)
     except Exception as e:
         if is_dev:
             logger.error("get_exif_data %s", e)
-    return params, info
+    return ImageGenerationInfo()
 
 
 def update_image_data(search_dirs: List[str], is_rebuild = False):
@@ -58,6 +48,7 @@ def update_image_data(search_dirs: List[str], is_rebuild = False):
         if not is_rebuild and not Folder.check_need_update(conn, folder_path):
             return
         print(f"Processing folder: {folder_path}")
+        parsed_params = ImageGenerationParams()
         for filename in os.listdir(folder_path):
             file_path = os.path.normpath(os.path.join(folder_path, filename))
             try:
@@ -68,11 +59,12 @@ def update_image_data(search_dirs: List[str], is_rebuild = False):
                 elif is_valid_media_path(file_path):
                     img = DbImg.get(conn, file_path)
                     if is_rebuild:
-                        parsed_params, info = get_exif_data(file_path)
+                        info = get_exif_data(file_path)
+                        parsed_params = info.params
                         if not img:
                             img = DbImg(
                                 file_path,
-                                info,
+                                info.raw_info,
                                 os.path.getsize(file_path),
                                 get_modified_date(file_path),
                             )
@@ -83,10 +75,11 @@ def update_image_data(search_dirs: List[str], is_rebuild = False):
                                 continue
                             else:
                                 DbImg.safe_batch_remove(conn=conn, image_ids=[img.id])
-                        parsed_params, info = get_exif_data(file_path)
+                        info = get_exif_data(file_path)
+                        parsed_params = info.params
                         img = DbImg(
                             file_path,
-                            info,
+                            info.raw_info,
                             os.path.getsize(file_path),
                             get_modified_date(file_path),
                         )
@@ -94,10 +87,10 @@ def update_image_data(search_dirs: List[str], is_rebuild = False):
 
                     if not parsed_params:
                         continue
-                    meta = parsed_params.get("meta", {})
-                    lora = parsed_params.get("lora", [])
-                    lyco = parsed_params.get("lyco", [])
-                    pos = parsed_params["pos_prompt"]
+                    meta = parsed_params.meta
+                    lora = parsed_params.extra.get("lora", [])
+                    lyco = parsed_params.extra.get("lyco", [])
+                    pos = parsed_params.pos_prompt
                     size_tag = Tag.get_or_create(
                         conn,
                         str(meta.get("Size-1", 0)) + " * " + str(meta.get("Size-2", 0)),

scripts/iib/tool.py

@@ -346,7 +346,7 @@ def is_img_created_by_comfyui(img: Image):
     return img.info.get('prompt') and img.info.get('workflow')
 
 def is_img_created_by_comfyui_with_webui_gen_info(img: Image):
-    return img.info.get('parameters')
+    return is_img_created_by_comfyui(img) and img.info.get('parameters')
 
 def get_comfyui_exif_data(img: Image):
     prompt = img.info.get('prompt')
@@ -569,3 +569,5 @@ def open_file_with_default_app(file_path):
         subprocess.call(['xdg-open', file_path])
     else:
         raise OSError(f'Unsupported operating system: {system}')
+def omit(d, keys):
+    return {k: v for k, v in d.items() if k not in keys}

scripts/iib/vendor/comfyui.py

@@ -0,0 +1,44 @@
+from PIL import Image
+
+from scripts.iib.tool import (
+    comfyui_exif_data_to_str,
+    is_img_created_by_comfyui,
+    is_img_created_by_comfyui_with_webui_gen_info,
+    get_comfyui_exif_data,
+    parse_generation_parameters,
+    read_sd_webui_gen_info_from_image,
+)
+from scripts.iib.vendor.model import ImageGenerationInfo, ImageGenerationParams
+
+
+class ComfyUIParser:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def parse(clz, img, file_path):
+        info = ""
+        params = None
+        if not clz.test(img):
+            raise Exception("image not matchd")
+        if is_img_created_by_comfyui_with_webui_gen_info(img):
+            info = read_sd_webui_gen_info_from_image(img, file_path)
+            params = parse_generation_parameters(info)
+        else:
+            params = get_comfyui_exif_data(img)
+            info = comfyui_exif_data_to_str(params)
+        return ImageGenerationInfo(
+            info,
+            ImageGenerationParams(
+                meta=params["meta"], pos_prompt=params["pos_prompt"], extra=params
+            ),
+        )
+
+    @classmethod
+    def test(clz, img: Image) -> bool:
+        try:
+            return is_img_created_by_comfyui(
+                img
+            ) or is_img_created_by_comfyui_with_webui_gen_info(img)
+        except Exception:
+            return False

scripts/iib/vendor/index.py

@@ -0,0 +1,13 @@
+from scripts.iib.vendor.comfyui import ComfyUIParser
+from scripts.iib.vendor.sd_webui import SdWebUIParser
+from scripts.iib.vendor.model import ImageGenerationInfo
+from PIL import Image
+
+
+def parse_image_info(image_path: str) -> ImageGenerationInfo:
+    parsers = [ComfyUIParser, SdWebUIParser]
+    with Image.open(image_path) as img:
+        for parser in parsers:
+            if parser.test(img):
+                return parser.parse(img, image_path)
+        raise Exception("matched parser is not found")

scripts/iib/vendor/model.py

@@ -0,0 +1,18 @@
+from scripts.iib.tool import omit
+
+
+class ImageGenerationParams:
+    def __init__(self, meta: dict = {}, pos_prompt: list = [], extra: dict = {}) -> None:
+        self.meta = meta
+        self.pos_prompt = pos_prompt
+        self.extra = omit(extra, ["meta", "por_prompt"])
+
+
+class ImageGenerationInfo:
+    def __init__(
+        self,
+        raw_info: str = "",
+        params: ImageGenerationParams = ImageGenerationParams(),
+    ):
+        self.raw_info = raw_info
+        self.params = params

scripts/iib/vendor/sd_webui.py

@@ -0,0 +1,32 @@
+from PIL import Image
+
+from scripts.iib.tool import (
+    parse_generation_parameters,
+    read_sd_webui_gen_info_from_image,
+)
+from scripts.iib.vendor.model import ImageGenerationInfo, ImageGenerationParams
+
+
+class SdWebUIParser:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def parse(clz, img: Image, file_path):
+        if not clz.test(img):
+            raise Exception("image not matchd")
+        info = read_sd_webui_gen_info_from_image(img, file_path)
+        params = parse_generation_parameters(info)
+        return ImageGenerationInfo(
+            info,
+            ImageGenerationParams(
+                meta=params["meta"], pos_prompt=params["pos_prompt"], extra=params
+            ),
+        )
+
+    @classmethod
+    def test(clz, img: Image):
+        try:
+            return True
+        except Exception as e:
+            return False