| 144 | |
| 145 | class CrossAttention(nn.Module): |
| 146 | def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.): |
| 147 | super().__init__() |
| 148 | inner_dim = dim_head * heads |
| 149 | context_dim = default(context_dim, query_dim) |
| 150 | |
| 151 | self.scale = dim_head ** -0.5 |
| 152 | self.heads = heads |
| 153 | |
| 154 | self.to_q = nn.Linear(query_dim, inner_dim, bias=False) |
| 155 | self.to_k = nn.Linear(context_dim, inner_dim, bias=False) |
| 156 | self.to_v = nn.Linear(context_dim, inner_dim, bias=False) |
| 157 | |
| 158 | self.to_out = nn.Sequential( |
| 159 | nn.Linear(inner_dim, query_dim), |
| 160 | nn.Dropout(dropout) |
| 161 | ) |
| 162 | |
| 163 | def forward(self, x, context=None, mask=None): |
| 164 | h = self.heads |