feat: add metadata to images

args_manager.py

@@ -20,7 +20,10 @@
                                 help="Prevent writing images and logs to hard drive.")
 
 args_parser.parser.add_argument("--disable-analytics", action='store_true',
-                                help="Disables analytics for Gradio", default=False)
+                                help="Disables analytics for Gradio.")
+
+args_parser.parser.add_argument("--disable-metadata", action='store_true',
+                                help="Disables saving metadata to images.")
 
 args_parser.parser.add_argument("--disable-preset-download", action='store_true',
                                 help="Disables downloading models for presets", default=False)

language/en.json

@@ -374,5 +374,12 @@
     "* Powered by Fooocus Inpaint Engine (beta)": "* Powered by Fooocus Inpaint Engine (beta)",
     "Fooocus Enhance": "Fooocus Enhance",
     "Fooocus Cinematic": "Fooocus Cinematic",
-    "Fooocus Sharp": "Fooocus Sharp"
+    "Fooocus Sharp": "Fooocus Sharp",
+    "Drag any image generated by Fooocus here": "Drag any image generated by Fooocus here",
+    "Metadata": "Metadata",
+    "Apply Metadata": "Apply Metadata",
+    "Metadata Scheme": "Metadata Scheme",
+    "Image Prompt parameters are not included. Use a1111 for compatibility with Civitai.": "Image Prompt parameters are not included. Use a1111 for compatibility with Civitai.",
+    "fooocus (json)": "fooocus (json)",
+    "a1111 (plain text)": "a1111 (plain text)"
 }

modules/async_worker.py

@@ -19,6 +19,7 @@ def __init__(self, args):
 def worker():
     global async_tasks
 
+    import os
     import traceback
     import math
     import numpy as np
@@ -39,13 +40,16 @@ def worker():
     import extras.ip_adapter as ip_adapter
     import extras.face_crop
     import fooocus_version
+    import args_manager
 
     from modules.sdxl_styles import apply_style, apply_wildcards, fooocus_expansion, apply_arrays
     from modules.private_logger import log
     from extras.expansion import safe_str
     from modules.util import remove_empty_str, HWC3, resize_image, \
         get_image_shape_ceil, set_image_shape_ceil, get_shape_ceil, resample_image, erode_or_dilate, ordinal_suffix
     from modules.upscaler import perform_upscale
+    from modules.flags import Performance
+    from modules.meta_parser import get_metadata_parser, MetadataScheme
 
     pid = os.getpid()
     print(f'Started worker with PID {pid}')
@@ -135,7 +139,7 @@ def handler(async_task):
         prompt = args.pop()
         negative_prompt = args.pop()
         style_selections = args.pop()
-        performance_selection = args.pop()
+        performance_selection = Performance(args.pop())
         aspect_ratios_selection = args.pop()
         image_number = args.pop()
         image_seed = args.pop()
@@ -153,6 +157,7 @@ def handler(async_task):
         inpaint_input_image = args.pop()
         inpaint_additional_prompt = args.pop()
         inpaint_mask_image_upload = args.pop()
+
         disable_preview = args.pop()
         disable_intermediate_results = args.pop()
         disable_seed_increment = args.pop()
@@ -190,8 +195,11 @@ def handler(async_task):
         invert_mask_checkbox = args.pop()
         inpaint_erode_or_dilate = args.pop()
 
+        save_metadata_to_images = args.pop() if not args_manager.args.disable_metadata else False
+        metadata_scheme = MetadataScheme(args.pop()) if not args_manager.args.disable_metadata else MetadataScheme.FOOOCUS
+
         cn_tasks = {x: [] for x in flags.ip_list}
-        for _ in range(4):
+        for _ in range(flags.controlnet_image_count):
             cn_img = args.pop()
             cn_stop = args.pop()
             cn_weight = args.pop()
@@ -216,17 +224,9 @@ def handler(async_task):
             print(f'Refiner disabled because base model and refiner are same.')
             refiner_model_name = 'None'
 
-        assert performance_selection in ['Speed', 'Quality', 'Extreme Speed']
-
-        steps = 30
-
-        if performance_selection == 'Speed':
-            steps = 30
-
-        if performance_selection == 'Quality':
-            steps = 60
+        steps = performance_selection.steps()
 
