├── README.md ├── .gitignore └── scripts └── stealth_pnginfo.py /README.md: -------------------------------------------------------------------------------- 1 | # sd-webui-stealth-pnginfo 2 | 3 | This extension embeds image generation metadata into the (unused) alpha channel of a PNG generated in [stable-diffusion-webui](https://github.com/AUTOMATIC1111/stable-diffusion-webui). 4 | 5 | **Will not work for JPEG images.** Untested on webp, should work in theory but who knows. 6 | 7 | ***Please note that as this data is embedded in the image pixels, it is not visible in an EXIF data viewer.*** 8 | 9 | It will show up in the PNG info tab in the webui and via bots like [Prompt Inspector](https://github.com/sALTaccount/PromptInspectorBot), but not via EXIF tools. 10 | 11 | Extension auto-activates and runs alongside standard webui PNG info functions, no extra settings are required. It can optionally be disabled in settings if desired. 12 | 13 | ### Acknowledgements 14 | 15 | Forked from [@ashen-sensored/sd_webui_stealth_pnginfo](https://github.com/ashen-sensored/sd_webui_stealth_pnginfo) - this was originally created to get around Discord stripping the text sections from PNGs, but they've reversed that change so Ashen is not maintaining that anymore. I find it convenient though (since it doesn't rely on that data being retained, and many websites/devices will strip it) so here it is. 16 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | share/python-wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | MANIFEST 28 | 29 | # PyInstaller 30 | # Usually these files are written by a python script from a template 31 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 32 | *.manifest 33 | *.spec 34 | 35 | # Installer logs 36 | pip-log.txt 37 | pip-delete-this-directory.txt 38 | 39 | # Unit test / coverage reports 40 | htmlcov/ 41 | .tox/ 42 | .nox/ 43 | .coverage 44 | .coverage.* 45 | .cache 46 | nosetests.xml 47 | coverage.xml 48 | *.cover 49 | *.py,cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | cover/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | .pybuilder/ 76 | target/ 77 | 78 | # Jupyter Notebook 79 | .ipynb_checkpoints 80 | 81 | # IPython 82 | profile_default/ 83 | ipython_config.py 84 | 85 | # pyenv 86 | # For a library or package, you might want to ignore these files since the code is 87 | # intended to run in multiple environments; otherwise, check them in: 88 | # .python-version 89 | 90 | # pipenv 91 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 92 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 93 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 94 | # install all needed dependencies. 95 | #Pipfile.lock 96 | 97 | # poetry 98 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 99 | # This is especially recommended for binary packages to ensure reproducibility, and is more 100 | # commonly ignored for libraries. 101 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 102 | #poetry.lock 103 | 104 | # pdm 105 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 106 | #pdm.lock 107 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 108 | # in version control. 109 | # https://pdm.fming.dev/#use-with-ide 110 | .pdm.toml 111 | 112 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 113 | __pypackages__/ 114 | 115 | # Celery stuff 116 | celerybeat-schedule 117 | celerybeat.pid 118 | 119 | # SageMath parsed files 120 | *.sage.py 121 | 122 | # Environments 123 | .env 124 | .venv 125 | env/ 126 | venv/ 127 | ENV/ 128 | env.bak/ 129 | venv.bak/ 130 | 131 | # Spyder project settings 132 | .spyderproject 133 | .spyproject 134 | 135 | # Rope project settings 136 | .ropeproject 137 | 138 | # mkdocs documentation 139 | /site 140 | 141 | # mypy 142 | .mypy_cache/ 143 | .dmypy.json 144 | dmypy.json 145 | 146 | # Pyre type checker 147 | .pyre/ 148 | 149 | # pytype static type analyzer 150 | .pytype/ 151 | 152 | # Cython debug symbols 153 | cython_debug/ 154 | 155 | # PyCharm 156 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 157 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 158 | # and can be added to the global gitignore or merged into this file. For a more nuclear 159 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 160 | .idea 161 | *.pt 162 | *.pth 163 | *.ckpt 164 | *.safetensors 165 | models/control_sd15_scribble.pth 166 | detected_maps/ -------------------------------------------------------------------------------- /scripts/stealth_pnginfo.py: -------------------------------------------------------------------------------- 1 | from modules import script_callbacks, shared, generation_parameters_copypaste 2 | from modules.script_callbacks import ImageSaveParams 3 | import gradio as gr 4 | from modules import images 5 | from PIL import Image 6 | from gradio import processing_utils 7 | import PIL 8 | import warnings 9 | import gzip 10 | 11 | 12 | def add_stealth_pnginfo(params: ImageSaveParams): 13 | stealth_pnginfo_enabled = shared.opts.data.get("stealth_pnginfo", True) 14 | stealth_pnginfo_mode = shared.opts.data.get('stealth_pnginfo_mode', 'alpha') 15 | stealth_pnginfo_compressed = shared.opts.data.get("stealth_pnginfo_compression", True) 16 | if not stealth_pnginfo_enabled: 17 | return 18 | if not params.filename.endswith('.png') or params.pnginfo is None: 19 | return 20 | if 'parameters' not in params.pnginfo: 21 | return 22 | add_data(params, stealth_pnginfo_mode, stealth_pnginfo_compressed) 23 | 24 | 25 | def prepare_data(params, mode='alpha', compressed=False): 26 | signature = f"stealth_{'png' if mode == 'alpha' else 'rgb'}{'info' if not compressed else 'comp'}" 27 | binary_signature = ''.join(format(byte, '08b') for byte in signature.encode('utf-8')) 28 | param = params.encode('utf-8') if not compressed else gzip.compress(bytes(params, 'utf-8')) 29 | binary_param = ''.join(format(byte, '08b') for byte in param) 30 | binary_param_len = format(len(binary_param), '032b') 31 | return binary_signature + binary_param_len + binary_param 32 | 33 | 34 | def add_data(params, mode='alpha', compressed=False): 35 | binary_data = prepare_data(params.pnginfo['parameters'], mode, compressed) 36 | if mode == 'alpha': 37 | params.image.putalpha(255) 38 | width, height = params.image.size 39 | pixels = params.image.load() 40 | index = 0 41 | end_write = False 42 | for x in range(width): 43 | for y in range(height): 44 | if index >= len(binary_data): 45 | end_write = True 46 | break 47 | values = pixels[x, y] 48 | if mode == 'alpha': 49 | r, g, b, a = values 50 | else: 51 | r, g, b = values 52 | if mode == 'alpha': 53 | a = (a & ~1) | int(binary_data[index]) 54 | index += 1 55 | else: 56 | r = (r & ~1) | int(binary_data[index]) 57 | if index + 1 < len(binary_data): 58 | g = (g & ~1) | int(binary_data[index + 1]) 59 | if index + 2 < len(binary_data): 60 | b = (b & ~1) | int(binary_data[index + 2]) 61 | index += 3 62 | pixels[x, y] = (r, g, b, a) if mode == 'alpha' else (r, g, b) 63 | if end_write: 64 | break 65 | 66 | 67 | def read_info_from_image_stealth(image): 68 | geninfo, items = original_read_info_from_image(image) 69 | # possible_sigs = {'stealth_pnginfo', 'stealth_pngcomp', 'stealth_rgbinfo', 'stealth_rgbcomp'} 70 | 71 | # respecting original pnginfo 72 | if geninfo is not None: 73 | return geninfo, items 74 | 75 | # trying to read stealth pnginfo 76 | width, height = image.size 77 | pixels = image.load() 78 | 79 | has_alpha = True if image.mode == 'RGBA' else False 80 | mode = None 81 | compressed = False 82 | binary_data = '' 83 | buffer_a = '' 84 | buffer_rgb = '' 85 | index_a = 0 86 | index_rgb = 0 87 | sig_confirmed = False 88 | confirming_signature = True 89 | reading_param_len = False 90 | reading_param = False 91 | read_end = False 92 | for x in range(width): 93 | for y in range(height): 94 | if has_alpha: 95 | r, g, b, a = pixels[x, y] 96 | buffer_a += str(a & 1) 97 | index_a += 1 98 | else: 99 | r, g, b = pixels[x, y] 100 | buffer_rgb += str(r & 1) 101 | buffer_rgb += str(g & 1) 102 | buffer_rgb += str(b & 1) 103 | index_rgb += 3 104 | if confirming_signature: 105 | if index_a == len('stealth_pnginfo') * 8: 106 | decoded_sig = bytearray(int(buffer_a[i:i + 8], 2) for i in 107 | range(0, len(buffer_a), 8)).decode('utf-8', errors='ignore') 108 | if decoded_sig in {'stealth_pnginfo', 'stealth_pngcomp'}: 109 | confirming_signature = False 110 | sig_confirmed = True 111 | reading_param_len = True 112 | mode = 'alpha' 113 | if decoded_sig == 'stealth_pngcomp': 114 | compressed = True 115 | buffer_a = '' 116 | index_a = 0 117 | else: 118 | read_end = True 119 | break 120 | elif index_rgb == len('stealth_pnginfo') * 8: 121 | decoded_sig = bytearray(int(buffer_rgb[i:i + 8], 2) for i in 122 | range(0, len(buffer_rgb), 8)).decode('utf-8', errors='ignore') 123 | if decoded_sig in {'stealth_rgbinfo', 'stealth_rgbcomp'}: 124 | confirming_signature = False 125 | sig_confirmed = True 126 | reading_param_len = True 127 | mode = 'rgb' 128 | if decoded_sig == 'stealth_rgbcomp': 129 | compressed = True 130 | buffer_rgb = '' 131 | index_rgb = 0 132 | elif reading_param_len: 133 | if mode == 'alpha': 134 | if index_a == 32: 135 | param_len = int(buffer_a, 2) 136 | reading_param_len = False 137 | reading_param = True 138 | buffer_a = '' 139 | index_a = 0 140 | else: 141 | if index_rgb == 33: 142 | pop = buffer_rgb[-1] 143 | buffer_rgb = buffer_rgb[:-1] 144 | param_len = int(buffer_rgb, 2) 145 | reading_param_len = False 146 | reading_param = True 147 | buffer_rgb = pop 148 | index_rgb = 1 149 | elif reading_param: 150 | if mode == 'alpha': 151 | if index_a == param_len: 152 | binary_data = buffer_a 153 | read_end = True 154 | break 155 | else: 156 | if index_rgb >= param_len: 157 | diff = param_len - index_rgb 158 | if diff < 0: 159 | buffer_rgb = buffer_rgb[:diff] 160 | binary_data = buffer_rgb 161 | read_end = True 162 | break 163 | else: 164 | # impossible 165 | read_end = True 166 | break 167 | if read_end: 168 | break 169 | if sig_confirmed and binary_data != '': 170 | # Convert binary string to UTF-8 encoded text 171 | byte_data = bytearray(int(binary_data[i:i + 8], 2) for i in range(0, len(binary_data), 8)) 172 | try: 173 | if compressed: 174 | decoded_data = gzip.decompress(bytes(byte_data)).decode('utf-8') 175 | else: 176 | decoded_data = byte_data.decode('utf-8', errors='ignore') 177 | geninfo = decoded_data 178 | except: 179 | pass 180 | return geninfo, items 181 | 182 | 183 | def send_rgb_image_and_dimension(x): 184 | if isinstance(x, Image.Image): 185 | img = x 186 | if img.mode == 'RGBA': 187 | img = img.convert('RGB') 188 | else: 189 | img = generation_parameters_copypaste.image_from_url_text(x) 190 | if img.mode == 'RGBA': 191 | img = img.convert('RGB') 192 | 193 | if shared.opts.send_size and isinstance(img, Image.Image): 194 | w = img.width 195 | h = img.height 196 | else: 197 | w = gr.update() 198 | h = gr.update() 199 | 200 | return img, w, h 201 | 202 | 203 | def on_ui_settings(): 204 | section = ('stealth_pnginfo', "Stealth PNGinfo") 205 | shared.opts.add_option("stealth_pnginfo", shared.OptionInfo( 206 | True, "Save Stealth PNGinfo", gr.Checkbox, {"interactive": True}, section=section)) 207 | shared.opts.add_option("stealth_pnginfo_mode", shared.OptionInfo( 208 | "alpha", "Stealth PNGinfo mode", gr.Dropdown, {"choices": ["alpha", "rgb"], "interactive": True}, 209 | section=section)) 210 | shared.opts.add_option("stealth_pnginfo_compression", shared.OptionInfo( 211 | True, "Stealth PNGinfo compression", gr.Checkbox, {"interactive": True}, section=section)) 212 | 213 | 214 | def custom_image_preprocess(self, x): 215 | if x is None: 216 | return x 217 | 218 | mask = "" 219 | if self.tool == "sketch" and self.source in ["upload", "webcam"]: 220 | assert isinstance(x, dict) 221 | x, mask = x["image"], x["mask"] 222 | 223 | assert isinstance(x, str) 224 | im = processing_utils.decode_base64_to_image(x) 225 | with warnings.catch_warnings(): 226 | warnings.simplefilter("ignore") 227 | im = im.convert(self.image_mode) 228 | if self.shape is not None: 229 | im = processing_utils.resize_and_crop(im, self.shape) 230 | if self.invert_colors: 231 | im = PIL.ImageOps.invert(im) 232 | if ( 233 | self.source == "webcam" 234 | and self.mirror_webcam is True 235 | and self.tool != "color-sketch" 236 | ): 237 | im = PIL.ImageOps.mirror(im) 238 | 239 | if self.tool == "sketch" and self.source in ["upload", "webcam"]: 240 | mask_im = None 241 | if mask is not None: 242 | mask_im = processing_utils.decode_base64_to_image(mask) 243 | 244 | return { 245 | "image": self._format_image(im), 246 | "mask": self._format_image(mask_im), 247 | } 248 | 249 | return self._format_image(im) 250 | 251 | 252 | def on_after_component_change_pnginfo_image_mode(component, **_kwargs): 253 | if type(component) is gr.State: 254 | return 255 | if type(component) is gr.Image and component.elem_id == 'pnginfo_image': 256 | component.image_mode = 'RGBA' 257 | 258 | def set_alpha_channel_to_zero(image): 259 | width, height = image.size 260 | pixels = image.load() 261 | 262 | for x in range(width): 263 | for y in range(height): 264 | r, g, b, a = pixels[x, y] 265 | pixels[x, y] = (r, g, b, 0) 266 | 267 | def clear_alpha(param): 268 | print('clear_alpha called') 269 | output_image = param['image'].convert('RGB') 270 | return output_image 271 | # set_alpha_channel_to_zero(input) 272 | # return input 273 | 274 | if type(component) is gr.Image and component.elem_id == 'img2maskimg': 275 | component.upload(clear_alpha, component, component) 276 | component.preprocess = custom_image_preprocess.__get__(component, gr.Image) 277 | 278 | 279 | def stealth_resize_image(resize_mode, im, width, height, upscaler_name=None): 280 | """ 281 | Resizes an image with the specified resize_mode, width, and height. 282 | 283 | Args: 284 | resize_mode: The mode to use when resizing the image. 285 | 0: Resize the image to the specified width and height. 286 | 1: Resize the image to fill the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, cropping the excess. 287 | 2: Resize the image to fit within the specified width and height, maintaining the aspect ratio, and then center the image within the dimensions, filling empty with data from image. 288 | im: The image to resize. 289 | width: The width to resize the image to. 290 | height: The height to resize the image to. 291 | upscaler_name: The name of the upscaler to use. If not provided, defaults to opts.upscaler_for_img2img. 292 | """ 293 | # convert to RGB 294 | if im.mode == 'RGBA': 295 | im = im.convert('RGB') 296 | 297 | return original_resize_image(resize_mode, im, width, height, upscaler_name) 298 | 299 | 300 | LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS) 301 | original_read_info_from_image = images.read_info_from_image 302 | images.read_info_from_image = read_info_from_image_stealth 303 | generation_parameters_copypaste.send_image_and_dimensions = send_rgb_image_and_dimension 304 | original_resize_image = images.resize_image 305 | images.resize_image = stealth_resize_image 306 | 307 | script_callbacks.on_ui_settings(on_ui_settings) 308 | script_callbacks.on_before_image_saved(add_stealth_pnginfo) 309 | script_callbacks.on_after_component(on_after_component_change_pnginfo_image_mode) 310 | --------------------------------------------------------------------------------