158 lines
5.8 KiB
Python
158 lines
5.8 KiB
Python
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
|