Update sub_quadratic_attention.py

backend/misc/sub_quadratic_attention.py

@@ -24,7 +24,8 @@ except ImportError:
 from torch import Tensor
 from typing import List
 
-from ldm_patched.modules import model_management
+from backend import memory_management
+
 
 def dynamic_slice(
         x: Tensor,
@@ -34,11 +35,13 @@ def dynamic_slice(
     slicing = [slice(start, start + size) for start, size in zip(starts, sizes)]
     return x[slicing]
 
+
 class AttnChunk(NamedTuple):
     exp_values: Tensor
     exp_weights_sum: Tensor
     max_score: Tensor
 
+
 class SummarizeChunk(Protocol):
     @staticmethod
     def __call__(
@@ -47,6 +50,7 @@ class SummarizeChunk(Protocol):
             value: Tensor,
     ) -> AttnChunk: ...
 
+
 class ComputeQueryChunkAttn(Protocol):
     @staticmethod
     def __call__(
@@ -55,6 +59,7 @@ class ComputeQueryChunkAttn(Protocol):
             value: Tensor,
     ) -> Tensor: ...
 
+
 def _summarize_chunk(
         query: Tensor,
         key_t: Tensor,
@@ -64,7 +69,7 @@ def _summarize_chunk(
         mask,
 ) -> AttnChunk:
     if upcast_attention:
-        with torch.autocast(enabled=False, device_type = 'cuda'):
+        with torch.autocast(enabled=False, device_type='cuda'):
             query = query.float()
             key_t = key_t.float()
             attn_weights = torch.baddbmm(
@@ -93,6 +98,7 @@ def _summarize_chunk(
     max_score = max_score.squeeze(-1)
     return AttnChunk(exp_values, exp_weights.sum(dim=-1), max_score)
 
+
 def _query_chunk_attention(
         query: Tensor,
         key_t: Tensor,
@@ -116,7 +122,7 @@ def _query_chunk_attention(
             (batch_x_heads, kv_chunk_size, v_channels_per_head)
         )
         if mask is not None:
-            mask = mask[:,:,chunk_idx:chunk_idx + kv_chunk_size]
+            mask = mask[:, :, chunk_idx:chunk_idx + kv_chunk_size]
 
         return summarize_chunk(query, key_chunk, value_chunk, mask=mask)
 
@@ -135,6 +141,7 @@ def _query_chunk_attention(
     all_weights = torch.unsqueeze(chunk_weights, -1).sum(dim=0)
     return all_values / all_weights
 
+
 # TODO: refactor CrossAttention#get_attention_scores to share code with this
 def _get_attention_scores_no_kv_chunking(
         query: Tensor,
@@ -145,7 +152,7 @@ def _get_attention_scores_no_kv_chunking(
         mask,
 ) -> Tensor:
     if upcast_attention:
-        with torch.autocast(enabled=False, device_type = 'cuda'):
+        with torch.autocast(enabled=False, device_type='cuda'):
             query = query.float()
             key_t = key_t.float()
             attn_scores = torch.baddbmm(
@@ -169,7 +176,7 @@ def _get_attention_scores_no_kv_chunking(
     try:
         attn_probs = attn_scores.softmax(dim=-1)
         del attn_scores
-    except model_management.OOM_EXCEPTION:
+    except memory_management.OOM_EXCEPTION:
         print("ran out of memory while running softmax in  _get_attention_scores_no_kv_chunking, trying slower in place softmax instead")
         attn_scores -= attn_scores.max(dim=-1, keepdim=True).values
         torch.exp(attn_scores, out=attn_scores)
@@ -180,10 +187,12 @@ def _get_attention_scores_no_kv_chunking(
     hidden_states_slice = torch.bmm(attn_probs.to(value.dtype), value)
     return hidden_states_slice
 
+
 class ScannedChunk(NamedTuple):
     chunk_idx: int
     attn_chunk: AttnChunk
 
+
 def efficient_dot_product_attention(
         query: Tensor,
         key_t: Tensor,
@@ -193,7 +202,7 @@ def efficient_dot_product_attention(
         kv_chunk_size_min: Optional[int] = None,
         use_checkpoint=True,
         upcast_attention=False,
-    mask = None,
+        mask=None,
 ):
     """Computes efficient dot-product attention given query, transposed key, and value.
       This is efficient version of attention presented in
@@ -234,7 +243,7 @@ def efficient_dot_product_attention(
         if mask is None:
             return None
         chunk = min(query_chunk_size, q_tokens)
-        return mask[:,chunk_idx:chunk_idx + chunk]
+        return mask[:, chunk_idx:chunk_idx + chunk]
 
     summarize_chunk: SummarizeChunk = partial(_summarize_chunk, scale=scale, upcast_attention=upcast_attention)
     summarize_chunk: SummarizeChunk = partial(checkpoint, summarize_chunk) if use_checkpoint else summarize_chunk