add free-u

CHANGELOG.md

@@ -4,9 +4,7 @@
 
 **TBD**: Candidates before release:  
 - Implement Lyco for *backend:diffusers*
-- Add FreeU for *backend:diffusers*  
-  for *backend:original* use extension: <https://github.com/ljleb/sd-webui-freeu>  
-- Add HyperTile: <https://github.com/tfernd/HyperTile>
+- Note: Free-U requires diffusers main branch, will be included in next release
 - Implement styles extra field
 
 This is a big one, with some major changes and new functionality...  
@@ -144,6 +142,11 @@ Upgrades are still possible and supported, but above is recommended for best exp
   - **GC**:  
     - custom garbage collect threshold to reduce vram memory usage, thanks @Disty0  
       see *settings -> compute -> gc*  
+- **Inference**  
+  - new section in **settings**  
+    - [Token Merging](https://github.com/dbolya/tomesd): not new, but updated  
+    - [Free-U](https://github.com/ChenyangSi/FreeU): new!  
+    - [HyperTile](https://github.com/tfernd/HyperTile): new!  
 - **General**  
   - **Startup**  
     - all main CLI parameters can now be set as environment variable as well  

html/locale_en.json

@@ -395,7 +395,7 @@
   {"id":"","label":"Enable upcast sampling","localized":"","hint":"Usually produces similar results to --no-half with better performance while using less memory"},
   {"id":"","label":"Enable upcast cross attention layer","localized":"","hint":""},
   {"id":"","label":"Disable NaN check in produced images/latent spaces","localized":"","hint":""},
-  {"id":"","label":"Attempt VAE roll back when produced NaN values (experimental)","localized":"","hint":"Requires Torch 2.1 and NaN check enabled"},
+  {"id":"","label":"Attempt VAE roll back when produced NaN values","localized":"","hint":"Requires Torch 2.1 and NaN check enabled"},
   {"id":"","label":"Use channels last as torch memory format","localized":"","hint":""},
   {"id":"","label":"Enable full-depth cuDNN benchmark feature","localized":"","hint":""},
   {"id":"","label":"Enable model compile","localized":"","hint":""},

javascript/style.css

@@ -41,6 +41,7 @@ tr { border-bottom: none !important; padding: 0.1em 0.5em !important; }
 .gradio-slider { max-width: 50%; margin-right: var(--spacing-sm) !important; }
 .gradio-slider input[type="number"] { width: 6em; margin-left: 0.5em; }
 .gradio-textbox { overflow: visible !important; }
+.gradio-radio { padding: 0 !important; }
 
 /* custom gradio elements */
 .accordion-compact { padding: 8px 0px 4px 0px !important; }
@@ -102,7 +103,8 @@ div#extras_scale_to_tab div.form{ flex-direction: row; }
 #quicksettings { width: fit-content; margin-top: 1em; }
 #quicksettings > button { padding: 0 1em 0 0; align-self: end; margin-bottom: var(--text-lg); }
 #settings { display: flex; gap: var(--layout-gap); }
-#settings div { border: none; gap: 0.5em; }
+#settings div { border: none; gap: 0; margin: 0 0 var(--layout-gap) 0px; padding: 0; }
+#settings .gr-group { max-width: 70em; }
 #settings > div.tab-content { flex: 10 0 75%; display: grid; }
 #settings > div.tab-content > div { border: none; padding: 0; }
 #settings > div.tab-content > div > div > div > div > div { flex-direction: unset; }

modules/processing.py

@@ -28,6 +28,7 @@ import modules.face_restoration
 import modules.images as images
 import modules.styles
 import modules.sd_hijack
+import modules.sd_hijack_freeu
 import modules.sd_samplers
 import modules.sd_samplers_common
 import modules.sd_models
@@ -639,6 +640,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
 
         if not shared.opts.cuda_compile:
             modules.sd_models.apply_token_merging(p.sd_model, p.get_token_merging_ratio())
+            modules.sd_hijack_freeu.apply_freeu(p.sd_model, shared.backend == shared.Backend.ORIGINAL)
 
         if shared.cmd_opts.profile:
             """

modules/sd_hijack_freeu.py

@@ -0,0 +1,151 @@
+import math
+import functools
+import torch
+from modules.shared import state, log, opts
+from modules.sd_hijack_unet import th
+
+# based on <https://github.com/ljleb/sd-webui-freeu/blob/main/lib_free_u/unet.py>
+# official params are b1,b2,s1,s2
+
+# extra params that can be made configurable if needed are:
+backbone_width = 0.5
+backbone_offset = 0.0
+skip_cutoff = 0.0
+skip_high_end_factor = 1.0
+start_ratio = 0.0
+stop_ratio = 1.0
+transition_smoothness = 0.0
+
+# internal state
+state_enabled = False
+cat_original = None
+
+
+def to_denoising_step(number, steps=None) -> int:
+    if steps is None:
+        steps = state.sampling_steps
+    if isinstance(number, float):
+        return int(number * steps)
+    return number
+
+
+def get_schedule_ratio():
+    start_step = to_denoising_step(start_ratio)
+    stop_step = to_denoising_step(stop_ratio)
+    if start_step == stop_step:
+        smooth_schedule_ratio = 0.0
+    elif state.sampling_step < start_step:
+        smooth_schedule_ratio = min(1.0, max(0.0, state.sampling_step / start_step))
+    else:
+        smooth_schedule_ratio = min(1.0, max(0.0, 1 + (state.sampling_step - start_step) / (start_step - stop_step)))
+    flat_schedule_ratio = 1.0 if start_step <= state.sampling_step < stop_step else 0.0
+    return lerp(flat_schedule_ratio, smooth_schedule_ratio, transition_smoothness)
+
+
+def lerp(a, b, r):
+    return (1-r)*a + r*b
+
+
+def free_u_cat_hijack(hs, *args, original_function, **kwargs):
+    print('HERE', state.sampling_step, state.sampling_steps)
+    if not opts.freeu_enabled:
+        return original_function(hs, *args, **kwargs)
+    schedule_ratio = get_schedule_ratio()
+    if schedule_ratio == 0:
+        return original_function(hs, *args, **kwargs)
+    try:
+        h, h_skip = hs
+        if list(kwargs.keys()) != ["dim"] or kwargs.get("dim", -1) != 1:
+            return original_function(hs, *args, **kwargs)
+    except ValueError:
+        return original_function(hs, *args, **kwargs)
+    dims = h.shape[1]
+    index = [1280, 640, 320].index(dims)
+    if index > 1: # not 1st or 2nd stage
+        return original_function([h, h_skip], *args, **kwargs)
+    region_begin, region_end, region_inverted = ratio_to_region(backbone_width, backbone_offset, dims)
+    mask = torch.arange(dims)
+    mask = (region_begin <= mask) & (mask <= region_end)
+    if region_inverted:
+        mask = ~mask
+    backbone_factor = opts.freeu_b1 if index == 0 else opts.freeu_b2
+    skip_factor = opts.freeu_s1 if index == 0 else opts.freeu_s2
+    h[:, mask] *= lerp(1, backbone_factor, schedule_ratio)
+    h_skip = filter_skip(h_skip, threshold=skip_cutoff, scale=lerp(1, skip_factor, schedule_ratio), scale_high=lerp(1, skip_high_end_factor, schedule_ratio))
+    return original_function([h, h_skip], *args, **kwargs)
+
+
+def no_gpu_complex_support():
+    mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
+    try:
+        import torch_directml
+    except ImportError:
+        dml_available = False
+    else:
+        dml_available = torch_directml.is_available()
+    return mps_available or dml_available
+
+
+def filter_skip(x, threshold, scale, scale_high):
+    if scale == 1 and scale_high == 1:
+        return x
+    fft_device = x.device
+    if no_gpu_complex_support():
+        fft_device = "cpu"
+    # FFT
+    x_freq = torch.fft.fftn(x.to(fft_device).float(), dim=(-2, -1)) # pylint: disable=E1102
+    x_freq = torch.fft.fftshift(x_freq, dim=(-2, -1)) # pylint: disable=E1102
+    B, C, H, W = x_freq.shape
+    mask = torch.full((B, C, H, W), float(scale_high), device=fft_device)
+    crow, ccol = H // 2, W // 2
+    threshold_row = max(1, math.floor(crow * threshold))
+    threshold_col = max(1, math.floor(ccol * threshold))
+    mask[..., crow - threshold_row:crow + threshold_row, ccol - threshold_col:ccol + threshold_col] = scale
+    x_freq *= mask
+    # IFFT
+    x_freq = torch.fft.ifftshift(x_freq, dim=(-2, -1)) # pylint: disable=E1102
+    x_filtered = torch.fft.ifftn(x_freq, dim=(-2, -1)).real.to(device=x.device, dtype=x.dtype) # pylint: disable=E1102
+    return x_filtered
+
+
+def ratio_to_region(width: float, offset: float, n: int):
+    if width < 0:
+        offset += width
+        width = -width
+    width = min(width, 1)
+    if offset < 0:
+        offset = 1 + offset - int(offset)
+    offset = math.fmod(offset, 1.0)
+    if width + offset <= 1:
+        inverted = False
+        start = offset * n
+        end = (width + offset) * n
+    else:
+        inverted = True
+        start = (width + offset - 1) * n
+        end = offset * n
+    return round(start), round(end), inverted
+
+
+def apply_freeu(model, backend_original):
+    global state_enabled # pylint: disable=global-statement
+    global cat_original # pylint: disable=global-statement
+    if backend_original:
+        if opts.freeu_enabled:
+            if not state_enabled: # otherwise already patched
+                cat_original = th.cat
+                th.cat = functools.partial(free_u_cat_hijack, original_function=th.cat)
+                state_enabled = True
+        else:
+            if cat_original is not None:
+                th.cat = cat_original
+                state_enabled = False
+    elif hasattr(model, 'enable_freeu'):
+        if opts.freeu_enabled:
+            model.enable_freeu(s1=opts.freeu_s1, s2=opts.freeu_s2, b1=opts.freeu_b1, b2=opts.freeu_b2)
+            state_enabled = True
+        elif state_enabled:
+            model.disable_freeu()
+            state_enabled = False
+    if opts.freeu_enabled:
+        log.info(f'Applying free-u: b1={opts.freeu_b1} b2={opts.freeu_b2} s1={opts.freeu_s1} s2={opts.freeu_s2}')

modules/sd_models.py

@@ -1249,7 +1249,7 @@ def apply_token_merging(sd_model, token_merging_ratio=0):
                 merge_crossattn=False,
                 merge_mlp=False
             )
-            shared.log.debug(f'Applying token merging: ratio={token_merging_ratio}')
+            shared.log.info(f'Applying token merging: ratio={token_merging_ratio}')
             sd_model.applied_token_merged_ratio = token_merging_ratio
         except Exception:
             shared.log.warning(f'Token merging not supported: pipeline={sd_model.__class__.__name__}')

modules/shared.py

@@ -402,23 +402,8 @@ options_templates.update(options_section(('sd', "Execution & Models"), {
     "sd_disable_ckpt": OptionInfo(False, "Disallow usage of models in ckpt format"),
 }))
 
-options_templates.update(options_section(('optimizations', "Optimizations"), {
-    "cross_attention_optimization": OptionInfo(cross_attention_optimization_default, "Cross-attention optimization method", gr.Radio, lambda: {"choices": shared_items.list_crossattention() }),
-    "cross_attention_options": OptionInfo([], "Cross-attention advanced options", gr.CheckboxGroup, lambda: {"choices": ['xFormers enable flash Attention', 'SDP disable memory attention']}),
-    "sub_quad_sep": OptionInfo("<h2>Sub-quadratic options</h2>", "", gr.HTML),
-    "sub_quad_q_chunk_size": OptionInfo(512, "cross-attention query chunk size", gr.Slider, {"minimum": 16, "maximum": 8192, "step": 8}),
-    "sub_quad_kv_chunk_size": OptionInfo(512, "cross-attention kv chunk size", gr.Slider, {"minimum": 0, "maximum": 8192, "step": 8}),
-    "sub_quad_chunk_threshold": OptionInfo(80, "cross-attention chunking threshold", gr.Slider, {"minimum": 0, "maximum": 100, "step": 1}),
-    "token_merging_sep": OptionInfo("<h2>Token Merging</h2>", "", gr.HTML),
-    "token_merging_ratio": OptionInfo(0.0, "Token merging ratio (txt2img)", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}),
-    "token_merging_ratio_img2img": OptionInfo(0.0, "Token merging ratio (img2img)", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}),
-    "token_merging_ratio_hr": OptionInfo(0.0, "Token merging ratio for (hires)", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}),
-    "inference_mode": OptionInfo("no-grad", "Torch inference mode", gr.Radio, lambda: {"choices": ["no-grad", "inference-mode", "none"]}),
-    "sd_vae_sliced_encode": OptionInfo(False, "VAE Slicing (original)"),
-}))
-
 options_templates.update(options_section(('cuda', "Compute Settings"), {
-    # "memmon_poll_rate": OptionInfo(2, "VRAM usage polls per second during generation", gr.Slider, {"minimum": 0, "maximum": 40, "step": 1}),
+    "math_sep": OptionInfo("<h2>Execution precision</h2>", "", gr.HTML),
     "precision": OptionInfo("Autocast", "Precision type", gr.Radio, lambda: {"choices": ["Autocast", "Full"]}),
     "cuda_dtype": OptionInfo("FP32" if sys.platform == "darwin" or cmd_opts.use_openvino else "BF16" if devices.backend == "ipex" else "FP16", "Device precision type", gr.Radio, lambda: {"choices": ["FP32", "FP16", "BF16"]}),
     "no_half": OptionInfo(False, "Use full precision for model (--no-half)", None, None, None),
@@ -427,23 +412,54 @@ options_templates.update(options_section(('cuda', "Compute Settings"), {
     "upcast_attn": OptionInfo(False, "Enable upcast cross attention layer"),
     "cuda_cast_unet": OptionInfo(False, "Use fixed UNet precision"),
     "disable_nan_check": OptionInfo(True, "Disable NaN check in produced images/latent spaces", gr.Checkbox, {"visible": False}),
-    "rollback_vae": OptionInfo(False, "Attempt VAE roll back when produced NaN values (experimental)"),
+    "rollback_vae": OptionInfo(False, "Attempt VAE roll back when produced NaN values"),
+
+    "cross_attention_sep": OptionInfo("<h2>Cross-attention</h2>", "", gr.HTML),
+    "cross_attention_optimization": OptionInfo(cross_attention_optimization_default, "Cross-attention optimization method", gr.Radio, lambda: {"choices": shared_items.list_crossattention() }),
+    "cross_attention_options": OptionInfo([], "Cross-attention advanced options", gr.CheckboxGroup, lambda: {"choices": ['xFormers enable flash Attention', 'SDP disable memory attention']}),
+    "sub_quad_sep": OptionInfo("<h3>Sub-quadratic options</h3>", "", gr.HTML),
+    "sub_quad_q_chunk_size": OptionInfo(512, "cross-attention query chunk size", gr.Slider, {"minimum": 16, "maximum": 8192, "step": 8}),
+    "sub_quad_kv_chunk_size": OptionInfo(512, "cross-attention kv chunk size", gr.Slider, {"minimum": 0, "maximum": 8192, "step": 8}),
+    "sub_quad_chunk_threshold": OptionInfo(80, "cross-attention chunking threshold", gr.Slider, {"minimum": 0, "maximum": 100, "step": 1}),
+
+    "other_sep": OptionInfo("<h2>Execution precision</h2>", "", gr.HTML),
     "opt_channelslast": OptionInfo(False, "Use channels last as torch memory format "),
     "cudnn_benchmark": OptionInfo(False, "Enable full-depth cuDNN benchmark feature"),
-    "ipex_optimize": OptionInfo(True if devices.backend == "ipex" else False, "Enable IPEX Optimize for Intel GPUs"),
-    "directml_memory_provider": OptionInfo(default_memory_provider, 'DirectML memory stats provider', gr.Radio, lambda: {"choices": memory_providers}),
     "torch_gc_threshold": OptionInfo(90, "VRAM usage threshold before running Torch GC to clear up VRAM", gr.Slider, {"minimum": 0, "maximum": 100, "step": 1}),
+
     "cuda_compile_sep": OptionInfo("<h2>Model Compile</h2>", "", gr.HTML),
     "cuda_compile": OptionInfo(True if cmd_opts.use_openvino else False, "Enable model compile"),
     "cuda_compile_backend": OptionInfo("openvino_fx" if cmd_opts.use_openvino else "none", "Model compile backend", gr.Radio, lambda: {"choices": ['none', 'inductor', 'cudagraphs', 'aot_ts_nvfuser', 'hidet', 'ipex', 'openvino_fx']}),
     "cuda_compile_mode": OptionInfo("default", "Model compile mode", gr.Radio, lambda: {"choices": ['default', 'reduce-overhead', 'max-autotune']}),
     "cuda_compile_fullgraph": OptionInfo(False, "Model compile fullgraph"),
     "cuda_compile_precompile": OptionInfo(False, "Model compile precompile"),
+    "cuda_compile_verbose": OptionInfo(False, "Model compile verbose mode"),
+    "cuda_compile_errors": OptionInfo(True, "Model compile suppress errors"),
+
+    "ipex_sep": OptionInfo("<h2>IPEX, DirectML and OpenVINO</h2>", "", gr.HTML),
+    "ipex_optimize": OptionInfo(True if devices.backend == "ipex" else False, "Enable IPEX Optimize for Intel GPUs"),
+    "directml_memory_provider": OptionInfo(default_memory_provider, 'DirectML memory stats provider', gr.Radio, lambda: {"choices": memory_providers}),
     "openvino_disable_model_caching": OptionInfo(False, "OpenVINO disable model caching"),
     "openvino_multi_gpu": OptionInfo(False, "OpenVINO use Multi GPU"),
     "openvino_remove_igpu_from_multi": OptionInfo(False, "OpenVINO remove iGPU from Multi GPU"),
-    "cuda_compile_verbose": OptionInfo(False, "Model compile verbose mode"),
-    "cuda_compile_errors": OptionInfo(True, "Model compile suppress errors"),
+}))
+
+options_templates.update(options_section(('advanced', "Inference Settings"), {
+    "token_merging_sep": OptionInfo("<h2>Token merging</h2>", "", gr.HTML),
+    "token_merging_ratio": OptionInfo(0.0, "Token merging ratio (txt2img)", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}),
+    "token_merging_ratio_img2img": OptionInfo(0.0, "Token merging ratio (img2img)", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}),
+    "token_merging_ratio_hr": OptionInfo(0.0, "Token merging ratio for (hires)", gr.Slider, {"minimum": 0.0, "maximum": 0.9, "step": 0.1}),
+
+    "freeu_sep": OptionInfo("<h2>FreeU</h2>", "", gr.HTML),
+    "freeu_enabled": OptionInfo(False, "FreeU enabled"),
+    "freeu_b1": OptionInfo(1.2, "1st stage backbone factor", gr.Slider, {"minimum": 1.0, "maximum": 2.0, "step": 0.01}),
+    "freeu_b2": OptionInfo(1.4, "2nd stage backbone factor", gr.Slider, {"minimum": 1.0, "maximum": 2.0, "step": 0.01}),
+    "freeu_s1": OptionInfo(0.9, "1st stage skip factor", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
+    "freeu_s2": OptionInfo(0.2, "2nd stage skip factor", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
+
+    "inference_mode_sep": OptionInfo("<h2>Inference mode</h2>", "", gr.HTML),
+    "inference_mode": OptionInfo("no-grad", "Torch inference mode", gr.Radio, lambda: {"choices": ["no-grad", "inference-mode", "none"]}),
+    "sd_vae_sliced_encode": OptionInfo(False, "VAE Slicing (original)"),
 }))
 
 options_templates.update(options_section(('diffusers', "Diffusers Settings"), {