Function attention_backend

src/diffusers/models/attention_dispatch.py:371–388

Context manager to set the active attention backend.

(backend: str | AttentionBackendName = AttentionBackendName.NATIVE)

@contextlib.contextmanager
def attention_backend(backend: str | AttentionBackendName = AttentionBackendName.NATIVE):
    """
    Context manager to set the active attention backend.
    """
    if backend not in _AttentionBackendRegistry._backends:
        raise ValueError(f"Backend {backend} is not registered.")

    backend = AttentionBackendName(backend)
    _check_attention_backend_requirements(backend)
    _maybe_download_kernel_for_backend(backend)

    old_backend = _AttentionBackendRegistry._active_backend
    _AttentionBackendRegistry.set_active_backend(backend)

    try:
        yield
    finally:
        _AttentionBackendRegistry.set_active_backend(old_backend)


def dispatch_attention_fn(