make style

Author: a-r-r-o-w

src/diffusers/models/_modeling_parallel.py

@@ -15,7 +15,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import contextlib
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union
 
@@ -25,8 +24,7 @@
 
 
 if TYPE_CHECKING:
-    from ..pipelines.pipeline_utils import DiffusionPipeline
-    from .modeling_utils import ModelMixin
+    pass
 
 
 logger = get_logger(__name__)  # pylint: disable=invalid-name

src/diffusers/models/attention_dispatch.py

@@ -235,9 +235,13 @@ def list_backends(cls):
         return list(cls._backends.keys())
 
     @classmethod
-    def _is_context_parallel_enabled(cls, backend: AttentionBackendName, parallel_config: Optional["ContextParallelConfig"]) -> bool:
+    def _is_context_parallel_enabled(
+        cls, backend: AttentionBackendName, parallel_config: Optional["ContextParallelConfig"]
+    ) -> bool:
         supports_context_parallel = backend in cls._supports_context_parallel
-        is_degree_greater_than_1 = parallel_config is not None and (parallel_config.ring_degree > 1 or parallel_config.ulysses_degree > 1)
+        is_degree_greater_than_1 = parallel_config is not None and (
+            parallel_config.ring_degree > 1 or parallel_config.ulysses_degree > 1
+        )
         return supports_context_parallel and is_degree_greater_than_1
 
 
@@ -285,9 +289,8 @@ def dispatch_attention_fn(
         backend_name = AttentionBackendName(backend)
         backend_fn = _AttentionBackendRegistry._backends.get(backend_name)
 
-    if (
-        parallel_config is not None
-        and not _AttentionBackendRegistry._is_context_parallel_enabled(backend_name, parallel_config)
+    if parallel_config is not None and not _AttentionBackendRegistry._is_context_parallel_enabled(
+        backend_name, parallel_config
     ):
         raise ValueError(
             f"Backend {backend_name} either does not support context parallelism or context parallelism "

src/diffusers/models/transformers/transformer_bria.py

@@ -162,7 +162,12 @@ def __call__(
             key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)
 
         hidden_states = dispatch_attention_fn(
-            query, key, value, attn_mask=attention_mask, backend=self._attention_backend, parallel_config=self._parallel_config
+            query,
+            key,
+            value,
+            attn_mask=attention_mask,
+            backend=self._attention_backend,
+            parallel_config=self._parallel_config,
         )
         hidden_states = hidden_states.flatten(2, 3)
         hidden_states = hidden_states.to(query.dtype)

src/diffusers/models/transformers/transformer_flux.py

@@ -116,7 +116,12 @@ def __call__(
             key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)
 
         hidden_states = dispatch_attention_fn(
-            query, key, value, attn_mask=attention_mask, backend=self._attention_backend, parallel_config=self._parallel_config
+            query,
+            key,
+            value,
+            attn_mask=attention_mask,
+            backend=self._attention_backend,
+            parallel_config=self._parallel_config,
         )
         hidden_states = hidden_states.flatten(2, 3)
         hidden_states = hidden_states.to(query.dtype)