add new attention dispatcher

Signed-off-by: Vladimir Mandic <mandic00@live.com>

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  

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}')

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:

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)
 

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("<h2>Attention Dispatcher</h2>", "", gr.HTML),
+        "hf_attention": OptionInfo('', "Attention dispatcher kernel", gr.Textbox),
     }))
 
     # --- Server Settings ---