Upload modules.py with huggingface_hub

modules.py

@@ -0,0 +1,365 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from fast_pytorch_kmeans import KMeans
+from torch import einsum
+import torch.distributed as dist
+from einops import rearrange
+
+
+def get_timestep_embedding(timesteps, embedding_dim):
+    """
+    This matches the implementation in Denoising Diffusion Probabilistic Models:
+    From Fairseq.
+    Build sinusoidal embeddings.
+    This matches the implementation in tensor2tensor, but differs slightly
+    from the description in Section 3.5 of "Attention Is All You Need".
+    """
+    assert len(timesteps.shape) == 1
+
+    half_dim = embedding_dim // 2
+    emb = math.log(10000) / (half_dim - 1)
+    emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
+    emb = emb.to(device=timesteps.device)
+    emb = timesteps.float()[:, None] * emb[None, :]
+    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+    if embedding_dim % 2 == 1:  # zero pad
+        emb = torch.nn.functional.pad(emb, (0,1,0,0))
+    return emb
+
+
+def nonlinearity(x):
+    # swish
+    return x*torch.sigmoid(x)
+
+
+def Normalize(in_channels):
+    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            self.conv = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, x):
+        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
+        if self.with_conv:
+            x = self.conv(x)
+        return x
+
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels, with_conv):
+        super().__init__()
+        self.with_conv = with_conv
+        if self.with_conv:
+            # no asymmetric padding in torch conv, must do it ourselves
+            self.conv = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=3,
+                                        stride=2,
+                                        padding=0)
+
+    def forward(self, x):
+        if self.with_conv:
+            pad = (0,1,0,1)
+            x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
+            x = self.conv(x)
+        else:
+            x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
+        return x
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False, dropout):
+        super().__init__()
+        self.in_channels = in_channels
+        out_channels = in_channels if out_channels is None else out_channels
+        self.out_channels = out_channels
+        self.use_conv_shortcut = conv_shortcut
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = torch.nn.Conv2d(in_channels,
+                                     out_channels,
+                                     kernel_size=3,
+                                     stride=1,
+                                     padding=1)
+        self.norm2 = Normalize(out_channels)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.conv2 = torch.nn.Conv2d(out_channels,
+                                     out_channels,
+                                     kernel_size=3,
+                                     stride=1,
+                                     padding=1)
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                self.conv_shortcut = torch.nn.Conv2d(in_channels,
+                                                     out_channels,
+                                                     kernel_size=3,
+                                                     stride=1,
+                                                     padding=1)
+            else:
+                self.nin_shortcut = torch.nn.Conv2d(in_channels,
+                                                    out_channels,
+                                                    kernel_size=1,
+                                                    stride=1,
+                                                    padding=0)
+
+    def forward(self, x):
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+
+        h = self.norm2(h)
+        h = nonlinearity(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            if self.use_conv_shortcut:
+                x = self.conv_shortcut(x)
+            else:
+                x = self.nin_shortcut(x)
+
+        return x+h
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels):
+        super().__init__()
+        self.in_channels = in_channels
+
+        self.norm = Normalize(in_channels)
+        self.q = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.k = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.v = torch.nn.Conv2d(in_channels,
+                                 in_channels,
+                                 kernel_size=1,
+                                 stride=1,
+                                 padding=0)
+        self.proj_out = torch.nn.Conv2d(in_channels,
+                                        in_channels,
+                                        kernel_size=1,
+                                        stride=1,
+                                        padding=0)
+
+
+    def forward(self, x):
+        h_ = x
+        h_ = self.norm(h_)
+        q = self.q(h_)
+        k = self.k(h_)
+        v = self.v(h_)
+
+        # compute attention
+        b,c,h,w = q.shape
+        q = q.reshape(b,c,h*w)
+        q = q.permute(0,2,1)   # b,hw,c
+        k = k.reshape(b,c,h*w) # b,c,hw
+        w_ = torch.bmm(q,k)     # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
+        w_ = w_ * (int(c)**(-0.5))
+        w_ = torch.nn.functional.softmax(w_, dim=2)
+
+        # attend to values
+        v = v.reshape(b,c,h*w)
+        w_ = w_.permute(0,2,1)   # b,hw,hw (first hw of k, second of q)
+        h_ = torch.bmm(v,w_)     # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
+        h_ = h_.reshape(b,c,h,w)
+
+        h_ = self.proj_out(h_)
+
+        return x+h_
+
+
+class Swish(nn.Module):
+    def forward(self, x):
+        return x * torch.sigmoid(x)
+
+
+class Encoder(nn.Module):
+    """
+    Encoder of VQ-GAN to map input batch of images to latent space.
+    Dimension Transformations:
+    3x256x256 --Conv2d--> 32x256x256
+    for loop:
+              --ResBlock--> 64x256x256 --DownBlock--> 64x128x128
+              --ResBlock--> 128x128x128 --DownBlock--> 128x64x64
+              --ResBlock--> 256x64x64  --DownBlock--> 256x32x32
+              --ResBlock--> 512x32x32
+    --ResBlock--> 512x32x32
+    --NonLocalBlock--> 512x32x32
+    --ResBlock--> 512x32x32
+    --GroupNorm-->
+    --Swish-->
+    --Conv2d-> 256x32x32
+    """
+
+    def __init__(self, in_channels=3, channels=[128, 128, 128, 256, 512, 512], attn_resolutions=[32], resolution=512, dropout=0.0, num_res_blocks=2, z_channels=256, **kwargs):
+        super(Encoder, self).__init__()
+        layers = [nn.Conv2d(in_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(num_res_blocks):
+                layers.append(ResnetBlock(in_channels=in_channels, out_channels=out_channels, dropout=0.0))
+                in_channels = out_channels
+                if resolution in attn_resolutions:
+                    layers.append(AttnBlock(in_channels))
+            if i < len(channels) - 2:
+                layers.append(Downsample(channels[i + 1], with_conv=True))
+                resolution //= 2
+        layers.append(ResnetBlock(in_channels=channels[-1], out_channels=channels[-1], dropout=0.0))
+        layers.append(AttnBlock(channels[-1]))
+        layers.append(ResnetBlock(in_channels=channels[-1], out_channels=channels[-1], dropout=0.0))
+        layers.append(Normalize(channels[-1]))
+        layers.append(Swish())
+        layers.append(nn.Conv2d(channels[-1], z_channels, 3, 1, 1))
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class Decoder(nn.Module):
+    def __init__(self, out_channels=3, channels=[128, 128, 128, 256, 512, 512], attn_resolutions=[32], resolution=512, dropout=0.0, num_res_blocks=2, z_channels=256, **kwargs):
+        super(Decoder, self).__init__()
+        ch_mult = channels[1:]
+        num_resolutions = len(ch_mult)
+        block_in = ch_mult[num_resolutions - 1]
+        curr_res = resolution// 2 ** (num_resolutions - 1)
+
+        layers = [nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1),
+                  ResnetBlock(in_channels=block_in, out_channels=block_in, dropout=0.0),
+                  AttnBlock(block_in),
+                  ResnetBlock(in_channels=block_in, out_channels=block_in, dropout=0.0)
+                  ]
+
+        for i in reversed(range(num_resolutions)):
+            block_out = ch_mult[i]
+            for i_block in range(num_res_blocks+1):
+                layers.append(ResnetBlock(in_channels=block_in, out_channels=block_out, dropout=0.0))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    layers.append(AttnBlock(block_in))
+            if i > 0:
+                layers.append(Upsample(block_in, with_conv=True))
+            curr_res = curr_res * 2
+
+        layers.append(Normalize(block_in))
+        layers.append(Swish())
+        layers.append(nn.Conv2d(block_in, out_channels, kernel_size=3, stride=1, padding=1))
+
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class Codebook(nn.Module):
+    """
+    Improved version over VectorQuantizer, can be used as a drop-in replacement. Mostly
+    avoids costly matrix multiplications and allows for post-hoc remapping of indices.
+    """
+    def __init__(self, codebook_size, codebook_dim, beta, init_steps=2000, reservoir_size=2e5):
+        super().__init__()
+        self.codebook_size = codebook_size
+        self.codebook_dim = codebook_dim
+        self.beta = beta
+
+        self.embedding = nn.Embedding(self.codebook_size, self.codebook_dim)
+        self.embedding.weight.data.uniform_(-1.0 / self.codebook_size, 1.0 / self.codebook_size)
+
+        self.q_start_collect, self.q_init, self.q_re_end, self.q_re_step = init_steps, init_steps * 3, init_steps * 30, init_steps // 2
+        self.q_counter = 0
+        self.reservoir_size = int(reservoir_size)
+        self.reservoir = None
+
+    def forward(self, z):
+        z = rearrange(z, 'b c h w -> b h w c').contiguous()
+        batch_size = z.size(0)
+        z_flattened = z.view(-1, self.codebook_dim)
+        if self.training:
+            self.q_counter += 1
+            # x_flat = x.permute(0, 2, 3, 1).reshape(-1, z.shape(1))
+            if self.q_counter > self.q_start_collect:
+                z_new = z_flattened.clone().detach().view(batch_size, -1, self.codebook_dim)
+                z_new = z_new[:, torch.randperm(z_new.size(1))][:, :10].reshape(-1, self.codebook_dim)
+                self.reservoir = z_new if self.reservoir is None else torch.cat([self.reservoir, z_new], dim=0)
+                self.reservoir = self.reservoir[torch.randperm(self.reservoir.size(0))[:self.reservoir_size]].detach()
+            if self.q_counter < self.q_init:
+                z_q = rearrange(z, 'b h w c -> b c h w').contiguous()
+                return z_q, z_q.new_tensor(0), None  # z_q, loss, min_encoding_indices
+            else:
+                # if self.q_counter < self.q_init + self.q_re_end:
+                if self.q_init <= self.q_counter < self.q_re_end:
+                    if (self.q_counter - self.q_init) % self.q_re_step == 0 or self.q_counter == self.q_init + self.q_re_end - 1:
+                        kmeans = KMeans(n_clusters=self.codebook_size)
+                        world_size = dist.get_world_size()
+                        print("Updating codebook from reservoir.")
+                        if world_size > 1:
+                            global_reservoir = [torch.zeros_like(self.reservoir) for _ in range(world_size)]
+                            dist.all_gather(global_reservoir, self.reservoir.clone())
+                            global_reservoir = torch.cat(global_reservoir, dim=0)
+                        else:
+                            global_reservoir = self.reservoir
+                        kmeans.fit_predict(global_reservoir)  # reservoir is 20k encoded latents
+                        self.embedding.weight.data = kmeans.centroids.detach()
+
+        d = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \
+            torch.sum(self.embedding.weight**2, dim=1) - 2 * \
+            torch.einsum('bd,dn->bn', z_flattened, rearrange(self.embedding.weight, 'n d -> d n'))
+
+        min_encoding_indices = torch.argmin(d, dim=1)
+        z_q = self.embedding(min_encoding_indices).view(z.shape)
+
+        # compute loss for embedding
+        loss = torch.mean((z_q.detach()-z)**2) + self.beta * torch.mean((z_q - z.detach()) ** 2)
+
+        # preserve gradients
+        z_q = z + (z_q - z).detach()
+
+        # reshape back to match original input shape
+        z_q = rearrange(z_q, 'b h w c -> b c h w').contiguous()
+
+        return z_q, loss, min_encoding_indices
+
+    def get_codebook_entry(self, indices, shape):
+        # get quantized latent vectors
+        z_q = self.embedding(indices)
+
+        if shape is not None:
+            z_q = z_q.view(shape)
+            # reshape back to match original input shape
+            z_q = z_q.permute(0, 3, 1, 2).contiguous()
+
+        return z_q
+
+
+if __name__ == '__main__':
+    enc = Encoder()
+    dec = Decoder()
+    print(sum([p.numel() for p in enc.parameters()]))
+    print(sum([p.numel() for p in dec.parameters()]))
+    x = torch.randn(1, 3, 512, 512)
+    res = enc(x)
+    print(res.shape)
+    res = dec(res)
+    print(res.shape)