-        if performance_selection == 'Extreme Speed':
+        if performance_selection == Performance.EXTREME_SPEED:
             print('Enter LCM mode.')
             progressbar(async_task, 1, 'Downloading LCM components ...')
             loras += [(modules.config.downloading_sdxl_lcm_lora(), 1.0)]
@@ -244,7 +244,6 @@ def handler(async_task):
             adm_scaler_positive = 1.0
             adm_scaler_negative = 1.0
             adm_scaler_end = 0.0
-            steps = 8
 
         print(f'[Parameters] Adaptive CFG = {adaptive_cfg}')
         print(f'[Parameters] Sharpness = {sharpness}')
@@ -305,16 +304,7 @@ def handler(async_task):
                     if 'fast' in uov_method:
                         skip_prompt_processing = True
                     else:
-                        steps = 18
-
-                        if performance_selection == 'Speed':
-                            steps = 18
-
-                        if performance_selection == 'Quality':
-                            steps = 36
-
-                        if performance_selection == 'Extreme Speed':
-                            steps = 8
+                        steps = performance_selection.steps_uov()
 
                     progressbar(async_task, 1, 'Downloading upscale models ...')
                     modules.config.downloading_upscale_model()
@@ -830,31 +820,50 @@ def callback(step, x0, x, total_steps, y):
 
                 img_paths = []
                 for x in imgs:
-                    d = [
-                        ('Prompt', task['log_positive_prompt']),
-                        ('Negative Prompt', task['log_negative_prompt']),
-                        ('Fooocus V2 Expansion', task['expansion']),
-                        ('Styles', str(raw_style_selections)),
-                        ('Performance', performance_selection),
-                        ('Resolution', str((width, height))),
-                        ('Sharpness', sharpness),
-                        ('Guidance Scale', guidance_scale),
-                        ('ADM Guidance', str((
-                            modules.patch.patch_settings[pid].positive_adm_scale,
-                            modules.patch.patch_settings[pid].negative_adm_scale,
-                            modules.patch.patch_settings[pid].adm_scaler_end))),
-                        ('Base Model', base_model_name),
-                        ('Refiner Model', refiner_model_name),
-                        ('Refiner Switch', refiner_switch),
-                        ('Sampler', sampler_name),
-                        ('Scheduler', scheduler_name),
-                        ('Seed', task['task_seed']),
-                    ]
+                    d = [('Prompt', 'prompt', task['log_positive_prompt']),
+                         ('Negative Prompt', 'negative_prompt', task['log_negative_prompt']),
+                         ('Fooocus V2 Expansion', 'prompt_expansion', task['expansion']),
+                         ('Styles', 'styles', str(raw_style_selections)),
+                         ('Performance', 'performance', performance_selection.value),
+                         ('Resolution', 'resolution', str((width, height))),
+                         ('Guidance Scale', 'guidance_scale', guidance_scale),
+                         ('Sharpness', 'sharpness', sharpness),
+                         ('ADM Guidance', 'adm_guidance', str((
+                             modules.patch.patch_settings[pid].positive_adm_scale,
+                             modules.patch.patch_settings[pid].negative_adm_scale,
+                             modules.patch.patch_settings[pid].adm_scaler_end))),
+                         ('Base Model', 'base_model', base_model_name),
+                         ('Refiner Model', 'refiner_model', refiner_model_name),
+                         ('Refiner Switch', 'refiner_switch', refiner_switch)]
+
+                    if refiner_model_name != 'None':
+                        if overwrite_switch > 0:
+                            d.append(('Overwrite Switch', 'overwrite_switch', overwrite_switch))
+                        if refiner_swap_method != flags.refiner_swap_method:
+                            d.append(('Refiner Swap Method', 'refiner_swap_method', refiner_swap_method))
+                    if modules.patch.patch_settings[pid].adaptive_cfg != modules.config.default_cfg_tsnr:
+                        d.append(('CFG Mimicking from TSNR', 'adaptive_cfg', modules.patch.patch_settings[pid].adaptive_cfg))
+
+                    d.append(('Sampler', 'sampler', sampler_name))
+                    d.append(('Scheduler', 'scheduler', scheduler_name))
+                    d.append(('Seed', 'seed', task['task_seed']))
+
+                    if freeu_enabled:
+                        d.append(('FreeU', 'freeu', str((freeu_b1, freeu_b2, freeu_s1, freeu_s2))))
+
+                    metadata_parser = None
+                    if save_metadata_to_images:
+                        metadata_parser = modules.meta_parser.get_metadata_parser(metadata_scheme)
+                        metadata_parser.set_data(task['log_positive_prompt'], task['positive'],
+                                                 task['log_negative_prompt'], task['negative'],
+                                                 steps, base_model_name, refiner_model_name, loras)
+
                     for li, (n, w) in enumerate(loras):
                         if n != 'None':
