Update vqgan_arch.py

vqgan_arch.py

@@ -1,66 +1,52 @@
-import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 from basicsr.utils.registry import ARCH_REGISTRY
 
+class VectorQuantizer(nn.Module):
+    def __init__(self, n_e, e_dim):
+        super().__init__()
+        self.n_e = n_e
+        self.e_dim = e_dim
+        self.embedding = nn.Parameter(torch.randn(n_e, e_dim))
+
+    def get_codebook_feat(self, indices, shape):
+        feat = self.embedding[indices]
+        feat = feat.view(shape[0], shape[1], shape[2], shape[3]).permute(0, 3, 1, 2)
+        return feat
+
+    def forward(self, z):
+        z_q = self.embedding[z]
+        loss = torch.mean((z_q.detach() - z) ** 2) + torch.mean((z_q - z.detach()) ** 2)
+        z_q = z + (z_q - z).detach()
+        return z_q, loss, None
 
 class ResBlock(nn.Module):
-    def __init__(self, in_channel, out_channel, downsample=False):
+    def __init__(self, in_channel, out_channel, downsample=False, downsample_method='nearest'):
         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)
+        self.downsample_method = downsample_method
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channel, out_channel, 3, 1, 1),
+            nn.ReLU(True),
+            nn.Conv2d(out_channel, out_channel, 3, 1, 1)
+        )
+        # Add a projection layer for the identity path if channels or spatial dimensions change
+        self.proj = nn.Conv2d(in_channel, out_channel, 1, 1, 0) if in_channel != out_channel else nn.Identity()
+        self.downsample_identity = nn.AvgPool2d(kernel_size=2, stride=2) if downsample else nn.Identity()
 
     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)
+        out = self.conv(x)
         if self.downsample:
-            identity = self.down(x)
+            out = F.interpolate(out, scale_factor=0.5, mode=self.downsample_method)
+        # Adjust the identity path to match out's dimensions
+        identity = self.proj(identity)  # Match channel dimensions
+        identity = self.downsample_identity(identity)  # Match spatial dimensions if downsampling
         out += identity
-        out = self.relu(out)
+        out = F.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):
@@ -71,7 +57,7 @@ class VQAutoEncoder(nn.Module):
         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))
+            self.encoder.append(ResBlock(curr_channel, next_channel, downsample=down, downsample_method=downsample_method))
             curr_channel = next_channel
         self.encoder.append(nn.Conv2d(curr_channel, z_channels, 3, 1, 1))
         self.quantize = VectorQuantizer(codebook_size, z_channels)