(self, query_dim, context_dim, n_heads, d_head)
| 178 | |
| 179 | class CrossAttention: |
| 180 | def __init__(self, query_dim, context_dim, n_heads, d_head): |
| 181 | self.to_q = Linear(query_dim, n_heads*d_head, bias=False) |
| 182 | self.to_k = Linear(context_dim, n_heads*d_head, bias=False) |
| 183 | self.to_v = Linear(context_dim, n_heads*d_head, bias=False) |
| 184 | self.scale = d_head ** -0.5 |
| 185 | self.num_heads = n_heads |
| 186 | self.head_size = d_head |
| 187 | self.to_out = [Linear(n_heads*d_head, query_dim)] |
| 188 | |
| 189 | def __call__(self, x, context=None): |
| 190 | context = x if context is None else context |
nothing calls this directly
no outgoing calls
no test coverage detected