From 5b99064dfbba98fb4846e4be750a7ab917b38fa9 Mon Sep 17 00:00:00 2001
From: Christian Heller <c.heller@plomlompom.de>
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