diff --git a/modules/images.py b/modules/images.py index d6488cec..4565a7ee 100644 --- a/modules/images.py +++ b/modules/images.py @@ -19,7 +19,7 @@ import string import json import hashlib -from modules import sd_samplers, shared, script_callbacks, errors +from modules import sd_samplers, shared, script_callbacks, errors, stealth_infotext from modules.paths_internal import roboto_ttf_file from modules.shared import opts @@ -264,6 +264,9 @@ def resize_image(resize_mode, im, width, height, upscaler_name=None, force_RGBA= upscaler_name: The name of the upscaler to use. If not provided, defaults to opts.upscaler_for_img2img. """ + if not force_RGBA and im.mode == 'RGBA': + im = im.convert('RGB') + upscaler_name = upscaler_name or opts.upscaler_for_img2img def resize(im, w, h): @@ -706,6 +709,8 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i pnginfo[pnginfo_section_name] = info params = script_callbacks.ImageSaveParams(image, p, fullfn, pnginfo) + if opts.enable_pnginfo: + stealth_infotext.add_stealth_pnginfo(params) script_callbacks.before_image_saved_callback(params) image = params.image @@ -782,44 +787,53 @@ IGNORED_INFO_KEYS = { def read_info_from_image(image: Image.Image) -> tuple[str | None, dict]: - items = (image.info or {}).copy() + """Read generation info from an image, checking standard metadata first, then stealth info if needed.""" - geninfo = items.pop('parameters', None) + def read_standard(): + items = (image.info or {}).copy() - if "exif" in items: - exif_data = items["exif"] - try: - exif = piexif.load(exif_data) - except OSError: - # memory / exif was not valid so piexif tried to read from a file - exif = None - 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") + geninfo = items.pop('parameters', None) - if exif_comment: - geninfo = exif_comment - elif "comment" in items: # for gif - if isinstance(items["comment"], bytes): - geninfo = items["comment"].decode('utf8', errors="ignore") - else: - geninfo = items["comment"] + if "exif" in items: + exif_data = items["exif"] + try: + exif = piexif.load(exif_data) + except OSError: + # memory / exif was not valid so piexif tried to read from a file + exif = None + 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") - for field in IGNORED_INFO_KEYS: - items.pop(field, None) + if exif_comment: + geninfo = exif_comment + elif "comment" in items: # for gif + if isinstance(items["comment"], bytes): + geninfo = items["comment"].decode('utf8', errors="ignore") + else: + geninfo = items["comment"] - if items.get("Software", None) == "NovelAI": - try: - json_info = json.loads(items["Comment"]) - sampler = sd_samplers.samplers_map.get(json_info["sampler"], "Euler a") + for field in IGNORED_INFO_KEYS: + items.pop(field, None) - geninfo = f"""{items["Description"]} -Negative prompt: {json_info["uc"]} -Steps: {json_info["steps"]}, Sampler: {sampler}, CFG scale: {json_info["scale"]}, Seed: {json_info["seed"]}, Size: {image.width}x{image.height}, Clip skip: 2, ENSD: 31337""" - except Exception: - errors.report("Error parsing NovelAI image generation parameters", exc_info=True) + if items.get("Software", None) == "NovelAI": + try: + json_info = json.loads(items["Comment"]) + sampler = sd_samplers.samplers_map.get(json_info["sampler"], "Euler a") + + geninfo = f"""{items["Description"]} + Negative prompt: {json_info["uc"]} + Steps: {json_info["steps"]}, Sampler: {sampler}, CFG scale: {json_info["scale"]}, Seed: {json_info["seed"]}, Size: {image.width}x{image.height}, Clip skip: 2, ENSD: 31337""" + except Exception: + errors.report("Error parsing NovelAI image generation parameters", exc_info=True) + + return geninfo, items + + geninfo, items = read_standard() + if geninfo is None: + geninfo = stealth_infotext.read_info_from_image_stealth(image) return geninfo, items diff --git a/modules/infotext_utils.py b/modules/infotext_utils.py index 3ad30438..cb703acd 100644 --- a/modules/infotext_utils.py +++ b/modules/infotext_utils.py @@ -197,10 +197,14 @@ def connect_paste_params_buttons(): def send_image_and_dimensions(x): if isinstance(x, Image.Image): img = x + if img.mode == 'RGBA': + img = img.convert('RGB') elif isinstance(x, list) and isinstance(x[0], tuple): img = x[0][0] else: img = image_from_url_text(x) + if img is not None and img.mode == 'RGBA': + img = img.convert('RGB') if shared.opts.send_size and isinstance(img, Image.Image): w = img.width diff --git a/modules/shared_options.py b/modules/shared_options.py index 06bd9930..ca96341e 100644 --- a/modules/shared_options.py +++ b/modules/shared_options.py @@ -358,6 +358,7 @@ Infotext is what this software calls the text that contains generation parameter It is displayed in UI below the image. To use infotext, paste it into the prompt and click the ↙️ paste button. """), "enable_pnginfo": OptionInfo(True, "Write infotext to metadata of the generated image"), + "stealth_pnginfo_option": OptionInfo("Alpha", "Stealth infotext mode", gr.Radio, {"choices": ["Alpha", "RGB", "None"]}).info("Ignored if infotext is disabled"), "save_txt": OptionInfo(False, "Create a text file with infotext next to every generated image"), "add_model_name_to_info": OptionInfo(True, "Add model name to infotext"), diff --git a/modules/stealth_infotext.py b/modules/stealth_infotext.py new file mode 100644 index 00000000..2204f104 --- /dev/null +++ b/modules/stealth_infotext.py @@ -0,0 +1,163 @@ +import gzip + +from modules.script_callbacks import ImageSaveParams +from modules import shared + + +def add_stealth_pnginfo(params: ImageSaveParams): + stealth_pnginfo_option = shared.opts.data.get('stealth_pnginfo_option', 'Alpha') + if not stealth_pnginfo_option or stealth_pnginfo_option == 'None': + return + if not params.filename.endswith('.png') or params.pnginfo is None: + return + if 'parameters' not in params.pnginfo: + return + add_data(params, str(stealth_pnginfo_option), True) + +def prepare_data(params, mode='Alpha', compressed=True): + signature = f"stealth_{'png' if mode == 'Alpha' else 'rgb'}{'info' if not compressed else 'comp'}" + binary_signature = ''.join(format(byte, '08b') for byte in signature.encode('utf-8')) + param = params.encode('utf-8') if not compressed else gzip.compress(bytes(params, 'utf-8')) + binary_param = ''.join(format(byte, '08b') for byte in param) + binary_param_len = format(len(binary_param), '032b') + return binary_signature + binary_param_len + binary_param + +def add_data(params, mode='Alpha', compressed=True): + binary_data = prepare_data(params.pnginfo['parameters'], mode, compressed) + if mode == 'Alpha': + params.image.putalpha(255) + width, height = params.image.size + pixels = params.image.load() + index = 0 + end_write = False + for x in range(width): + for y in range(height): + if index >= len(binary_data): + end_write = True + break + values = pixels[x, y] + if mode == 'Alpha': + r, g, b, a = values + else: + r, g, b = values + if mode == 'Alpha': + a = (a & ~1) | int(binary_data[index]) + index += 1 + else: + r = (r & ~1) | int(binary_data[index]) + if index + 1 < len(binary_data): + g = (g & ~1) | int(binary_data[index + 1]) + if index + 2 < len(binary_data): + b = (b & ~1) | int(binary_data[index + 2]) + index += 3 + pixels[x, y] = (r, g, b, a) if mode == 'Alpha' else (r, g, b) + if end_write: + break + +def read_info_from_image_stealth(image): + geninfo = None + 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 + 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 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 + reading_param_len = True + mode = 'alpha' + if decoded_sig == 'stealth_pngcomp': + compressed = True + buffer_a = '' + index_a = 0 + 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'}: + confirming_signature = False + sig_confirmed = 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 + break + if read_end: + break + 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 geninfo diff --git a/modules/ui.py b/modules/ui.py index b6865f15..1bc85331 100644 --- a/modules/ui.py +++ b/modules/ui.py @@ -927,7 +927,7 @@ def create_ui(): with gr.Blocks(analytics_enabled=False) as pnginfo_interface: with ResizeHandleRow(equal_height=False): with gr.Column(variant='panel'): - image = gr.Image(elem_id="pnginfo_image", label="Source", source="upload", interactive=True, type="pil", height="50vh") + image = gr.Image(elem_id="pnginfo_image", label="Source", source="upload", interactive=True, type="pil", height="50vh", image_mode="RGBA") with gr.Column(variant='panel'): html = gr.HTML()