From 191f8148e431f5018fd026a165afb9c175e7653d Mon Sep 17 00:00:00 2001 From: Manuel Schmid Date: Mon, 15 Jan 2024 20:47:26 +0100 Subject: [PATCH] feat: add support for A1111 metadata schema https://github.com/AUTOMATIC1111/stable-diffusion-webui/blob/cf2772fab0af5573da775e7437e6acdca424f26e/modules/processing.py#L672 --- modules/async_worker.py | 127 +++++++++++++++++++++++----------------- modules/config.py | 5 ++ modules/flags.py | 5 ++ webui.py | 7 ++- 4 files changed, 90 insertions(+), 54 deletions(-) diff --git a/modules/async_worker.py b/modules/async_worker.py index 3e0c1bac..98eaad4b 100644 --- a/modules/async_worker.py +++ b/modules/async_worker.py @@ -143,6 +143,7 @@ def worker(): inpaint_additional_prompt = args.pop() inpaint_mask_image_upload = args.pop() save_metadata_to_images = args.pop() if not args_manager.args.disable_metadata else False + metadata_schema = args.pop() if not args_manager.args.disable_metadata else 'fooocus' cn_tasks = {x: [] for x in flags.ip_list} for _ in range(4): @@ -780,69 +781,89 @@ 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 + metadata_string = '' + if save_metadata_to_images and metadata_schema == 'fooocus': + 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, + 'denoising_strength': denoising_strength, + 'freeu': advanced_parameters.freeu_enabled, + 'img2img': input_image_checkbox, + 'prompt_expansion': task['expansion'] } - 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: + if advanced_parameters.freeu_enabled: 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 + '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 'cn' in goals: - metadata |= { - 'canny_low_threshold': advanced_parameters.canny_low_threshold, 'canny_high_threshold': advanced_parameters.canny_high_threshold, - } + if 'vary' in goals: + metadata |= { + 'uov_method': uov_method + #'uov_input_image': raw_uov_input_image + } - 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 + 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 |= { - # 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 - } + '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 } - cn_task_index += 1 - metadata |= {'software': f'Fooocus v{fooocus_version.version}'} - metadata_string = json.dumps(metadata, ensure_ascii=False) + 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 |= { + 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) + elif save_metadata_to_images and metadata_schema == 'a1111': + generation_params = { + "Steps": steps, + "Sampler": sampler_name, + "CFG scale": cfg_scale, + "Seed": task['task_seed'], + "Size": f"{width}x{height}", + #"Model hash": p.sd_model_hash if opts.add_model_hash_to_info else None, + "Model": base_model_name, + "Denoising strength": denoising_strength, + "Version": f'Fooocus v{fooocus_version.version}', + "User": 'mashb1t', + } + + generation_params_text = ", ".join([k if k == v else f'{k}: {v}' for k, v in generation_params.items() if v is not None]) + negative_prompt_text = f"\nNegative prompt: {raw_negative_prompt}" if raw_negative_prompt else "" + metadata_string = f"{raw_prompt}{raw_negative_prompt}\n{generation_params_text}".strip() for x in imgs: d = [ diff --git a/modules/config.py b/modules/config.py index 7d760edf..28104859 100644 --- a/modules/config.py +++ b/modules/config.py @@ -320,6 +320,11 @@ default_save_metadata_to_images = get_config_item_or_set_default( default_value=False, validator=lambda x: isinstance(x, bool) ) +default_metadata_schema = get_config_item_or_set_default( + key='default_metadata_schema', + default_value='fooocus', + validator=lambda x: x in [y[1] for y in modules.flags.metadata_schema if y[1] == x] +) example_inpaint_prompts = [[x] for x in example_inpaint_prompts] diff --git a/modules/flags.py b/modules/flags.py index 27f2d716..873736df 100644 --- a/modules/flags.py +++ b/modules/flags.py @@ -32,6 +32,11 @@ default_parameters = { cn_ip: (0.5, 0.6), cn_ip_face: (0.9, 0.75), cn_canny: (0.5, 1.0), cn_cpds: (0.5, 1.0) } # stop, weight +metadata_schema =[ + ('Fooocus (json)', 'fooocus'), + ('A1111 (plain text)', 'a1111'), +] + inpaint_engine_versions = ['None', 'v1', 'v2.5', 'v2.6'] performance_selections = ['Speed', 'Quality', 'Extreme Speed'] diff --git a/webui.py b/webui.py index e3814ff9..1a88a87b 100644 --- a/webui.py +++ b/webui.py @@ -385,6 +385,11 @@ with shared.gradio_root: if not args_manager.args.disable_metadata: save_metadata_to_images = gr.Checkbox(label='Save Metadata to Images', value=modules.config.default_save_metadata_to_images, info='Adds parameters to generated images allowing manual regeneration.') + metadata_schema = gr.Radio(label='Metadata Schema', choices=flags.metadata_schema, value=modules.config.default_metadata_schema, + info='Use A1111 for compatibility with Civitai.') + + save_metadata_to_images.change(lambda x: gr.update(visible=not x), inputs=[save_metadata_to_images], outputs=[metadata_schema], + queue=False, show_progress=False) with gr.Tab(label='Control'): debugging_cn_preprocessor = gr.Checkbox(label='Debug Preprocessors', value=False, @@ -534,7 +539,7 @@ with shared.gradio_root: ctrls += [outpaint_selections, inpaint_input_image, inpaint_additional_prompt, inpaint_mask_image] if not args_manager.args.disable_metadata: - ctrls += [save_metadata_to_images] + ctrls += [save_metadata_to_images, metadata_schema] ctrls += ip_ctrls