From 5b99064dfbba98fb4846e4be750a7ab917b38fa9 Mon Sep 17 00:00:00 2001 From: Christian Heller Date: Fri, 13 Sep 2024 03:34:47 +0200 Subject: [PATCH] Switch from EXIF to PNG chunk tEXt metadata. --- browser.py | 26 +++++++++++++------------- requirements.txt | 1 - stable.py | 8 +++++--- stable/core.py | 10 +++++----- 4 files changed, 23 insertions(+), 22 deletions(-) diff --git a/browser.py b/browser.py index 1f25d2d..b38cf10 100755 --- a/browser.py +++ b/browser.py @@ -3,7 +3,8 @@ from json import dump as json_dump, load as json_load from os.path import exists as path_exists, join as path_join, abspath from argparse import ArgumentParser -from exiftool import ExifToolHelper # type: ignore +from PIL import Image +from PIL.PngImagePlugin import PngImageFile import gi # type: ignore gi.require_version('Gtk', '4.0') gi.require_version('Gdk', '4.0') @@ -95,14 +96,15 @@ class ImgItem(FileItem): for k in cached.keys(): setattr(self, k, cached[k]) - def set_metadata(self, exif_tool, cache): - """Set instance attributes from 'Comment' EXIF tag, write to cache.""" - for d in exif_tool.get_tags([self.full_path], ['Comment']): - for k, v in d.items(): - if k.endswith('Comment'): - gen_params = GenParams.from_str(v) - for k, v_ in gen_params.as_dict.items(): - setattr(self, k, v_) + def set_metadata(self, cache): + """Set instance attributes from 'image file's GenParams PNG chunk.""" + img = Image.open(self.full_path) + if isinstance(img, PngImageFile): + gen_params_as_str = img.text.get('generation_parameters', '') + if gen_params_as_str: + gen_params = GenParams.from_str(gen_params_as_str) + for k, v_ in gen_params.as_dict.items(): + setattr(self, k, v_) cached = {} for k in (k.lower() for k in GEN_PARAMS): cached[k] = getattr(self, k) @@ -119,7 +121,6 @@ class MainWindow(Gtk.Window): recurse_dirs: bool per_row: int metadata: Gtk.TextBuffer - sort_order: list sort_store: Gtk.ListStore sort_selection: Gtk.SingleSelection prev_key: list @@ -622,9 +623,8 @@ class MainWindow(Gtk.Window): if '' == item.model: to_set_metadata_on += [item] self.gallery_store.append(item) - with ExifToolHelper() as et: - for item in to_set_metadata_on: - item.set_metadata(et, cache) + for item in to_set_metadata_on: + item.set_metadata(cache) old_selection = self.gallery_selection.props.selected_item self.block_file_selection_updates = True diff --git a/requirements.txt b/requirements.txt index e17a459..e5365cf 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,4 @@ # for stable.py -pyexiftool pillow torch diffusers diff --git a/stable.py b/stable.py index 694b250..351add1 100755 --- a/stable.py +++ b/stable.py @@ -20,9 +20,9 @@ def parse_args(): help='model filename (-P will pre prefixed, but may ' 'also be full path on its own)') parser.add_argument('-o', '--output', - help='output filename or path; if -q > 1, will insert ' - 'incremented counter number; if no image file ' - 'extension included, defaults to .png') + help='output filename or path; if -q > 1 or name ' + 'pre-existing, will insert incremented counter ' + 'number; will append .png extension if not provided') parser.add_argument('-p', '--prompt', help='textual guidance to image generation') parser.add_argument('-q', '--quantity', default=1, type=int, @@ -105,6 +105,8 @@ def run(): dir_path = dirname(args.output) if dirname(args.output) else '.' filename = basename(args.output) filename_sans_ext, ext = splitext(filename) + if ext not in {'', '.png', '.PNG'}: + raise Exception('Can only export to PNG.') ext = ext if ext else '.png' filename_with_ext = f'{filename_sans_ext}{ext}' start_at_idx = 0 diff --git a/stable/core.py b/stable/core.py index 3d8d4a0..4d0b05a 100644 --- a/stable/core.py +++ b/stable/core.py @@ -3,7 +3,7 @@ from logging import (Formatter as LogFormatter, captureWarnings, from diffusers import StableDiffusionPipeline from diffusers.utils import logging from torch import Generator, float16 -from exiftool import ExifToolHelper # type: ignore +from PIL.PngImagePlugin import PngInfo SAFETY_CHECKER_WARNING_PATTERN = 'You have disabled the safety checker' @@ -51,6 +51,7 @@ class ImageMaker: return [s.__name__ for s in self.pipe.scheduler.compatibles] def gen_image_to(self, path): + """Create image and write as file with metadata to path.""" if None in {self.gen_params.seed, self.gen_params.prompt, self.gen_params.guidance, self.gen_params.height, self.gen_params.width, self.gen_params.n_steps}: @@ -63,7 +64,6 @@ class ImageMaker: width=self.gen_params.width, num_inference_steps=self.gen_params.n_steps, ).images[0] - image.save(path) - with ExifToolHelper() as et: - et.set_tags([path], tags={'Comment': self.gen_params.to_str}, - params=['-overwrite_original']) + png_info = PngInfo() + png_info.add_text('generation_parameters', self.gen_params.to_str) + image.save(path, pnginfo=png_info) -- 2.30.2