Create vqgan_arch.py

vqgan_arch.py

@@ -0,0 +1,102 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from basicsr.utils.registry import ARCH_REGISTRY
+
+
+class ResBlock(nn.Module):
+    def __init__(self, in_channel, out_channel, downsample=False):
+        super().__init__()
+        self.downsample = downsample
+        self.conv1 = nn.Conv2d(in_channel, out_channel, 3, 1, 1, bias=False)
+        self.conv2 = nn.Conv2d(out_channel, out_channel, 3, 1, 1, bias=False)
+        self.norm1 = nn.BatchNorm2d(out_channel)
+        self.norm2 = nn.BatchNorm2d(out_channel)
+        self.relu = nn.LeakyReLU(0.2, inplace=True)
+        if downsample:
+            self.down = nn.Conv2d(in_channel, out_channel, 3, 2, 1)
+
+    def forward(self, x):
+        identity = x
+        out = self.conv1(x)
+        out = self.norm1(out)
+        out = self.relu(out)
+        out = self.conv2(out)
+        out = self.norm2(out)
+        if self.downsample:
+            identity = self.down(x)
+        out += identity
+        out = self.relu(out)
+        return out
+
+
+class VectorQuantizer(nn.Module):
+    def __init__(self, n_e, e_dim, beta=0.25):
+        super().__init__()
+        self.n_e = n_e
+        self.e_dim = e_dim
+        self.beta = beta
+        self.embedding = nn.Embedding(self.n_e, self.e_dim)
+        self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e)
+
+    def forward(self, z):
+        z = z.permute(0, 2, 3, 1).contiguous()
+        z_flattened = z.view(-1, self.e_dim)
+        d = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \
+            torch.sum(self.embedding.weight**2, dim=1) - 2 * \
+            torch.matmul(z_flattened, self.embedding.weight.t())
+        min_encoding_indices = torch.argmin(d, dim=1)
+        z_q = self.embedding(min_encoding_indices).view(z.shape)
+        loss = torch.mean((z_q.detach() - z) ** 2) + self.beta * torch.mean((z_q - z.detach()) ** 2)
+        z_q = z + (z_q - z).detach()
+        z_q = z_q.permute(0, 3, 1, 2).contiguous()
+        return z_q, loss, min_encoding_indices
+
+    def get_codebook_feat(self, indices, shape):
+        z_q = self.embedding(indices)
+        if len(z_q.shape) == 2:
+            z_q = z_q.view(shape[0], shape[1], shape[2], shape[3])
+        z_q = z_q.permute(0, 3, 1, 2).contiguous()
+        return z_q
+
+
+@ARCH_REGISTRY.register()
+class VQAutoEncoder(nn.Module):
+    def __init__(self, in_channel, channel, down_factor, downsample_method, downsample_steps, z_channels, codebook_size):
+        super().__init__()
+        self.encoder = nn.ModuleList()
+        self.decoder = nn.ModuleList()
+        curr_channel = in_channel
+        for i in range(downsample_steps):
+            next_channel = channel * down_factor[i]
+            down = i < len(down_factor) - 1
+            self.encoder.append(ResBlock(curr_channel, next_channel, downsample=down))
+            curr_channel = next_channel
+        self.encoder.append(nn.Conv2d(curr_channel, z_channels, 3, 1, 1))
+        self.quantize = VectorQuantizer(codebook_size, z_channels)
+        self.decoder.append(nn.Conv2d(z_channels, curr_channel, 3, 1, 1))
+        for i in range(downsample_steps - 1, -1, -1):
+            next_channel = channel * down_factor[i]
+            up = i > 0
+            self.decoder.append(ResBlock(curr_channel, next_channel, downsample=False))
+            if up:
+                self.decoder.append(nn.Upsample(scale_factor=down_factor[i], mode=downsample_method))
+            curr_channel = next_channel
+        self.decoder.append(nn.Conv2d(curr_channel, in_channel, 3, 1, 1))
+
+    def encode(self, x):
+        for module in self.encoder:
+            x = module(x)
+        return x
+
+    def decode(self, z):
+        for module in self.decoder:
+            z = module(z)
+        return z
+
+    def forward(self, x):
+        z = self.encode(x)
+        z_q, quant_loss, _ = self.quantize(z)
+        out = self.decode(z_q)
+        return out, quant_loss