-                            d.append((f'LoRA {li + 1}', f'{n} : {w}'))
-                    d.append(('Version', 'v' + fooocus_version.version))
-                    img_paths.append(log(x, d))
+                            d.append((f'LoRA {li + 1}', f'lora_combined_{li + 1}', f'{n} : {w}'))
+
+                    d.append(('Version', 'version', 'Fooocus v' + fooocus_version.version))
+                    img_paths.append(log(x, d, metadata_parser))
 
                 yield_result(async_task, img_paths, do_not_show_finished_images=len(tasks) == 1 or disable_intermediate_results)
             except ldm_patched.modules.model_management.InterruptProcessingException as e:

modules/config.py

@@ -8,7 +8,7 @@
 
 from modules.model_loader import load_file_from_url
 from modules.util import get_files_from_folder, makedirs_with_log
-
+from modules.flags import Performance, MetadataScheme
 
 config_path = os.path.abspath("./config.txt")
 config_example_path = os.path.abspath("config_modification_tutorial.txt")
@@ -293,8 +293,8 @@ def get_config_item_or_set_default(key, default_value, validator, disable_empty_
 )
 default_performance = get_config_item_or_set_default(
     key='default_performance',
-    default_value='Speed',
-    validator=lambda x: x in modules.flags.performance_selections
+    default_value=Performance.SPEED.value,
+    validator=lambda x: x in Performance.list()
 )
 default_advanced_checkbox = get_config_item_or_set_default(
     key='default_advanced_checkbox',
@@ -369,6 +369,21 @@ def get_config_item_or_set_default(key, default_value, validator, disable_empty_
     ],
     validator=lambda x: isinstance(x, list) and all(isinstance(v, str) for v in x)
 )
+default_save_metadata_to_images = get_config_item_or_set_default(
+    key='default_save_metadata_to_images',
+    default_value=False,
+    validator=lambda x: isinstance(x, bool)
+)
+default_metadata_scheme = get_config_item_or_set_default(
+    key='default_metadata_scheme',
+    default_value=MetadataScheme.FOOOCUS.value,
+    validator=lambda x: x in [y[1] for y in modules.flags.metadata_scheme if y[1] == x]
+)
+metadata_created_by = get_config_item_or_set_default(
+    key='metadata_created_by',
+    default_value='',
+    validator=lambda x: isinstance(x, str)
+)
 
 example_inpaint_prompts = [[x] for x in example_inpaint_prompts]
 
@@ -391,6 +406,7 @@ def get_config_item_or_set_default(key, default_value, validator, disable_empty_
     "default_prompt_negative",
     "default_styles",
     "default_aspect_ratio",
+    "default_save_metadata_to_images",
     "checkpoint_downloads",
     "embeddings_downloads",
     "lora_downloads",

modules/flags.py

@@ -1,3 +1,5 @@
+from enum import IntEnum, Enum
+
 disabled = 'Disabled'
 enabled = 'Enabled'
 subtle_variation = 'Vary (Subtle)'
@@ -10,16 +12,49 @@
     disabled, subtle_variation, strong_variation, upscale_15, upscale_2, upscale_fast
 ]
 
