diff --git a/language/en.json b/language/en.json index ec7ef09d..17bffa17 100644 --- a/language/en.json +++ b/language/en.json @@ -18,6 +18,7 @@ "Art/Anime": "Art/Anime", "Describe this Image into Prompt": "Describe this Image into Prompt", "Image Size and Recommended Size": "Image Size and Recommended Size", + "Upscale or Variation": "Upscale or Variation", "Upscale or Variation:": "Upscale or Variation:", "Disabled": "Disabled", "Vary (Subtle)": "Vary (Subtle)", diff --git a/modules/async_worker.py b/modules/async_worker.py index 6ab89eb3..2d051210 100644 --- a/modules/async_worker.py +++ b/modules/async_worker.py @@ -116,6 +116,7 @@ class AsyncTask: self.enhance_input_image = args.pop() self.enhance_checkbox = args.pop() self.enhance_uov_method = args.pop() + self.enhance_uov_processing_order = args.pop() self.enhance_ctrls = [] for _ in range(modules.config.default_enhance_tabs): enhance_enabled = args.pop() @@ -1018,6 +1019,45 @@ def worker(): del task_enhance['c'], task_enhance['uc'] # Save memory return current_progress, imgs[0] + def enhance_upscale(all_steps, async_task, base_progress, callback, controlnet_canny_path, controlnet_cpds_path, + current_task_id, denoising_strength, done_steps_inpainting, done_steps_upscaling, enhance_steps, + final_scheduler_name, height, img, preparation_steps, switch, tiled, total_count, use_expansion, + use_style, use_synthetic_refiner, width): + # reset inpaint worker to prevent tensor size issues and not mix upscale and inpainting + inpaint_worker.current_task = None + + current_task_id += 1 + current_progress = int(base_progress + (100 - preparation_steps) / float(all_steps) * (done_steps_upscaling + done_steps_inpainting)) + goals_enhance = [] + img, skip_prompt_processing, steps = prepare_upscale( + async_task, goals_enhance, img, async_task.enhance_uov_method, async_task.performance_selection, + enhance_steps, current_progress) + steps, _, _, _ = apply_overrides(async_task, steps, height, width) + exception_result = '' + if len(goals_enhance) > 0: + try: + current_progress, img = process_enhance( + all_steps, async_task, callback, controlnet_canny_path, + controlnet_cpds_path, current_progress, current_task_id, denoising_strength, False, + 'None', 0.0, 0.0, async_task.negative_prompt, async_task.prompt, final_scheduler_name, + goals_enhance, height, img, None, preparation_steps, steps, switch, tiled, total_count, + use_expansion, use_style, use_synthetic_refiner, width) + + except ldm_patched.modules.model_management.InterruptProcessingException: + if async_task.last_stop == 'skip': + print('User skipped') + async_task.last_stop = False + # also skip all enhance steps for this image, but add the steps to the progress bar + if async_task.enhance_uov_processing_order == flags.enhancement_uov_before: + done_steps_inpainting += len(async_task.enhance_ctrls) * enhance_steps + exception_result = 'continue' + else: + print('User stopped') + exception_result = 'break' + finally: + done_steps_upscaling += steps + return current_task_id, done_steps_inpainting, done_steps_upscaling, img, exception_result + @torch.no_grad() @torch.inference_mode() def handler(async_task: AsyncTask): @@ -1277,39 +1317,16 @@ def worker(): for img in images_to_enhance: enhancement_image_start_time = time.perf_counter() - # upscale if not disabled or already in goals - if should_process_enhance_uov: - current_task_id += 1 - current_progress = int(base_progress + (100 - preparation_steps) / float(all_steps) * (done_steps_upscaling + done_steps_inpainting)) - goals_enhance = [] - img, skip_prompt_processing, steps = prepare_upscale(async_task, goals_enhance, img, - async_task.enhance_uov_method, - async_task.performance_selection, - enhance_steps, current_progress) - - steps, _, _, _ = apply_overrides(async_task, steps, height, width) - - if len(goals_enhance) > 0: - try: - current_progress, img = process_enhance( - all_steps, async_task, callback, controlnet_canny_path, - controlnet_cpds_path, current_progress, current_task_id, denoising_strength, False, - 'None', 0.0, 0.0, async_task.negative_prompt, async_task.prompt, final_scheduler_name, - goals_enhance, height, img, None, preparation_steps, steps, switch, tiled, total_count, - use_expansion, use_style, use_synthetic_refiner, width) - - except ldm_patched.modules.model_management.InterruptProcessingException: - if async_task.last_stop == 'skip': - print('User skipped') - async_task.last_stop = False - # also skip all enhance steps for this image, but add the steps to the progress bar - done_steps_inpainting += len(async_task.enhance_ctrls) * enhance_steps - continue - else: - print('User stopped') - break - finally: - done_steps_upscaling += steps + if should_process_enhance_uov and async_task.enhance_uov_processing_order == flags.enhancement_uov_before: + current_task_id, done_steps_inpainting, done_steps_upscaling, img, exception_result = enhance_upscale( + all_steps, async_task, base_progress, callback, controlnet_canny_path, controlnet_cpds_path, + current_task_id, denoising_strength, done_steps_inpainting, done_steps_upscaling, enhance_steps, + final_scheduler_name, height, img, preparation_steps, switch, tiled, total_count, use_expansion, + use_style, use_synthetic_refiner, width) + if exception_result == 'continue': + continue + elif exception_result == 'break': + break # inpaint for all other tabs for enhance_mask_dino_prompt_text, enhance_prompt, enhance_negative_prompt, enhance_mask_model, enhance_mask_sam_model, enhance_mask_text_threshold, enhance_mask_box_threshold, enhance_mask_sam_max_detections, enhance_inpaint_disable_initial_latent, enhance_inpaint_engine, enhance_inpaint_strength, enhance_inpaint_respective_field, enhance_inpaint_erode_or_dilate, enhance_mask_invert in async_task.enhance_ctrls: @@ -1379,6 +1396,17 @@ def worker(): enhancement_task_time = time.perf_counter() - enhancement_task_start_time print(f'Enhancement time: {enhancement_task_time:.2f} seconds') + if should_process_enhance_uov and async_task.enhance_uov_processing_order == flags.enhancement_uov_after: + current_task_id, done_steps_inpainting, done_steps_upscaling, img, exception_result = enhance_upscale( + all_steps, async_task, base_progress, callback, controlnet_canny_path, controlnet_cpds_path, + current_task_id, denoising_strength, done_steps_inpainting, done_steps_upscaling, enhance_steps, + final_scheduler_name, height, img, preparation_steps, switch, tiled, total_count, use_expansion, + use_style, use_synthetic_refiner, width) + if exception_result == 'continue': + continue + elif exception_result == 'break': + break + enhancement_image_time = time.perf_counter() - enhancement_image_start_time print(f'Enhancement image time: {enhancement_image_time:.2f} seconds') diff --git a/modules/config.py b/modules/config.py index b0dd7d43..d985f118 100644 --- a/modules/config.py +++ b/modules/config.py @@ -515,6 +515,12 @@ default_enhance_tabs = get_config_item_or_set_default( validator=lambda x: isinstance(x, int) and 1 <= x <= 5, expected_type=int ) +enhance_uov_processing_order = get_config_item_or_set_default( + key='default_enhance_uov_processing_order', + default_value=modules.flags.enhancement_uov_before, + validator=lambda x: x in modules.flags.enhancement_uov_processing_order, + expected_type=int +) default_sam_max_detections = get_config_item_or_set_default( key='default_sam_max_detections', default_value=0, diff --git a/modules/flags.py b/modules/flags.py index a75eed98..4050b21b 100644 --- a/modules/flags.py +++ b/modules/flags.py @@ -12,6 +12,15 @@ uov_list = [ disabled, subtle_variation, strong_variation, upscale_15, upscale_2, upscale_fast ] +enhancement_uov_before = "before" +enhancement_uov_before_label = "Before First Enhancement" +enhancement_uov_after = "after" +enhancement_uov_after_label = "After Last Enhancement" +enhancement_uov_processing_order = [ + (enhancement_uov_before_label, enhancement_uov_before), + (enhancement_uov_after_label, enhancement_uov_after) +] + CIVITAI_NO_KARRAS = ["euler", "euler_ancestral", "heun", "dpm_fast", "dpm_adaptive", "ddim", "uni_pc"] # fooocus: a1111 (Civitai) diff --git a/webui.py b/webui.py index c26666e4..8c7a4fca 100644 --- a/webui.py +++ b/webui.py @@ -353,14 +353,18 @@ with shared.gradio_root: with gr.Row(visible=False) as enhance_input_panel: with gr.Tabs(): - with gr.TabItem(label='#1'): + with gr.TabItem(label='Upscale or Variation'): with gr.Row(): with gr.Column(): enhance_uov_method = gr.Radio(label='Upscale or Variation:', choices=flags.uov_list, value=flags.disabled) - gr.HTML('\U0001F4D4 Document') + enhance_uov_processing_order = gr.Radio(label='Order of Processing', + info='Before is slower (larger area to enhance), but might enhance the overall image quality, whereas after is faster but may apply changes to the image which were already fixed by enhance. Use after when enhancing large areas.', + choices=flags.enhancement_uov_processing_order, + value=modules.config.enhance_uov_processing_order) + # gr.HTML('\U0001F4D4 Document') enhance_ctrls = [] for index in range(modules.config.default_enhance_tabs): - with gr.TabItem(label=f'#{index + 2}') as enhance_tab_item: + 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) @@ -937,7 +941,7 @@ with shared.gradio_root: ctrls += ip_ctrls ctrls += [debugging_dino, dino_erode_or_dilate, debugging_enhance_masks_checkbox, - enhance_input_image, enhance_checkbox, enhance_uov_method] + enhance_input_image, enhance_checkbox, enhance_uov_method, enhance_uov_processing_order] ctrls += enhance_ctrls def parse_meta(raw_prompt_txt, is_generating):