feat: add option to process uov either before or after enhance

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)",

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 and async_task.enhance_uov_processing_order == flags.enhancement_uov_before:
-            if should_process_enhance_uov:
+                current_task_id, done_steps_inpainting, done_steps_upscaling, img, exception_result = enhance_upscale(
-                current_task_id += 1
+                    all_steps, async_task, base_progress, callback, controlnet_canny_path, controlnet_cpds_path,
-                current_progress = int(base_progress + (100 - preparation_steps) / float(all_steps) * (done_steps_upscaling + done_steps_inpainting))
+                    current_task_id, denoising_strength, done_steps_inpainting, done_steps_upscaling, enhance_steps,
-                goals_enhance = []
+                    final_scheduler_name, height, img, preparation_steps, switch, tiled, total_count, use_expansion,
-                img, skip_prompt_processing, steps = prepare_upscale(async_task, goals_enhance, img,
+                    use_style, use_synthetic_refiner, width)
-                                                                     async_task.enhance_uov_method,
+                if exception_result == 'continue':
-                                                                     async_task.performance_selection,
+                    continue
-                                                                     enhance_steps, current_progress)
+                elif exception_result == 'break':
-
+                    break
-                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
 
             # 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')
 

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,

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)

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('<a href="https://github.com/lllyasviel/Fooocus/discussions/390" target="_blank">\U0001F4D4 Document</a>')
+                                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('<a href="https://github.com/lllyasviel/Fooocus/discussions/390" target="_blank">\U0001F4D4 Document</a>')
                     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):