(model, controller)
| 172 | |
| 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 | |
| 213 | def __call__(self, *args): |
| 214 | return args[0] |
| 215 | |
| 216 | def __init__(self): |
| 217 | self.num_att_layers = 0 |
| 218 | |
| 219 | if controller is None: |
| 220 | controller = DummyController() |
| 221 | |
| 222 | def register_recr(net_, count, place_in_unet): |
| 223 | if net_.__class__.__name__ == 'CrossAttention': |
| 224 | net_.forward = ca_forward(net_, place_in_unet) |
| 225 | return count + 1 |
| 226 | elif hasattr(net_, 'children'): |
| 227 | for net__ in net_.children(): |
| 228 | count = register_recr(net__, count, place_in_unet) |
| 229 | return count |
| 230 | |
| 231 | cross_att_count = 0 |
no test coverage detected