refactor: add step before parsing to set data in parser

- add constructor for MetadataSchema class
- remove showable and copyable from log output
- add functional hash cache (model hashing takes about 5 seconds, only required once per model, using hash lazy loading)

modules/async_worker.py

@@ -42,9 +42,10 @@ def worker():
     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, calculate_sha256
+        get_image_shape_ceil, set_image_shape_ceil, get_shape_ceil, resample_image, erode_or_dilate
     from modules.upscaler import perform_upscale
-    from modules.flags import Performance, MetadataScheme, lora_count
+    from modules.flags import Performance, lora_count
+    from modules.metadata import get_metadata_parser, MetadataScheme
 
     try:
         async_gradio_app = shared.gradio_root
@@ -193,18 +194,6 @@ def worker():
             modules.patch.negative_adm_scale = advanced_parameters.adm_scaler_negative = 1.0
             modules.patch.adm_scaler_end = advanced_parameters.adm_scaler_end = 0.0
 
-        # TODO move hashing to metadata mapper as this slows down the generation process
-        base_model_path = os.path.join(modules.config.path_checkpoints, base_model_name)
-        base_model_hash = calculate_sha256(base_model_path)
-
-        refiner_model_path = os.path.join(modules.config.path_checkpoints, refiner_model_name)
-        refiner_model_hash = calculate_sha256(refiner_model_path) if refiner_model_name != 'None' else ''
-
-        lora_hashes = []
-        for (n, w) in loras:
-            lora_path = os.path.join(modules.config.path_loras, n) if n != 'None' else ''
-            lora_hashes.append(calculate_sha256(lora_path) if n != 'None' else '')
-
         modules.patch.adaptive_cfg = advanced_parameters.adaptive_cfg
         print(f'[Parameters] Adaptive CFG = {modules.patch.adaptive_cfg}')
 
@@ -777,61 +766,54 @@ def worker():
                     imgs = [inpaint_worker.current_task.post_process(x) for x in imgs]
 
                 for x in imgs:
-                    d = [('Prompt', 'prompt', task['log_positive_prompt'], True, True),
+                    d = [('Prompt', 'prompt', task['log_positive_prompt']),
-                         ('Full Positive Prompt', 'full_prompt', task['positive'], False, False),
+                         ('Negative Prompt', 'negative_prompt', task['log_negative_prompt']),
-                         ('Negative Prompt', 'negative_prompt', task['log_negative_prompt'], True, True),
+                         ('Fooocus V2 Expansion', 'prompt_expansion', task['expansion']),
-                         ('Full Negative Prompt', 'full_negative_prompt', task['negative'], False, False),
+                         ('Styles', 'styles', str(raw_style_selections)),
-                         ('Fooocus V2 Expansion', 'prompt_expansion', task['expansion'], True, True),
+                         ('Performance', 'performance', performance_selection.value),
-                         ('Styles', 'styles', str(raw_style_selections), True, True),
+                         ('Resolution', 'resolution', str((width, height))),
-                         ('Performance', 'performance', performance_selection.value, True, True),
+                         ('Guidance Scale', 'guidance_scale', guidance_scale),
-                         ('Steps', 'steps', steps, False, False),
+                         ('Sharpness', 'sharpness', sharpness),
-                         ('Resolution', 'resolution', str((width, height)), True, True),
-                         ('Guidance Scale', 'guidance_scale', guidance_scale, True, True),
-                         ('Sharpness', 'sharpness', sharpness, True, True),
                          ('ADM Guidance', 'adm_guidance', str((
                              modules.patch.positive_adm_scale,
                              modules.patch.negative_adm_scale,
-                             modules.patch.adm_scaler_end)), True, True),
+                             modules.patch.adm_scaler_end))),
-                         ('Base Model', 'base_model', base_model_name, True, True),
+                         ('Base Model', 'base_model', base_model_name),
-                         ('Base Model Hash', 'base_model_hash', base_model_hash, False, False), # TODO move to metadata and use cache
+                         ('Refiner Model', 'refiner_model', refiner_model_name),
-                         ('Refiner Model', 'refiner_model', refiner_model_name, True, True),
+                         ('Refiner Switch', 'refiner_switch', refiner_switch)]
-                         ('Refiner Model Hash', 'refiner_model_hash', refiner_model_hash, False, False), # TODO move to metadata and use cache
-                         ('Refiner Switch', 'refiner_switch', refiner_switch, True, True)]
 
                     # TODO evaluate if this should always be added
                     if refiner_model_name != 'None':
                         if advanced_parameters.overwrite_switch > 0:
