From ef9fd293ff33f83338153a1a93c64014d067e3ce Mon Sep 17 00:00:00 2001 From: Manuel Schmid Date: Sun, 16 Jun 2024 12:43:23 +0200 Subject: [PATCH] wip: rename stage2 to enhance --- modules/async_worker.py | 88 +++++++++---------- modules/config.py | 14 +-- webui.py | 187 ++++++++++++++++++++-------------------- 3 files changed, 145 insertions(+), 144 deletions(-) diff --git a/modules/async_worker.py b/modules/async_worker.py index 04a2705a..8d6004f8 100644 --- a/modules/async_worker.py +++ b/modules/async_worker.py @@ -111,23 +111,23 @@ class AsyncTask: self.debugging_dino = args.pop() self.dino_erode_or_dilate = args.pop() - self.stage2_ctrls = [] - for _ in range(modules.config.default_stage2_tabs): - stage2_enabled = args.pop() - # stage2_mode = args.pop() - stage2_mask_dino_prompt_text = args.pop() - stage2_mask_box_threshold = args.pop() - stage2_mask_text_threshold = args.pop() - stage2_mask_sam_max_num_boxes = args.pop() - stage2_mask_sam_model = args.pop() - if stage2_enabled: - self.stage2_ctrls.append([ - # stage2_mode, - stage2_mask_dino_prompt_text, - stage2_mask_box_threshold, - stage2_mask_text_threshold, - stage2_mask_sam_max_num_boxes, - stage2_mask_sam_model, + self.enhance_ctrls = [] + for _ in range(modules.config.default_enhance_tabs): + enhance_enabled = args.pop() + # enhance_mode = args.pop() + enhance_mask_dino_prompt_text = args.pop() + enhance_mask_box_threshold = args.pop() + enhance_mask_text_threshold = args.pop() + enhance_mask_sam_max_num_boxes = args.pop() + enhance_mask_sam_model = args.pop() + if enhance_enabled: + self.enhance_ctrls.append([ + # enhance_mode, + enhance_mask_dino_prompt_text, + enhance_mask_box_threshold, + enhance_mask_text_threshold, + enhance_mask_sam_max_num_boxes, + enhance_mask_sam_model, ]) @@ -1038,24 +1038,24 @@ def worker(): current_task_id, denoising_strength, final_scheduler_name, goals, initial_latent, switch, task, tasks, tiled, use_expansion, width, height) - # stage2 - progressbar(async_task, current_progress, 'Processing stage2 ...') + # enhance + progressbar(async_task, current_progress, 'Processing enhance ...') final_unet = pipeline.final_unet - if len(async_task.stage2_ctrls) == 0 or 'inpaint' in goals: - print(f'[Stage2] Skipping, preconditions aren\'t met') + if len(async_task.enhance_ctrls) == 0 or 'inpaint' in goals: + print(f'[Enhance] Skipping, preconditions aren\'t met') continue for img in imgs: - for stage2_mask_dino_prompt_text, stage2_mask_box_threshold, stage2_mask_text_threshold, stage2_mask_sam_max_num_boxes, stage2_mask_sam_model in async_task.stage2_ctrls: - print(f'[Stage2] Searching for "{stage2_mask_dino_prompt_text}"') + for enhance_mask_dino_prompt_text, enhance_mask_box_threshold, enhance_mask_text_threshold, enhance_mask_sam_max_num_boxes, enhance_mask_sam_model in async_task.enhance_ctrls: + print(f'[Enhance] Searching for "{enhance_mask_dino_prompt_text}"') mask, dino_detection_count, sam_detection_count, sam_detection_on_mask_count = generate_mask_from_image(img, sam_options=SAMOptions( - dino_prompt=stage2_mask_dino_prompt_text, - dino_box_threshold=stage2_mask_box_threshold, - dino_text_threshold=stage2_mask_text_threshold, + dino_prompt=enhance_mask_dino_prompt_text, + dino_box_threshold=enhance_mask_box_threshold, + dino_text_threshold=enhance_mask_text_threshold, dino_erode_or_dilate=async_task.dino_erode_or_dilate, dino_debug=async_task.debugging_dino, - max_num_boxes=stage2_mask_sam_max_num_boxes, - model_type=stage2_mask_sam_model + max_num_boxes=enhance_mask_sam_max_num_boxes, + model_type=enhance_mask_sam_model )) mask = mask[:, :, 0] @@ -1065,12 +1065,12 @@ def worker(): do_not_show_finished_images=len( tasks) == 1 or async_task.disable_intermediate_results) - print(f'[Stage2] {dino_detection_count} boxes detected') - print(f'[Stage2] {sam_detection_count} segments detected in boxes') - print(f'[Stage2] {sam_detection_on_mask_count} segments applied to mask') + print(f'[Enhance] {dino_detection_count} boxes detected') + print(f'[Enhance] {sam_detection_count} segments detected in boxes') + print(f'[Enhance] {sam_detection_on_mask_count} segments applied to mask') if dino_detection_count == 0 or not async_task.debugging_dino and sam_detection_on_mask_count == 0: - print(f'[Stage2] No "{stage2_mask_dino_prompt_text}" detected, skipping') + print(f'[Enhance] No "{enhance_mask_dino_prompt_text}" detected, skipping') continue # TODO make configurable @@ -1094,21 +1094,21 @@ def worker(): # patch_samplers(async_task) # defaults from inpaint mode improve details - denoising_strength_stage2 = 0.5 - inpaint_respective_field_stage2 = 0.0 - inpaint_head_model_path_stage2 = None - inpaint_parameterized_stage2 = False # inpaint_engine = None, improve detail + denoising_strength_enhance = 0.5 + inpaint_respective_field_enhance = 0.0 + inpaint_head_model_path_enhance = None + inpaint_parameterized_enhance = False # inpaint_engine = None, improve detail - goals_stage2 = ['inpaint'] - denoising_strength_stage2, initial_latent_stage2, width_stage2, height_stage2 = apply_inpaint( - async_task, None, inpaint_head_model_path_stage2, img, mask, - inpaint_parameterized_stage2, denoising_strength_stage2, - inpaint_respective_field_stage2, switch, current_progress, True) + goals_enhance = ['inpaint'] + denoising_strength_enhance, initial_latent_enhance, width_enhance, height_enhance = apply_inpaint( + async_task, None, inpaint_head_model_path_enhance, img, mask, + inpaint_parameterized_enhance, denoising_strength_enhance, + inpaint_respective_field_enhance, switch, current_progress, True) imgs2, img_paths, current_progress = process_task(all_steps, async_task, callback, controlnet_canny_path, controlnet_cpds_path, - current_task_id, denoising_strength_stage2, final_scheduler_name, goals_stage2, - initial_latent_stage2, switch, task, tasks, tiled, use_expansion, width_stage2, - height_stage2) + current_task_id, denoising_strength_enhance, final_scheduler_name, goals_enhance, + initial_latent_enhance, switch, task, tasks, tiled, use_expansion, width_enhance, + height_enhance) # reset and prepare next iteration img = imgs2[0] diff --git a/modules/config.py b/modules/config.py index 3c5389b9..929dd9ce 100644 --- a/modules/config.py +++ b/modules/config.py @@ -502,16 +502,16 @@ example_inpaint_prompts = get_config_item_or_set_default( validator=lambda x: isinstance(x, list) and all(isinstance(v, str) for v in x), expected_type=list ) -example_stage2_prompts = get_config_item_or_set_default( - key='example_stage2_prompts', +example_enhance_prompts = get_config_item_or_set_default( + key='example_enhance_prompts', default_value=[ 'face', 'eye', 'mouth', 'hair', 'hand', 'body' ], validator=lambda x: isinstance(x, list) and all(isinstance(v, str) for v in x), expected_type=list ) -default_stage2_tabs = get_config_item_or_set_default( - key='default_stage2_tabs', +default_enhance_tabs = get_config_item_or_set_default( + key='default_enhance_tabs', default_value=3, validator=lambda x: isinstance(x, int) and 1 <= x <= 5, expected_type=int @@ -548,7 +548,7 @@ metadata_created_by = get_config_item_or_set_default( ) example_inpaint_prompts = [[x] for x in example_inpaint_prompts] -example_stage2_prompts = [[x] for x in example_stage2_prompts] +example_enhance_prompts = [[x] for x in example_enhance_prompts] default_inpaint_mask_model = get_config_item_or_set_default( key='default_inpaint_mask_model', @@ -557,8 +557,8 @@ default_inpaint_mask_model = get_config_item_or_set_default( expected_type=str ) -default_stage2_inpaint_mask_model = get_config_item_or_set_default( - key='default_stage2_inpaint_mask_model', +default_enhance_inpaint_mask_model = get_config_item_or_set_default( + key='default_enhance_inpaint_mask_model', default_value='sam', validator=lambda x: x in modules.flags.inpaint_mask_models, expected_type=str diff --git a/webui.py b/webui.py index 8cf5a087..4043eddf 100644 --- a/webui.py +++ b/webui.py @@ -175,7 +175,7 @@ with shared.gradio_root: skip_button.click(skip_clicked, inputs=currentTask, outputs=currentTask, queue=False, show_progress=False) with gr.Row(elem_classes='advanced_check_row'): input_image_checkbox = gr.Checkbox(label='Input Image', value=False, container=False, elem_classes='min_check') - stage2_checkbox = gr.Checkbox(label='Stage2', value=False, container=False, elem_classes='min_check') + enhance_checkbox = gr.Checkbox(label='Enhance', value=False, container=False, elem_classes='min_check') advanced_checkbox = gr.Checkbox(label='Advanced', value=modules.config.default_advanced_checkbox, container=False, elem_classes='min_check') with gr.Row(visible=False) as image_input_panel: with gr.Tabs(): @@ -329,118 +329,119 @@ with shared.gradio_root: metadata_input_image.upload(trigger_metadata_preview, inputs=metadata_input_image, outputs=metadata_json, queue=False, show_progress=True) - with gr.Row(visible=False) as stage2_input_panel: + with gr.Row(visible=False) as enhance_input_panel: with gr.Tabs(): - stage2_ctrls = [] - for index in range(modules.config.default_stage2_tabs): - with gr.TabItem(label=f'#{index + 1}') as stage2_tab_item: - stage2_enabled = gr.Checkbox(label='Enable', value=False, elem_classes='min_check', - container=False) - gr.HTML('DISCLAIMER: Stage2 does not work with Inpaint or Outpaint and will be skipped.') + enhance_ctrls = [] + for index in range(modules.config.default_enhance_tabs): + with gr.TabItem(label=f'#{index + 1}') as enhance_tab_item: + enhance_enabled = gr.Checkbox(label='Enable', value=False, elem_classes='min_check', + container=False) + gr.HTML( + 'DISCLAIMER: The enhance feature does not work with Inpaint or Outpaint and will be skipped.') - stage2_mask_dino_prompt_text = gr.Textbox(label='Detection prompt', - info='Use singular whenever possible', - interactive=True, - visible=modules.config.default_stage2_inpaint_mask_model == 'sam') - example_stage2_mask_dino_prompt_text = gr.Dataset( - samples=modules.config.example_stage2_prompts, + enhance_mask_dino_prompt_text = gr.Textbox(label='Detection prompt', + info='Use singular whenever possible', + interactive=True, + visible=modules.config.default_enhance_inpaint_mask_model == 'sam') + example_enhance_mask_dino_prompt_text = gr.Dataset( + samples=modules.config.example_enhance_prompts, label='Additional Prompt Quick List', - components=[stage2_mask_dino_prompt_text], - visible=modules.config.default_stage2_inpaint_mask_model == 'sam') - example_stage2_mask_dino_prompt_text.click(lambda x: x[0], - inputs=example_stage2_mask_dino_prompt_text, - outputs=stage2_mask_dino_prompt_text, - show_progress=False, queue=False) + components=[enhance_mask_dino_prompt_text], + visible=modules.config.default_enhance_inpaint_mask_model == 'sam') + example_enhance_mask_dino_prompt_text.click(lambda x: x[0], + inputs=example_enhance_mask_dino_prompt_text, + outputs=enhance_mask_dino_prompt_text, + show_progress=False, queue=False) - stage2_prompt = gr.Textbox(label="Enhancement positive prompt", - placeholder="Uses original prompt instead if empty.") - stage2_negative_prompt = gr.Textbox(label="Enhancement negative prompt", - placeholder="Uses original negative prompt instead if empty.") + enhance_prompt = gr.Textbox(label="Enhancement positive prompt", + placeholder="Uses original prompt instead if empty.") + enhance_negative_prompt = gr.Textbox(label="Enhancement negative prompt", + placeholder="Uses original negative prompt instead if empty.") with gr.Accordion("Detection", open=False): # TODO check if limiting to SAM is better - stage2_mask_inpaint_mask_model = gr.Dropdown(label='Mask generation model', - choices=flags.inpaint_mask_models, - value=modules.config.default_stage2_inpaint_mask_model) - stage2_mask_inpaint_mask_cloth_category = gr.Dropdown(label='Cloth category', - choices=flags.inpaint_mask_cloth_category, - value=modules.config.default_inpaint_mask_cloth_category, - visible=modules.config.default_stage2_inpaint_mask_model == 'u2net_cloth_seg', - interactive=True) + enhance_mask_model = gr.Dropdown(label='Mask generation model', + choices=flags.inpaint_mask_models, + value=modules.config.default_enhance_inpaint_mask_model) + enhance_mask_cloth_category = gr.Dropdown(label='Cloth category', + choices=flags.inpaint_mask_cloth_category, + value=modules.config.default_inpaint_mask_cloth_category, + visible=modules.config.default_enhance_inpaint_mask_model == 'u2net_cloth_seg', + interactive=True) with gr.Accordion("SAM Options", - visible=modules.config.default_stage2_inpaint_mask_model == 'sam', + visible=modules.config.default_enhance_inpaint_mask_model == 'sam', open=False) as sam_options: - stage2_mask_sam_model = gr.Dropdown(label='SAM model', - choices=flags.inpaint_mask_sam_model, - value=modules.config.default_inpaint_mask_sam_model, - interactive=True) - stage2_mask_box_threshold = gr.Slider(label="Box Threshold", minimum=0.0, - maximum=1.0, value=0.3, step=0.05, - interactive=True) - stage2_mask_text_threshold = gr.Slider(label="Text Threshold", minimum=0.0, - maximum=1.0, value=0.25, step=0.05, + enhance_mask_sam_model = gr.Dropdown(label='SAM model', + choices=flags.inpaint_mask_sam_model, + value=modules.config.default_inpaint_mask_sam_model, + interactive=True) + enhance_mask_box_threshold = gr.Slider(label="Box Threshold", minimum=0.0, + maximum=1.0, value=0.3, step=0.05, interactive=True) - stage2_mask_sam_max_num_boxes = gr.Slider(label="Maximum number of box detections", - minimum=1, maximum=5, - value=modules.config.default_sam_max_num_boxes, - step=1, interactive=True) + enhance_mask_text_threshold = gr.Slider(label="Text Threshold", minimum=0.0, + maximum=1.0, value=0.25, step=0.05, + interactive=True) + enhance_mask_sam_max_num_boxes = gr.Slider(label="Maximum number of box detections", + minimum=1, maximum=5, + value=modules.config.default_sam_max_num_boxes, + step=1, interactive=True) with gr.Accordion("Inpaint", visible=True, open=False): - stage2_mask_inpaint_mode = gr.Dropdown(choices=modules.flags.inpaint_options, - value=modules.flags.inpaint_option_default, - # TODO test - label='Method', interactive=True) - stage2_mask_inpaint_disable_initial_latent = gr.Checkbox( + enhance_inpaint_mode = gr.Dropdown(choices=modules.flags.inpaint_options, + value=modules.flags.inpaint_option_default, + # TODO test + label='Method', interactive=True) + enhance_inpaint_disable_initial_latent = gr.Checkbox( label='Disable initial latent in inpaint', value=False) - stage2_mask_inpaint_engine = gr.Dropdown(label='Inpaint Engine', - value=modules.config.default_inpaint_engine_version, - choices=flags.inpaint_engine_versions, - info='Version of Fooocus inpaint model') - stage2_mask_inpaint_strength = gr.Slider(label='Inpaint Denoising Strength', - minimum=0.0, maximum=1.0, step=0.001, - value=1.0, - info='Same as the denoising strength in A1111 inpaint. ' - 'Only used in inpaint, not used in outpaint. ' - '(Outpaint always use 1.0)') - stage2_mask_inpaint_respective_field = gr.Slider(label='Inpaint Respective Field', - minimum=0.0, maximum=1.0, step=0.001, - value=0.618, - info='The area to inpaint. ' - 'Value 0 is same as "Only Masked" in A1111. ' - 'Value 1 is same as "Whole Image" in A1111. ' - 'Only used in inpaint, not used in outpaint. ' - '(Outpaint always use 1.0)') + enhance_inpaint_engine = gr.Dropdown(label='Inpaint Engine', + value=modules.config.default_inpaint_engine_version, + choices=flags.inpaint_engine_versions, + info='Version of Fooocus inpaint model') + enhance_inpaint_strength = gr.Slider(label='Inpaint Denoising Strength', + minimum=0.0, maximum=1.0, step=0.001, + value=1.0, + info='Same as the denoising strength in A1111 inpaint. ' + 'Only used in inpaint, not used in outpaint. ' + '(Outpaint always use 1.0)') + enhance_inpaint_respective_field = gr.Slider(label='Inpaint Respective Field', + minimum=0.0, maximum=1.0, step=0.001, + value=0.618, + info='The area to inpaint. ' + 'Value 0 is same as "Only Masked" in A1111. ' + 'Value 1 is same as "Whole Image" in A1111. ' + 'Only used in inpaint, not used in outpaint. ' + '(Outpaint always use 1.0)') - stage2_ctrls += [ - stage2_enabled, - stage2_mask_dino_prompt_text, - stage2_prompt, - stage2_negative_prompt, - stage2_mask_sam_model, - stage2_mask_text_threshold, - stage2_mask_box_threshold, - stage2_mask_sam_max_num_boxes, - stage2_mask_inpaint_disable_initial_latent, - stage2_mask_inpaint_engine, - stage2_mask_inpaint_strength, - stage2_mask_inpaint_respective_field + enhance_ctrls += [ + enhance_enabled, + enhance_mask_dino_prompt_text, + enhance_prompt, + enhance_negative_prompt, + enhance_mask_sam_model, + enhance_mask_text_threshold, + enhance_mask_box_threshold, + enhance_mask_sam_max_num_boxes, + enhance_inpaint_disable_initial_latent, + enhance_inpaint_engine, + enhance_inpaint_strength, + enhance_inpaint_respective_field ] - stage2_mask_inpaint_mode.input(inpaint_mode_change, inputs=stage2_mask_inpaint_mode, outputs=[ + enhance_inpaint_mode.input(inpaint_mode_change, inputs=enhance_inpaint_mode, outputs=[ inpaint_additional_prompt, outpaint_selections, example_inpaint_prompts, - stage2_mask_inpaint_disable_initial_latent, stage2_mask_inpaint_engine, - stage2_mask_inpaint_strength, stage2_mask_inpaint_respective_field + enhance_inpaint_disable_initial_latent, enhance_inpaint_engine, + enhance_inpaint_strength, enhance_inpaint_respective_field ], show_progress=False, queue=False) - stage2_mask_inpaint_mask_model.change( + enhance_mask_model.change( lambda x: [gr.update(visible=x == 'u2net_cloth_seg')] + [gr.update(visible=x == 'sam')] * 2 + [gr.Dataset.update(visible=x == 'sam', - samples=modules.config.example_stage2_prompts)], - inputs=stage2_mask_inpaint_mask_model, - outputs=[stage2_mask_inpaint_mask_cloth_category, stage2_mask_dino_prompt_text, sam_options, - example_stage2_mask_dino_prompt_text], + samples=modules.config.example_enhance_prompts)], + inputs=enhance_mask_model, + outputs=[enhance_mask_cloth_category, enhance_mask_dino_prompt_text, sam_options, + example_enhance_mask_dino_prompt_text], queue=False, show_progress=False) switch_js = "(x) => {if(x){viewer_to_bottom(100);viewer_to_bottom(500);}else{viewer_to_top();} return x;}" @@ -457,8 +458,8 @@ with shared.gradio_root: desc_tab.select(lambda: 'desc', outputs=current_tab, queue=False, _js=down_js, show_progress=False) metadata_tab.select(lambda: 'metadata', outputs=current_tab, queue=False, _js=down_js, show_progress=False) - stage2_checkbox.change(lambda x: gr.update(visible=x), inputs=stage2_checkbox, - outputs=stage2_input_panel, queue=False, show_progress=False, _js=switch_js) + enhance_checkbox.change(lambda x: gr.update(visible=x), inputs=enhance_checkbox, + outputs=enhance_input_panel, queue=False, show_progress=False, _js=switch_js) with gr.Column(scale=1, visible=modules.config.default_advanced_checkbox) as advanced_column: with gr.Tab(label='Settings'): @@ -894,7 +895,7 @@ with shared.gradio_root: ctrls += [save_metadata_to_images, metadata_scheme] ctrls += ip_ctrls - ctrls += [debugging_dino, dino_erode_or_dilate] + stage2_ctrls + ctrls += [debugging_dino, dino_erode_or_dilate] + enhance_ctrls def parse_meta(raw_prompt_txt, is_generating): loaded_json = None