Add native "Stealth" infotext support (#279)

Author: catboxanon

modules/gradio_extensons.py

@@ -1,4 +1,7 @@
 import gradio as gr
+from gradio import processing_utils
+import warnings
+import PIL.ImageOps
 
 from modules import scripts, ui_tempdir, patches
 
@@ -74,10 +77,75 @@ def Blocks_get_config_file(self, *args, **kwargs):
     return config
 
 
+def Image_upload_handler(self, x):
+    """Handles conversion of uploaded images to RGB"""
+    if isinstance(x, dict) and 'image' in x:
+        output_image = x['image'].convert('RGB')
+        return output_image
+    return x
+
+def Image_custom_preprocess(self, x):
+    """Custom preprocessing for images with masks"""
+    if x is None:
+        return x
+        
+    mask = ""
+    if self.tool == "sketch" and self.source in ["upload", "webcam"]:
+        if isinstance(x, dict):
+            x, mask = x["image"], x["mask"]
+            
+    if not isinstance(x, str):
+        return x
+        
+    im = processing_utils.decode_base64_to_image(x)
+    
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        im = im.convert(self.image_mode)
+        
+    if self.shape is not None:
+        im = processing_utils.resize_and_crop(im, self.shape)
+        
+    if self.invert_colors:
+        im = PIL.ImageOps.invert(im)
+        
+    if (self.source == "webcam" 
+        and self.mirror_webcam is True 
+        and self.tool != "color-sketch"):
+        im = PIL.ImageOps.mirror(im)
+        
+    if self.tool == "sketch" and self.source in ["upload", "webcam"]:
+        mask_im = None
+        if mask is not None:
+            mask_im = processing_utils.decode_base64_to_image(mask)
+        return {
+            "image": self._format_image(im),
+            "mask": self._format_image(mask_im)
+        }
+        
+    return self._format_image(im)
+
+def Image_init_extension(self, *args, **kwargs):
+    """Extended initialization for Image components"""
+    res = original_Image_init(self, *args, **kwargs)
+    
+    # Only apply to inpaint with mask component for now
+    if getattr(self, 'elem_id', None) == 'img2maskimg':
+        self.upload(
+            fn=Image_upload_handler.__get__(self, gr.Image),
+            inputs=self,
+            outputs=self
+        )
+        self.preprocess = Image_custom_preprocess.__get__(self, gr.Image)
+    
+    return res
+
+
 original_IOComponent_init = patches.patch(__name__, obj=gr.components.IOComponent, field="__init__", replacement=IOComponent_init)
 original_Block_get_config = patches.patch(__name__, obj=gr.blocks.Block, field="get_config", replacement=Block_get_config)
 original_BlockContext_init = patches.patch(__name__, obj=gr.blocks.BlockContext, field="__init__", replacement=BlockContext_init)
 original_Blocks_get_config_file = patches.patch(__name__, obj=gr.blocks.Blocks, field="get_config_file", replacement=Blocks_get_config_file)
+original_Image_init = patches.patch(__name__, obj=gr.components.Image, field="__init__", replacement=Image_init_extension)
 
 
 ui_tempdir.install_ui_tempdir_override()

modules/images.py

@@ -21,7 +21,7 @@
 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
 
@@ -276,6 +276,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):
@@ -715,6 +718,9 @@ 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
@@ -787,44 +793,53 @@ def _atomically_save_image(image_to_save, filename_without_extension, extension)
 
 
 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")
-
-        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"]
+        geninfo = items.pop('parameters', None)
+
+        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")
+
+            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"]
 
-    for field in IGNORED_INFO_KEYS:
-        items.pop(field, None)
+        for field in IGNORED_INFO_KEYS:
+            items.pop(field, None)
 
-    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)
+        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

@@ -188,8 +188,12 @@ 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')
     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

@@ -431,6 +431,7 @@
 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,162 @@
+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):
+    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

@@ -881,7 +881,7 @@ def select_img2img_tab(tab):
     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")
+                image = gr.Image(elem_id="pnginfo_image", label="Source", source="upload", interactive=True, type="pil", image_mode="RGBA")
 
             with gr.Column(variant='panel'):
                 html = gr.HTML()