| 94 | |
| 95 | class CrossAttention(nn.Module): |
| 96 | def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0): |
| 97 | super().__init__() |
| 98 | inner_dim = dim_head * heads |
| 99 | context_dim = default(context_dim, query_dim) |
| 100 | |
| 101 | self.scale = dim_head**-0.5 |
| 102 | self.heads = heads |
| 103 | |
| 104 | self.to_q = nn.Linear(query_dim, inner_dim, bias=False) |
| 105 | self.to_k = nn.Linear(context_dim, inner_dim, bias=False) |
| 106 | self.to_v = nn.Linear(context_dim, inner_dim, bias=False) |
| 107 | |
| 108 | self.to_out = nn.Sequential( |
| 109 | nn.Linear(inner_dim, query_dim), nn.Dropout(dropout) |
| 110 | ) |
| 111 | |
| 112 | def forward(self, x, text, mask=None): |
| 113 | B, L, C = x.shape |