encoder.py

import torch.nn as nn
from helper import ResidualBlock, NonLocalBlock, DownSampleBlock, UpSampleBlock, GroupNorm, Swish, LinearCombo


"""
Code for class Encoder adapted from https://github.com/dome272/VQGAN-pytorch/blob/main/encoder.py
"""
class Encoder(nn.Module):
    def __init__(self, args):
        super(Encoder, self).__init__()
        channels = args.encoder_channels
        resolution = args.encoder_start_resolution
        layers = [nn.Conv2d(args.image_channels, channels[0], 3, 1, 1)]
        for i in range(len(channels)-1):
            in_channels = channels[i]
            out_channels = channels[i + 1]
            for j in range(args.encoder_num_res_blocks):
                layers.append(ResidualBlock(in_channels, out_channels))
                in_channels = out_channels
                if resolution in args.encoder_attn_resolutions:
                    layers.append(NonLocalBlock(in_channels))
            if i != len(channels)-2:
                layers.append(DownSampleBlock(channels[i+1]))
                resolution //= 2
        layers.append(ResidualBlock(channels[-1], channels[-1]))
        layers.append(NonLocalBlock(channels[-1]))
        layers.append(ResidualBlock(channels[-1], channels[-1]))
        layers.append(GroupNorm(channels[-1]))
        layers.append(Swish())
        layers.append(nn.Conv2d(channels[-1], args.latent_dim, 3, 1, 1))
        self.model = nn.Sequential(*layers)

    def forward(self, x):
        return self.model(x)
    

# Encoder for CVQGAN option
class CondEncoder(nn.Module):
    def __init__(self, args):
        super(CondEncoder, self).__init__()
        self.c_fmap_dim = args.c_fmap_dim
        self.model = nn.Sequential(
            LinearCombo(args.c_input_dim, args.c_hidden_dim),
            LinearCombo(args.c_hidden_dim, args.c_hidden_dim),
            nn.Linear(args.c_hidden_dim, args.c_latent_dim*args.c_fmap_dim**2)
        )
    
    def forward(self, x):
        encoded = self.model(x)
        s = encoded.shape
        return encoded.view(s[0], s[1]//self.c_fmap_dim**2, self.c_fmap_dim, self.c_fmap_dim)