(self, x)
| 492 | self.output_ch = 2*z_channels if double_z else z_channels |
| 493 | |
| 494 | def forward(self, x): |
| 495 | #assert x.shape[2] == x.shape[3] == self.resolution, "{}, {}, {}".format(x.shape[2], x.shape[3], self.resolution) |
| 496 | |
| 497 | # timestep embedding |
| 498 | temb = None |
| 499 | |
| 500 | # downsampling |
| 501 | hs = [self.conv_in(x)] |
| 502 | for i_level in range(self.num_resolutions): |
| 503 | for i_block in range(self.num_res_blocks): |
| 504 | h = self.down[i_level].block[i_block](hs[-1], temb) |
| 505 | if len(self.down[i_level].attn) > 0: |
| 506 | h = self.down[i_level].attn[i_block](h) |
| 507 | hs.append(h) |
| 508 | if i_level != self.num_resolutions-1: |
| 509 | hs.append(self.down[i_level].downsample(hs[-1])) |
| 510 | |
| 511 | # middle |
| 512 | h = hs[-1] |
| 513 | h = self.mid.block_1(h, temb) |
| 514 | h = self.mid.attn_1(h) |
| 515 | h = self.mid.block_2(h, temb) |
| 516 | |
| 517 | # end |
| 518 | h = self.norm_out(h) |
| 519 | h = nonlinearity(h) |
| 520 | h = self.conv_out(h) |
| 521 | return h |
| 522 | |
| 523 | |
| 524 | class DummyDecoder(nn.Module): |
nothing calls this directly
no test coverage detected