diff --git a/scripts/iib/api.py b/scripts/iib/api.py index b57f63b..4dda336 100644 --- a/scripts/iib/api.py +++ b/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] diff --git a/scripts/iib/db/update_image_data.py b/scripts/iib/db/update_image_data.py index ecbf4fd..a45bb23 100644 --- a/scripts/iib/db/update_image_data.py +++ b/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)), diff --git a/scripts/iib/tool.py b/scripts/iib/tool.py index cd80ad5..482af76 100644 --- a/scripts/iib/tool.py +++ b/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} diff --git a/scripts/iib/vendor/comfyui.py b/scripts/iib/vendor/comfyui.py new file mode 100644 index 0000000..cf46a87 --- /dev/null +++ b/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 diff --git a/scripts/iib/vendor/fooocus.py b/scripts/iib/vendor/fooocus.py new file mode 100644 index 0000000..e69de29 diff --git a/scripts/iib/vendor/index.py b/scripts/iib/vendor/index.py new file mode 100644 index 0000000..b99c08a --- /dev/null +++ b/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") diff --git a/scripts/iib/vendor/model.py b/scripts/iib/vendor/model.py new file mode 100644 index 0000000..86e2e03 --- /dev/null +++ b/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 diff --git a/scripts/iib/vendor/novelai.py b/scripts/iib/vendor/novelai.py new file mode 100644 index 0000000..e69de29 diff --git a/scripts/iib/vendor/sd_webui.py b/scripts/iib/vendor/sd_webui.py new file mode 100644 index 0000000..a773f3d --- /dev/null +++ b/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 diff --git a/scripts/iib/vendor/stable_swarm_ui.py b/scripts/iib/vendor/stable_swarm_ui.py new file mode 100644 index 0000000..e69de29