wip: rename stage2 to enhance

This commit is contained in:
Manuel Schmid 2024-06-16 12:43:23 +02:00
parent a3dcd6e1fd
commit ef9fd293ff
No known key found for this signature in database
GPG Key ID: 32C4F7569B40B84B
3 changed files with 145 additions and 144 deletions

View File

@ -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]

View File

@ -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

187
webui.py
View File

@ -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