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

Code before:
            backend (`str`):
                The name of the backend to set. Must be one of the available backends defined in
                `AttentionBackendName`. Available backends can be found in
                `diffusers.attention_dispatch.AttentionBackendName`. Defaults to torch native scaled dot product
                attention as backend.
        """
        from .attention import AttentionModuleMixin
        from .attention_dispatch import (
            AttentionBackendName,
            _AttentionBackendRegistry,
            _check_attention_backend_requirements,
            _maybe_download_kernel_for_backend,
        )

        # TODO: the following will not be required when everything is refactored to AttentionModuleMixin
        from .attention_processor import Attention, MochiAttention

        logger.warning("Attention backends are an experimental feature and the API may be subject to change.")

        backend = backend.lower()
        available_backends = {x.value for x in AttentionBackendName.__members__.values()}
        if backend not in available_backends:
            raise ValueError(f"`{backend=}` must be one of the following: " + ", ".join(available_backends))

        backend = AttentionBackendName(backend)
        if not _AttentionBackendRegistry._is_context_parallel_available(backend):
            compatible_backends = sorted(_AttentionBackendRegistry._supports_context_parallel)
            raise ValueError(
                f"Context parallelism is enabled but backend '{backend.value}' "
                f"which does not support context parallelism. "
                f"Please set a compatible attention backend: {compatible_backends} using `model.set_attention_backend()` before "
                f"calling `model.enable_parallelism()`."
            )

        _check_attention_backend_requirements(backend)
        _maybe_download_kernel_for_backend(backend)

        attention_classes = (Attention, MochiAttention, AttentionModuleMixin)
        for module in self.modules():
            if not isinstance(module, attention_classes):
                continue
            processor = module.processor
            if processor is None or not hasattr(processor, "_attention_backend"):
                continue
            processor._attention_backend = backend

        # Important to set the active backend so that it propagates gracefully throughout.
        _AttentionBackendRegistry.set_active_backend(backend)

    def reset_attention_backend(self) -> None:
        """

Review comment (suggestion):
This will raise an error if the backend doesn't support context parallel (regardless of whether you want to use CP or not) 

We need to first check if the attention processor has a `_parallel_config` set and then check if the backend supports context parallel. 


Code after:
                `diffusers.attention_dispatch.AttentionBackendName`. Defaults to torch native scaled dot product
                attention as backend.
        """
        from .attention import AttentionModuleMixin
        from .attention_dispatch import (
            AttentionBackendName,
            _AttentionBackendRegistry,
            _check_attention_backend_requirements,
            _maybe_download_kernel_for_backend,
        )

        # TODO: the following will not be required when everything is refactored to AttentionModuleMixin
        from .attention_processor import Attention, MochiAttention

        logger.warning("Attention backends are an experimental feature and the API may be subject to change.")
        attention_classes = (Attention, MochiAttention, AttentionModuleMixin)

        parallel_config_set = False
        for module in self.modules():
            if not isinstance(module, attention_classes):
                continue
            processor = module.processor
            if getattr(processor, "_parallel_config", None) is not None:
                parallel_config_set = True
                break

        backend = backend.lower()
        available_backends = {x.value for x in AttentionBackendName.__members__.values()}
        if backend not in available_backends:
            raise ValueError(f"`{backend=}` must be one of the following: " + ", ".join(available_backends))

        backend = AttentionBackendName(backend)
        if parallel_config_set and not _AttentionBackendRegistry._is_context_parallel_available(backend):
            compatible_backends = sorted(_AttentionBackendRegistry._supports_context_parallel)
            raise ValueError(
                f"Context parallelism is enabled but current attention backend '{backend.value}' "
                f"does not support context parallelism. "
                f"Please set a compatible attention backend: {compatible_backends} using `model.set_attention_backend()`."
            )

        _check_attention_backend_requirements(backend)
        _maybe_download_kernel_for_backend(backend)

        for module in self.modules():
            if not isinstance(module, attention_classes):
                continue
            processor = module.processor
            if processor is None or not hasattr(processor, "_attention_backend"):
                continue
            processor._attention_backend = backend