feat: add exif data processing (saving and loading)

modules/meta_parser.py

@@ -7,6 +7,7 @@ from pathlib import Path
 import gradio as gr
 from PIL import Image
 
+import fooocus_version
 import modules.config
 import modules.sdxl_styles
 from modules.flags import MetadataScheme, Performance, Steps
@@ -518,16 +519,28 @@ def get_metadata_parser(metadata_scheme: MetadataScheme) -> MetadataParser:
             raise NotImplementedError
 
 
-def read_info_from_image(filepath) -> tuple[str | None, dict, MetadataScheme | None]:
+def read_info_from_image(filepath) -> tuple[str | None, MetadataScheme | None]:
     with Image.open(filepath) as image:
         items = (image.info or {}).copy()
 
     parameters = items.pop('parameters', None)
+    metadata_scheme = items.pop('fooocus_scheme', None)
+    exif = items.pop('exif', None)
+
     if parameters is not None and is_json(parameters):
         parameters = json.loads(parameters)
+    elif exif is not None:
+        exif = image.getexif()
+        # 0x9286 = UserComment
+        parameters = exif.get(0x9286, None)
+        # 0x927C = MakerNote
+        metadata_scheme = exif.get(0x927C, None)
+
+        if is_json(parameters):
+            parameters = json.loads(parameters)
 
     try:
-        metadata_scheme = MetadataScheme(items.pop('fooocus_scheme', None))
+        metadata_scheme = MetadataScheme(metadata_scheme)
     except ValueError:
         metadata_scheme = None
 
@@ -538,4 +551,16 @@ def read_info_from_image(filepath) -> tuple[str | None, dict, MetadataScheme | N
         if isinstance(parameters, str):
             metadata_scheme = MetadataScheme.A1111
 
-    return parameters, items, metadata_scheme
+    return parameters, metadata_scheme
+
+
+def get_exif(metadata: str | None, metadata_scheme: str):
+    exif = Image.Exif()
+    # tags see see https://github.com/python-pillow/Pillow/blob/9.2.x/src/PIL/ExifTags.py
+    # 0x9286 = UserComment
+    exif[0x9286] = metadata
+    # 0x0131 = Software
+    exif[0x0131] = 'Fooocus v' + fooocus_version.version
+    # 0x927C = MakerNote
+    exif[0x927C] = metadata_scheme
+    return exif

modules/private_logger.py

@@ -6,9 +6,8 @@ import urllib.parse
 
 from PIL import Image
 from PIL.PngImagePlugin import PngInfo
-from fooocus_version import version
 from modules.util import generate_temp_filename
-from modules.meta_parser import MetadataParser
+from modules.meta_parser import MetadataParser, get_exif
 
 log_cache = {}
 
@@ -27,7 +26,8 @@ def log(img, metadata, metadata_parser: MetadataParser | None = None, output_for
     date_string, local_temp_filename, only_name = generate_temp_filename(folder=path_outputs, extension=output_format)
     os.makedirs(os.path.dirname(local_temp_filename), exist_ok=True)
 
-    parsed_parameters = metadata_parser.parse_string(metadata) if metadata_parser else None
+    parsed_parameters = metadata_parser.parse_string(metadata) if metadata_parser is not None else ''
+    image = Image.fromarray(img)
 
     if output_format == 'png':
         if parsed_parameters != '':
@@ -36,14 +36,13 @@ def log(img, metadata, metadata_parser: MetadataParser | None = None, output_for
             pnginfo.add_text('fooocus_scheme', metadata_parser.get_scheme().value)
         else:
             pnginfo = None
-
-        Image.fromarray(img).save(local_temp_filename, pnginfo=pnginfo)
+        image.save(local_temp_filename, pnginfo=pnginfo)
     elif output_format == 'jpg':
-        Image.fromarray(img).save(local_temp_filename, quality=95, optimize=True, progressive=True)
+        image.save(local_temp_filename, quality=95, optimize=True, progressive=True, exif=get_exif(parsed_parameters, metadata_parser.get_scheme().value) if metadata_parser else Image.Exif())
     elif output_format == 'webp':
-        Image.fromarray(img).save(local_temp_filename, quality=95, lossless=False)
+        image.save(local_temp_filename, quality=95, lossless=False, exif=get_exif(parsed_parameters, metadata_parser.get_scheme().value) if metadata_parser else Image.Exif())
     else:
-        Image.fromarray(img).save(local_temp_filename)
+        image.save(local_temp_filename)
 
     if args_manager.args.disable_image_log:
         return local_temp_filename