Repository: huggingface/diffusers
File: src/diffusers/models/attention_dispatch.py (Python)

Code before:

if TYPE_CHECKING:
    from ._modeling_parallel import ParallelConfig

_REQUIRED_FLASH_VERSION = "2.6.3"
_REQUIRED_AITER_VERSION = "0.1.5"
_REQUIRED_SAGE_VERSION = "2.1.1"
_REQUIRED_FLEX_VERSION = "2.5.0"
_REQUIRED_XLA_VERSION = "2.2"
_REQUIRED_XFORMERS_VERSION = "0.0.29"

_CAN_USE_FLASH_ATTN = is_flash_attn_available() and is_flash_attn_version(">=", _REQUIRED_FLASH_VERSION)
_CAN_USE_FLASH_ATTN_3 = is_flash_attn_3_available()
_CAN_USE_AITER_ATTN = is_aiter_available() and is_aiter_version(">=", _REQUIRED_AITER_VERSION)
_CAN_USE_SAGE_ATTN = is_sageattention_available() and is_sageattention_version(">=", _REQUIRED_SAGE_VERSION)
_CAN_USE_FLEX_ATTN = is_torch_version(">=", _REQUIRED_FLEX_VERSION)
_CAN_USE_NPU_ATTN = is_torch_npu_available()
_CAN_USE_XLA_ATTN = is_torch_xla_available() and is_torch_xla_version(">=", _REQUIRED_XLA_VERSION)
_CAN_USE_XFORMERS_ATTN = is_xformers_available() and is_xformers_version(">=", _REQUIRED_XFORMERS_VERSION)


if _CAN_USE_FLASH_ATTN:
    try:
        from flash_attn import flash_attn_func, flash_attn_varlen_func
        from flash_attn.flash_attn_interface import _wrapped_flash_attn_backward, _wrapped_flash_attn_forward
    except (ImportError, OSError) as e:
        # Handle ABI mismatch or other import failures gracefully.
        # This can happen when flash_attn was compiled against a different PyTorch version.
        _flash_attn_logger = get_logger(__name__)
        _flash_attn_logger.warning(
            f"flash_attn is installed but failed to import: {e}. "
            f"Falling back to native PyTorch attention."
        )
        _CAN_USE_FLASH_ATTN = False
        flash_attn_func = None
        flash_attn_varlen_func = None
        _wrapped_flash_attn_backward = None
        _wrapped_flash_attn_forward = None
else:
    flash_attn_func = None
    flash_attn_varlen_func = None
    _wrapped_flash_attn_backward = None
    _wrapped_flash_attn_forward = None


if _CAN_USE_FLASH_ATTN_3:
    try:
        from flash_attn_interface import flash_attn_func as flash_attn_3_func
        from flash_attn_interface import flash_attn_varlen_func as flash_attn_3_varlen_func
    except (ImportError, OSError) as e:
        _flash_attn_3_logger = get_logger(__name__)

Review comment (suggestion):
Think we can include `RuntimeError` in the exceptions list as well. 

Code after:

if TYPE_CHECKING:
    from ._modeling_parallel import ParallelConfig

_REQUIRED_FLASH_VERSION = "2.6.3"
_REQUIRED_AITER_VERSION = "0.1.5"
_REQUIRED_SAGE_VERSION = "2.1.1"
_REQUIRED_FLEX_VERSION = "2.5.0"
_REQUIRED_XLA_VERSION = "2.2"
_REQUIRED_XFORMERS_VERSION = "0.0.29"

logger = get_logger(__name__)  # pylint: disable=invalid-name

_CAN_USE_FLASH_ATTN = is_flash_attn_available() and is_flash_attn_version(">=", _REQUIRED_FLASH_VERSION)
_CAN_USE_FLASH_ATTN_3 = is_flash_attn_3_available()
_CAN_USE_AITER_ATTN = is_aiter_available() and is_aiter_version(">=", _REQUIRED_AITER_VERSION)
_CAN_USE_SAGE_ATTN = is_sageattention_available() and is_sageattention_version(">=", _REQUIRED_SAGE_VERSION)
_CAN_USE_FLEX_ATTN = is_torch_version(">=", _REQUIRED_FLEX_VERSION)
_CAN_USE_NPU_ATTN = is_torch_npu_available()
_CAN_USE_XLA_ATTN = is_torch_xla_available() and is_torch_xla_version(">=", _REQUIRED_XLA_VERSION)
_CAN_USE_XFORMERS_ATTN = is_xformers_available() and is_xformers_version(">=", _REQUIRED_XFORMERS_VERSION)


if _CAN_USE_FLASH_ATTN:
    try:
        from flash_attn import flash_attn_func, flash_attn_varlen_func
        from flash_attn.flash_attn_interface import _wrapped_flash_attn_backward, _wrapped_flash_attn_forward
    except (ImportError, OSError, RuntimeError) as e:
        # Handle ABI mismatch or other import failures gracefully.
        # This can happen when flash_attn was compiled against a different PyTorch version.
        logger.warning(f"flash_attn is installed but failed to import: {e}. Falling back to native PyTorch attention.")
        _CAN_USE_FLASH_ATTN = False
        flash_attn_func = None
        flash_attn_varlen_func = None
        _wrapped_flash_attn_backward = None
        _wrapped_flash_attn_forward = None
else:
    flash_attn_func = None
    flash_attn_varlen_func = None
    _wrapped_flash_attn_backward = None
    _wrapped_flash_attn_forward = None


if _CAN_USE_FLASH_ATTN_3:
    try:
        from flash_attn_interface import flash_attn_func as flash_attn_3_func
        from flash_attn_interface import flash_attn_varlen_func as flash_attn_3_varlen_func
    except (ImportError, OSError, RuntimeError) as e:
        logger.warning(f"flash_attn_3 failed to import: {e}. Falling back to native attention.")
        _CAN_USE_FLASH_ATTN_3 = False
        flash_attn_3_func = None