up

Author: sayakpaul

src/diffusers/models/modeling_utils.py

@@ -611,14 +611,24 @@ def set_attention_backend(self, backend: str) -> None:
         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 not _AttentionBackendRegistry._is_context_parallel_available(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 backend '{backend.value}' "
@@ -630,7 +640,6 @@ def set_attention_backend(self, backend: str) -> None:
         _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