Add AutoEncoder (VQVAE)

models/blocks.py

@@ -398,3 +398,111 @@ class UpBlockUnet(nn.Module):
                 out = out + out_attn
 
         return out
+
+
+class UpBlock(nn.Module):
+    r"""
+    Up conv block with attention.
+    Sequence of following blocks
+    1. Upsample
+    1. Concatenate Down block output
+    2. Resnet block with time embedding
+    3. Attention Block
+    """
+
+    def __init__(self, in_channels, out_channels, t_emb_dim,
+                 up_sample, num_heads, num_layers, attn, norm_channels):
+        super().__init__()
+        self.num_layers = num_layers
+        self.up_sample = up_sample
+        self.t_emb_dim = t_emb_dim
+        self.attn = attn
+        self.resnet_conv_first = nn.ModuleList(
+            [
+                nn.Sequential(
+                    nn.GroupNorm(norm_channels, in_channels if i ==
+                                 0 else out_channels),
+                    nn.SiLU(),
+                    nn.Conv2d(in_channels if i == 0 else out_channels, out_channels, kernel_size=3, stride=1,
+                              padding=1),
+                )
+                for i in range(num_layers)
+            ]
+        )
+
+        if self.t_emb_dim is not None:
+            self.t_emb_layers = nn.ModuleList([
+                nn.Sequential(
+                    nn.SiLU(),
+                    nn.Linear(t_emb_dim, out_channels)
+                )
+                for _ in range(num_layers)
+            ])
+
+        self.resnet_conv_second = nn.ModuleList(
+            [
+                nn.Sequential(
+                    nn.GroupNorm(norm_channels, out_channels),
+                    nn.SiLU(),
+                    nn.Conv2d(out_channels, out_channels,
+                              kernel_size=3, stride=1, padding=1),
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        if self.attn:
+            self.attention_norms = nn.ModuleList(
+                [
+                    nn.GroupNorm(norm_channels, out_channels)
+                    for _ in range(num_layers)
+                ]
+            )
+
+            self.attentions = nn.ModuleList(
+                [
+                    nn.MultiheadAttention(
+                        out_channels, num_heads, batch_first=True)
+                    for _ in range(num_layers)
+                ]
+            )
+
+        self.residual_input_conv = nn.ModuleList(
+            [
+                nn.Conv2d(in_channels if i == 0 else out_channels,
+                          out_channels, kernel_size=1)
+                for i in range(num_layers)
+            ]
+        )
+        self.up_sample_conv = nn.ConvTranspose2d(in_channels, in_channels,
+                                                 4, 2, 1) \
+            if self.up_sample else nn.Identity()
+
+    def forward(self, x, out_down=None, t_emb=None):
+        # Upsample
+        x = self.up_sample_conv(x)
+
+        # Concat with Downblock output
+        if out_down is not None:
+            x = torch.cat([x, out_down], dim=1)
+
+        out = x
+        for i in range(self.num_layers):
+            # Resnet Block
+            resnet_input = out
+            out = self.resnet_conv_first[i](out)
+            if self.t_emb_dim is not None:
+                out = out + self.t_emb_layers[i](t_emb)[:, :, None, None]
+            out = self.resnet_conv_second[i](out)
+            out = out + self.residual_input_conv[i](resnet_input)
+
+            # Self Attention
+            if self.attn:
+                batch_size, channels, h, w = out.shape
+                in_attn = out.reshape(batch_size, channels, h * w)
+                in_attn = self.attention_norms[i](in_attn)
+                in_attn = in_attn.transpose(1, 2)
+                out_attn, _ = self.attentions[i](in_attn, in_attn, in_attn)
+                out_attn = out_attn.transpose(1, 2).reshape(
+                    batch_size, channels, h, w)
+                out = out + out_attn
+        return out

models/vqvae.py

@@ -0,0 +1,156 @@
+import enum
+from sre_compile import dis
+import torch
+import torch.nn as nn
+from models.blocks import DownBlock, UpBlock, MidBlock
+
+
+class VQVAE(nn.Module):
+    def __init__(self, im_channels, model_config):
+        super().__init__()
+        self.down_channels = model_config['down_channels']
+        self.mid_channels = model_config['mid_channels']
+        self.down_sample = model_config['down_sample']
+        self.num_down_layers = model_config['num_down_layers']
+        self.num_up_layers = model_config['num_up_layers']
+        self.num_mid_layers = model_config['num_mid_layers']
+
+        # To disable attn in encoder and decoder blocks
+        self.attn = model_config['attn']
+
+        # Latent Dimension
+        self.z_channels = model_config["z_channels"]
+        self.codebook_size = model_config["codebook_size"]
+        self.norm_channels = model_config["norm_channels"]
+        self.num_heads = model_config["num_heads"]
+
+        assert self.mid_channels[0] == self.down_channels[-1]
+        assert self.mid_channels[-1] == self.down_channels[-1]
+        assert len(self.down_sample) == len(self.down_channels) - 1
+        assert len(self.attns) == len(self.down_channels) - 1
+
+        self.upsample = list(reversed(self.down_sample))
+
+        # Encoder
+        self.encoder_conv_one = nn.Conv2d(
+            im_channels, self.down_channels[0], kernel_size=3, padding=1, stride=1)
+
+        self.encoder_layers = nn.ModuleList([])
+        for i in range(len(self.down_channels) - 1):
+            self.encoder_layers.append(DownBlock(self.down_channels[i], self.down_channels[i+1],
+                                                 t_emd_dim=None, down_sample=self.down_sample[i],
+                                                 num_heads=self.num_heads, num_layers=self.num_down_layers,
+                                                 attn=self.attns[i], norm_channels=self.norm_channels))
+        self.encode_mid_blocks = nn.ModuleList([])
+        for i in range(len(self.down_channels)-1):
+            self.encode_mid_blocks.append(MidBlock(self.down_channels[i], self.down_channels[i+1],
+                                          t_emb_dim=None, num_heads=self.num_heads, num_layers=self.num_mid_layers,
+                                          norm_dim=self.norm_channels))
+        self.encoder_norm_out = nn.GroupNorm(
+            self.norm_channels, self.down_channels[-1])
+        self.encoder_conv_out = nn.Conv2d(
+            self.down_channels[-1], self.z_channels, kernel_size=3, padding=1)
+
+        # Pre-Quantization Convolution (Before comparing to code blocks to get embedding matrix)
+        self.pre_quant_conv = nn.Conv2d(
+            self.z_channels, self.z_channels, kernel_size=1)
+
+        # Code book
+        self.embedding = nn.Embedding(self.codebook_size, self.z_channels)
+
+        # Decoder
+        self.post_quant_conv = nn.Conv2d(
+            self.z_channels, self.z_channels, kernel_size=1)
+        self.decoder_conv_out = nn.Conv2d(
+            self.z_channels, self.mid_channels[-1], kernel_size=3, padding=1)
+
+        # Midblock + UpBlock
+        self.decode_mids = nn.ModuleList([])
+        for i in reversed(range(1, len(self.mid_channels))):
+            self.decode_mids.append(MidBlock(self.mid_channels[i], self.mid_channels[i-1],
+                                             t_emb_dim=None, num_heads=self.num_heads,
+                                             num_layers=self.num_mid_layers,
+                                             norm_dim=self.norm_channels))
+        self.decoder_layers = nn.ModuleList([])
+        for i in reversed(range(1, len(self.down_channels))):
+            self.decoder_layers.append(UpBlock(self.down_channels[i], self.down_channels[i-1],
+                                       t_emb_dim=None, up_sample=self.down_sample[i-1], num_heads=self.num_heads,
+                                       num_layers=self.num_up_layers,
+                                       attn=self.attn[i-1],
+                                       norm_channels=self.norm_channels))
+
+        self.decoder_norm_out = nn.GroupNorm(
+            self.norm_channels, self.down_channels[0])
+        self.decoder_conv_out = nn.Conv2d(
+            self.down_channels[0], im_channels, kernel_size=3, padding=1)
+
+    def quantize(self, x):
+        B, C, H, W = x.shape,
+
+        # B,C,H,W -> B,H,W,C
+        x = x.permute(0, 2, 3, 1)
+
+        # B,H,W,C -> B, H*W, C
+        x = x.reshape(x.size(0), -1, x.size(-1))
+
+        # Find nearest neighbours/codebook vectors
+        # Distance between  B,H*W,C and B,K,C
+        dist = torch.cdist(
+            x, self.embedding.weight[None, :].repeat((x.size(0), 1, 1)))
+
+        min_encoding_indices = torch.argmin(dist, dim=-1)
+
+        # Replace encoder output with codebook vector
+        quant_out = torch.index_select(
+            self.embedding.weight, 0, min_encoding_indices.view(-1))
+
+        # x -> B*H*W,C
+        x = x.reshape((-1, x.size(-1)))
+        commitment_loss = torch.mean((quant_out.detach() - x) ** 2)
+        codebook_loss = torch.mean((quant_out - x.detach()) ** 2)
+        quantize_loss = {
+            "codebook_loss": codebook_loss,
+            "commitment_loss": commitment_loss
+        }
+
+        # Straight through estimation
+        quant_out = x - (quant_out - x).detach()
+
+        # quant_out -> B,C,H,W
+        quant_out = quant_out.reshape((B, H, W, C)).permute(0, 3, 1, 2)
+        min_encoding_indices = min_encoding_indices.reshape(
+            (-1, quant_out.size(-2), quant_out.size(-1)))
+
+        return quant_out, quantize_loss, min_encoding_indices
+
+    def encode(self, x):
+        out = self.encoder_conv_one(x)
+        for _, down in enumerate(self.encoder_layers):
+            out = down(out)
+        for mid in self.encode_mid_blocks:
+            out = mid(out)
+        out = self.encoder_norm_out(out)
+        out = nn.SiLU()(out)
+        out = self.encoder_conv_out(out)
+        out = self.pre_quant_conv(out)
+        out, quant_losses, _ = self.quantize(out)
+        return out, quant_losses
+
+    def decode(self, z):
+        out = z
+        out = self.post_quant_conv(out)
+        out = self.decoder_conv_in(out)
+        for mid in self.decode_mids:
+            out = mid(out)
+        for up in self.decoder_layers:
+            out = up(out)
+
+        out = self.decoder_norm_out(out)
+        out = nn.SiLU(out)
+        out = self.decoder_conv_out(out)
+        return out
+
+    def forward(self, x):
+        z, quant_losses = self.encode(x)
+        out = self.decode(z)
+        return out, z, quant_losses