diff --git a/modules/async_worker.py b/modules/async_worker.py index ff12f5ec..0de7daa4 100644 --- a/modules/async_worker.py +++ b/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), - ('Full Positive Prompt', 'full_prompt', task['positive'], False, False), - ('Negative Prompt', 'negative_prompt', task['log_negative_prompt'], True, True), - ('Full Negative Prompt', 'full_negative_prompt', task['negative'], False, False), - ('Fooocus V2 Expansion', 'prompt_expansion', task['expansion'], True, True), - ('Styles', 'styles', str(raw_style_selections), True, True), - ('Performance', 'performance', performance_selection.value, True, True), - ('Steps', 'steps', steps, False, False), - ('Resolution', 'resolution', str((width, height)), True, True), - ('Guidance Scale', 'guidance_scale', guidance_scale, True, True), - ('Sharpness', 'sharpness', sharpness, True, True), + 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.positive_adm_scale, modules.patch.negative_adm_scale, - modules.patch.adm_scaler_end)), True, True), - ('Base Model', 'base_model', base_model_name, True, True), - ('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, True, True), - ('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)] + modules.patch.adm_scaler_end))), + ('Base Model', 'base_model', base_model_name), + ('Refiner Model', 'refiner_model', refiner_model_name), + ('Refiner Switch', 'refiner_switch', refiner_switch)] # 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(('Scheduler', 'scheduler', scheduler_name, True, True)) - d.append(('Seed', 'seed', task['task_seed'], True, True)) + d.append(('Sampler', 'sampler', sampler_name)) + d.append(('Scheduler', 'scheduler', scheduler_name)) + 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} 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((f'LoRA {li + 1}', f'lora_combined_{li + 1}', f'{n} : {w}')) - 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 != '': - d.append(('Created By', 'created_by', modules.config.metadata_created_by, False, False)) - - log(x, d, save_metadata_to_images, metadata_scheme) + log(x, d, metadata_parser) yield_result(async_task, imgs, do_not_show_finished_images=len(tasks) == 1) except ldm_patched.modules.model_management.InterruptProcessingException as e: diff --git a/modules/metadata.py b/modules/metadata.py index 868c2b48..6c03edc5 100644 --- a/modules/metadata.py +++ b/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: - data[f'lora_combined_{li + 1}'] = f'{filename} : {weight}' + if lora_name == path.stem: + 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']) - - 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) + width, height = eval(data['resolution']) 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_hash']: data['refiner_model_hash'] + self.fooocus_to_a1111['refiner_model']: self.refiner_model_name, + 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 - negative_prompt_resolved = ', '.join(data['full_negative_prompt']) #TODO add loras to negative prompt if even possible + positive_prompt_resolved = ', '.join(self.full_prompt) # TODO add loras to positive 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: - # 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: - value = Path(value).stem - metadata[li] = (label, key, value, show_in_log, copy_in_log) + for li, (label, key, value) in enumerate(metadata): + # remove model folder paths from metadata + if key.startswith('lora_combined_'): + name, weight = value.split(' : ') + name = Path(name).stem + value = f'{name} : {weight}' + metadata[li] = (label, key, value) - 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): diff --git a/modules/private_logger.py b/modules/private_logger.py index 1afcaa55..61264da0 100644 --- a/modules/private_logger.py +++ b/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: - metadata_parser = modules.metadata.get_metadata_parser(metadata_scheme) + pnginfo = None + 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) - else: - pnginfo = None + pnginfo.add_text('fooocus_scheme', metadata_parser.get_scheme().value) + 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"
\n" item += f"" item += ""
{only_name}
" - for label, key, value, showable, copyable in metadata: - if showable: - value_txt = str(value).replace('\n', '
') - item += f"\n" + for label, key, value in metadata: + value_txt = str(value).replace('\n', '
') + item += f"\n" item += "
{label}{value_txt}
{label}{value_txt}
" - 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"
" item += "