From 83b4c65757289093e63d0fe9b685468862e93419 Mon Sep 17 00:00:00 2001 From: zanllp Date: Wed, 3 Jul 2024 11:13:34 +0800 Subject: [PATCH] Implement a basic plugin mechanism --- .gitignore | 5 +++-- README.md | 18 ++++++++-------- plugins/.gitkeep | 0 scripts/iib/api.py | 2 +- scripts/iib/db/update_image_data.py | 18 +++++++++++++--- scripts/iib/parsers/index.py | 3 ++- scripts/iib/plugin.py | 33 +++++++++++++++++++++++++++++ scripts/iib/tool.py | 14 +++++++++++- 8 files changed, 76 insertions(+), 17 deletions(-) create mode 100644 plugins/.gitkeep create mode 100644 scripts/iib/plugin.py diff --git a/.gitignore b/.gitignore index b2686c9..d1e745c 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,3 @@ - *.log __pycache__ iib.db @@ -16,4 +15,6 @@ out/**/* venv/**/* zip_temp/*.zip iib_db_backup -db_migrate_temp.db \ No newline at end of file +db_migrate_temp.db +plugins/* +!plugins/.gitkeep \ No newline at end of file diff --git a/README.md b/README.md index a003c75..b65b3c5 100644 --- a/README.md +++ b/README.md @@ -14,16 +14,16 @@ https://github.com/zanllp/sd-webui-infinite-image-browsing/assets/25872019/807b8 ### Software Support and Development Progress Overview +| Software | Support | Development Progress | Provided by | +| ---------------------- | ---------------- | -------------------- | ----------- | +| Stable Diffusion web UI| Supported | -- | Built-in | +| ComfyUI | Partially supported | -- | Built-in | +| Fooocus | Supported | -- | Built-in | +| NovelAI | Supported | -- | Built-in | +| StableSwarmUI | Supported | -- | Built-in | +| Pixiv | Supported | -- | [Plugin](https://github.com/zanllp/pixiv_iib_plugin) | -| Software | Support | Development Progress | -| ---------------------- | ---------------- | -------------------- | -| Stable Diffusion web UI| Supported | -- | -| ComfyUI | Partially supported | -- | -| Fooocus | Supported | -- | -| NovelAI | Supported | -- | -| StableSwarmUI | Supported | -- | - -You can add your own parser through https://github.com/zanllp/sd-webui-infinite-image-browsing/tree/main/scripts/iib/parsers +You can add your own parser through [parsers](https://github.com/zanllp/sd-webui-infinite-image-browsing/tree/main/scripts/iib/parsers), You can also implement parsing for more formats by writing plugins. For reference, see [pixiv_iib_plugin](https://github.com/zanllp/pixiv_iib_plugin) ## Key Features diff --git a/plugins/.gitkeep b/plugins/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/scripts/iib/api.py b/scripts/iib/api.py index 56eb61b..49871a4 100644 --- a/scripts/iib/api.py +++ b/scripts/iib/api.py @@ -62,7 +62,7 @@ from scripts.iib.seq import seq import urllib.parse from scripts.iib.fastapi_video import range_requests_response, close_video_file_reader from scripts.iib.parsers.index import parse_image_info - +import scripts.iib.plugin try: import pillow_avif diff --git a/scripts/iib/db/update_image_data.py b/scripts/iib/db/update_image_data.py index 621d20b..c9b5a88 100644 --- a/scripts/iib/db/update_image_data.py +++ b/scripts/iib/db/update_image_data.py @@ -14,6 +14,7 @@ from scripts.iib.tool import ( from scripts.iib.parsers.model import ImageGenerationInfo, ImageGenerationParams from scripts.iib.logger import logger from scripts.iib.parsers.index import parse_image_info +from scripts.iib.plugin import plugin_inst_map # 定义一个函数来获取图片文件的EXIF数据 def get_exif_data(file_path): @@ -109,6 +110,16 @@ def rebuild_image_index(search_dirs: List[str]): conn.commit() update_image_data(search_dirs=search_dirs, is_rebuild=True) + +def get_extra_meta_keys_from_plugins(source_identifier: str): + try: + plugin = plugin_inst_map.get(source_identifier) + if plugin: + return plugin.extra_convert_to_tag_meta_keys + except Exception as e: + logger.error("get_extra_meta_keys_from_plugins %s", e) + return [] + def build_single_img_idx(conn, file_path, is_rebuild, safe_save_img_tag): img = DbImg.get(conn, file_path) parsed_params = None @@ -153,7 +164,7 @@ def build_single_img_idx(conn, file_path, is_rebuild, safe_save_img_tag): safe_save_img_tag(ImageTag(img.id, size_tag.id)) media_type_tag = Tag.get_or_create(conn, "Image" if is_image_file(file_path) else "Video", 'Media Type') safe_save_img_tag(ImageTag(img.id, media_type_tag.id)) - for k in [ + keys = [ "Model", "Sampler", "Source Identifier", @@ -162,8 +173,9 @@ def build_single_img_idx(conn, file_path, is_rebuild, safe_save_img_tag): "Size", "Refiner", "Hires upscaler" - ]: - + ] + keys += get_extra_meta_keys_from_plugins(meta.get("Source Identifier", "")) + for k in keys: v = case_insensitive_get(meta, k) if not v: continue diff --git a/scripts/iib/parsers/index.py b/scripts/iib/parsers/index.py index 9e3912f..93dcdd0 100644 --- a/scripts/iib/parsers/index.py +++ b/scripts/iib/parsers/index.py @@ -6,11 +6,12 @@ from scripts.iib.parsers.model import ImageGenerationInfo from scripts.iib.parsers.stable_swarm_ui import StableSwarmUIParser from scripts.iib.logger import logger from PIL import Image +from scripts.iib.plugin import plugin_insts import traceback def parse_image_info(image_path: str) -> ImageGenerationInfo: - parsers = [ + parsers = plugin_insts + [ ComfyUIParser, FooocusParser, NovelAIParser, diff --git a/scripts/iib/plugin.py b/scripts/iib/plugin.py new file mode 100644 index 0000000..eff4d69 --- /dev/null +++ b/scripts/iib/plugin.py @@ -0,0 +1,33 @@ +import os +import importlib.util +import sys + +def load_plugins(plugin_dir): + plugins = [] + for filename in os.listdir(plugin_dir): + main_module_path = os.path.join(plugin_dir, filename, 'main.py') + if not os.path.exists(main_module_path): + continue + spec = importlib.util.spec_from_file_location('main', main_module_path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + plugins.append(module) + return plugins + +plugin_insts = [] +plugin_inst_map = {} +# 使用插件 +try: + plugin_dir = os.path.normpath(os.path.join(os.path.dirname(__file__), '../../', 'plugins')) + plugins = load_plugins(plugin_dir) + sys.path.append(plugin_dir) + for plugin in plugins: + try: + res = plugin.Main() + plugin_insts.append(res) + plugin_inst_map[res.source_identifier] = res + print(f'IIB loaded plugin: {res.name}') + except Exception as e: + print(f'Error running plugin {plugin.__class__.__name__}: {e}') +except Exception as e: + print(f'Error loading plugins: {e}') \ No newline at end of file diff --git a/scripts/iib/tool.py b/scripts/iib/tool.py index 1099f45..251a5fd 100644 --- a/scripts/iib/tool.py +++ b/scripts/iib/tool.py @@ -659,4 +659,16 @@ def case_insensitive_get(d, key, default=None): for k, v in d.items(): if k.lower() == key.lower(): return v - return default \ No newline at end of file + return default + +def build_sd_webui_style_img_gen_info(prompt, negative_prompt = 'None', meta = {}): + res = f"{prompt}\nNegative prompt: {negative_prompt}\n" + for k, v in meta.items(): + res += f"{k}: {v}, " + return res + +def map_dict_keys(value_dict, map_dict=None): + if map_dict is None: + return value_dict + else: + return {map_dict.get(key, key): value for key, value in value_dict.items()} \ No newline at end of file