Add native "Stealth" infotext support (#2684)

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
 

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

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"),

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

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()