| 173 | |
| 174 | def register_attention_control(model, controller): |
| 175 | def ca_forward(self, place_in_unet): |
| 176 | to_out = self.to_out |
| 177 | if type(to_out) is torch.nn.modules.container.ModuleList: |
| 178 | to_out = self.to_out[0] |
| 179 | else: |
| 180 | to_out = self.to_out |
| 181 | |
| 182 | def forward(x, context=None, mask=None): |
| 183 | batch_size, sequence_length, dim = x.shape |
| 184 | h = self.heads |
| 185 | q = self.to_q(x) |
| 186 | is_cross = context is not None |
| 187 | context = context if is_cross else x |
| 188 | k = self.to_k(context) |
| 189 | v = self.to_v(context) |
| 190 | q = self.reshape_heads_to_batch_dim(q) |
| 191 | k = self.reshape_heads_to_batch_dim(k) |
| 192 | v = self.reshape_heads_to_batch_dim(v) |
| 193 | |
| 194 | sim = torch.einsum("b i d, b j d -> b i j", q, k) * self.scale |
| 195 | |
| 196 | if mask is not None: |
| 197 | mask = mask.reshape(batch_size, -1) |
| 198 | max_neg_value = -torch.finfo(sim.dtype).max |
| 199 | mask = mask[:, None, :].repeat(h, 1, 1) |
| 200 | sim.masked_fill_(~mask, max_neg_value) |
| 201 | |
| 202 | # attention, what we cannot get enough of |
| 203 | attn = sim.softmax(dim=-1) |
| 204 | attn = controller(attn, is_cross, place_in_unet) |
| 205 | out = torch.einsum("b i j, b j d -> b i d", attn, v) |
| 206 | out = self.reshape_batch_dim_to_heads(out) |
| 207 | return to_out(out) |
| 208 | |
| 209 | return forward |
| 210 | |
| 211 | class DummyController: |
| 212 | |