add more developer debug tools (#603)

* add more developer debug tools

* add more developer debug tools

fooocus_extras/preprocessors.py

@@ -1,12 +1,13 @@
 import cv2
 import numpy as np
+import modules.advanced_parameters as advanced_parameters
 
 
 def centered_canny(x: np.ndarray):
     assert isinstance(x, np.ndarray)
     assert x.ndim == 2 and x.dtype == np.uint8
 
-    y = cv2.Canny(x, 64, 128)
+    y = cv2.Canny(x, int(advanced_parameters.canny_low_threshold), int(advanced_parameters.canny_high_threshold))
     y = y.astype(np.float32) / 255.0
     return y
 

fooocus_version.py

@@ -1 +1 @@
-version = '2.1.27'
+version = '2.1.28'

modules/advanced_parameters.py

@@ -2,7 +2,8 @@ adm_scaler_positive, adm_scaler_negative, adm_scaler_end, adaptive_cfg, sampler_
     scheduler_name, overwrite_step, overwrite_switch, overwrite_width, overwrite_height, \
     overwrite_vary_strength, overwrite_upscale_strength, \
     mixing_image_prompt_and_vary_upscale, mixing_image_prompt_and_inpaint, \
-    debugging_cn_preprocessor, controlnet_softness = [None] * 16
+    debugging_cn_preprocessor, controlnet_softness, canny_low_threshold, canny_high_threshold, \
+    freeu_enabled, freeu_b1, freeu_b2, freeu_s1, freeu_s2 = [None] * 23
 
 
 def set_all_advanced_parameters(*args):
@@ -10,12 +11,14 @@ def set_all_advanced_parameters(*args):
         scheduler_name, overwrite_step, overwrite_switch, overwrite_width, overwrite_height, \
         overwrite_vary_strength, overwrite_upscale_strength, \
         mixing_image_prompt_and_vary_upscale, mixing_image_prompt_and_inpaint, \
-        debugging_cn_preprocessor, controlnet_softness
+        debugging_cn_preprocessor, controlnet_softness, canny_low_threshold, canny_high_threshold, \
+        freeu_enabled, freeu_b1, freeu_b2, freeu_s1, freeu_s2
 
     adm_scaler_positive, adm_scaler_negative, adm_scaler_end, adaptive_cfg, sampler_name, \
         scheduler_name, overwrite_step, overwrite_switch, overwrite_width, overwrite_height, \
         overwrite_vary_strength, overwrite_upscale_strength, \
         mixing_image_prompt_and_vary_upscale, mixing_image_prompt_and_inpaint, \
-        debugging_cn_preprocessor, controlnet_softness = args
+        debugging_cn_preprocessor, controlnet_softness, canny_low_threshold, canny_high_threshold, \
+        freeu_enabled, freeu_b1, freeu_b2, freeu_s1, freeu_s2 = args
 
     return

modules/async_worker.py

@@ -449,6 +449,16 @@ def worker():
             if len(cn_tasks[flags.cn_ip]) > 0:
                 pipeline.final_unet = ip_adapter.patch_model(pipeline.final_unet, cn_tasks[flags.cn_ip])
 
+        if advanced_parameters.freeu_enabled:
+            print(f'FreeU is enabled!')
+            pipeline.final_unet = core.apply_freeu(
+                pipeline.final_unet,
+                advanced_parameters.freeu_b1,
+                advanced_parameters.freeu_b2,
+                advanced_parameters.freeu_s1,
+                advanced_parameters.freeu_s2
+            )
+
         results = []
         all_steps = steps * image_number
 

modules/core.py

@@ -20,6 +20,7 @@ import comfy.samplers
 from comfy.sd import load_checkpoint_guess_config
 from nodes import VAEDecode, EmptyLatentImage, VAEEncode, VAEEncodeTiled, VAEDecodeTiled, VAEEncodeForInpaint, \
     ControlNetApplyAdvanced
+from comfy_extras.nodes_freelunch import FreeU
 from comfy.sample import prepare_mask
 from modules.patch import patched_sampler_cfg_function, patched_model_function_wrapper
 from comfy.lora import model_lora_keys_unet, model_lora_keys_clip, load_lora
@@ -32,6 +33,7 @@ opVAEDecodeTiled = VAEDecodeTiled()
 opVAEEncodeTiled = VAEEncodeTiled()
 opVAEEncodeForInpaint = VAEEncodeForInpaint()
 opControlNetApplyAdvanced = ControlNetApplyAdvanced()
+opFreeU = FreeU()
 
 
 class StableDiffusionModel:
@@ -42,6 +44,12 @@ class StableDiffusionModel:
         self.clip_vision = clip_vision
 
 
+@torch.no_grad()
+@torch.inference_mode()
+def apply_freeu(model, b1, b2, s1, s2):
+    return opFreeU.patch(model=model, b1=b1, b2=b2, s1=s1, s2=s2)[0]
+
+
 @torch.no_grad()
 @torch.inference_mode()
 def load_controlnet(ckpt_filename):

webui.py

@@ -208,7 +208,7 @@ with shared.gradio_root:
                 dev_mode = gr.Checkbox(label='Developer Debug Mode', value=False, container=False)
 
                 with gr.Column(visible=False) as dev_tools:
-                    with gr.Tab(label='Developer Control and Debug Tools'):
+                    with gr.Tab(label='Developer Debug Tools'):
                         adm_scaler_positive = gr.Slider(label='Positive ADM Guidance Scaler', minimum=0.1, maximum=3.0,
                                                         step=0.001, value=1.5, info='The scaler multiplied to positive ADM (use 1.0 to disable). ')
                         adm_scaler_negative = gr.Slider(label='Negative ADM Guidance Scaler', minimum=0.1, maximum=3.0,
@@ -241,19 +241,38 @@ with shared.gradio_root:
                                                                minimum=-1, maximum=1.0, step=0.001, value=-1,
                                                                info='Set as negative number to disable. For developer debugging.')
 
-                        mixing_image_prompt_and_vary_upscale = gr.Checkbox(label='Mixing Image Prompt and Vary/Upscale', value=False)
+                    with gr.Tab(label='Control Debug'):
-                        mixing_image_prompt_and_inpaint = gr.Checkbox(label='Mixing Image Prompt and Inpaint', value=False)
-
                         debugging_cn_preprocessor = gr.Checkbox(label='Debug Preprocessor of ControlNets', value=False)
 
+                        mixing_image_prompt_and_vary_upscale = gr.Checkbox(label='Mixing Image Prompt and Vary/Upscale',
+                                                                           value=False)
+                        mixing_image_prompt_and_inpaint = gr.Checkbox(label='Mixing Image Prompt and Inpaint',
+                                                                      value=False)
+
                         controlnet_softness = gr.Slider(label='Softness of ControlNet', minimum=0.0, maximum=1.0,
-                                                        step=0.001, value=0.25, info='Similar to the Control Mode in A1111 (use 0.0 to disable). ')
+                                                        step=0.001, value=0.25,
+                                                        info='Similar to the Control Mode in A1111 (use 0.0 to disable). ')
+
+                        with gr.Tab(label='Canny'):
+                            canny_low_threshold = gr.Slider(label='Canny Low Threshold', minimum=1, maximum=255,
+                                                            step=1, value=64)
+                            canny_high_threshold = gr.Slider(label='Canny High Threshold', minimum=1, maximum=255,
+                                                             step=1, value=128)
+
+                    with gr.Tab(label='FreeU'):
+                        freeu_enabled = gr.Checkbox(label='Enabled', value=False)
+                        freeu_b1 = gr.Slider(label='B1', minimum=0, maximum=2, step=0.01, value=1.01)
+                        freeu_b2 = gr.Slider(label='B2', minimum=0, maximum=2, step=0.01, value=1.02)
+                        freeu_s1 = gr.Slider(label='S1', minimum=0, maximum=4, step=0.01, value=0.99)
+                        freeu_s2 = gr.Slider(label='S2', minimum=0, maximum=4, step=0.01, value=0.95)
+                        freeu_ctrls = [freeu_enabled, freeu_b1, freeu_b2, freeu_s1, freeu_s2]
 
                 adps = [adm_scaler_positive, adm_scaler_negative, adm_scaler_end, adaptive_cfg, sampler_name,
                         scheduler_name, overwrite_step, overwrite_switch, overwrite_width, overwrite_height,
                         overwrite_vary_strength, overwrite_upscale_strength,
                         mixing_image_prompt_and_vary_upscale, mixing_image_prompt_and_inpaint,
-                                debugging_cn_preprocessor, controlnet_softness]
+                        debugging_cn_preprocessor, controlnet_softness, canny_low_threshold, canny_high_threshold]
+                adps += freeu_ctrls
 
                 def dev_mode_checked(r):
                     return gr.update(visible=r)