Add support for StableSwarmUI

README.md

@@ -19,7 +19,7 @@ https://github.com/zanllp/sd-webui-infinite-image-browsing/assets/25872019/807b8
 | ComfyUI                | Partially supported | --                |
 | Fooocus                | Supported        | --                   |
 | NovelAI                | Supported    | --            |
-| StableSwarmUI          | Not supported    | Planned              |
+| StableSwarmUI          | Supported    | --              |
 
 You can add your own parser through https://github.com/zanllp/sd-webui-infinite-image-browsing/tree/main/scripts/iib/parsers
 

scripts/iib/db/datamodel.py

@@ -622,6 +622,12 @@ class Folder:
         with closing(conn.cursor()) as cur:
             cur.execute("DELETE FROM folders WHERE path = ?", (folder_path,))
 
+    @classmethod
+    def remove_all(cls, conn: Connection):
+        with closing(conn.cursor()) as cur:
+            cur.execute("DELETE FROM folders")
+            conn.commit()
+
 
 class ExtraPathType(Enum):
     scanned = "scanned"

scripts/iib/db/update_image_data.py

@@ -8,7 +8,8 @@ from scripts.iib.tool import (
     get_video_type,
     is_dev,
     get_modified_date,
-    is_image_file
+    is_image_file,
+    case_insensitive_get
 )
 from scripts.iib.parsers.model import ImageGenerationInfo, ImageGenerationParams
 from scripts.iib.logger import logger
@@ -30,6 +31,9 @@ def update_image_data(search_dirs: List[str], is_rebuild = False):
     conn = DataBase.get_conn()
     tag_incr_count_rec: Dict[int, int] = {}
 
+    if is_rebuild:
+        Folder.remove_all(conn)
+
     def safe_save_img_tag(img_tag: ImageTag):
         tag_incr_count_rec[img_tag.tag_id] = (
             tag_incr_count_rec.get(img_tag.tag_id, 0) + 1
@@ -38,7 +42,7 @@ def update_image_data(search_dirs: List[str], is_rebuild = False):
 
     # 递归处理每个文件夹
     def process_folder(folder_path: str):
-        if not is_rebuild and not Folder.check_need_update(conn, folder_path):
+        if not Folder.check_need_update(conn, folder_path):
             return
         print(f"Processing folder: {folder_path}")
         for filename in os.listdir(folder_path):
@@ -160,7 +164,7 @@ def build_single_img_idx(conn, file_path, is_rebuild, safe_save_img_tag):
         "Hires upscaler"
     ]:
         
-        v = meta.get(k)
+        v = case_insensitive_get(meta, k)
         if not v:
             continue
         

scripts/iib/parsers/index.py

@@ -3,13 +3,20 @@ from scripts.iib.parsers.sd_webui import SdWebUIParser
 from scripts.iib.parsers.fooocus import FooocusParser
 from scripts.iib.parsers.novelai import NovelAIParser
 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
 import traceback
 
 
 def parse_image_info(image_path: str) -> ImageGenerationInfo:
-    parsers = [ComfyUIParser, FooocusParser, NovelAIParser, SdWebUIParser]
+    parsers = [
+        ComfyUIParser,
+        FooocusParser,
+        NovelAIParser,
+        StableSwarmUIParser,
+        SdWebUIParser,
+    ]
     with Image.open(image_path) as img:
         for parser in parsers:
             if parser.test(img, image_path):

scripts/iib/parsers/novelai.py

@@ -3,11 +3,10 @@ from PIL import Image
 
 from scripts.iib.tool import (
     parse_generation_parameters,
+    replace_punctuation
 )
 from scripts.iib.parsers.model import ImageGenerationInfo, ImageGenerationParams
 
-def replace_punctuation(input_string):
-    return input_string.replace(',', ' ').replace('\n', ' ')
 
 class NovelAIParser:
     def __init__(self):

scripts/iib/parsers/stable_swarm_ui.py

@@ -0,0 +1,61 @@
+from PIL import Image
+
+import piexif
+import piexif.helper
+from scripts.iib.tool import parse_generation_parameters, replace_punctuation
+from scripts.iib.parsers.model import ImageGenerationInfo, ImageGenerationParams
+from PIL.ExifTags import TAGS
+import json
+
+
+class StableSwarmUIParser:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def get_exif_data(clz, image: Image) -> str:
+        items = image.info or {}
+
+        if "exif" in items:
+            exif = piexif.load(items["exif"])
+            exif_bytes = (
+                (exif or {}).get("Exif", {}).get(piexif.ExifIFD.UserComment, b"")
+            )
+
+        unicode_start = exif_bytes.find(b"UNICODE")
+        if unicode_start == -1:
+            raise ValueError("'UNICODE' markup isn't found")
+
+        unicode_data = exif_bytes[unicode_start + len("UNICODE") + 1 :]
+        geninfo = unicode_data.decode("utf-16")
+        return geninfo
+
+    @classmethod
+    def parse(clz, img: Image, file_path):
+        if not clz.test(img, file_path):
+            raise Exception("The input image does not match the current parser.")
+        exif_data = json.loads(clz.get_exif_data(img))["sui_image_params"]
+        prompt = exif_data.pop("prompt")
+        negativeprompt = exif_data.pop("negativeprompt")
+        steps = exif_data.pop("steps")
+        meta_kv = [f"Steps: {steps}", "Source Identifier: StableSwarmUI"]
+        for key, value in exif_data.items():
+            value = replace_punctuation(str(value))
+            meta_kv.append(f"{key}: {value}")
+        meta = ", ".join(meta_kv)
+        info = "\n".join([prompt, f"Negative prompt: {negativeprompt}", meta])
+        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, file_path: str):
+        try:
+            exif = clz.get_exif_data(img)
+            return exif.find("sui_image_params") != -1
+        except Exception as e:
+            return False

scripts/iib/tool.py

@@ -424,7 +424,6 @@ def read_sd_webui_gen_info_from_image(image: Image, path="") -> str:
     """
     items = image.info or {}
     geninfo = items.pop("parameters", None)
-
     if "exif" in items:
         exif = piexif.load(items["exif"])
         exif_comment = (exif or {}).get("Exif", {}).get(piexif.ExifIFD.UserComment, b"")
@@ -624,4 +623,15 @@ def get_current_tag():
         else:
             return None
     except subprocess.CalledProcessError:
-        return None
+        return None
+    
+    
+def replace_punctuation(input_string):
+    return input_string.replace(',', ' ').replace('\n', ' ')
+
+
+def case_insensitive_get(d, key, default=None):
+    for k, v in d.items():
+        if k.lower() == key.lower():
+            return v
+    return default