(self, x, context=None)
| 319 | self.use_linear = use_linear |
| 320 | |
| 321 | def forward(self, x, context=None): |
| 322 | # note: if no context is given, cross-attention defaults to self-attention |
| 323 | if not isinstance(context, list): |
| 324 | context = [context] |
| 325 | b, c, h, w = x.shape |
| 326 | x_in = x |
| 327 | x = self.norm(x) |
| 328 | if not self.use_linear: |
| 329 | x = self.proj_in(x) |
| 330 | x = rearrange(x, 'b c h w -> b (h w) c').contiguous() |
| 331 | if self.use_linear: |
| 332 | x = self.proj_in(x) |
| 333 | for i, block in enumerate(self.transformer_blocks): |
| 334 | x = block(x, context=context[i]) |
| 335 | if self.use_linear: |
| 336 | x = self.proj_out(x) |
| 337 | x = rearrange(x, 'b (h w) c -> b c h w', h=h, w=w).contiguous() |
| 338 | if not self.use_linear: |
| 339 | x = self.proj_out(x) |
| 340 | return x + x_in |
| 341 |
nothing calls this directly
no outgoing calls
no test coverage detected