-                            d.append(('Overwrite Switch', 'overwrite_switch', advanced_parameters.overwrite_switch, True, True))
+                            d.append(('Overwrite Switch', 'overwrite_switch', advanced_parameters.overwrite_switch))
                         if refiner_swap_method != flags.refiner_swap_method:
-                            d.append(('Refiner Swap Method', 'refiner_swap_method', refiner_swap_method, True, True))
+                            d.append(('Refiner Swap Method', 'refiner_swap_method', refiner_swap_method))
                     if advanced_parameters.adaptive_cfg != modules.config.default_cfg_tsnr:
-                        d.append(('CFG Mimicking from TSNR', 'adaptive_cfg', advanced_parameters.adaptive_cfg, True, True))
+                        d.append(('CFG Mimicking from TSNR', 'adaptive_cfg', advanced_parameters.adaptive_cfg))
 
-                    d.append(('Sampler', 'sampler', sampler_name, True, True))
+                    d.append(('Sampler', 'sampler', sampler_name))
-                    d.append(('Scheduler', 'scheduler', scheduler_name, True, True))
+                    d.append(('Scheduler', 'scheduler', scheduler_name))
-                    d.append(('Seed', 'seed', task['task_seed'], True, True))
+                    d.append(('Seed', 'seed', task['task_seed']))
 
                     if advanced_parameters.freeu_enabled:
                         d.append(('FreeU', 'freeu', str((
                             advanced_parameters.freeu_b1,
                             advanced_parameters.freeu_b2,
                             advanced_parameters.freeu_s1,
-                            advanced_parameters.freeu_s2)), True, True))
+                            advanced_parameters.freeu_s2))))
+
+                    metadata_parser = None
+                    if save_metadata_to_images:
+                        metadata_parser = modules.metadata.get_metadata_parser(metadata_scheme)
+                        metadata_parser.set_data(task['positive'], 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'lora_combined_{li + 1}', f'{n} : {w}', True, True))
+                            d.append((f'LoRA {li + 1}', f'lora_combined_{li + 1}', f'{n} : {w}'))
-                            d.append((f'LoRA {li + 1} Name', f'lora_name_{li + 1}', n, False, False))
-                            d.append((f'LoRA {li + 1} Weight', f'lora_weight_{li + 1}', w, False, False))
-                            # TODO move hashes to metadata handling
-                            d.append((f'LoRA {li + 1} Hash', f'lora_hash_{li + 1}', lora_hashes[li], False, False))
 
-                    d.append(('Version', 'version', 'Fooocus v' + fooocus_version.version, True, True))
+                    d.append(('Version', 'version', 'Fooocus v' + fooocus_version.version))
 
-                    if modules.config.metadata_created_by != '':
+                    log(x, d, metadata_parser)
-                        d.append(('Created By', 'created_by', modules.config.metadata_created_by, False, False))
-
-                    log(x, d, save_metadata_to_images, metadata_scheme)
 
                 yield_result(async_task, imgs, do_not_show_finished_images=len(tasks) == 1)
             except ldm_patched.modules.model_management.InterruptProcessingException as e:

modules/metadata.py

@@ -1,11 +1,12 @@
 import json
+import os
 import re
 from abc import ABC, abstractmethod
 from pathlib import Path
 from PIL import Image
 
 import modules.config
-from modules.flags import MetadataScheme, Performance, Steps, lora_count_with_lcm
+from modules.flags import MetadataScheme, Performance, Steps
 from modules.util import quote, unquote, extract_styles_from_prompt, is_json, calculate_sha256
 
 re_param_code = r'\s*(\w[\w \-/]+):\s*("(?:\\.|[^\\"])+"|[^,]*)(?:,|$)'
