diff --git a/README.md b/README.md index 809cef3..c24cb6c 100644 --- a/README.md +++ b/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 diff --git a/scripts/iib/db/datamodel.py b/scripts/iib/db/datamodel.py index b7bbb5a..6254238 100644 --- a/scripts/iib/db/datamodel.py +++ b/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" diff --git a/scripts/iib/db/update_image_data.py b/scripts/iib/db/update_image_data.py index 2fdd4cd..621d20b 100644 --- a/scripts/iib/db/update_image_data.py +++ b/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 diff --git a/scripts/iib/parsers/index.py b/scripts/iib/parsers/index.py index 3281e84..9e3912f 100644 --- a/scripts/iib/parsers/index.py +++ b/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): diff --git a/scripts/iib/parsers/novelai.py b/scripts/iib/parsers/novelai.py index 7b69178..680e1cd 100644 --- a/scripts/iib/parsers/novelai.py +++ b/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): diff --git a/scripts/iib/parsers/stable_swarm_ui.py b/scripts/iib/parsers/stable_swarm_ui.py index e69de29..8e41e04 100644 --- a/scripts/iib/parsers/stable_swarm_ui.py +++ b/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 diff --git a/scripts/iib/tool.py b/scripts/iib/tool.py index 2fced86..6f49a0f 100644 --- a/scripts/iib/tool.py +++ b/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 \ No newline at end of file + 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 \ No newline at end of file