(self, x, emb)
| 254 | |
| 255 | |
| 256 | def _forward(self, x, emb): |
| 257 | if self.updown: |
| 258 | in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1] |
| 259 | h = in_rest(x) |
| 260 | h = self.h_upd(h) |
| 261 | x = self.x_upd(x) |
| 262 | h = in_conv(h) |
| 263 | else: |
| 264 | h = self.in_layers(x) |
| 265 | emb_out = self.emb_layers(emb).type(h.dtype) |
| 266 | while len(emb_out.shape) < len(h.shape): |
| 267 | emb_out = emb_out[..., None] |
| 268 | if self.use_scale_shift_norm: |
| 269 | out_norm, out_rest = self.out_layers[0], self.out_layers[1:] |
| 270 | scale, shift = th.chunk(emb_out, 2, dim=1) |
| 271 | h = out_norm(h) * (1 + scale) + shift |
| 272 | h = out_rest(h) |
| 273 | else: |
| 274 | h = h + emb_out |
| 275 | h = self.out_layers(h) |
| 276 | return self.skip_connection(x) + h |
| 277 | |
| 278 | |
| 279 | class AttentionBlock(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected