157 lines
6.2 KiB
Python
157 lines
6.2 KiB
Python
import json
|
|
import gzip
|
|
import piexif
|
|
from PIL import Image
|
|
|
|
# modules/images.py from Stable Diffusion WebUI
|
|
def read_info_from_image(image: Image.Image) -> tuple[str | None, dict]:
|
|
items = (image.info or {}).copy()
|
|
|
|
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'')
|
|
try:
|
|
exif_comment = piexif.helper.UserComment.load(exif_comment)
|
|
except ValueError:
|
|
exif_comment = exif_comment.decode('utf8', errors="ignore")
|
|
|
|
if exif_comment:
|
|
items['exif comment'] = exif_comment
|
|
geninfo = exif_comment
|
|
|
|
if items.get("Software", None) == "NovelAI":
|
|
try:
|
|
json_info = json.loads(items["Comment"])
|
|
geninfo = f"""{items["Description"]}
|
|
Negative prompt: {json_info["uc"]}
|
|
Steps: {json_info["steps"]}, CFG scale: {json_info["scale"]}, Seed: {json_info["seed"]}, Size: {image.width}x{image.height}, Clip skip: 2, ENSD: 31337"""
|
|
except Exception:
|
|
print("Failed to parse NovelAI info")
|
|
|
|
return geninfo, items
|
|
|
|
# https://github.com/ashen-sensored/sd_webui_stealth_pnginfo/blob/main/scripts/stealth_pnginfo.py
|
|
def read_info_from_image_stealth(image):
|
|
# if tensor, convert to PIL image
|
|
if hasattr(image, 'cpu'):
|
|
image = image.cpu().numpy() #((1, 1, 1280, 3), '<f4')
|
|
image = image[0].astype('uint8') #((1, 1280, 3), 'uint8')
|
|
image = Image.fromarray(image)
|
|
# trying to read stealth pnginfo
|
|
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
|
|
never_confirmed = True
|
|
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
|
|
if confirming_signature:
|
|
if x * height + y > 120 and never_confirmed:
|
|
return ''
|
|
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'}:
|
|
#print(f"Found signature at {x}, {y}")
|
|
confirming_signature = False
|
|
sig_confirmed = True
|
|
reading_param_len = True
|
|
mode = 'alpha'
|
|
if decoded_sig == 'stealth_pngcomp':
|
|
compressed = True
|
|
buffer_a = ''
|
|
index_a = 0
|
|
never_confirmed = False
|
|
else:
|
|
read_end = True
|
|
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'}:
|
|
#print(f"Found signature at {x}, {y}")
|
|
confirming_signature = False
|
|
sig_confirmed = True
|
|
reading_param_len = True
|
|
mode = 'rgb'
|
|
if decoded_sig == 'stealth_rgbcomp':
|
|
compressed = True
|
|
buffer_rgb = ''
|
|
index_rgb = 0
|
|
never_confirmed = False
|
|
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
|
|
break
|
|
if read_end:
|
|
break
|
|
geninfo = ''
|
|
if sig_confirmed and binary_data != '':
|
|
# 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 str(geninfo)
|
|
|