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

Code before:
            num_blocks_per_group=num_blocks_per_group,
            non_blocking=non_blocking,
            use_stream=use_stream,
            record_stream=record_stream,
            low_cpu_mem_usage=low_cpu_mem_usage,
            offload_to_disk_path=offload_to_disk_path,
        )

    def set_attention_backend(self, backend: str) -> None:
        """
        Set the attention backend for the model.

        Args:
            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

        # 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)
        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

    def reset_attention_backend(self) -> None:
        """
        Resets the attention backend for the model. Following calls to `forward` will use the environment default or
        the torch native scaled dot product attention.
        """
        from .attention import AttentionModuleMixin
        from .attention_processor import Attention, MochiAttention


Review comment (refactor):
cc @DN6 Added this here because things are probably going to moving a little too fast for now. #11941 introduces several changes to simplify the interface for operating with the templated implementations. For custom parameters that the backends can accept, it is very easy for a user to register a custom dispatch function (i.e. we won't support all parameter customizability on our end and just use the defaults)

Code after:
            num_blocks_per_group=num_blocks_per_group,
            non_blocking=non_blocking,
            use_stream=use_stream,
            record_stream=record_stream,
            low_cpu_mem_usage=low_cpu_mem_usage,
            offload_to_disk_path=offload_to_disk_path,
        )

    def set_attention_backend(self, backend: str) -> None:
        """
        Set the attention backend for the model.

        Args:
            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, _check_attention_backend_requirements

        # 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)
        _check_attention_backend_requirements(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

    def reset_attention_backend(self) -> None:
        """
        Resets the attention backend for the model. Following calls to `forward` will use the environment default or
        the torch native scaled dot product attention.
        """
        from .attention import AttentionModuleMixin
        from .attention_processor import Attention, MochiAttention