add more developer debug tools (#603)

* add more developer debug tools

* add more developer debug tools

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):