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.save(local_temp_filename, pnginfo=pnginfo)
-        Image.fromarray(img).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

webui.py

@@ -278,15 +278,12 @@ with shared.gradio_root:
                             metadata_import_button = gr.Button(value='Apply Metadata')
 
                         def trigger_metadata_preview(filepath):
-                            parameters, items, metadata_scheme = modules.meta_parser.read_info_from_image(filepath)
+                            parameters, metadata_scheme = modules.meta_parser.read_info_from_image(filepath)
 
                             results = {}
                             if parameters is not None:
                                 results['parameters'] = parameters
 
-                            if items:
-                                results['items'] = items
-
                             if isinstance(metadata_scheme, flags.MetadataScheme):
                                 results['metadata_scheme'] = metadata_scheme.value
 
@@ -326,8 +323,7 @@ with shared.gradio_root:
 
                 output_format = gr.Radio(label='Output Format',
                                             choices=modules.flags.output_formats,
-                                            value=modules.config.default_output_format,
+                                            value=modules.config.default_output_format)
-                                            info='Metadata support has only been implemented for png.')
 
                 negative_prompt = gr.Textbox(label='Negative Prompt', show_label=True, placeholder="Type prompt here.",
                                              info='Describing what you do not want to see.', lines=2,
@@ -743,7 +739,7 @@ with shared.gradio_root:
         load_parameter_button.click(modules.meta_parser.load_parameter_button_click, inputs=[prompt, state_is_generating], outputs=load_data_outputs, queue=False, show_progress=False)
 
         def trigger_metadata_import(filepath, state_is_generating):
-            parameters, items, metadata_scheme = modules.meta_parser.read_info_from_image(filepath)
+            parameters, metadata_scheme = modules.meta_parser.read_info_from_image(filepath)
             if parameters is None:
                 print('Could not find metadata in the image!')
                 parsed_parameters = {}