Method
__init__
(
self,
query_dim,
context_dim=None,
heads=8,
dim_head=64,
dropout=0.0,
backend=None,
)
Source from the content-addressed store, hash-verified
| 173 | |
| 174 | class CrossAttention(nn.Module): |
| 175 | def __init__( |
| 176 | self, |
| 177 | query_dim, |
| 178 | context_dim=None, |
| 179 | heads=8, |
| 180 | dim_head=64, |
| 181 | dropout=0.0, |
| 182 | backend=None, |
| 183 | ): |
| 184 | super().__init__() |
| 185 | inner_dim = dim_head * heads |
| 186 | context_dim = default(context_dim, query_dim) |
| 187 | |
| 188 | self.scale = dim_head**-0.5 |
| 189 | self.heads = heads |
| 190 | |
| 191 | self.to_q = nn.Linear(query_dim, inner_dim, bias=False) |
| 192 | self.to_k = nn.Linear(context_dim, inner_dim, bias=False) |
| 193 | self.to_v = nn.Linear(context_dim, inner_dim, bias=False) |
| 194 | |
| 195 | self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout)) |
| 196 | self.backend = backend |
| 197 | |
| 198 | def forward( |
| 199 | self, |
Callers
nothing calls this directly
Tested by
no test coverage detected