Class _AttentionBackendRegistry

src/diffusers/models/attention_dispatch.py:267–313

class _AttentionBackendRegistry:
    _backends = {}
    _constraints = {}
    _supported_arg_names = {}
    _supports_context_parallel = set()
    _active_backend = AttentionBackendName(DIFFUSERS_ATTN_BACKEND)
    _checks_enabled = DIFFUSERS_ATTN_CHECKS

    @classmethod
    def register(
        cls,
        backend: AttentionBackendName,
        constraints: list[Callable] | None = None,
        supports_context_parallel: bool = False,
    ):
        logger.debug(f"Registering attention backend: {backend} with constraints: {constraints}")

        def decorator(func):
            cls._backends[backend] = func
            cls._constraints[backend] = constraints or []
            cls._supported_arg_names[backend] = set(inspect.signature(func).parameters.keys())
            if supports_context_parallel:
                cls._supports_context_parallel.add(backend.value)

            return func

        return decorator

    @classmethod
    def get_active_backend(cls):
        return cls._active_backend, cls._backends[cls._active_backend]

    @classmethod
    def set_active_backend(cls, backend: str):
        cls._active_backend = backend

    @classmethod
    def list_backends(cls):
        return list(cls._backends.keys())

    @classmethod
    def _is_context_parallel_available(
        cls,
        backend: AttentionBackendName,
    ) -> bool:
        supports_context_parallel = backend.value in cls._supports_context_parallel
        return supports_context_parallel


@dataclass