@@ -25,6 +26,16 @@ def get_sha256(filepath):
 
 
 class MetadataParser(ABC):
+    def __init__(self):
+        self.full_prompt: str = ''
+        self.full_negative_prompt: str = ''
+        self.steps: int = 30
+        self.base_model_name: str = ''
+        self.base_model_hash: str = ''
+        self.refiner_model_name: str = ''
+        self.refiner_model_hash: str = ''
+        self.loras: list = []
+
     @abstractmethod
     def get_scheme(self) -> MetadataScheme:
         raise NotImplementedError
@@ -37,6 +48,27 @@ class MetadataParser(ABC):
     def parse_string(self, metadata: dict) -> str:
         raise NotImplementedError
 
+    def set_data(self, full_prompt, full_negative_prompt, steps, base_model_name, refiner_model_name, loras):
+        self.full_prompt = full_prompt
+        self.full_negative_prompt = full_negative_prompt
+        self.steps = steps
+        self.base_model_name = Path(base_model_name).stem
+
+        base_model_path = os.path.join(modules.config.path_checkpoints, base_model_name)
+        self.base_model_hash = get_sha256(base_model_path)
+
+        if refiner_model_name not in ['', 'None']:
+            self.refiner_model_name = Path(refiner_model_name).stem
+            refiner_model_path = os.path.join(modules.config.path_checkpoints, refiner_model_name)
+            self.refiner_model_hash = get_sha256(refiner_model_path)
+
+        self.loras = []
+        for (lora_name, lora_weight) in loras:
+            if lora_name != 'None':
+                lora_path = os.path.join(modules.config.path_loras, lora_name)
+                lora_hash = get_sha256(lora_path)
+                self.loras.append((Path(lora_name).stem, lora_weight, lora_hash))
+
 
 class A1111MetadataParser(MetadataParser):
     def get_scheme(self) -> MetadataScheme:
@@ -63,6 +95,7 @@ class A1111MetadataParser(MetadataParser):
         'refiner_model_hash': 'Refiner hash',
         'lora_hashes': 'Lora hashes',
         'lora_weights': 'Lora weights',
+        'created_by': 'User',
         'version': 'Version'
     }
 
@@ -127,65 +160,64 @@ class A1111MetadataParser(MetadataParser):
             lora_filenames = modules.config.lora_filenames.copy()
             lora_filenames.remove(modules.config.downloading_sdxl_lcm_lora())
             for li, lora in enumerate(data['lora_hashes'].split(', ')):
-                name, _, weight = lora.split(': ')
+                lora_name, lora_hash, lora_weight = lora.split(': ')
                 for filename in lora_filenames:
                     path = Path(filename)
-                    if name == path.stem:
+                    if lora_name == path.stem:
-                        data[f'lora_combined_{li + 1}'] = f'{filename} : {weight}'
+                        data[f'lora_combined_{li + 1}'] = f'{filename} : {lora_weight}'
                         break
 
         return data
 
     def parse_string(self, metadata: dict) -> str:
-        data = {k: v for _, k, v, _, _ in metadata}
+        data = {k: v for _, k, v in metadata}
 