-KSAMPLER_NAMES = ["euler", "euler_ancestral", "heun", "heunpp2","dpm_2", "dpm_2_ancestral",
-                  "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu",
-                  "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm"]
+CIVITAI_NO_KARRAS = ["euler", "euler_ancestral", "heun", "dpm_fast", "dpm_adaptive", "ddim", "uni_pc"]
+
+# fooocus: a1111 (Civitai)
+KSAMPLER = {
+    "euler": "Euler",
+    "euler_ancestral": "Euler a",
+    "heun": "Heun",
+    "heunpp2": "",
+    "dpm_2": "DPM2",
+    "dpm_2_ancestral": "DPM2 a",
+    "lms": "LMS",
+    "dpm_fast": "DPM fast",
+    "dpm_adaptive": "DPM adaptive",
+    "dpmpp_2s_ancestral": "DPM++ 2S a",
+    "dpmpp_sde": "DPM++ SDE",
+    "dpmpp_sde_gpu": "DPM++ SDE",
+    "dpmpp_2m": "DPM++ 2M",
+    "dpmpp_2m_sde": "DPM++ 2M SDE",
+    "dpmpp_2m_sde_gpu": "DPM++ 2M SDE",
+    "dpmpp_3m_sde": "",
+    "dpmpp_3m_sde_gpu": "",
+    "ddpm": "",
+    "lcm": "LCM"
+}
+
+SAMPLER_EXTRA = {
+    "ddim": "DDIM",
+    "uni_pc": "UniPC",
+    "uni_pc_bh2": ""
+}
+
+SAMPLERS = KSAMPLER | SAMPLER_EXTRA
+
+KSAMPLER_NAMES = list(KSAMPLER.keys())
 
 SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform", "lcm", "turbo"]
-SAMPLER_NAMES = KSAMPLER_NAMES + ["ddim", "uni_pc", "uni_pc_bh2"]
+SAMPLER_NAMES = KSAMPLER_NAMES + list(SAMPLER_EXTRA.keys())
 
 sampler_list = SAMPLER_NAMES
 scheduler_list = SCHEDULER_NAMES
 
+refiner_swap_method = 'joint'
+
 cn_ip = "ImagePrompt"
 cn_ip_face = "FaceSwap"
 cn_canny = "PyraCanny"
@@ -33,12 +68,56 @@
 }  # stop, weight
 
 inpaint_engine_versions = ['None', 'v1', 'v2.5', 'v2.6']
-performance_selections = ['Speed', 'Quality', 'Extreme Speed']
-
 inpaint_option_default = 'Inpaint or Outpaint (default)'
 inpaint_option_detail = 'Improve Detail (face, hand, eyes, etc.)'
 inpaint_option_modify = 'Modify Content (add objects, change background, etc.)'
 inpaint_options = [inpaint_option_default, inpaint_option_detail, inpaint_option_modify]
 
 desc_type_photo = 'Photograph'
 desc_type_anime = 'Art/Anime'
+
+
+class MetadataScheme(Enum):
+    FOOOCUS = 'fooocus'
+    A1111 = 'a1111'
+
+
+metadata_scheme = [
+    (f'{MetadataScheme.FOOOCUS.value} (json)', MetadataScheme.FOOOCUS.value),
+    (f'{MetadataScheme.A1111.value} (plain text)', MetadataScheme.A1111.value),
+]
+
+lora_count = 5
+
+controlnet_image_count = 4
+
+
+class Steps(IntEnum):
+    QUALITY = 60
+    SPEED = 30
+    EXTREME_SPEED = 8
+
+
+class StepsUOV(IntEnum):
+    QUALITY = 36
+    SPEED = 18
+    EXTREME_SPEED = 8
+
+
+class Performance(Enum):
+    QUALITY = 'Quality'
+    SPEED = 'Speed'
+    EXTREME_SPEED = 'Extreme Speed'
+
+    @classmethod
+    def list(cls) -> list:
+        return list(map(lambda c: c.value, cls))
+
+    def steps(self) -> int | None:
+        return Steps[self.name].value if Steps[self.name] else None
+
+    def steps_uov(self) -> int | None:
+        return StepsUOV[self.name].value if Steps[self.name] else None
+
+
+performance_selections = Performance.list()