diff --git a/.env.example b/.env.example index 36aae60..8c71100 100644 --- a/.env.example +++ b/.env.example @@ -39,3 +39,11 @@ IIB_ACCESS_CONTROL=auto # It can be set to a string value that represents a specific permission or set of permissions, such as "read-only", "write-only", "read-write", or "no-access". # This variable can be used to restrict access to certain API endpoints or data sources based on the permissions required by the user. # IIB_ACCESS_CONTROL_PERMISSION=read-write + + + +# ---------------------------- PARSER_CONFIG ---------------------------- +# This attribute is used to control whether to enable SdWebUIStealthParser. +# Due to the high performance cost of parsing this type of file, it is disabled by default. +# Set to 'true' to enable it. +IIB_ENABLE_SD_WEBUI_STEALTH_PARSER=false diff --git a/scripts/iib/parsers/index.py b/scripts/iib/parsers/index.py index 66fb5db..86f2e61 100644 --- a/scripts/iib/parsers/index.py +++ b/scripts/iib/parsers/index.py @@ -1,3 +1,4 @@ +import os from scripts.iib.parsers.comfyui import ComfyUIParser from scripts.iib.parsers.sd_webui import SdWebUIParser from scripts.iib.parsers.fooocus import FooocusParser @@ -5,6 +6,7 @@ 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.parsers.invoke_ai import InvokeAIParser +from scripts.iib.parsers.sd_webui_stealth import SdWebUIStealthParser from scripts.iib.logger import logger from PIL import Image from scripts.iib.plugin import plugin_insts @@ -12,14 +14,19 @@ import traceback def parse_image_info(image_path: str) -> ImageGenerationInfo: + enable_stealth_parser = os.getenv('IIB_ENABLE_SD_WEBUI_STEALTH_PARSER', 'false').lower() == 'true' parsers = plugin_insts + [ ComfyUIParser, FooocusParser, NovelAIParser, InvokeAIParser, StableSwarmUIParser, - SdWebUIParser, ] + + if enable_stealth_parser: + parsers.append(SdWebUIStealthParser) + + parsers.append(SdWebUIParser) with Image.open(image_path) as img: for parser in parsers: if parser.test(img, image_path): diff --git a/scripts/iib/parsers/sd_webui_stealth.py b/scripts/iib/parsers/sd_webui_stealth.py new file mode 100644 index 0000000..668b716 --- /dev/null +++ b/scripts/iib/parsers/sd_webui_stealth.py @@ -0,0 +1,157 @@ +from PIL import Image +import gzip +from scripts.iib.tool import ( + parse_generation_parameters, + read_sd_webui_gen_info_from_image, +) +from scripts.iib.parsers.model import ImageGenerationInfo, ImageGenerationParams +# https://github.com/neggles/sd-webui-stealth-pnginfo/blob/main/scripts/stealth_pnginfo.py +def read_info_from_image_stealth(image, fast_check=False): + width, height = image.size + pixels = image.load() + + has_alpha = True if image.mode == 'RGBA' else False + mode = None + compressed = False + binary_data = '' + buffer_a = '' + buffer_rgb = '' + index_a = 0 + index_rgb = 0 + sig_confirmed = False + confirming_signature = True + reading_param_len = False + reading_param = False + read_end = False + cyc_count = 0 + for x in range(width): + for y in range(height): + if has_alpha: + r, g, b, a = pixels[x, y] + buffer_a += str(a & 1) + index_a += 1 + else: + r, g, b = pixels[x, y] + buffer_rgb += str(r & 1) + buffer_rgb += str(g & 1) + buffer_rgb += str(b & 1) + index_rgb += 3 + cyc_count += 1 + if fast_check and cyc_count> 256: + return False + if confirming_signature: + if index_a == len('stealth_pnginfo') * 8: + decoded_sig = bytearray(int(buffer_a[i:i + 8], 2) for i in + range(0, len(buffer_a), 8)).decode('utf-8', errors='ignore') + if decoded_sig in {'stealth_pnginfo', 'stealth_pngcomp'}: + confirming_signature = False + sig_confirmed = True + if fast_check: + return True + reading_param_len = True + mode = 'alpha' + if decoded_sig == 'stealth_pngcomp': + compressed = True + buffer_a = '' + index_a = 0 + else: + read_end = True + if fast_check: + return False + break + elif index_rgb == len('stealth_pnginfo') * 8: + decoded_sig = bytearray(int(buffer_rgb[i:i + 8], 2) for i in + range(0, len(buffer_rgb), 8)).decode('utf-8', errors='ignore') + if decoded_sig in {'stealth_rgbinfo', 'stealth_rgbcomp'}: + confirming_signature = False + sig_confirmed = True + if fast_check: + return True + reading_param_len = True + mode = 'rgb' + if decoded_sig == 'stealth_rgbcomp': + compressed = True + buffer_rgb = '' + index_rgb = 0 + elif reading_param_len: + if mode == 'alpha': + if index_a == 32: + param_len = int(buffer_a, 2) + reading_param_len = False + reading_param = True + buffer_a = '' + index_a = 0 + else: + if index_rgb == 33: + pop = buffer_rgb[-1] + buffer_rgb = buffer_rgb[:-1] + param_len = int(buffer_rgb, 2) + reading_param_len = False + reading_param = True + buffer_rgb = pop + index_rgb = 1 + elif reading_param: + if mode == 'alpha': + if index_a == param_len: + binary_data = buffer_a + read_end = True + break + else: + if index_rgb >= param_len: + diff = param_len - index_rgb + if diff < 0: + buffer_rgb = buffer_rgb[:diff] + binary_data = buffer_rgb + read_end = True + break + else: + # impossible + read_end = True + if fast_check: + return False + break + if read_end: + break + if sig_confirmed and binary_data != '': + if fast_check: + return True + # Convert binary string to UTF-8 encoded text + byte_data = bytearray(int(binary_data[i:i + 8], 2) for i in range(0, len(binary_data), 8)) + try: + if compressed: + decoded_data = gzip.decompress(bytes(byte_data)).decode('utf-8') + else: + decoded_data = byte_data.decode('utf-8', errors='ignore') + geninfo = decoded_data + except: + pass + return geninfo + + + +class SdWebUIStealthParser: + def __init__(self): + pass + + @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.") + info = read_info_from_image_stealth(img) + if not info: + return ImageGenerationInfo() + info += ", Source Identifier: Stable Diffusion web UI(Stealth)" + 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: + return bool(read_info_from_image_stealth(img, fast_check=True)) + except Exception as e: + return False