Implement a basic plugin mechanism

.gitignore

@@ -1,4 +1,3 @@
-
 *.log
 __pycache__
 iib.db
@@ -17,3 +16,5 @@ venv/**/*
 zip_temp/*.zip
 iib_db_backup
 db_migrate_temp.db
+plugins/*
+!plugins/.gitkeep

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 |
+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)
-| ---------------------- | ---------------- | -------------------- |
-| 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
 
 ## Key Features
 

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

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

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,

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}')

scripts/iib/tool.py

@@ -660,3 +660,15 @@ def case_insensitive_get(d, key, default=None):
         if k.lower() == key.lower():
             return v
     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()}