code format fixes

Author: Bissmella

src/diffusers/models/attention_dispatch.py

@@ -44,7 +44,6 @@
     is_xformers_version,
 )
 from ..utils.constants import DIFFUSERS_ATTN_BACKEND, DIFFUSERS_ATTN_CHECKS
-from ..utils import is_torch_version
 
 
 if TYPE_CHECKING:
@@ -1076,7 +1075,7 @@ def _all_to_all_dim_exchange(x: torch.Tensor, scatter_idx: int = 2, gather_idx:
 
         # B, S_LOCAL, H, D -> group_world_size, S_LOCAL, B, H_LOCAL, D
         x_temp = x.reshape(batch_size, seq_len_local, group_world_size, num_heads_local, head_dim).transpose(0, 2).contiguous()
-        
+
 
         if group_world_size >1:
             out = _all_to_all_single(x_temp, group=group)
@@ -1365,7 +1364,7 @@ def TemplatedUnifiedAttention(
     forward_op,
     backward_op,
     _parallel_config: Optional["ParallelConfig"] = None,
-    scatter_idx: int =2, 
+    scatter_idx: int =2,
     gather_idx: int =1,
     ):
     """

tests/others/test_unified_sp_attention.py

@@ -1,13 +1,19 @@
 import math
+import os
+
 import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
+
 from diffusers.models.attention_dispatch import TemplatedUnifiedAttention
-import os
+from diffusers.models._modeling_parallel import (
+        ParallelConfig,
+        ContextParallelConfig
+    )
 
 def run(rank, world_size):
     dist.init_process_group(
-        backend="gloo", 
+        backend="gloo",
         rank=rank,
         world_size=world_size
     )
@@ -21,10 +27,7 @@ def run(rank, world_size):
 
     q.requires_grad_(True)
 
-    from diffusers.models._modeling_parallel import (
-        ParallelConfig, 
-        ContextParallelConfig
-    )
+    
 
     pc = ParallelConfig(
         context_parallel_config=ContextParallelConfig(
@@ -126,4 +129,4 @@ def dummy_backward_op(ctx, grad_out, *args, **kwargs):
     world_size = 4
     os.environ["MASTER_ADDR"] = "localhost"
     os.environ["MASTER_PORT"] = "12355"
-    mp.spawn(run, args=(world_size,), nprocs=world_size)
+    mp.spawn(run, args=(world_size,), nprocs=world_size)