| 258 | class MemoryEfficientCrossAttention(nn.Module): |
| 259 | # https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223 |
| 260 | def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0, **kwargs): |
| 261 | super().__init__() |
| 262 | print( |
| 263 | f"Setting up {self.__class__.__name__}. Query dim is {query_dim}, context_dim is {context_dim} and using " |
| 264 | f"{heads} heads with a dimension of {dim_head}." |
| 265 | ) |
| 266 | inner_dim = dim_head * heads |
| 267 | context_dim = default(context_dim, query_dim) |
| 268 | |
| 269 | self.heads = heads |
| 270 | self.dim_head = dim_head |
| 271 | |
| 272 | self.to_q = nn.Linear(query_dim, inner_dim, bias=False) |
| 273 | self.to_k = nn.Linear(context_dim, inner_dim, bias=False) |
| 274 | self.to_v = nn.Linear(context_dim, inner_dim, bias=False) |
| 275 | |
| 276 | self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout)) |
| 277 | self.attention_op: Optional[Any] = None |
| 278 | |
| 279 | def forward( |
| 280 | self, |