| 177 | return ret |
| 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 |
| 191 | q,k,v = self.to_q(x), self.to_k(context), self.to_v(context) |
| 192 | q = q.reshape(x.shape[0], -1, self.num_heads, self.head_size).permute(0,2,1,3) # (bs, num_heads, time, head_size) |
| 193 | k = k.reshape(x.shape[0], -1, self.num_heads, self.head_size).permute(0,2,3,1) # (bs, num_heads, head_size, time) |
| 194 | v = v.reshape(x.shape[0], -1, self.num_heads, self.head_size).permute(0,2,1,3) # (bs, num_heads, time, head_size) |
| 195 | |
| 196 | score = q.dot(k) * self.scale |
| 197 | weights = score.softmax() # (bs, num_heads, time, time) |
| 198 | attention = weights.dot(v).permute(0,2,1,3) # (bs, time, num_heads, head_size) |
| 199 | |
| 200 | h_ = attention.reshape(shape=(x.shape[0], -1, self.num_heads * self.head_size)) |
| 201 | return h_.sequential(self.to_out) |
| 202 | |
| 203 | class GEGLU: |
| 204 | def __init__(self, dim_in, dim_out): |