(self, x)
| 406 | |
| 407 | |
| 408 | def forward(self, x): |
| 409 | #assert x.shape[2] == x.shape[3] == self.resolution, "{}, {}, {}".format(x.shape[2], x.shape[3], self.resolution) |
| 410 | |
| 411 | # timestep embedding |
| 412 | temb = None |
| 413 | |
| 414 | # downsampling |
| 415 | hs = [self.conv_in(x)] |
| 416 | for i_level in range(self.num_resolutions): |
| 417 | for i_block in range(self.num_res_blocks): |
| 418 | h = self.down[i_level].block[i_block](hs[-1], temb) |
| 419 | if len(self.down[i_level].attn) > 0: |
| 420 | h = self.down[i_level].attn[i_block](h) |
| 421 | hs.append(h) |
| 422 | if i_level != self.num_resolutions-1: |
| 423 | hs.append(self.down[i_level].downsample(hs[-1])) |
| 424 | |
| 425 | # middle |
| 426 | h = hs[-1] |
| 427 | h = self.mid.block_1(h, temb) |
| 428 | h = self.mid.attn_1(h) |
| 429 | h = self.mid.block_2(h, temb) |
| 430 | |
| 431 | # end |
| 432 | h = self.norm_out(h) |
| 433 | h = nonlinearity(h) |
| 434 | h = self.conv_out(h) |
| 435 | return h |
| 436 | |
| 437 | class DummyDecoder(nn.Module): |
| 438 | def __init__(self, **ignorekwargs): |
nothing calls this directly
no test coverage detected