feat: add metadata logging for images

inspired by https://github.com/MoonRide303/Fooocus-MRE

modules/async_worker.py

@@ -16,6 +16,7 @@ def worker():
 
     import traceback
     import math
+    import json
     import numpy as np
     import torch
     import time
@@ -354,6 +355,9 @@ def worker():
 
         progressbar(async_task, 1, 'Initializing ...')
 
+        raw_prompt = prompt
+        raw_negative_prompt = negative_prompt
+
         if not skip_prompt_processing:
 
             prompts = remove_empty_str([safe_str(p) for p in prompt.splitlines()], default='')
@@ -774,6 +778,70 @@ def worker():
                 if inpaint_worker.current_task is not None:
                     imgs = [inpaint_worker.current_task.post_process(x) for x in imgs]
 
+                metadata = {
+                    'prompt': raw_prompt, 'negative_prompt': raw_negative_prompt, 'styles': str(raw_style_selections),
+                    'real_prompt': task['log_positive_prompt'], 'real_negative_prompt': task['log_negative_prompt'],
+                    'seed': task['task_seed'], 'width': width, 'height': height,
+                    'sampler': sampler_name, 'scheduler': scheduler_name, 'performance': performance_selection,
+                    'steps': steps, 'refiner_switch': refiner_switch, 'sharpness': sharpness, 'cfg': cfg_scale,
+                    'base_model': base_model_name, 'refiner_model': refiner_model_name,
+                    'freeu': advanced_parameters.freeu_enabled,
+                    'img2img': input_image_checkbox,
+                    'prompt_expansion': task['expansion']
+                }
+
+                if advanced_parameters.freeu_enabled:
+                    metadata |= {
+                        'freeu_b1': advanced_parameters.freeu_b1, 'freeu_b2': advanced_parameters.freeu_b2, 'freeu_s1': advanced_parameters.freeu_s1, 'freeu_s2': advanced_parameters.freeu_s2
+                    }
+
+                if 'vary' in goals:
+                    metadata |= {
+                        'uov_method': uov_method, 'denoising_strength': denoising_strength,
+                        #'uov_input_image': raw_uov_input_image
+                    }
+
+                if 'upscale' in goals:
+                    metadata |= {
+                        'uov_method': uov_method, 'scale': f, 
+                        #'uov_input_image': uov_input_image
+                    }
+
+                if 'inpaint' in goals:
+                    if len(outpaint_selections) > 0:
+                        metadata |= {
+                            'outpaint_selections': outpaint_selections
+                        }
+                    else:
+                        metadata |= {
+                            'inpaint_additional_prompt': inpaint_additional_prompt, 'inpaint_mask_upload': advanced_parameters.inpaint_mask_upload_checkbox, 'invert_mask': advanced_parameters.invert_mask_checkbox,
+                            'inpaint_disable_initial_latent': advanced_parameters.inpaint_disable_initial_latent, 'inpaint_engine': advanced_parameters.inpaint_engine,
+                            'inpaint_strength': advanced_parameters.inpaint_strength, 'inpaint_respective_field': advanced_parameters.inpaint_respective_field,
+                            #'inpaint_image': inpaint_image, 'inpaint_mask': inpaint_mask
+                        }
+
+                if 'cn' in goals:
+                    metadata |= {
+                        'canny_low_threshold': advanced_parameters.canny_low_threshold, 'canny_high_threshold': advanced_parameters.canny_high_threshold,
+                    }
+
+                    ip_list = {x: [] for x in flags.ip_list}
+                    cn_task_index = 1
+                    for cn_type in ip_list:
+                        for cn_task in cn_tasks[cn_type]:
+                            cn_img, cn_stop, cn_weight = cn_task
+                            metadata |= {
+                                # TODO check (A1111) compatibility
+                                f'image_prompt_{cn_task_index}': {
+                                    'cn_type': cn_type, 'cn_stop': cn_stop, 'cn_weight': cn_weight, 
+                                    #'cn_image': cn_img
+                                }
+                            }
+                            cn_task_index += 1
+
+                metadata |= {'software': f'Fooocus v{fooocus_version.version}'}
+                metadata_string = json.dumps(metadata, ensure_ascii=False)
+
                 for x in imgs:
                     d = [
                         ('Prompt', task['log_positive_prompt']),
@@ -799,7 +867,7 @@ def worker():
                         if n != 'None':
                             d.append((f'LoRA {li + 1}', f'{n} : {w}'))
                     d.append(('Version', 'v' + fooocus_version.version))
-                    log(x, d)
+                    log(x, d, metadata_string, True)
 
                 yield_result(async_task, imgs, do_not_show_finished_images=len(tasks) == 1)
             except ldm_patched.modules.model_management.InterruptProcessingException as e:

modules/private_logger.py

@@ -5,6 +5,7 @@ import json
 import urllib.parse
 
 from PIL import Image
+from PIL.PngImagePlugin import PngInfo
 from modules.util import generate_temp_filename
 
 
@@ -18,13 +19,20 @@ def get_current_html_path():
     return html_name
 
 
-def log(img, dic):
+def log(img, dic, metadata=None, save_metadata_to_image=False):
     if args_manager.args.disable_image_log:
         return
 
     date_string, local_temp_filename, only_name = generate_temp_filename(folder=modules.config.path_outputs, extension='png')
     os.makedirs(os.path.dirname(local_temp_filename), exist_ok=True)
-    Image.fromarray(img).save(local_temp_filename)
+
+    if save_metadata_to_image:
+        pnginfo = PngInfo()
+        pnginfo.add_text("Comment", metadata)
+    else:
+        pnginfo = None
+    Image.fromarray(img).save(local_temp_filename, pnginfo=pnginfo)
+
     html_name = os.path.join(os.path.dirname(local_temp_filename), 'log.html')
 
     css_styles = (