-        width, heigth = eval(data['resolution'])
+        width, height = eval(data['resolution'])
-
-        lora_hashes = []
-        for index in range(lora_count_with_lcm):
-            key = f'lora_name_{index + 1}'
-            if key in data:
-                lora_name = Path(data[f'lora_name_{index + 1}']).stem
-                lora_weight = data[f'lora_weight_{index + 1}']
-                lora_hash = data[f'lora_hash_{index + 1}']
-                # workaround for Fooocus not knowing LoRA name in LoRA metadata
-                lora_hashes.append(f'{lora_name}: {lora_hash}: {lora_weight}')
-        lora_hashes_string = ', '.join(lora_hashes)
 
         generation_params = {
             self.fooocus_to_a1111['performance']: data['performance'],
-            self.fooocus_to_a1111['steps']: data['steps'],
+            self.fooocus_to_a1111['steps']: self.steps,
             self.fooocus_to_a1111['sampler']: data['sampler'],
             self.fooocus_to_a1111['seed']: data['seed'],
-            self.fooocus_to_a1111['resolution']: f'{width}x{heigth}',
+            self.fooocus_to_a1111['resolution']: f'{width}x{height}',
             self.fooocus_to_a1111['guidance_scale']: data['guidance_scale'],
             self.fooocus_to_a1111['sharpness']: data['sharpness'],
             self.fooocus_to_a1111['adm_guidance']: data['adm_guidance'],
-            # TODO load model by name / hash
             self.fooocus_to_a1111['base_model']: Path(data['base_model']).stem,
-            self.fooocus_to_a1111['base_model_hash']: data['base_model_hash']
+            self.fooocus_to_a1111['base_model_hash']: self.base_model_hash,
         }
 
-        if 'refiner_model' in data and data['refiner_model'] != 'None' and 'refiner_model_hash' in data:
+        # TODO evaluate if this should always be added
+        if self.refiner_model_name not in ['', 'None']:
             generation_params |= {
-                self.fooocus_to_a1111['refiner_model']: Path(data['refiner_model']).stem,
+                self.fooocus_to_a1111['refiner_model']: self.refiner_model_name,
-                self.fooocus_to_a1111['refiner_model_hash']: data['refiner_model_hash']
+                self.fooocus_to_a1111['refiner_model_hash']: self.refiner_model_hash
             }
 
         for key in ['adaptive_cfg', 'overwrite_switch', 'refiner_swap_method', 'freeu']:
             if key in data:
                 generation_params[self.fooocus_to_a1111[key]] = data[key]
 
+        lora_hashes = []
+        for index, (lora_name, lora_weight, lora_hash) in enumerate(self.loras):
+            # workaround for Fooocus not knowing LoRA name in LoRA metadata
+            lora_hashes.append(f'{lora_name}: {lora_hash}: {lora_weight}')
+        lora_hashes_string = ', '.join(lora_hashes)
+
         generation_params |= {
             self.fooocus_to_a1111['lora_hashes']: lora_hashes_string,
             self.fooocus_to_a1111['version']: data['version']
         }
 
+        if modules.config.metadata_created_by != '':
+            generation_params[self.fooocus_to_a1111['created_by']] = modules.config.metadata_created_by
+
         generation_params_text = ", ".join(
             [k if k == v else f'{k}: {quote(v)}' for k, v in generation_params.items() if v is not None])
         # TODO check if multiline positive prompt is correctly processed
-        positive_prompt_resolved = ', '.join(data['full_prompt']) #TODO add loras to positive prompt if even possible
+        positive_prompt_resolved = ', '.join(self.full_prompt)  # TODO add loras to positive prompt if even possible
-        negative_prompt_resolved = ', '.join(data['full_negative_prompt']) #TODO add loras to negative prompt if even possible
+        negative_prompt_resolved = ', '.join(
+            self.full_negative_prompt)  # TODO add loras to negative prompt if even possible
         negative_prompt_text = f"\nNegative prompt: {negative_prompt_resolved}" if negative_prompt_resolved else ""
         return f"{positive_prompt_resolved}{negative_prompt_text}\n{generation_params_text}".strip()
 
@@ -200,11 +232,11 @@ class FooocusMetadataParser(MetadataParser):
         lora_filenames.remove(modules.config.downloading_sdxl_lcm_lora())
 
         for key, value in metadata.items():
-            if value == '' or value == 'None':
+            if value in ['', 'None']:
                 continue
             if key in ['base_model', 'refiner_model']:
                 metadata[key] = self.replace_value_with_filename(key, value, model_filenames)
-            elif key.startswith(('lora_combined_', 'lora_name_')):
+            elif key.startswith('lora_combined_'):
                 metadata[key] = self.replace_value_with_filename(key, value, lora_filenames)
             else:
                 continue
