Spaces:
Runtime error
Runtime error
| import torch | |
| from torch import nn | |
| from typing import Optional, Any | |
| try: | |
| import xformers | |
| import xformers.ops | |
| XFORMERS_IS_AVAILBLE = True | |
| except: | |
| XFORMERS_IS_AVAILBLE = False | |
| class MemoryEfficientCrossAttention(nn.Module): | |
| # https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223 | |
| def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0): | |
| super().__init__() | |
| # print(f"Setting up {self.__class__.__name__}. Query dim is {query_dim}, context_dim is {context_dim} and using {heads} heads.") | |
| inner_dim = dim_head * heads | |
| context_dim = context_dim or query_dim | |
| self.heads = heads | |
| self.dim_head = dim_head | |
| self.to_q = nn.Linear(query_dim, inner_dim, bias=False) | |
| self.to_k = nn.Linear(context_dim, inner_dim, bias=False) | |
| self.to_v = nn.Linear(context_dim, inner_dim, bias=False) | |
| self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout)) | |
| self.attention_op: Optional[Any] = None | |
| def forward(self, x, context=None, mask=None): | |
| q = self.to_q(x) | |
| context = x if context is None else context | |
| k = self.to_k(context) | |
| v = self.to_v(context) | |
| b, _, _ = q.shape | |
| q, k, v = map( | |
| lambda t: t.unsqueeze(3) | |
| .reshape(b, t.shape[1], self.heads, self.dim_head) | |
| .permute(0, 2, 1, 3) | |
| .reshape(b * self.heads, t.shape[1], self.dim_head) | |
| .contiguous(), | |
| (q, k, v), | |
| ) | |
| out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=self.attention_op) | |
| if mask is not None: | |
| raise NotImplementedError | |
| out = ( | |
| out.unsqueeze(0) | |
| .reshape(b, self.heads, out.shape[1], self.dim_head) | |
| .permute(0, 2, 1, 3) | |
| .reshape(b, out.shape[1], self.heads * self.dim_head) | |
| ) | |
| return self.to_out(out) |