From a0531c70fdbd9d6e3a040d1be4ea0341d518305f Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Fri, 29 May 2026 21:10:53 +0200 Subject: [PATCH] add new attention dispatcher Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 6 ++++ modules/attention.py | 63 ++++++++++++++++++++++++++++++++++++++- modules/loader.py | 1 + modules/sd_models.py | 4 ++- modules/ui_definitions.py | 3 ++ 5 files changed, 75 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 01ca1ec4d..257976dd5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -57,6 +57,12 @@ Plus continued work on modernization of codebase: UI is now fully TypeScript bas - **Gallery** add clear cache button to folder menu - **UV** much updated `--uv` support for fast installs now also supports global `uv` if present in the system + - **Attention Dispatcher** new attention backends dispatcher + in *settings -> attention Dispatcher* + allows to use pluggable kernels defines either in packages or in new [kernels](https://huggingface.co/docs/kernels/index) library + see [backends](https://huggingface.co/docs/diffusers/optimization/attention_backends#available-backends) for list of available attention backends + *note* compatiblity matrix between torch backend, torch version and model specifics is relatively small at the moment + *note* does not replace existing *attention* settings - **Changes** - all **Guidance** params are now set to *-1* by default to allow using model defaults and avoid confusion with different model behaviour log will print default values used by model if not set by user diff --git a/modules/attention.py b/modules/attention.py index 282718e58..a8d93e25d 100644 --- a/modules/attention.py +++ b/modules/attention.py @@ -1,6 +1,6 @@ from functools import wraps import torch -from modules import rocm +from modules import rocm, errors from modules.logger import log from installer import install, installed, torch_info @@ -255,3 +255,64 @@ def set_diffusers_attention(pipe, quiet = False): log.debug(f"Torch attention: slicing={shared.opts.attention_slicing}") pipe.current_attn_name = shared.opts.cross_attention_optimization + + +orig_get_kernel = None +def get_kernel_hijack(repo_id, revision=None, version=None, backend=None, user_agent=None, trust_remote_code: bool | list[str] = False): # pylint: disable=unused-argument + log.debug(f'Attention dispatcher hub: repo="{repo_id}" revision={revision} version={version} backend={backend}') + user_agent = 'kernels/0.14.1' + module = None + try: + module = orig_get_kernel(repo_id, revision=revision, version=version, backend=backend, user_agent=user_agent, trust_remote_code=True) + except Exception as e: + log.error(f'Attention dispatcher hub: {e}') + errors.display(e, 'kernels') + return module + + +def get_hf_api_hijack(user_agent = None): # pylint: disable=unused-argument + from huggingface_hub import HfApi + return HfApi(library_name="kernels", user_agent="donottrack") + + +def set_attention_dispatcher(pipe): + global orig_get_kernel # pylint: disable=global-statement + from modules import shared + attn = shared.opts.hf_attention.strip().lower() + if pipe is None or not hasattr(pipe, 'transformer') or not hasattr(pipe.transformer, 'set_attention_backend'): + return + + from diffusers.models import attention_dispatch as a + backends = [b.value for b in a._AttentionBackendRegistry.list_backends()] # pylint: disable=protected-access + # https://huggingface.co/docs/kernels/index + # https://huggingface.co/docs/diffusers/optimization/attention_backends#available-backends + + if 'hub' in attn: + try: + install('kernels==0.14.1') + import kernels + import kernels.utils + log.debug(f'Attention dispatcher: kernels={kernels.__version__}') + if orig_get_kernel is None: + orig_get_kernel = kernels.get_kernel + kernels.get_kernel = get_kernel_hijack + kernels.utils._get_hf_api = get_hf_api_hijack # pylint: disable=protected-access + from diffusers.utils import import_utils + import_utils._kernels_available = True # pylint: disable=protected-access + import_utils._kernels_version = kernels.__version__ # pylint: disable=protected-access + except Exception as e: + log.error(f'Attention dispatcher kernels: {e}') + return + + prev = a._AttentionBackendRegistry.get_active_backend() # pylint: disable=protected-access + if attn in backends: + try: + pipe.transformer.set_attention_backend(attn) + except Exception as e: + log.error(f'Attention dispatcher: target={attn} {e}') + current = a._AttentionBackendRegistry.get_active_backend() # pylint: disable=protected-access + log.debug(f'Attention dispatcher: target={attn} previous={prev[0].value} active={current[0]} list={backends}') + elif len(attn) > 0: + log.warning(f'Attention dispatcher: active={prev[0].value} list={backends} target={attn} not found') + else: + log.debug(f'Attention dispatcher: active={prev[0].value} list={backends}') diff --git a/modules/loader.py b/modules/loader.py index d57bfd3b7..3f018e964 100644 --- a/modules/loader.py +++ b/modules/loader.py @@ -208,6 +208,7 @@ try: except Exception as e: log.error(f'Loader: diffusers=={diffusers.__version__ if "diffusers" in sys.modules else None} {e}') log.error('Please restart re-run the installer') + # errors.display(e, 'diffusers') sys.exit(1) try: diff --git a/modules/sd_models.py b/modules/sd_models.py index aefdb2915..cc36d57ea 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -197,7 +197,7 @@ def set_diffuser_options(sd_model, vae=None, op:str='model', offload:bool=True, shared.opts.sdnq_use_quantized_matmul = False if module.quantization_config.use_quantized_matmul != shared.opts.sdnq_use_quantized_matmul: from modules.sdnq.loader import apply_sdnq_options_to_model - log.debug(f'Setting {op} {module_name}: sdnq_use_quantized_matmul={shared.opts.sdnq_use_quantized_matmul}') + # log.debug(f'Setting {op} {module_name}: sdnq_use_quantized_matmul={shared.opts.sdnq_use_quantized_matmul}') module = apply_sdnq_options_to_model(module, use_quantized_matmul=shared.opts.sdnq_use_quantized_matmul) setattr(sd_model, module_name, module) @@ -968,6 +968,7 @@ def load_diffuser(checkpoint_info=None, op='model', revision=None): # pylint: di prompt_parser_diffusers.cache.clear() set_diffuser_options(sd_model, vae, op, offload=False) + attention.set_attention_dispatcher(sd_model) sd_model = model_quant.do_post_load_quant(sd_model, allow=allow_post_quant) # run this before move model so it can be compressed in CPU timer.load.record("options") @@ -1018,6 +1019,7 @@ def load_diffuser(checkpoint_info=None, op='model', revision=None): # pylint: di if shared.opts.diffusers_offload_mode != 'balanced': devices.torch_gc(force=True, reason='load') + if sd_model is not None: script_callbacks.model_loaded_callback(sd_model) diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index fd37a6abd..63cc1907e 100644 --- a/modules/ui_definitions.py +++ b/modules/ui_definitions.py @@ -230,6 +230,9 @@ def create_settings(cmd_opts): "xformers_options": OptionInfo(['Flash attention'], "xFormers options", gr.CheckboxGroup, {"choices": ['Flash attention'] }), "dynamic_attention_slice_rate": OptionInfo(0.5, "Dynamic Attention slicing rate", gr.Slider, {"minimum": 0.01, "maximum": max(gpu_memory,4), "step": 0.01}), "dynamic_attention_trigger_rate": OptionInfo(1, "Dynamic Attention trigger rate", gr.Slider, {"minimum": 0.01, "maximum": max(gpu_memory,4)*2, "step": 0.01}), + + "hf_attention_sep": OptionInfo("

Attention Dispatcher

", "", gr.HTML), + "hf_attention": OptionInfo('', "Attention dispatcher kernel", gr.Textbox), })) # --- Server Settings ---