selector ¶

def backend_name_to_enum(backend_name: str) -> AttentionBackendEnum | None:
    """
    Convert a string backend name to a _Backend enum value.

    Returns:
    * _Backend: enum value if backend_name is a valid in-tree type
    * None: otherwise it's an invalid in-tree type or an out-of-tree platform is
            loaded.
    """
    assert backend_name is not None
    return AttentionBackendEnum[backend_name] if backend_name in AttentionBackendEnum.__members__ else \
          None

def coerce_attn_backend(attn_backend: AttentionBackendEnum | str | None, ) -> AttentionBackendEnum | None:
    """Normalize an explicit backend selection.

    Environment-variable parsing remains permissive via
    :func:`backend_name_to_enum`, but typed/config-driven call sites should
    fail fast on typos instead of silently falling back to another backend.
    """
    if attn_backend is None or isinstance(attn_backend, AttentionBackendEnum):
        return attn_backend
    if not isinstance(attn_backend, str) or not attn_backend.strip():
        raise ValueError("attention backend must be a non-empty string, "
                         f"an AttentionBackendEnum, or None; got {attn_backend!r}")

    backend_name = attn_backend.strip().upper()
    backend = backend_name_to_enum(backend_name)
    if backend is None:
        raise ValueError(f"Unknown attention backend {attn_backend!r}. "
                         f"Expected one of {sorted(AttentionBackendEnum.__members__)}")
    return backend

def get_env_variable_attn_backend() -> AttentionBackendEnum | None:
    '''
    Get the backend override specified by the FastVideo attention
    backend environment variable, if one is specified.

    Returns:

    * _Backend enum value if an override is specified
    * None otherwise
    '''
    backend_name = os.environ.get(STR_BACKEND_ENV_VAR)
    return (None if backend_name is None else backend_name_to_enum(backend_name))

def get_global_forced_attn_backend() -> AttentionBackendEnum | None:
    '''
    Get the currently-forced choice of attention backend,
    or None if auto-selection is currently enabled.
    '''
    return forced_attn_backend

def global_force_attn_backend(attn_backend: AttentionBackendEnum | None) -> None:
    '''
    Force all attention operations to use a specified backend.

    Passing `None` for the argument re-enables automatic
    backend selection.,

    Arguments:

    * attn_backend: backend selection (None to revert to auto)
    '''
    global forced_attn_backend
    forced_attn_backend = attn_backend
    # Backend selection is cached by tensor shape/dtype, while the global
    # override is intentionally not part of that cache key. Invalidate cached
    # resolutions whenever the override changes so independently constructed
    # role models can bind different attention implementations.
    _cached_get_attn_backend.cache_clear()

@contextmanager
def global_force_attn_backend_context_manager(attn_backend: AttentionBackendEnum) -> Generator[None, None, None]:
    '''
    Globally force a FastVideo attention backend override within a
    context manager, reverting the global attention backend
    override to its prior state upon exiting the context
    manager.

    Arguments:

    * attn_backend: attention backend to force

    Returns:

    * Generator
    '''

    # Save the current state of the global backend override (if any)
    original_value = get_global_forced_attn_backend()

    # Globally force the new backend override
    global_force_attn_backend(attn_backend)

    # Yield control back to the enclosed code block
    try:
        yield
    finally:
        # Revert the original global backend override, if any
        global_force_attn_backend(original_value)