(self, x)
| 430 | padding=1) |
| 431 | |
| 432 | def forward(self, x): |
| 433 | # timestep embedding |
| 434 | temb = None |
| 435 | |
| 436 | # print(f'encoder-input={x.shape}') |
| 437 | # downsampling |
| 438 | hs = [self.conv_in(x)] |
| 439 | # print(f'encoder-conv in feat={hs[0].shape}') |
| 440 | for i_level in range(self.num_resolutions): |
| 441 | for i_block in range(self.num_res_blocks): |
| 442 | h = self.down[i_level].block[i_block](hs[-1], temb) |
| 443 | # print(f'encoder-down feat={h.shape}') |
| 444 | if len(self.down[i_level].attn) > 0: |
| 445 | h = self.down[i_level].attn[i_block](h) |
| 446 | hs.append(h) |
| 447 | if i_level != self.num_resolutions-1: |
| 448 | # print(f'encoder-downsample (input)={hs[-1].shape}') |
| 449 | hs.append(self.down[i_level].downsample(hs[-1])) |
| 450 | # print(f'encoder-downsample (output)={hs[-1].shape}') |
| 451 | |
| 452 | # middle |
| 453 | h = hs[-1] |
| 454 | h = self.mid.block_1(h, temb) |
| 455 | # print(f'encoder-mid1 feat={h.shape}') |
| 456 | h = self.mid.attn_1(h) |
| 457 | h = self.mid.block_2(h, temb) |
| 458 | # print(f'encoder-mid2 feat={h.shape}') |
| 459 | |
| 460 | # end |
| 461 | h = self.norm_out(h) |
| 462 | h = nonlinearity(h) |
| 463 | h = self.conv_out(h) |
| 464 | # print(f'end feat={h.shape}') |
| 465 | return h |
| 466 | |
| 467 | |
| 468 | class Decoder(nn.Module): |
nothing calls this directly
no test coverage detected