@@ -212,20 +244,33 @@ class FooocusMetadataParser(MetadataParser):
         return metadata
 
     def parse_string(self, metadata: list) -> str:
+        for li, (label, key, value) in enumerate(metadata):
             # remove model folder paths from metadata
-        for li, (label, key, value, show_in_log, copy_in_log) in enumerate(metadata):
-            if value == '' or value == 'None':
-                continue
-            if key in ['base_model', 'refiner_model'] or key.startswith(('lora_combined_', 'lora_name_')):
             if key.startswith('lora_combined_'):
                 name, weight = value.split(' : ')
                 name = Path(name).stem
                 value = f'{name} : {weight}'
-                else:
+                metadata[li] = (label, key, value)
-                    value = Path(value).stem
-                metadata[li] = (label, key, value, show_in_log, copy_in_log)
 
-        return json.dumps({k: v for _, k, v, _, _ in metadata})
+        res = {k: v for _, k, v in metadata}
+
+        res['full_prompt'] = self.full_prompt
+        res['full_negative_prompt'] = self.full_negative_prompt
+        res['steps'] = self.steps
+        res['base_model'] = self.base_model_name
+        res['base_model_hash'] = self.base_model_hash
+
+        # TODO evaluate if this should always be added
+        if self.refiner_model_name not in ['', 'None']:
+            res['refiner_model'] = self.refiner_model_name
+            res['refiner_model_hash'] = self.refiner_model_hash
+
+        res['loras'] = self.loras
+
+        if modules.config.metadata_created_by != '':
+            res['created_by'] = modules.config.metadata_created_by
+
+        return json.dumps(res)
 
     @staticmethod
     def replace_value_with_filename(key, value, filenames):

modules/private_logger.py

@@ -7,7 +7,7 @@ import urllib.parse
 from PIL import Image
 from PIL.PngImagePlugin import PngInfo
 from modules.util import generate_temp_filename
-from modules.metadata import MetadataScheme
+from modules.metadata import MetadataParser
 
 
 log_cache = {}
@@ -20,22 +20,21 @@ def get_current_html_path():
     return html_name
 
 
-def log(img, metadata, save_metadata_to_image=False, metadata_scheme: MetadataScheme = MetadataScheme.FOOOCUS):
+def log(img, metadata, metadata_parser: MetadataParser | None = None):
     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)
 
-    if save_metadata_to_image:
+    pnginfo = None
-        metadata_parser = modules.metadata.get_metadata_parser(metadata_scheme)
+    if metadata_parser is not None:
         parsed_parameters = metadata_parser.parse_string(metadata)
 
         pnginfo = PngInfo()
         pnginfo.add_text('parameters', parsed_parameters)
-        pnginfo.add_text('fooocus_scheme', metadata_scheme.value)
+        pnginfo.add_text('fooocus_scheme', metadata_parser.get_scheme().value)
-    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')
@@ -98,13 +97,12 @@ def log(img, metadata, save_metadata_to_image=False, metadata_scheme: MetadataSc
     item = f"<div id=\"{div_name}\" class=\"image-container\"><hr><table><tr>\n"
     item += f"<td><a href=\"{only_name}\" target=\"_blank\"><img src='{only_name}' onerror=\"this.closest('.image-container').style.display='none';\" loading='lazy'></img></a><div>{only_name}</div></td>"
     item += "<td><table class='metadata'>"
-    for label, key, value, showable, copyable in metadata:
+    for label, key, value in metadata:
-        if showable:
         value_txt = str(value).replace('\n', ' </br> ')
         item += f"<tr><td class='label'>{label}</td><td class='value'>{value_txt}</td></tr>\n"
     item += "</table>"
 
-    js_txt = urllib.parse.quote(json.dumps({k: v for _, k, v, _, copyable in metadata if copyable}, indent=0), safe='')
+    js_txt = urllib.parse.quote(json.dumps({k: v for _, k, v in metadata}, indent=0), safe='')
     item += f"</br><button onclick=\"to_clipboard('{js_txt}')\">Copy to Clipboard</button>"
 
     item += "</td>"