diff --git a/CHANGELOG.md b/CHANGELOG.md index 2b5888d2c..2325c721e 100644 --- a/CHANGELOG.md +++ b/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: -- Add 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 diff --git a/extensions-builtin/sd-extension-system-info b/extensions-builtin/sd-extension-system-info index 2bb95bead..2fd1fcd9d 160000 --- a/extensions-builtin/sd-extension-system-info +++ b/extensions-builtin/sd-extension-system-info @@ -1 +1 @@ -Subproject commit 2bb95beadeef813ca667573325b05ebbfd307e6b +Subproject commit 2fd1fcd9d2353809f875fca1b6164630100f8087 diff --git a/html/locale_en.json b/html/locale_en.json index ba2a63285..92b41bf39 100644 --- a/html/locale_en.json +++ b/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":""}, diff --git a/javascript/style.css b/javascript/style.css index c03b6eb60..afdf698ba 100644 --- a/javascript/style.css +++ b/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; } diff --git a/modules/processing.py b/modules/processing.py index 5547808a2..d40f1d33e 100644 --- a/modules/processing.py +++ b/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: """ diff --git a/modules/sd_hijack_freeu.py b/modules/sd_hijack_freeu.py new file mode 100644 index 000000000..51448561a --- /dev/null +++ b/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 +# 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}') diff --git a/modules/sd_models.py b/modules/sd_models.py index cc51b2a11..7120271d1 100644 --- a/modules/sd_models.py +++ b/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__}') diff --git a/modules/shared.py b/modules/shared.py index 7c02b2aa4..15960445b 100644 --- a/modules/shared.py +++ b/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("

Sub-quadratic options

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

Token Merging

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

Execution precision

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

Cross-attention

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

Sub-quadratic options

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

Execution precision

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

Model Compile

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

IPEX, DirectML and OpenVINO

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

Token merging

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

FreeU

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

Inference mode

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