feat: Stealth Infotext
zanllp
parent
bece389db9
commit
b9f993d777
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
Loading…
Reference in New Issue