| 780 | return x |
| 781 | |
| 782 | class FirstStagePostProcessor(nn.Module): |
| 783 | |
| 784 | def __init__(self, ch_mult:list, in_channels, |
| 785 | pretrained_model:nn.Module=None, |
| 786 | reshape=False, |
| 787 | n_channels=None, |
| 788 | dropout=0., |
| 789 | pretrained_config=None): |
| 790 | super().__init__() |
| 791 | if pretrained_config is None: |
| 792 | assert pretrained_model is not None, 'Either "pretrained_model" or "pretrained_config" must not be None' |
| 793 | self.pretrained_model = pretrained_model |
| 794 | else: |
| 795 | assert pretrained_config is not None, 'Either "pretrained_model" or "pretrained_config" must not be None' |
| 796 | self.instantiate_pretrained(pretrained_config) |
| 797 | |
| 798 | self.do_reshape = reshape |
| 799 | |
| 800 | if n_channels is None: |
| 801 | n_channels = self.pretrained_model.encoder.ch |
| 802 | |
| 803 | self.proj_norm = Normalize(in_channels,num_groups=in_channels//2) |
| 804 | self.proj = nn.Conv2d(in_channels,n_channels,kernel_size=3, |
| 805 | stride=1,padding=1) |
| 806 | |
| 807 | blocks = [] |
| 808 | downs = [] |
| 809 | ch_in = n_channels |
| 810 | for m in ch_mult: |
| 811 | blocks.append(ResnetBlock(in_channels=ch_in,out_channels=m*n_channels,dropout=dropout)) |
| 812 | ch_in = m * n_channels |
| 813 | downs.append(Downsample(ch_in, with_conv=False)) |
| 814 | |
| 815 | self.model = nn.ModuleList(blocks) |
| 816 | self.downsampler = nn.ModuleList(downs) |
| 817 | |
| 818 | |
| 819 | def instantiate_pretrained(self, config): |
| 820 | model = instantiate_from_config(config) |
| 821 | self.pretrained_model = model.eval() |
| 822 | # self.pretrained_model.train = False |
| 823 | for param in self.pretrained_model.parameters(): |
| 824 | param.requires_grad = False |
| 825 | |
| 826 | |
| 827 | @torch.no_grad() |
| 828 | def encode_with_pretrained(self,x): |
| 829 | c = self.pretrained_model.encode(x) |
| 830 | if isinstance(c, DiagonalGaussianDistribution): |
| 831 | c = c.mode() |
| 832 | return c |
| 833 | |
| 834 | def forward(self,x): |
| 835 | z_fs = self.encode_with_pretrained(x) |
| 836 | z = self.proj_norm(z_fs) |
| 837 | z = self.proj(z) |
| 838 | z = nonlinearity(z) |
| 839 |
nothing calls this directly
no outgoing calls
no test coverage detected