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