(self, x, use_cp=False)
| 385 | ) |
| 386 | |
| 387 | def forward(self, x, use_cp=False): |
| 388 | # assert x.shape[2] == x.shape[3] == self.resolution, "{}, {}, {}".format(x.shape[2], x.shape[3], self.resolution) |
| 389 | # timestep embedding |
| 390 | temb = None |
| 391 | |
| 392 | # downsampling |
| 393 | hs = [self.conv_in(x)] |
| 394 | for i_level in range(self.num_resolutions): |
| 395 | for i_block in range(self.num_res_blocks): |
| 396 | h = self.down[i_level].block[i_block](hs[-1], temb) |
| 397 | if len(self.down[i_level].attn) > 0: |
| 398 | h = self.down[i_level].attn[i_block](h) |
| 399 | hs.append(h) |
| 400 | if i_level != self.num_resolutions - 1: |
| 401 | hs.append(self.down[i_level].downsample(hs[-1])) |
| 402 | |
| 403 | # middle |
| 404 | h = hs[-1] |
| 405 | h = self.mid.block_1(h, temb) |
| 406 | # h = self.mid.attn_1(h) |
| 407 | h = self.mid.block_2(h, temb) |
| 408 | |
| 409 | # end |
| 410 | h = self.norm_out(h) |
| 411 | h = nonlinearity(h) |
| 412 | h = self.conv_out(h) |
| 413 | return h |
nothing calls this directly
no test coverage detected