initiate

app.py

@@ -0,0 +1,194 @@
+import gradio as gr
+from io import BytesIO
+import os
+import sys
+
+import matplotlib.pyplot as plt
+import matplotlib
+import numpy as np
+from PIL import Image
+from omegaconf import OmegaConf
+
+import torch
+from torchvision import transforms as T
+
+from optvq.models.quantizer import sinkhorn
+from optvq.utils.init import seed_everything
+seed_everything(42)
+from optvq.models.vqgan_hf import VQModelHF
+matplotlib.rcParams['font.family'] = 'Times New Roman'
+
+#################
+N_data = 50
+N_code = 20
+dim = 2
+handler = None
+device = torch.device("cpu")
+#################
+
+def nearest(src, trg):
+    dis_mat = torch.cdist(src, trg)
+    min_idx = torch.argmin(dis_mat, dim=-1)
+    return min_idx
+
+def normalize(A, dim, mode="all"):
+    if mode == "all":
+        A = (A - A.mean()) / (A.std() + 1e-6)
+        A = A - A.min()
+    elif mode == "dim":
+        A = A / dim
+    elif mode == "null":
+        pass
+    return A
+
+def draw_NN(data, code):
+    # nearest neighbor method
+    indices = nearest(data, code)
+    data = data.numpy()
+    code = code.numpy()
+    
+    plt.figure(figsize=(3, 2.5), dpi=400)
+    # draw arrows in blue color, alpha=0.5
+    for i in range(data.shape[0]):
+        idx = indices[i].item()
+        start = data[i]
+        end = code[idx]
+        plt.arrow(start[0], start[1], end[0] - start[0], end[1] - start[1],
+                    head_width=0.05, head_length=0.05, fc='red', ec='red', alpha=0.6,
+                    ls="-", lw=0.5)
+    plt.scatter(data[:, 0], data[:, 1], s=10, marker="o", c="gray", label="Data")
+    plt.scatter(code[:, 0], code[:, 1], s=25, marker="*", c="blue", label="Code")
+    plt.legend(loc="lower right")
+    plt.grid(color="gray", alpha=0.8, ls="-.", lw=0.5)
+    plt.title("Nearest neighbor")
+
+    buf = BytesIO()
+    plt.savefig(buf, format="png")
+    buf.seek(0)
+    image = Image.open(buf)
+    return image
+
+def draw_optvq(data, code):
+    cost = torch.cdist(data, code, p=2.0)
+    cost = normalize(cost, dim, mode="all")
+    Q = sinkhorn(cost, n_iters=5, epsilon=10, is_distributed=False)
+    indices = torch.argmax(Q, dim=-1)
+    data = data.numpy()
+    code = code.numpy()
+
+    plt.figure(figsize=(3, 2.5), dpi=400)
+    # draw arrows in blue color, alpha=0.5
+    for i in range(data.shape[0]):
+        idx = indices[i].item()
+        start = data[i]
+        end = code[idx]
+        plt.arrow(start[0], start[1], end[0] - start[0], end[1] - start[1],
+                    head_width=0.05, head_length=0.05, fc='green', ec='green', alpha=0.6,
+                    ls="-", lw=0.5)
+    plt.scatter(data[:, 0], data[:, 1], s=10, marker="o", c="gray", label="Data")
+    plt.scatter(code[:, 0], code[:, 1], s=25, marker="*", c="blue", label="Code")
+    plt.legend(loc="lower right")
+    plt.grid(color="gray", alpha=0.8, ls="-.", lw=0.5)
+    plt.title("Optimal Transport (OptVQ)")
+
+    buf = BytesIO()
+    plt.savefig(buf, format="png")
+    buf.seek(0)
+    image = Image.open(buf)
+    return image
+
+def draw_process(x, y, std):
+    data = torch.randn(N_data, dim)
+    code = torch.randn(N_code, dim) * std
+    code[:, 0] += x
+    code[:, 1] += y
+
+    image_NN = draw_NN(data, code)
+    image_optvq = draw_optvq(data, code)
+    
+    return image_NN, image_optvq
+
+class Handler:
+    def __init__(self, device):
+        self.transform = T.Compose([
+            T.Resize(256),
+            T.CenterCrop(256),
+            T.ToTensor()
+        ])
+        self.device = device
+
+        
+        self.basevq = VQModelHF.from_pretrained("BorelTHU/basevq-16x16x4")
+        self.basevq.to(self.device)
+        self.basevq.eval()
+
+        self.vqgan = VQModelHF.from_pretrained("BorelTHU/vqgan-16x16")
+        self.vqgan.to(self.device)
+        self.vqgan.eval()
+
+        self.optvq = VQModelHF.from_pretrained("BorelTHU/optvq-16x16x4")
+        self.optvq.to(self.device)
+        self.optvq.eval()
+
+    def tensor_to_image(self, tensor):
+        img = tensor.squeeze(0).cpu().permute(1, 2, 0).numpy()
+        img = (img + 1) / 2 * 255
+        img = img.astype("uint8")
+        return img
+
+    def process_image(self, img: np.ndarray):
+        img = Image.fromarray(img.astype("uint8"))
+        img = self.transform(img)
+        img = img.unsqueeze(0).to(self.device)
+        with torch.no_grad():
+            img = 2 * img - 1
+            # basevq
+            quant, *_ = self.basevq.encode(img)
+            basevq_rec = self.basevq.decode(quant)
+            # vqgan
+            quant, *_ = self.vqgan.encode(img)
+            vqgan_rec = self.vqgan.decode(quant)
+            # optvq
+            quant, *_ = self.optvq.encode(img)
+            optvq_rec = self.optvq.decode(quant)
+        
+        # tensor to PIL image
+        img = self.tensor_to_image(img)
+        basevq_rec = self.tensor_to_image(basevq_rec)
+        vqgan_rec = self.tensor_to_image(vqgan_rec)
+        optvq_rec = self.tensor_to_image(optvq_rec)
+        return img, basevq_rec, vqgan_rec, optvq_rec
+
+if __name__ == "__main__":
+    # create the model handler
+    handler = Handler(device=device)
+
+    # create the interface
+    with gr.Blocks() as demo:
+        gr.Textbox(value="This demo shows the image reconstruction comparison between OptVQ and other methods. The input image is resized to 256 x 256 and then fed into the models. The output images are the reconstructed images from the latent codes.", label="Demo 1: Image reconstruction results")
+        with gr.Row():
+            with gr.Column():
+                image_input = gr.Image(label="Input data", image_mode="RGB", type="numpy")
+                btn_demo1 = gr.Button(value="Run reconstruction")
+            image_basevq = gr.Image(label="BaseVQ rec.")
+            image_vqgan = gr.Image(label="VQGAN rec.")
+            image_optvq = gr.Image(label="OptVQ rec.")
+        btn_demo1.click(fn=handler.process_image, inputs=[image_input], outputs=[image_input, image_basevq, image_vqgan, image_optvq])
+
+        gr.Textbox(value="This demo shows the 2D visualizations of nearest neighbor and optimal transport (OptVQ) methods. The data points are randomly generated from a normal distribution, and the matching results are shown as arrows with different colors.", label="Demo 2: 2D visualizations of matching results")
+        with gr.Row():
+            with gr.Column():
+                input_x = gr.Slider(label="x", value=0, minimum=-10, maximum=10, step=0.1)
+                input_y = gr.Slider(label="y", value=0, minimum=-10, maximum=10, step=0.1)
+                input_std = gr.Slider(label="std", value=1, minimum=0, maximum=5, step=0.1)
+                btn_demo2 = gr.Button(value="Run 2D example")
+            output_nn = gr.Image(label="NN")
+            output_optvq = gr.Image(label="OptVQ")
+        
+        # set the function
+        input_x.change(fn=draw_process, inputs=[input_x, input_y, input_std], outputs=[output_nn, output_optvq])
+        input_y.change(fn=draw_process, inputs=[input_x, input_y, input_std], outputs=[output_nn, output_optvq])
+        input_std.change(fn=draw_process, inputs=[input_x, input_y, input_std], outputs=[output_nn, output_optvq])
+        btn_demo2.click(fn=draw_process, inputs=[input_x, input_y, input_std], outputs=[output_nn, output_optvq])
+
+    demo.launch()

optvq/data/dataloader.py

@@ -0,0 +1,50 @@
+# ------------------------------------------------------------------------------
+# OptVQ: Preventing Local Pitfalls in Vector Quantization via Optimal Transport
+# Copyright (c) 2024 Borui Zhang. All Rights Reserved.
+# Licensed under the MIT License [see LICENSE for details]
+# ------------------------------------------------------------------------------
+
+from typing import Union
+
+import torch
+from torch.utils.data import Subset
+
+def maybe_get_subset(dataset, subset_size: Union[int, float] = None, num_data_repeat: int = None):
+    """
+    num_data_repeat is aimed at avoiding the consuming loading time for the small dataset
+    """
+    if subset_size is None:
+        return dataset
+    else:
+        if subset_size < 1.0:
+            subset_size = len(dataset) * subset_size
+        subset_size = int(subset_size)
+        selected_indices = torch.randperm(len(dataset))[:subset_size]
+        if num_data_repeat is not None:
+            selected_indices = selected_indices.repeat(num_data_repeat)
+        return Subset(dataset, selected_indices)
+
+class LoaderWrapper:
+    """
+    write a dataloader class, given total steps, recursively loading data
+    """
+    def __init__(self, loader, total_iterations: int = None):
+        self.loader = loader
+        self.total_iterations = total_iterations if total_iterations is not None else len(loader)
+    
+    def __iter__(self):
+        self.generator = iter(self.loader)
+        self.counter = 0
+        return self
+    
+    def __next__(self):
+        if self.counter >= self.total_iterations:
+            self.counter = 0
+            raise StopIteration
+        else:
+            self.counter += 1
+            try:
+                return next(self.generator)
+            except StopIteration:
+                self.generator = iter(self.loader)
+                return next(self.generator)

optvq/data/dataset.py

@@ -0,0 +1,62 @@
+# ------------------------------------------------------------------------------
+# OptVQ: Preventing Local Pitfalls in Vector Quantization via Optimal Transport
+# Copyright (c) 2024 Borui Zhang. All Rights Reserved.
+# Licensed under the MIT License [see LICENSE for details]
+# ------------------------------------------------------------------------------
+
+import os
+from PIL import Image
+import numpy as np
+
+import torch
+from torch.utils.data import Dataset
+from torchvision import transforms
+
+from .preprocessor import normalize_params
+
+class ImageNetDataset(Dataset):
+    def __init__(self, root, transform=None, convert_to_numpy: bool = True, post_normalize: str = "plain"):
+        self.root = root
+        self.transform = transform
+        self.convert_to_numpy = convert_to_numpy
+        self.post_normalize = transforms.Normalize(
+            **normalize_params[post_normalize]
+        )
+
+        # find classes
+        classes = sorted(entry.name for entry in os.scandir(root) if entry.is_dir())
+        class_to_idx = {cls_name: i for i, cls_name in enumerate(classes)}
+
+        # make dataset
+        self.samples = []
+        self.extensions = []
+        for target_class in sorted(class_to_idx.keys()):
+            class_idx = class_to_idx[target_class]
+            target_dir = os.path.join(root, target_class)
+            if not os.path.isdir(target_dir):
+                continue
+            for fname in sorted(os.listdir(target_dir)):
+                path = os.path.join(target_dir, fname)
+                item = (path, class_idx)
+                self.samples.append(item)
+                ext = path.split(".")[-1]
+                if ext not in self.extensions:
+                    self.extensions.append(ext)
+    
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        path, label = self.samples[index]
+        image = Image.open(path)
+        if not image.mode == "RGB":
+            image = image.convert("RGB")
+        if self.convert_to_numpy:
+            image = np.array(image).astype("uint8")
+        # image augmentation
+        image = self.transform(image=image)["image"]
+        # to tensor and normalize
+        image = (image / 255).astype(np.float32)
+        image = torch.from_numpy(image).permute(2, 0, 1)
+        image = self.post_normalize(image)
+        return image, label

optvq/data/preprocessor.py

@@ -0,0 +1,71 @@
+# ------------------------------------------------------------------------------
+# OptVQ: Preventing Local Pitfalls in Vector Quantization via Optimal Transport
+# Copyright (c) 2024 Borui Zhang. All Rights Reserved.
+# Licensed under the MIT License [see LICENSE for details]
+# ------------------------------------------------------------------------------
+
+from typing import Optional
+import numpy as np
+
+from torchvision import transforms
+import albumentations as A
+
+BICUBIC = transforms.InterpolationMode.BICUBIC
+
+normalize_params = {
+    "plain": {"mean": (0.5,), "std": (0.5,)},
+    "cnn": {"mean": (0.485, 0.456, 0.406), "std": (0.229, 0.224, 0.225)},
+    "clip": {"mean": (0.48145466, 0.4578275, 0.40821073), "std": (0.26862954, 0.26130258, 0.27577711)}
+}
+
+recover_map_dict = {
+    "plain": transforms.Normalize(
+        mean=(-1,), std=(2,)
+    ),
+    "cnn": transforms.Normalize(
+        mean=(-0.485/0.229, -0.456/0.224, -0.406/0.225),
+        std=(1/0.229, 1/0.224, 1/0.225)
+    ),
+    "clip": transforms.Normalize(
+        mean=(-0.48145466/0.26862954, -0.4578275/0.26130258, -0.40821073/0.27577711),
+        std=(1/0.26862954, 1/0.26130258, 1/0.27577711)
+    )
+}
+
+def get_recover_map(name: str):
+    return recover_map_dict[name]
+
+###########################################
+# Preprocessor
+###########################################
+
+def plain_preprocessor(resize: Optional[int] = 32):
+    return transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize((0.5,), (0.5,)),
+        transforms.Resize(resize),
+    ])
+
+def imagenet_preprocessor(resize: Optional[int] = 256, is_train: bool = True):
+    if is_train:
+        # augmentation v1
+        # transform = A.Compose([
+        #     A.SmallestMaxSize(max_size=resize),
+        #     A.RandomCrop(height=resize, width=resize),
+        #     A.HorizontalFlip(p=0.5),
+        # ])
+
+        # augmentation v2
+        transform = A.Compose([
+            A.SmallestMaxSize(max_size=resize),
+            A.RandomResizedCrop(width=resize, height=resize, scale=(0.2, 1.0)),
+            A.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4, hue=0.1, p=0.8),
+            A.GaussianBlur(blur_limit=7, p=0.5),
+            A.HorizontalFlip(p=0.5),
+        ])
+    else:
+        transform = A.Compose([
+            A.SmallestMaxSize(max_size=resize),
+            A.CenterCrop(height=resize, width=resize),
+        ])
+    return transform

optvq/losses/aeloss.py

@@ -0,0 +1,76 @@
+# ------------------------------------------------------------------------------
+# OptVQ: Preventing Local Pitfalls in Vector Quantization via Optimal Transport
+# Copyright (c) 2024 Borui Zhang. All Rights Reserved.
+# Licensed under the MIT License [see LICENSE for details]
+# ------------------------------------------------------------------------------
+# Modified from [thuanz123/enhancing-transformers](https://github.com/thuanz123/enhancing-transformers)
+# Copyright (c) 2022 Thuan H. Nguyen. All Rights Reserved.
+# ------------------------------------------------------------------------------
+# Modified from [CompVis/taming-transformers](https://github.com/CompVis/taming-transformers)
+# Copyright (c) 2020 Patrick Esser and Robin Rombach and Björn Ommer. All Rights Reserved.
+# ------------------------------------------------------------------------------
+
+from typing import Tuple
+
+import lpips
+
+import torch
+import torch.nn as nn
+
+class AELoss(nn.Module):
+    """
+    Args:
+        loss_q_weight (float): weight for quantization loss
+        loss_l1_weight (float): weight for L1 loss (loglaplace)
+        loss_l2_weight (float): weight for L2 loss (loggaussian)
+        loss_p_weight (float): weight for perceptual loss
+    """
+    def __init__(self, loss_q_weight: float = 1.0,
+                 loss_l1_weight: float = 1.0,
+                 loss_l2_weight: float = 1.0,
+                 loss_p_weight: float = 1.0) -> None:
+        super().__init__()
+        self.loss_type = ["aeloss"]
+
+        self.perceptual_loss = lpips.LPIPS(net="vgg", verbose=False)
+        # freeze the perceptual loss
+        for param in self.perceptual_loss.parameters():
+            param.requires_grad = False
+
+        self.loss_q_weight = loss_q_weight 
+        self.loss_l1_weight = loss_l1_weight 
+        self.loss_l2_weight = loss_l2_weight
+        self.loss_p_weight = loss_p_weight 
+
+    @torch.autocast(device_type="cuda", enabled=False)
+    def forward(self, q_loss: torch.FloatTensor, 
+                x: torch.FloatTensor, 
+                x_rec: torch.FloatTensor, *args, **kwargs) -> Tuple:
+        x = x.float()
+        x_rec = x_rec.float()
+        
+        # compute l1 loss
+        loss_l1 = (x_rec - x).abs().mean() if self.loss_l1_weight > 0.0 else torch.tensor(0.0, device=x.device)
+
+        # compute l2 loss
+        loss_l2 = (x_rec - x).pow(2).mean() if self.loss_l2_weight > 0.0 else torch.tensor(0.0, device=x.device)
+
+        # compute perceptual loss
+        loss_p = self.perceptual_loss(x, x_rec).mean() if self.loss_p_weight > 0.0 else torch.tensor(0.0, device=x.device)
+
+        # compute total loss
+        loss = self.loss_p_weight * loss_p + \
+               self.loss_l1_weight * loss_l1 + \
+               self.loss_l2_weight * loss_l2
+        loss += self.loss_q_weight * q_loss
+
+        # get the log
+        log = {
+            "loss": loss.detach(),
+            "loss_p": loss_p.detach(),
+            "loss_l1": loss_l1.detach(),
+            "loss_l2": loss_l2.detach(),
+            "loss_q": q_loss.detach()
+        }
+
+        return loss, log

optvq/losses/aeloss_disc.py

@@ -0,0 +1,177 @@
+# ------------------------------------------------------------------------------
+# OptVQ: Preventing Local Pitfalls in Vector Quantization via Optimal Transport
+# Copyright (c) 2024 Borui Zhang. All Rights Reserved.
+# Licensed under the MIT License [see LICENSE for details]
+# ------------------------------------------------------------------------------
+# Modified from [thuanz123/enhancing-transformers](https://github.com/thuanz123/enhancing-transformers)
+# Copyright (c) 2022 Thuan H. Nguyen. All Rights Reserved.
+# ------------------------------------------------------------------------------
+# Modified from [CompVis/taming-transformers](https://github.com/CompVis/taming-transformers)
+# Copyright (c) 2020 Patrick Esser and Robin Rombach and Björn Ommer. All Rights Reserved.
+# ------------------------------------------------------------------------------
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import lpips
+
+from optvq.models.discriminator import NLayerDiscriminator, weights_init
+
+class DummyLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+
+def hinge_d_loss(logits_real, logits_fake):
+    loss_real = torch.mean(F.relu(1. - logits_real))
+    loss_fake = torch.mean(F.relu(1. + logits_fake))
+    d_loss = 0.5 * (loss_real + loss_fake)
+    return d_loss
+
+
+def vanilla_d_loss(logits_real, logits_fake):
+    d_loss = 0.5 * (
+        torch.mean(torch.nn.functional.softplus(-logits_real)) +
+        torch.mean(torch.nn.functional.softplus(logits_fake)))
+    return d_loss
+
+
+class AELossWithDisc(nn.Module):
+    def __init__(self, 
+                 disc_start, 
+                 pixelloss_weight=1.0,
+                 disc_in_channels=3, 
+                 disc_num_layers=3, 
+                 use_actnorm=False, 
+                 disc_ndf=64, 
+                 disc_conditional=False,
+                 disc_loss="hinge",
+                 loss_l1_weight: float = 1.0,
+                 loss_l2_weight: float = 1.0,
+                 loss_p_weight: float = 1.0,
+                 loss_q_weight: float = 1.0,
+                 loss_g_weight: float = 1.0,
+                 loss_d_weight: float = 1.0
+        ):
+        super(AELossWithDisc, self).__init__()
+        assert disc_loss in ["hinge", "vanilla"]
+        
+        self.pixel_weight = pixelloss_weight
+        self.perceptual_loss = lpips.LPIPS(net="vgg", verbose=False).eval()
+
+        self.loss_l1_weight = loss_l1_weight
+        self.loss_l2_weight = loss_l2_weight
+        self.loss_p_weight = loss_p_weight
+        self.loss_q_weight = loss_q_weight
+        self.loss_g_weight = loss_g_weight
+        self.loss_d_weight = loss_d_weight
+
+        self.discriminator = NLayerDiscriminator(input_nc=disc_in_channels,
+                                                 n_layers=disc_num_layers,
+                                                 use_actnorm=use_actnorm,
+                                                 ndf=disc_ndf
+                                                 ).apply(weights_init)
+        self.discriminator_iter_start = disc_start
+        if disc_loss == "hinge":
+            self.disc_loss = hinge_d_loss
+        elif disc_loss == "vanilla":
+            self.disc_loss = vanilla_d_loss
+        else:
+            raise ValueError(f"Unknown GAN loss '{disc_loss}'.")
+        print(f"VQLPIPSWithDiscriminator running with {disc_loss} loss.")
+
+        self.disc_conditional = disc_conditional
+
+    def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer=None):
+        nll_grads = torch.autograd.grad(nll_loss, last_layer, retain_graph=True)[0]
+        g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]
+
+        g_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4)
+        g_weight = torch.clamp(g_weight, 0.0, 1e4).detach()
+        g_weight = g_weight * self.loss_g_weight
+        
+        # detection nan
+        if torch.isnan(g_weight).any():
+            g_weight = torch.tensor(0.0, device=g_weight.device)
+        return g_weight
+
+    @torch.autocast(device_type="cuda", enabled=False)
+    def forward(self, codebook_loss, inputs, reconstructions, mode, last_layer=None, cond=None, global_step=0):
+        x = inputs.contiguous().float()
+        x_rec = reconstructions.contiguous().float()
+
+        # compute q loss
+        loss_q = codebook_loss.mean()
+
+        # compute l1 loss
+        loss_l1 = (x_rec - x).abs().mean() if self.loss_l1_weight > 0.0 else torch.tensor(0.0, device=x.device)
+
+        # compute l2 loss
+        loss_l2 = (x_rec - x).pow(2).mean() if self.loss_l2_weight > 0.0 else torch.tensor(0.0, device=x.device)
+
+        # compute perceptual loss
+        loss_p = self.perceptual_loss(x, x_rec).mean() if self.loss_p_weight > 0.0 else torch.tensor(0.0, device=x.device)
+
+        # intigrate reconstruction loss
+        loss_rec = loss_l1 * self.loss_l1_weight + \
+                   loss_l2 * self.loss_l2_weight + \
+                   loss_p * self.loss_p_weight
+
+        # setup the factor_disc
+        if global_step < self.discriminator_iter_start:
+            factor_disc = 0.0
+        else:
+            factor_disc = 1.0
+
+        # now the GAN part
+        if mode == 0:
+            # generator update
+            if cond is None:
+                assert not self.disc_conditional
+                logits_fake = self.discriminator(x_rec)
+            else:
+                assert self.disc_conditional
+                logits_fake = self.discriminator(torch.cat((x_rec, cond), dim=1))
+
+            # compute g loss
+            loss_g = - logits_fake.mean()
+
+            try:
+                loss_g_weight = self.calculate_adaptive_weight(loss_rec, loss_g, last_layer=last_layer)
+            except RuntimeError:
+                # assert not self.training
+                loss_g_weight = torch.tensor(0.0)
+
+            loss = loss_g * loss_g_weight * factor_disc  + \
+                   loss_q * self.loss_q_weight + \
+                   loss_rec
+
+            log = {"total_loss": loss.item(),
+                   "loss_q": loss_q.item(),
+                   "loss_rec": loss_rec.item(),
+                   "loss_l1": loss_l1.item(),
+                   "loss_l2": loss_l2.item(),
+                   "loss_p": loss_p.item(),
+                   "loss_g": loss_g.item(),
+                   "loss_g_weight": loss_g_weight.item(),
+                   "factor_disc": factor_disc,
+            }
+            return loss, log
+
+        if mode == 1:
+            # second pass for discriminator update
+            if cond is None:
+                logits_real = self.discriminator(x.detach())
+                logits_fake = self.discriminator(x_rec.detach())
+            else:
+                logits_real = self.discriminator(torch.cat((x.detach(), cond), dim=1))
+                logits_fake = self.discriminator(torch.cat((x_rec.detach(), cond), dim=1))
+
+            loss_d = self.disc_loss(logits_real, logits_fake).mean()
+            loss = loss_d * self.loss_d_weight
+
+            log = {"loss_d": loss_d.item(),
+                   "logits_real": logits_real.mean().item(),
+                   "logits_fake": logits_fake.mean().item()
+            }
+            return loss, log

optvq/models/backbone/diffusion.py

@@ -0,0 +1,787 @@
+# ------------------------------------------------------------------------------
+# Copy from [CompVis/taming-transformers](https://github.com/CompVis/taming-transformers)
+# Copyright (c) 2020 Patrick Esser and Robin Rombach and Björn Ommer. All Rights Reserved.
+# ------------------------------------------------------------------------------
+
+# pytorch_diffusion + derived encoder decoder
+import math
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+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, temb_channels=512):
+        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)
+        if temb_channels > 0:
+            self.temb_proj = torch.nn.Linear(temb_channels,
+                                             out_channels)
+        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, temb):
+        h = x
+        h = self.norm1(h)
+        h = nonlinearity(h)
+        h = self.conv1(h)
+
+        if temb is not None:
+            h = h + self.temb_proj(nonlinearity(temb))[:,:,None,None]
+
+        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 Model(nn.Module):
+    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
+                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
+                 resolution, use_timestep=True):
+        super().__init__()
+        self.ch = ch
+        self.temb_ch = self.ch*4
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        self.use_timestep = use_timestep
+        if self.use_timestep:
+            # timestep embedding
+            self.temb = nn.Module()
+            self.temb.dense = nn.ModuleList([
+                torch.nn.Linear(self.ch,
+                                self.temb_ch),
+                torch.nn.Linear(self.temb_ch,
+                                self.temb_ch),
+            ])
+
+        # downsampling
+        self.conv_in = torch.nn.Conv2d(in_channels,
+                                       self.ch,
+                                       kernel_size=3,
+                                       stride=1,
+                                       padding=1)
+
+        curr_res = resolution
+        in_ch_mult = (1,)+tuple(ch_mult)
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch*in_ch_mult[i_level]
+            block_out = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions-1:
+                down.downsample = Downsample(block_in, resamp_with_conv)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch*ch_mult[i_level]
+            skip_in = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks+1):
+                if i_block == self.num_res_blocks:
+                    skip_in = ch*in_ch_mult[i_level]
+                block.append(ResnetBlock(in_channels=block_in+skip_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in, resamp_with_conv)
+                curr_res = curr_res * 2
+            self.up.insert(0, up) # prepend to get consistent order
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        out_ch,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+
+    def forward(self, x, t=None):
+        #assert x.shape[2] == x.shape[3] == self.resolution
+
+        if self.use_timestep:
+            # timestep embedding
+            assert t is not None
+            temb = get_timestep_embedding(t, self.ch)
+            temb = self.temb.dense[0](temb)
+            temb = nonlinearity(temb)
+            temb = self.temb.dense[1](temb)
+        else:
+            temb = None
+
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1], temb)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions-1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks+1):
+                h = self.up[i_level].block[i_block](
+                    torch.cat([h, hs.pop()], dim=1), temb)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class Encoder(nn.Module):
+    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
+                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
+                 resolution, z_channels, double_z=True, **ignore_kwargs):
+        super().__init__()
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+
+        # downsampling
+        self.conv_in = torch.nn.Conv2d(in_channels,
+                                       self.ch,
+                                       kernel_size=3,
+                                       stride=1,
+                                       padding=1)
+
+        curr_res = resolution
+        in_ch_mult = (1,)+tuple(ch_mult)
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch*in_ch_mult[i_level]
+            block_out = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions-1:
+                down.downsample = Downsample(block_in, resamp_with_conv)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        2*z_channels if double_z else z_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    @property
+    def hidden_dim(self):
+        return self.conv_out.out_channels
+
+    def forward(self, x):
+        #assert x.shape[2] == x.shape[3] == self.resolution, "{}, {}, {}".format(x.shape[2], x.shape[3], self.resolution)
+
+        # timestep embedding
+        temb = None
+
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1], temb)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions-1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class Decoder(nn.Module):
+    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
+                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
+                 resolution, z_channels, give_pre_end=False, **ignorekwargs):
+        super().__init__()
+        self.ch = ch
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.give_pre_end = give_pre_end
+
+        # compute in_ch_mult, block_in and curr_res at lowest res
+        in_ch_mult = (1,)+tuple(ch_mult)
+        block_in = ch*ch_mult[self.num_resolutions-1]
+        curr_res = resolution // 2**(self.num_resolutions-1)
+        self.z_shape = (1,z_channels,curr_res,curr_res)
+        print("Working with z of shape {} = {} dimensions.".format(
+            self.z_shape, np.prod(self.z_shape)))
+
+        # z to block_in
+        self.conv_in = torch.nn.Conv2d(z_channels,
+                                       block_in,
+                                       kernel_size=3,
+                                       stride=1,
+                                       padding=1)
+
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks+1):
+                block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in, resamp_with_conv)
+                curr_res = curr_res * 2
+            self.up.insert(0, up) # prepend to get consistent order
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        out_ch,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    @property
+    def hidden_dim(self):
+        return self.conv_in.in_channels
+
+    def forward(self, z):
+        #assert z.shape[1:] == self.z_shape[1:]
+        self.last_z_shape = z.shape
+
+        # timestep embedding
+        temb = None
+
+        # z to block_in
+        h = self.conv_in(z)
+
+        # middle
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks+1):
+                h = self.up[i_level].block[i_block](h, temb)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        if self.give_pre_end:
+            return h
+
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class VUNet(nn.Module):
+    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
+                 attn_resolutions, dropout=0.0, resamp_with_conv=True,
+                 in_channels, c_channels,
+                 resolution, z_channels, use_timestep=False, **ignore_kwargs):
+        super().__init__()
+        self.ch = ch
+        self.temb_ch = self.ch*4
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        self.resolution = resolution
+
+        self.use_timestep = use_timestep
+        if self.use_timestep:
+            # timestep embedding
+            self.temb = nn.Module()
+            self.temb.dense = nn.ModuleList([
+                torch.nn.Linear(self.ch,
+                                self.temb_ch),
+                torch.nn.Linear(self.temb_ch,
+                                self.temb_ch),
+            ])
+
+        # downsampling
+        self.conv_in = torch.nn.Conv2d(c_channels,
+                                       self.ch,
+                                       kernel_size=3,
+                                       stride=1,
+                                       padding=1)
+
+        curr_res = resolution
+        in_ch_mult = (1,)+tuple(ch_mult)
+        self.down = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = ch*in_ch_mult[i_level]
+            block_out = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level != self.num_resolutions-1:
+                down.downsample = Downsample(block_in, resamp_with_conv)
+                curr_res = curr_res // 2
+            self.down.append(down)
+
+        self.z_in = torch.nn.Conv2d(z_channels,
+                                    block_in,
+                                    kernel_size=1,
+                                    stride=1,
+                                    padding=0)
+        # middle
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=2*block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in,
+                                       out_channels=block_in,
+                                       temb_channels=self.temb_ch,
+                                       dropout=dropout)
+
+        # upsampling
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = ch*ch_mult[i_level]
+            skip_in = ch*ch_mult[i_level]
+            for i_block in range(self.num_res_blocks+1):
+                if i_block == self.num_res_blocks:
+                    skip_in = ch*in_ch_mult[i_level]
+                block.append(ResnetBlock(in_channels=block_in+skip_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level != 0:
+                up.upsample = Upsample(block_in, resamp_with_conv)
+                curr_res = curr_res * 2
+            self.up.insert(0, up) # prepend to get consistent order
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        out_ch,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+
+    def forward(self, x, z):
+        #assert x.shape[2] == x.shape[3] == self.resolution
+
+        if self.use_timestep:
+            # timestep embedding
+            assert t is not None
+            temb = get_timestep_embedding(t, self.ch)
+            temb = self.temb.dense[0](temb)
+            temb = nonlinearity(temb)
+            temb = self.temb.dense[1](temb)
+        else:
+            temb = None
+
+        # downsampling
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1], temb)
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level != self.num_resolutions-1:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        # middle
+        h = hs[-1]
+        z = self.z_in(z)
+        h = torch.cat((h,z),dim=1)
+        h = self.mid.block_1(h, temb)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h, temb)
+
+        # upsampling
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks+1):
+                h = self.up[i_level].block[i_block](
+                    torch.cat([h, hs.pop()], dim=1), temb)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level != 0:
+                h = self.up[i_level].upsample(h)
+
+        # end
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h
+
+
+class SimpleDecoder(nn.Module):
+    def __init__(self, in_channels, out_channels, *args, **kwargs):
+        super().__init__()
+        self.model = nn.ModuleList([nn.Conv2d(in_channels, in_channels, 1),
+                                     ResnetBlock(in_channels=in_channels,
+                                                 out_channels=2 * in_channels,
+                                                 temb_channels=0, dropout=0.0),
+                                     ResnetBlock(in_channels=2 * in_channels,
+                                                out_channels=4 * in_channels,
+                                                temb_channels=0, dropout=0.0),
+                                     ResnetBlock(in_channels=4 * in_channels,
+                                                out_channels=2 * in_channels,
+                                                temb_channels=0, dropout=0.0),
+                                     nn.Conv2d(2*in_channels, in_channels, 1),
+                                     Upsample(in_channels, with_conv=True)])
+        # end
+        self.norm_out = Normalize(in_channels)
+        self.conv_out = torch.nn.Conv2d(in_channels,
+                                        out_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, x):
+        for i, layer in enumerate(self.model):
+            if i in [1,2,3]:
+                x = layer(x, None)
+            else:
+                x = layer(x)
+
+        h = self.norm_out(x)
+        h = nonlinearity(h)
+        x = self.conv_out(h)
+        return x
+
+
+class UpsampleDecoder(nn.Module):
+    def __init__(self, in_channels, out_channels, ch, num_res_blocks, resolution,
+                 ch_mult=(2,2), dropout=0.0):
+        super().__init__()
+        # upsampling
+        self.temb_ch = 0
+        self.num_resolutions = len(ch_mult)
+        self.num_res_blocks = num_res_blocks
+        block_in = in_channels
+        curr_res = resolution // 2 ** (self.num_resolutions - 1)
+        self.res_blocks = nn.ModuleList()
+        self.upsample_blocks = nn.ModuleList()
+        for i_level in range(self.num_resolutions):
+            res_block = []
+            block_out = ch * ch_mult[i_level]
+            for i_block in range(self.num_res_blocks + 1):
+                res_block.append(ResnetBlock(in_channels=block_in,
+                                         out_channels=block_out,
+                                         temb_channels=self.temb_ch,
+                                         dropout=dropout))
+                block_in = block_out
+            self.res_blocks.append(nn.ModuleList(res_block))
+            if i_level != self.num_resolutions - 1:
+                self.upsample_blocks.append(Upsample(block_in, True))
+                curr_res = curr_res * 2
+
+        # end
+        self.norm_out = Normalize(block_in)
+        self.conv_out = torch.nn.Conv2d(block_in,
+                                        out_channels,
+                                        kernel_size=3,
+                                        stride=1,
+                                        padding=1)
+
+    def forward(self, x):
+        # upsampling
+        h = x
+        for k, i_level in enumerate(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.res_blocks[i_level][i_block](h, None)
+            if i_level != self.num_resolutions - 1:
+                h = self.upsample_blocks[k](h)
+        h = self.norm_out(h)
+        h = nonlinearity(h)
+        h = self.conv_out(h)
+        return h

optvq/models/backbone/simple_cnn.py

@@ -0,0 +1,90 @@
+# ------------------------------------------------------------------------------
+# OptVQ: Preventing Local Pitfalls in Vector Quantization via Optimal Transport
+# Copyright (c) 2024 Borui Zhang. All Rights Reserved.
+# Licensed under the MIT License [see LICENSE for details]
+# ------------------------------------------------------------------------------
+
+import torch.nn as nn
+
+class PlainCNNEncoder(nn.Module):
+    def __init__(self, in_dim: int = 3):
+        super(PlainCNNEncoder, self).__init__()
+
+        self.in_dim = in_dim
+
+        self.in_fc = nn.Conv2d(in_channels=in_dim, out_channels=16,
+                                kernel_size=3, stride=1, padding=1, bias=True)
+        self.act0 = nn.ReLU(inplace=True)
+
+        self.down1 = nn.Conv2d(in_channels=16, out_channels=16, kernel_size=2, stride=2)
+        self.conv1 = nn.Conv2d(in_channels=16, out_channels=16, 
+                               kernel_size=3, stride=1, padding=1, bias=True)
+        self.act1 = nn.ReLU(inplace=True)
+
+        self.down2 = nn.Conv2d(in_channels=16, out_channels=16, kernel_size=2, stride=2)
+        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, 
+                               kernel_size=3, stride=1, padding=1, bias=True)
+        self.act2 = nn.ReLU(inplace=True)
+
+        self.out_fc = nn.Conv2d(in_channels=32, out_channels=32,
+                                kernel_size=3, stride=1, padding=1, bias=True)
+    
+    @property
+    def hidden_dim(self):
+        return 32
+
+    def forward(self, x):
+        x = self.in_fc(x)
+        x = self.act0(x)
+
+        x = self.down1(x)
+        x = self.conv1(x)
+        x = self.act1(x)
+
+        x = self.down2(x)
+        x = self.conv2(x)
+        x = self.act2(x)
+
+        x = self.out_fc(x)
+        return x
+
+class PlainCNNDecoder(nn.Module):
+    def __init__(self, out_dim: int = 3):
+        super(PlainCNNDecoder, self).__init__()
+        self.out_dim = out_dim
+
+        self.in_fc = nn.Conv2d(in_channels=32, out_channels=32,
+                               kernel_size=3, stride=1, padding=1, bias=True)
+        
+        self.act1 = nn.ReLU(inplace=True)
+        self.up1 = nn.ConvTranspose2d(in_channels=32, out_channels=16, kernel_size=2, stride=2)
+        self.conv1 = nn.Conv2d(in_channels=16, out_channels=16, 
+                               kernel_size=3, stride=1, padding=1, bias=True)
+        
+        self.act2 = nn.ReLU(inplace=True)
+        self.up2 = nn.ConvTranspose2d(in_channels=16, out_channels=16, kernel_size=2, stride=2)
+        self.conv2 = nn.Conv2d(in_channels=16, out_channels=16, 
+                               kernel_size=3, stride=1, padding=1, bias=True)
+        
+        self.act3 = nn.ReLU(inplace=True)
+        self.out_fc = nn.Conv2d(in_channels=16, out_channels=out_dim,
+                               kernel_size=3, stride=1, padding=1, bias=True)
+    
+    @property
+    def hidden_dim(self):
+        return 32
+    
+    def forward(self, x):
+        x = self.in_fc(x)
+
+        x = self.act1(x)
+        x = self.up1(x)
+        x = self.conv1(x)
+
+        x = self.act2(x)
+        x = self.up2(x)
+        x = self.conv2(x)
+        
+        x = self.act3(x)
+        x = self.out_fc(x)
+        return x

optvq/models/discriminator.py

@@ -0,0 +1,155 @@
+# ------------------------------------------------------------------------------
+# Copy from [CompVis/taming-transformers](https://github.com/CompVis/taming-transformers)
+# Copyright (c) 2020 Patrick Esser and Robin Rombach and Björn Ommer. All Rights Reserved.
+# ------------------------------------------------------------------------------
+
+import functools
+import torch.nn as nn
+import torch
+
+class ActNorm(nn.Module):
+    def __init__(self, num_features, logdet=False, affine=True,
+                 allow_reverse_init=False):
+        assert affine
+        super().__init__()
+        self.logdet = logdet
+        self.loc = nn.Parameter(torch.zeros(1, num_features, 1, 1))
+        self.scale = nn.Parameter(torch.ones(1, num_features, 1, 1))
+        self.allow_reverse_init = allow_reverse_init
+
+        self.register_buffer('initialized', torch.tensor(0, dtype=torch.uint8))
+
+    def initialize(self, input):
+        with torch.no_grad():
+            flatten = input.permute(1, 0, 2, 3).contiguous().view(input.shape[1], -1)
+            mean = (
+                flatten.mean(1)
+                .unsqueeze(1)
+                .unsqueeze(2)
+                .unsqueeze(3)
+                .permute(1, 0, 2, 3)
+            )
+            std = (
+                flatten.std(1)
+                .unsqueeze(1)
+                .unsqueeze(2)
+                .unsqueeze(3)
+                .permute(1, 0, 2, 3)
+            )
+
+            self.loc.data.copy_(-mean)
+            self.scale.data.copy_(1 / (std + 1e-6))
+
+    def forward(self, input, reverse=False):
+        if reverse:
+            return self.reverse(input)
+        if len(input.shape) == 2:
+            input = input[:,:,None,None]
+            squeeze = True
+        else:
+            squeeze = False
+
+        _, _, height, width = input.shape
+
+        if self.training and self.initialized.item() == 0:
+            self.initialize(input)
+            self.initialized.fill_(1)
+
+        h = self.scale * (input + self.loc)
+
+        if squeeze:
+            h = h.squeeze(-1).squeeze(-1)
+
+        if self.logdet:
+            log_abs = torch.log(torch.abs(self.scale))
+            logdet = height*width*torch.sum(log_abs)
+            logdet = logdet * torch.ones(input.shape[0]).to(input)
+            return h, logdet
+
+        return h
+
+    def reverse(self, output):
+        if self.training and self.initialized.item() == 0:
+            if not self.allow_reverse_init:
+                raise RuntimeError(
+                    "Initializing ActNorm in reverse direction is "
+                    "disabled by default. Use allow_reverse_init=True to enable."
+                )
+            else:
+                self.initialize(output)
+                self.initialized.fill_(1)
+
+        if len(output.shape) == 2:
+            output = output[:,:,None,None]
+            squeeze = True
+        else:
+            squeeze = False
+
+        h = output / self.scale - self.loc
+
+        if squeeze:
+            h = h.squeeze(-1).squeeze(-1)
+        return h
+
+
+
+def weights_init(m):
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        nn.init.normal_(m.weight.data, 0.0, 0.02)
+    elif classname.find('BatchNorm') != -1:
+        nn.init.normal_(m.weight.data, 1.0, 0.02)
+        nn.init.constant_(m.bias.data, 0)
+
+
+class NLayerDiscriminator(nn.Module):
+    """Defines a PatchGAN discriminator as in Pix2Pix
+        --> see https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py
+    """
+    def __init__(self, input_nc=3, ndf=64, n_layers=3, use_actnorm=False):
+        """Construct a PatchGAN discriminator
+        Parameters:
+            input_nc (int)  -- the number of channels in input images
+            ndf (int)       -- the number of filters in the last conv layer
+            n_layers (int)  -- the number of conv layers in the discriminator
+            norm_layer      -- normalization layer
+        """
+        super(NLayerDiscriminator, self).__init__()
+        if not use_actnorm:
+            norm_layer = nn.BatchNorm2d
+        else:
+            norm_layer = ActNorm
+        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
+            use_bias = norm_layer.func != nn.BatchNorm2d
+        else:
+            use_bias = norm_layer != nn.BatchNorm2d
+
+        kw = 4
+        padw = 1
+        sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
+        nf_mult = 1
+        nf_mult_prev = 1
+        for n in range(1, n_layers):  # gradually increase the number of filters
+            nf_mult_prev = nf_mult
+            nf_mult = min(2 ** n, 8)
+            sequence += [
+                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
+                norm_layer(ndf * nf_mult),
+                nn.LeakyReLU(0.2, True)
+            ]
+
+        nf_mult_prev = nf_mult
+        nf_mult = min(2 ** n_layers, 8)
+        sequence += [
+            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
+            norm_layer(ndf * nf_mult),
+            nn.LeakyReLU(0.2, True)
+        ]
+
+        sequence += [
+            nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]  # output 1 channel prediction map
+        self.main = nn.Sequential(*sequence)
+
+    def forward(self, input):
+        """Standard forward."""
+        return self.main(input)

optvq/models/quantizer.py

@@ -0,0 +1,260 @@
+# ------------------------------------------------------------------------------
+# OptVQ: Preventing Local Pitfalls in Vector Quantization via Optimal Transport
+# Copyright (c) 2024 Borui Zhang. All Rights Reserved.
+# Licensed under the MIT License [see LICENSE for details]
+# ------------------------------------------------------------------------------
+
+import math
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+import torch.distributed as dist
+
+import optvq.utils.logger as L
+
+class VectorQuantizer(nn.Module):
+    def __init__(self, n_e: int = 1024, e_dim: int = 128, 
+                 beta: float = 1.0, use_norm: bool = False,
+                 use_proj: bool = True, fix_codes: bool = False,
+                 loss_q_type: str = "ce",
+                 num_head: int = 1,
+                 start_quantize_steps: int = None):
+        super(VectorQuantizer, self).__init__()
+        self.n_e = n_e
+        self.e_dim = e_dim
+        self.beta = beta
+        self.loss_q_type = loss_q_type
+        self.num_head = num_head
+        self.start_quantize_steps = start_quantize_steps
+        self.code_dim = self.e_dim // self.num_head
+
+        self.norm = lambda x: F.normalize(x, p=2.0, dim=-1, eps=1e-6) if use_norm else x
+        assert not use_norm, f"use_norm=True is no longer supported! Because the norm operation without theorectical analysis may cause unpredictable unstability."
+        self.use_proj = use_proj
+
+        self.embedding = nn.Embedding(num_embeddings=n_e, embedding_dim=self.code_dim)
+        if use_proj:
+            self.proj = nn.Linear(self.code_dim, self.code_dim)
+            torch.nn.init.normal_(self.proj.weight, std=self.code_dim ** -0.5)
+        if fix_codes:
+            self.embedding.weight.requires_grad = False
+
+    def reshape_input(self, x: Tensor):
+        """
+        (B, C, H, W) / (B, T, C) -> (B, T, C)
+        """
+        if x.ndim == 4:
+            _, C, H, W = x.size()
+            x = x.permute(0, 2, 3, 1).contiguous().view(-1, H * W, C)
+            return x, {"size": (H, W)}
+        elif x.ndim == 3:
+            return x, None
+        else:
+            raise ValueError("Invalid input shape!")
+
+    def recover_output(self, x: Tensor, info):
+        if info is not None:
+            H, W = info["size"]
+            if x.ndim == 3: # features (B, T, C) -> (B, C, H, W)
+                C = x.size(2)
+                return x.view(-1, H, W, C).permute(0, 3, 1, 2).contiguous()
+            elif x.ndim == 2: # indices (B, T) -> (B, H, W)
+                return x.view(-1, H, W)
+            else:
+                raise ValueError("Invalid input shape!")
+        else: # features (B, T, C) or indices (B, T)
+            return x
+    
+    def get_codebook(self, return_numpy: bool = True):
+        embed = self.proj(self.embedding.weight) if self.use_proj else self.embedding.weight
+        if return_numpy:
+            return embed.data.cpu().numpy()
+        else:
+            return embed.data
+
+    def quantize_input(self, query, reference):
+        # compute the distance matrix
+        query2ref = torch.cdist(query, reference, p=2.0) # (B1, B2)
+
+        # find the nearest embedding
+        indices = torch.argmin(query2ref, dim=-1) # (B1,)
+        nearest_ref = reference[indices] # (B1, C)
+            
+        return indices, nearest_ref, query2ref
+
+    def compute_codebook_loss(self, query, indices, nearest_ref, beta: float, query2ref):
+        # compute the loss
+        if self.loss_q_type == "l2":
+            loss = torch.mean((query - nearest_ref.detach()).pow(2)) + \
+                   torch.mean((nearest_ref - query.detach()).pow(2)) * beta
+        elif self.loss_q_type == "l1":
+            loss = torch.mean((query - nearest_ref.detach()).abs()) + \
+                   torch.mean((nearest_ref - query.detach()).abs()) * beta
+        elif self.loss_q_type == "ce":
+            loss = F.cross_entropy(- query2ref, indices)
+
+        return loss
+    
+    def compute_quantized_output(self, x, x_q):
+        if self.start_quantize_steps is not None:
+            if self.training and L.log.total_steps < self.start_quantize_steps:
+                L.log.add_scalar("params/quantize_ratio", 0.0)
+                return x
+            else:
+                L.log.add_scalar("params/quantize_ratio", 1.0)
+                return x + (x_q - x).detach()
+        else:
+            L.log.add_scalar("params/quantize_ratio", 1.0)
+            return x + (x_q - x).detach()  
+
+    @torch.autocast(device_type="cuda", enabled=False)
+    def forward(self, x: Tensor):
+        """
+        Quantize the input tensor x with the embedding table.
+
+        Args:
+            x (Tensor): input tensor with shape (B, C, H, W) or (B, T, C)
+        Returns:
+            (tuple) containing: (x_q, loss, indices)
+        """
+        x = x.float()
+        x, info = self.reshape_input(x)
+        B, T, C = x.size()
+        x = x.view(-1, C) # (B * T, C)
+        embed = self.proj(self.embedding.weight) if self.use_proj else self.embedding.weight
+
+        # split the x if multi-head is used
+        if self.num_head > 1:
+            x = x.view(-1, self.code_dim) # (B * T * nH, dC)
+
+        # compute the distance between x and each embedding
+        x, embed = self.norm(x), self.norm(embed)
+
+        # compute losses
+        indices, x_q, query2ref = self.quantize_input(x, embed)
+        loss = self.compute_codebook_loss(
+            query=x, indices=indices, nearest_ref=x_q, 
+            beta=self.beta, query2ref=query2ref
+        )
+
+        # compute statistics
+        if self.training and L.GET_STATS:
+            with torch.no_grad():
+                num_unique = torch.unique(indices).size(0)
+                x_norm_mean = torch.mean(x.norm(dim=-1))
+                embed_norm_mean = torch.mean(embed.norm(dim=-1))
+                diff_norm_mean = torch.mean((x_q - x).norm(dim=-1))
+                x2e_mean = query2ref.mean()
+                L.log.add_scalar("params/num_unique", num_unique)
+                L.log.add_scalar("params/x_norm", x_norm_mean.item())
+                L.log.add_scalar("params/embed_norm", embed_norm_mean.item())
+                L.log.add_scalar("params/diff_norm", diff_norm_mean.item())
+                L.log.add_scalar("params/x2e_mean", x2e_mean.item())
+    
+        # compute the final x_q
+        x_q = self.compute_quantized_output(x, x_q).view(B, T, C)
+        indices = indices.view(B, T, self.num_head)
+
+        # for output
+        x_q = self.recover_output(x_q, info)
+        indices = self.recover_output(indices, info)
+        
+        return x_q, loss, indices
+
+def sinkhorn(cost: Tensor, n_iters: int = 3, epsilon: float = 1, is_distributed: bool = False):
+    """
+    Sinkhorn algorithm.
+    Args:
+        cost (Tensor): shape with (B, K)
+    """
+    Q = torch.exp(- cost * epsilon).t() # (K, B)
+    if is_distributed:
+        B = Q.size(1) * dist.get_world_size()
+    else:
+        B = Q.size(1)
+    K = Q.size(0)
+
+    # make the matrix sums to 1
+    sum_Q = torch.sum(Q)
+    if is_distributed:
+        dist.all_reduce(sum_Q)
+    Q /= (sum_Q + 1e-8)
+
+    for _ in range(n_iters):
+        # normalize each row: total weight per prototype must be 1/K
+        sum_of_rows = torch.sum(Q, dim=1, keepdim=True)
+        if is_distributed:
+            dist.all_reduce(sum_of_rows)
+        Q /= (sum_of_rows + 1e-8)
+        Q /= K
+
+        # normalize each column: total weight per sample must be 1/B
+        Q /= (torch.sum(Q, dim=0, keepdim=True) + 1e-8)
+        Q /= B
+    
+    Q *= B # the columns must sum to 1 so that Q is an assignment
+    return Q.t() # (B, K)
+
+class VectorQuantizerSinkhorn(VectorQuantizer):
+    def __init__(self, epsilon: float = 10.0, n_iters: int = 5, 
+                 normalize_mode: str = "all", use_prob: bool = True,
+                 *args, **kwargs):
+        super(VectorQuantizerSinkhorn, self).__init__(*args, **kwargs)
+        self.epsilon = epsilon
+        self.n_iters = n_iters
+        self.normalize_mode = normalize_mode
+        self.use_prob = use_prob
+
+    def normalize(self, A, dim, mode="all"):
+        if mode == "all":
+            A = (A - A.mean()) / (A.std() + 1e-6)
+            A = A - A.min()
+        elif mode == "dim":
+            A = A / math.sqrt(dim)
+        elif mode == "null":
+            pass
+        return A
+
+    def quantize_input(self, query, reference):
+        # compute the distance matrix
+        query2ref = torch.cdist(query, reference, p=2.0) # (B1, B2)
+        
+        # compute the assignment matrix
+        with torch.no_grad():
+            is_distributed = dist.is_initialized() and dist.get_world_size() > 1
+            normalized_cost = self.normalize(query2ref, dim=reference.size(1), mode=self.normalize_mode)
+            Q = sinkhorn(normalized_cost, n_iters=self.n_iters, epsilon=self.epsilon, is_distributed=is_distributed)
+                
+        if self.use_prob:
+            # avoid the zero value problem
+            max_q_id = torch.argmax(Q, dim=-1)
+            Q[torch.arange(Q.size(0)), max_q_id] += 1e-8
+            indices = torch.multinomial(Q, num_samples=1).squeeze()
+        else:
+            indices = torch.argmax(Q, dim=-1)
+        nearest_ref = reference[indices]
+
+        if self.training and L.GET_STATS:
+            if L.log.total_steps % 1000 == 0:
+                L.log.add_histogram("params/normalized_cost", normalized_cost)
+        
+        return indices, nearest_ref, query2ref
+
+class Identity(VectorQuantizer):
+    @torch.autocast(device_type="cuda", enabled=False)
+    def forward(self, x: Tensor):
+        x = x.float()
+        loss_q = torch.tensor(0.0, device=x.device, dtype=x.dtype)
+
+        # compute statistics
+        if self.training and L.GET_STATS:
+            with torch.no_grad():
+                x_flatten, _ = self.reshape_input(x)
+                x_norm_mean = torch.mean(x_flatten.norm(dim=-1))
+                L.log.add_scalar("params/x_norm", x_norm_mean.item())
+        
+        return x, loss_q, None

optvq/models/vqgan.py

@@ -0,0 +1,168 @@
+# ------------------------------------------------------------------------------
+# OptVQ: Preventing Local Pitfalls in Vector Quantization via Optimal Transport
+# Copyright (c) 2024 Borui Zhang. All Rights Reserved.
+# Licensed under the MIT License [see LICENSE for details]
+# ------------------------------------------------------------------------------
+# Modified from [CompVis/taming-transformers](https://github.com/CompVis/taming-transformers)
+# Copyright (c) 2020 Patrick Esser and Robin Rombach and Björn Ommer. All Rights Reserved.
+# ------------------------------------------------------------------------------
+
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+
+import optvq.utils.logger as L
+
+class Identity(nn.Module):
+    def forward(self, x):
+        return x
+
+class VQModel(nn.Module):
+    def __init__(self,
+            encoder: nn.Module,
+            decoder: nn.Module,
+            loss: nn.Module,
+            quantize: nn.Module,
+            ckpt_path: str = None,
+            ignore_keys=[],
+            image_key="image",
+            colorize_nlabels=None,
+            monitor=None,
+            use_connector: bool = True,
+        ):
+        super(VQModel, self).__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.loss = loss
+        self.quantize = quantize
+        self.use_connector = use_connector
+
+        encoder_dim = self.encoder.hidden_dim
+        decoder_dim = self.decoder.hidden_dim
+        embed_dim = self.quantize.e_dim
+        
+        if not use_connector:
+            self.quant_conv = Identity()
+            self.post_quant_conv = Identity()
+            assert encoder_dim == embed_dim, f"{encoder_dim} != {embed_dim}"
+            assert decoder_dim == embed_dim, f"{decoder_dim} != {embed_dim}"
+        else:
+            self.quant_conv = torch.nn.Conv2d(encoder_dim, embed_dim, 1)
+            self.post_quant_conv = torch.nn.Conv2d(embed_dim, decoder_dim, 1)
+
+        if ckpt_path is not None:
+            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
+        self.image_key = image_key
+        if colorize_nlabels is not None:
+            assert type(colorize_nlabels)==int
+            self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
+        if monitor is not None:
+            self.monitor = monitor
+
+    def init_from_ckpt(self, path, ignore_keys=list()):
+        sd = torch.load(path, map_location="cpu")["state_dict"]
+        keys = list(sd.keys())
+        for k in keys:
+            for ik in ignore_keys:
+                if k.startswith(ik):
+                    print("Deleting key {} from state_dict.".format(k))
+                    del sd[k]
+        self.load_state_dict(sd, strict=False)
+        print(f"Restored from {path}")
+
+    def encode(self, x):
+        h = self.encoder(x)
+        h = self.quant_conv(h)
+        quant, emb_loss, indices = self.quantize(h)
+        return quant, emb_loss, indices
+
+    def decode(self, quant):
+        quant = self.post_quant_conv(quant)
+        dec = self.decoder(quant)
+        return dec
+
+    def decode_code(self, code_b):
+        quant_b = self.quantize.embed_code(code_b)
+        dec = self.decode(quant_b)
+        return dec
+
+    def forward(self, x, mode: int = 0, global_step: int = None):
+        """
+        Args:
+            x (torch.Tensor): input tensor
+            mode (int): 0 for autoencoder, 1 for discriminator
+            global_step (int): global step for adaptive discriminator weight
+        """
+        global_step = global_step if global_step is not None else L.log.total_steps
+        quant, qloss, indices = self.encode(x)
+        xrec = self.decode(quant)
+        if mode == 0:
+            # compute the autoencoder loss
+            loss, log_dict = self.loss(qloss, x, xrec, mode, last_layer=self.get_last_layer(), global_step=global_step)
+        elif mode == 1:
+            # compute the discriminator loss
+            loss, log_dict = self.loss(qloss, x, xrec, mode, last_layer=self.get_last_layer(), global_step=global_step)
+        elif mode == 2:
+            # compute the hidden embedding
+            h = self.encoder(x)
+            h = self.quant_conv(h)
+            return h
+        return loss, log_dict, indices
+
+    def get_input(self, batch, k):
+        x = batch[k]
+        if len(x.shape) == 3:
+            x = x[..., None]
+        x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format)
+        return x.float()
+
+    def get_last_layer(self):
+        if hasattr(self.decoder, "conv_out"):
+            return self.decoder.conv_out.weight
+        elif hasattr(self.decoder, "out_fc"):
+            return self.decoder.out_fc.weight
+        elif hasattr(self.decoder, "inv_conv"):
+            return self.decoder.inv_conv.weight
+        else:
+            raise NotImplementedError(f"Cannot find last layer in decoder")
+
+    def log_images(self, batch, **kwargs):
+        log = dict()
+        x = self.get_input(batch, self.image_key)
+        x = x.to(self.device)
+        xrec, _ = self(x)
+        if x.shape[1] > 3:
+            # colorize with random projection
+            assert xrec.shape[1] > 3
+            x = self.to_rgb(x)
+            xrec = self.to_rgb(xrec)
+        log["inputs"] = x
+        log["reconstructions"] = xrec
+        return log
+
+    def to_rgb(self, x):
+        assert self.image_key == "segmentation"
+        if not hasattr(self, "colorize"):
+            self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
+        x = F.conv2d(x, weight=self.colorize)
+        x = 2.*(x-x.min())/(x.max()-x.min()) - 1.
+        return x
+
+    # The functions below are deprecated
+
+    def validation_step(self, batch, batch_idx):
+        x = self.get_input(batch, self.image_key)
+        xrec, qloss = self(x)
+        aeloss, log_dict_ae = self.loss(qloss, x, xrec, 0, self.global_step,
+                                            last_layer=self.get_last_layer(), split="val")
+
+        discloss, log_dict_disc = self.loss(qloss, x, xrec, 1, self.global_step,
+                                            last_layer=self.get_last_layer(), split="val")
+        rec_loss = log_dict_ae["val/rec_loss"]
+        self.log("val/rec_loss", rec_loss,
+                   prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log("val/aeloss", aeloss,
+                   prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
+        self.log_dict(log_dict_ae)
+        self.log_dict(log_dict_disc)
+        return self.log_dict

optvq/models/vqgan_hf.py

@@ -0,0 +1,69 @@
+# ------------------------------------------------------------------------------
+# OptVQ: Preventing Local Pitfalls in Vector Quantization via Optimal Transport
+# Copyright (c) 2024 Borui Zhang. All Rights Reserved.
+# Licensed under the MIT License [see LICENSE for details]
+# ------------------------------------------------------------------------------
+
+# Convert a Pytorch model to a Hugging Face model
+
+import torch.nn as nn
+
+from huggingface_hub import PyTorchModelHubMixin
+
+from optvq.models.backbone.diffusion import Encoder, Decoder
+from optvq.models.quantizer import VectorQuantizer, VectorQuantizerSinkhorn
+from optvq.losses.aeloss_disc import AELossWithDisc
+from optvq.models.vqgan import VQModel
+
+class VQModelHF(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, 
+            encoder: dict = {},
+            decoder: dict = {},
+            loss: dict = {},
+            quantize: dict = {},
+            quantize_type: str = "optvq",
+            ckpt_path: str = None,
+            ignore_keys=[],
+            image_key="image",
+            colorize_nlabels=None,
+            monitor=None,
+            use_connector: bool = True,
+        ):
+        super(VQModelHF, self).__init__()
+        encoder = Encoder(**encoder)
+        decoder = Decoder(**decoder)
+        quantizer = self.setup_quantizer(quantize, quantize_type)
+        loss = AELossWithDisc(**loss)
+
+        self.model = VQModel(
+            encoder=encoder,
+            decoder=decoder,
+            loss=loss,
+            quantize=quantizer,
+            ckpt_path=ckpt_path,
+            ignore_keys=ignore_keys,
+            image_key=image_key,
+            colorize_nlabels=colorize_nlabels,
+            monitor=monitor,
+            use_connector=use_connector,
+        )
+    
+    def setup_quantizer(self, quantizer_config, quantize_type):
+        if quantize_type == "optvq":
+            quantizer = VectorQuantizerSinkhorn(**quantizer_config)
+        elif quantize_type == "basevq":
+            quantizer = VectorQuantizer(**quantizer_config)
+        else:
+            raise ValueError(f"Unknown quantizer type: {quantize_type}")
+        return quantizer
+    
+    def encode(self, x):
+        return self.model.encode(x)
+
+    def decode(self, x):
+        return self.model.decode(x)
+
+    def forward(self, x):
+        quant, *_ = self.encode(x)
+        rec = self.decode(quant)
+        return quant, rec

optvq/trainer/arguments.py

@@ -0,0 +1,53 @@
+# ------------------------------------------------------------------------------
+# OptVQ: Preventing Local Pitfalls in Vector Quantization via Optimal Transport
+# Copyright (c) 2024 Borui Zhang. All Rights Reserved.
+# Licensed under the MIT License [see LICENSE for details]
+# ------------------------------------------------------------------------------
+
+import argparse
+
+def get_parser():
+    parser = argparse.ArgumentParser()
+    # arguments with high priority
+    parser.add_argument("--seed", type=int, default=42,
+                        help="The random seed.")
+    parser.add_argument("--gpu", type=int, nargs="+", default=None,
+                        help="The GPU ids to use.")
+    parser.add_argument("--is_distributed", action="store_true", default=False)
+    parser.add_argument("--config", type=str, default=None,
+                        help="The path to the configuration file.")
+    parser.add_argument("--resume", type=str, default=None,
+                        help="The path to the checkpoint to resume.")
+    parser.add_argument("--device_rank", type=int, default=0)
+
+    # arguments for the training
+    parser.add_argument("--log_dir", type=str, default=None,
+                        help="The path to the log directory.")
+    parser.add_argument("--mode", type=str, default="train",
+                        help="options: train, test")
+    parser.add_argument("--use_initiate", type=str, default=None, 
+                        help="Options: random, kmeans")
+    parser.add_argument("--epochs", type=int, default=None)
+    parser.add_argument("--enterpoint", type=str, default=None)
+    parser.add_argument("--code_path", type=str, default=None)
+    parser.add_argument("--embed_path", type=str, default=None)
+    
+    # arguments for the data
+    parser.add_argument("--use_train_subset", type=float, default=None,
+                        help="The size of the training subset. None means using the full training set.")
+    parser.add_argument("--use_train_repeat", type=int, default=None,
+                        help="The number of times to repeat the training set.")
+    parser.add_argument("--use_val_subset", type=float, default=None,
+                        help="The size of the validation subset. None means using the full validation set.")
+    parser.add_argument("--batch_size", type=int, default=None)
+    parser.add_argument("--gradient_accumulate", type=int, default=None)
+
+    # arguments for the optimizer
+    parser.add_argument("--lr", type=float, default=None)
+    parser.add_argument("--mul_lr", type=float, default=None)
+
+    # arguments for the model
+    parser.add_argument("--num_codes", type=int, default=None,
+                        help="The number of codes.")
+    
+    return parser

optvq/trainer/pipeline.py

@@ -0,0 +1,362 @@
+# ------------------------------------------------------------------------------
+# OptVQ: Preventing Local Pitfalls in Vector Quantization via Optimal Transport
+# Copyright (c) 2024 Borui Zhang. All Rights Reserved.
+# Licensed under the MIT License [see LICENSE for details]
+# ------------------------------------------------------------------------------
+
+from typing import Callable
+import argparse
+import os
+from omegaconf import OmegaConf
+from functools import partial
+from torchinfo import summary
+
+import torch
+import torch.nn as nn
+from torch.utils.data.dataloader import DataLoader
+
+from optvq.utils.init import initiate_from_config_recursively
+from optvq.data.dataloader import maybe_get_subset
+import optvq.utils.logger as L
+
+def setup_config(opt: argparse.Namespace):
+    L.log.info("\n\n### Setting up the configurations. ###")
+
+    # load the config files
+    config = OmegaConf.load(opt.config)
+
+    # overwrite the certain arguments according to the config.args mapping
+    for key, value in config.args_map.items():
+        if hasattr(opt, key) and getattr(opt, key) is not None:
+            msg = f"config.{value} = opt.{key}"
+            L.log.info(f"Overwrite the config: {msg}")
+            exec(msg)
+    
+    return config
+
+def setup_dataloader(data, batch_size, is_distributed: bool = True, is_train: bool = True, num_workers: int = 8):
+    if is_train:
+        if is_distributed:
+            # setup the sampler
+            sampler = torch.utils.data.distributed.DistributedSampler(data, shuffle=True, drop_last=True)
+            # setup the dataloader
+            loader = DataLoader(
+                dataset=data, batch_size=batch_size, num_workers=num_workers,
+                drop_last=True, sampler=sampler, persistent_workers=True, pin_memory=True
+            )
+        else:
+            # setup the dataloader
+            loader = DataLoader(
+                dataset=data, batch_size=batch_size, num_workers=num_workers,
+                drop_last=True, shuffle=True, persistent_workers=True, pin_memory=True
+            )
+    else:
+        if is_distributed:
+            # setup the sampler
+            sampler = torch.utils.data.distributed.DistributedSampler(data, shuffle=False, drop_last=False)
+            # setup the dataloader
+            loader = DataLoader(
+                dataset=data, batch_size=batch_size, num_workers=num_workers,
+                drop_last=False, sampler=sampler, persistent_workers=True, pin_memory=True
+            )
+        else:
+            # setup the dataloader
+            loader = DataLoader(
+                dataset=data, batch_size=batch_size, num_workers=num_workers,
+                drop_last=False, shuffle=False, persistent_workers=True, pin_memory=True
+            )
+    
+    return loader
+
+def setup_dataset(config: OmegaConf):
+    L.log.info("\n\n### Setting up the datasets. ###")
+
+    # setup the training dataset
+    train_data = initiate_from_config_recursively(config.data.train)
+    if config.data.use_train_subset is not None:
+        train_data = maybe_get_subset(train_data, subset_size=config.data.use_train_subset, num_data_repeat=config.data.use_train_repeat)
+    L.log.info(f"Training dataset size: {len(train_data)}")
+
+    # setup the validation dataset
+    val_data = initiate_from_config_recursively(config.data.val)
+    if config.data.use_val_subset is not None:
+        val_data = maybe_get_subset(val_data, subset_size=config.data.use_val_subset)
+    L.log.info(f"Validation dataset size: {len(val_data)}")
+
+    return train_data, val_data
+
+def setup_model(config: OmegaConf, device):
+    L.log.info("\n\n### Setting up the models. ###")
+
+    # setup the model
+    model = initiate_from_config_recursively(config.model.autoencoder)
+    if config.is_distributed:
+        # apply syncBN
+        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
+        # model to devices
+        model = model.to(device)
+        find_unused_parameters = True
+        model = torch.nn.parallel.DistributedDataParallel(
+            module=model, device_ids=[config.gpu],
+            find_unused_parameters=find_unused_parameters
+        )
+        model_ori = model.module
+    else:
+        model = model.to(device)
+        model_ori = model
+    
+    input_size = config.data.train.params.transform.params.resize
+    in_channels = getattr(model_ori.encoder, "in_dim", 3)
+    sout = summary(model_ori, (1, in_channels, input_size, input_size), device="cuda", verbose=0)
+    L.log.info(sout)
+
+    # count the total number of parameters
+    for name, module in model_ori.named_children():
+        num_params = sum(p.numel() for p in module.parameters())
+        num_trainable = sum(p.numel() for p in module.parameters() if p.requires_grad)
+        L.log.info(f"Module: {name}, Total params: {num_params}, Trainable params: {num_trainable}")
+
+    return model
+
+### factory functions
+
+def get_setup_optimizers(config):
+    name = config.train.pipeline
+    func_name = "setup_optimizers_" + name
+    return globals()[func_name]
+
+def get_pipeline(config):
+    name = config.train.pipeline
+    func_name = "pipeline_" + name
+    return globals()[func_name]
+
+def _forward_backward(
+    config,
+    x: torch.Tensor,
+    forward: Callable,
+    model: nn.Module,
+    optimizer: torch.optim.Optimizer,
+    scheduler: torch.optim.lr_scheduler._LRScheduler,
+    scaler: torch.cuda.amp.GradScaler,
+):
+    with torch.autocast(device_type="cuda", dtype=torch.bfloat16, 
+                        enabled=config.use_amp):
+        # forward pass
+        loss, *output = forward(x)
+        loss_acc = loss / config.data.gradient_accumulate
+    scaler.scale(loss_acc).backward()
+    # gradient accumulate
+    if L.log.total_steps % config.data.gradient_accumulate == 0:
+        scaler.unscale_(optimizer)
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        scaler.step(optimizer)
+        optimizer.zero_grad()
+        scaler.update()
+    
+    if scheduler is not None:
+        scheduler.step()
+    return loss, output
+
+### autoencoder version
+def _find_weight_decay_id(modules: list, params_ids: list, 
+                          include_class: tuple = (nn.Linear, nn.Conv2d, 
+                                                  nn.ConvTranspose2d,
+                                                  nn.MultiheadAttention), 
+                          include_name: list = ["weight"]):
+    for mod in modules:
+        for sub_mod in mod.modules():
+            if isinstance(sub_mod, include_class):
+                for name, param in sub_mod.named_parameters():
+                    if any([k in name for k in include_name]):
+                        params_ids.append(id(param))
+    params_ids = list(set(params_ids))
+    return params_ids
+
+def set_weight_decay(modules: list):
+    weight_decay_ids = _find_weight_decay_id(modules, [])
+    wd_params, wd_names, no_wd_params, no_wd_names = [], [], [], []
+    for mod in modules:
+        for name, param in mod.named_parameters():
+            if id(param) in weight_decay_ids:
+                wd_params.append(param)
+                wd_names.append(name)
+            else:
+                no_wd_params.append(param)
+                no_wd_names.append(name)
+    return wd_params, wd_names, no_wd_params, no_wd_names
+
+def setup_optimizers_ae(config: OmegaConf, model: nn.Module, total_steps: int):
+    L.log.info("\n\n### Setting up the optimizers and schedulers. ###")
+    
+    # compute the total batch size and the learning rate
+    total_batch_size = config.data.batch_size * config.world_size * config.data.gradient_accumulate
+    total_learning_rate = config.train.learning_rate * total_batch_size
+    multipled_learning_rate = total_learning_rate * config.train.mul_learning_rate
+    L.log.info(f"Total batch size: {total_batch_size} = {config.data.batch_size} * {config.world_size} * {config.data.gradient_accumulate}")
+    L.log.info(f"Total learning rate: {total_learning_rate} = {config.train.learning_rate} * {total_batch_size}")
+    L.log.info(f"Multipled learning rate: {multipled_learning_rate} = {total_learning_rate} * {config.train.mul_learning_rate}")
+
+    # setup the optimizers
+    param_group = []
+    ## base learning rate
+    wd_params, wd_names, no_wd_params, no_wd_names = set_weight_decay([model.encoder, model.decoder, model.quant_conv, model.post_quant_conv])
+    param_group.append({
+        "params": wd_params, "lr": total_learning_rate, "eps": 1e-7,
+        "weight_decay": config.train.weight_decay, "beta": (0.9, 0.999),
+    })
+    param_group.append({
+        "params": no_wd_params, "lr": total_learning_rate, "eps": 1e-7,
+        "weight_decay": 0.0, "beta": (0.9, 0.999),
+    })
+    ## multipled learning rate
+    wd_params, wd_names, no_wd_params, no_wd_names = set_weight_decay([model.quantize])
+    param_group.append({
+        "params": wd_params, "lr": multipled_learning_rate, "eps": 1e-7,
+        "weight_decay": config.train.weight_decay, "beta": (0.9, 0.999),
+    })
+    param_group.append({
+        "params": no_wd_params, "lr": multipled_learning_rate, "eps": 1e-7,
+        "weight_decay": 0.0, "beta": (0.9, 0.999),
+    })
+
+    optimizer_ae = torch.optim.AdamW(param_group)
+    optimizer_dict = {"optimizer_ae": optimizer_ae}
+
+    # setup the schedulers
+    scheduler_ae = torch.optim.lr_scheduler.OneCycleLR(
+        optimizer=optimizer_ae, max_lr=[total_learning_rate, total_learning_rate, multipled_learning_rate, multipled_learning_rate],
+        total_steps=total_steps, pct_start=0.01, anneal_strategy="cos"
+    )
+    scheduler_dict = {"scheduler_ae": scheduler_ae}
+
+    # setup the scalers
+    scaler_dict = {"scaler_ae": torch.GradScaler(enabled=config.use_amp)}
+    L.log.info(f"Enable AMP: {config.use_amp}")
+    return optimizer_dict, scheduler_dict, scaler_dict
+
+def pipeline_ae(
+    config,
+    x: torch.Tensor,
+    model: nn.Module,
+    optimizers: dict,
+    schedulers: dict,
+    scalers: dict,
+):
+    assert "optimizer_ae" in optimizers
+    assert "scheduler_ae" in schedulers
+    assert "scaler_ae" in scalers
+
+    optimizer = optimizers["optimizer_ae"]
+    scheduler = schedulers["scheduler_ae"]
+    scaler = scalers["scaler_ae"]
+
+    forward = partial(model, mode=0)
+    _, (loss_ae_dict, indices) = _forward_backward(config, x, forward, model, optimizer, scheduler, scaler)
+    
+    log_per_step = loss_ae_dict
+    log_per_epoch = {"indices": indices}
+    return log_per_step, log_per_epoch
+
+### autoencoder + disc version
+
+def setup_optimizers_ae_disc(config: OmegaConf, model: nn.Module, total_steps: int):
+    L.log.info("\n\n### Setting up the optimizers and schedulers. ###")
+    
+    # compute the total batch size and the learning rate
+    total_batch_size = config.data.batch_size * config.world_size * config.data.gradient_accumulate
+    total_learning_rate = config.train.learning_rate * total_batch_size
+    multipled_learning_rate = total_learning_rate * config.train.mul_learning_rate
+    L.log.info(f"Total batch size: {total_batch_size} = {config.data.batch_size} * {config.world_size} * {config.data.gradient_accumulate}")
+    L.log.info(f"Total learning rate: {total_learning_rate} = {config.train.learning_rate} * {total_batch_size}")
+    L.log.info(f"Multipled learning rate: {multipled_learning_rate} = {total_learning_rate} * {config.train.mul_learning_rate}")
+
+    # setup the optimizers
+    param_group = []
+    ## base learning rate
+    wd_params, wd_names, no_wd_params, no_wd_names = set_weight_decay([model.encoder, model.decoder, model.quant_conv, model.post_quant_conv])
+    param_group.append({
+        "params": wd_params, "lr": total_learning_rate, "eps": 1e-7,
+        "weight_decay": config.train.weight_decay, "beta": (0.9, 0.999),
+    })
+    param_group.append({
+        "params": no_wd_params, "lr": total_learning_rate, "eps": 1e-7,
+        "weight_decay": 0.0, "beta": (0.9, 0.999),
+    })
+    ## multipled learning rate
+    wd_params, wd_names, no_wd_params, no_wd_names = set_weight_decay([model.quantize])
+    param_group.append({
+        "params": wd_params, "lr": multipled_learning_rate, "eps": 1e-7,
+        "weight_decay": config.train.weight_decay, "beta": (0.9, 0.999),
+    })
+    param_group.append({
+        "params": no_wd_params, "lr": multipled_learning_rate, "eps": 1e-7,
+        "weight_decay": 0.0, "beta": (0.9, 0.999),
+    })
+    optimizer_ae = torch.optim.AdamW(param_group)
+
+    param_group = []
+    wd_params, wd_names, no_wd_params, no_wd_names = set_weight_decay([model.loss.discriminator])
+    param_group.append({
+        "params": wd_params, "lr": total_learning_rate, "eps": 1e-7,
+        "weight_decay": config.train.weight_decay, "beta": (0.9, 0.999),
+    })
+    param_group.append({
+        "params": no_wd_params, "lr": total_learning_rate, "eps": 1e-7,
+        "weight_decay": 0.0, "beta": (0.9, 0.999),
+    })
+    optimizer_disc = torch.optim.AdamW(param_group)
+    optimizer_dict = {"optimizer_ae": optimizer_ae, "optimizer_disc": optimizer_disc}
+
+    # setup the schedulers
+    scheduler_ae = torch.optim.lr_scheduler.OneCycleLR(
+        optimizer=optimizer_ae, max_lr=[total_learning_rate, total_learning_rate, multipled_learning_rate, multipled_learning_rate],
+        total_steps=total_steps, pct_start=0.01, anneal_strategy="cos"
+    )
+    scheduler_disc = torch.optim.lr_scheduler.OneCycleLR(
+        optimizer=optimizer_disc, max_lr=[total_learning_rate, total_learning_rate],
+        total_steps=total_steps, pct_start=0.01, anneal_strategy="cos"
+    )
+    scheduler_dict = {"scheduler_ae": scheduler_ae, "scheduler_disc": scheduler_disc}
+
+    # setup the scalers
+    scaler_dict = {"scaler_ae": torch.GradScaler(enabled=config.use_amp), 
+                   "scaler_disc": torch.GradScaler(enabled=config.use_amp)}
+    L.log.info(f"Enable AMP: {config.use_amp}")
+    return optimizer_dict, scheduler_dict, scaler_dict
+
+def pipeline_ae_disc(
+    config, 
+    x: torch.Tensor,
+    model: nn.Module,
+    optimizers: dict,
+    schedulers: dict,
+    scalers: dict,
+):
+    # autoencoder step
+    assert "optimizer_ae" in optimizers
+    assert "scheduler_ae" in schedulers
+    assert "scaler_ae" in scalers
+
+    optimizer = optimizers["optimizer_ae"]
+    scheduler = schedulers["scheduler_ae"]
+    scaler = scalers["scaler_ae"]
+
+    forward = partial(model, mode=0)
+    _, (loss_ae_dict, indices) = _forward_backward(config, x, forward, model, optimizer, scheduler, scaler)
+    
+    log_per_step = loss_ae_dict
+    log_per_epoch = {"indices": indices}
+
+    # discriminator step
+    assert "optimizer_disc" in optimizers
+    assert "scheduler_disc" in schedulers
+    assert "scaler_disc" in scalers
+
+    optimizer = optimizers["optimizer_disc"]
+    scheduler = schedulers["scheduler_disc"]
+    scaler = scalers["scaler_disc"]
+
+    forward = partial(model, mode=1)
+    _, (loss_disc_dict, _) = _forward_backward(config, x, forward, model, optimizer, scheduler, scaler)
+    log_per_step.update(loss_disc_dict)
+    return log_per_step, log_per_epoch

optvq/utils/func.py

@@ -0,0 +1,27 @@
+# ------------------------------------------------------------------------------
+# OptVQ: Preventing Local Pitfalls in Vector Quantization via Optimal Transport
+# Copyright (c) 2024 Borui Zhang. All Rights Reserved.
+# Licensed under the MIT License [see LICENSE for details]
+# ------------------------------------------------------------------------------
+
+from typing import Tuple, Union, Iterable
+from omegaconf import OmegaConf
+import os
+
+import torch
+import torch.distributed as dist
+
+def dist_all_gather(x):
+    tensor_list = [torch.zeros_like(x) for _ in range(dist.get_world_size())]
+    dist.all_gather(tensor_list, x)
+    x = torch.cat(tensor_list, dim=0)
+    return x
+
+def any_2tuple(data: Union[int, Tuple[int]]) -> Tuple[int]:
+    if isinstance(data, int):
+        return (data, data)
+    elif isinstance(data, Iterable):
+        assert len(data) == 2, "target size must be tuple of (w, h)"
+        return tuple(data)
+    else:
+        raise ValueError("target size must be int or tuple of (w, h)") 

optvq/utils/init.py

@@ -0,0 +1,36 @@
+# ------------------------------------------------------------------------------
+# OptVQ: Preventing Local Pitfalls in Vector Quantization via Optimal Transport
+# Copyright (c) 2024 Borui Zhang. All Rights Reserved.
+# Licensed under the MIT License [see LICENSE for details]
+# ------------------------------------------------------------------------------
+
+import importlib
+import random
+import numpy as np
+from typing import Mapping, Optional
+
+import torch
+
+def initiate_from_config(config: Mapping):
+    assert "target" in config, f"Expected key `target` to initialize!"
+    module, cls = config["target"].rsplit(".", 1)
+    meta_class = getattr(importlib.import_module(module, package=None), cls)
+    return meta_class(**config.get("params", dict()))
+
+def initiate_from_config_recursively(config: Mapping):
+    assert "target" in config, f"Expected key `target` to initialize!"
+    update_config = {"target": config["target"], "params": {}}
+    for k, v in config["params"].items():
+        if isinstance(v, Mapping) and "target" in v:
+            sub_instance = initiate_from_config_recursively(v)
+            update_config["params"][k] = sub_instance
+        else:
+            update_config["params"][k] = v
+    return initiate_from_config(update_config)
+
+def seed_everything(seed: Optional[int] = None):
+    seed = int(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True

optvq/utils/logger.py

@@ -0,0 +1,386 @@
+# ------------------------------------------------------------------------------
+# OptVQ: Preventing Local Pitfalls in Vector Quantization via Optimal Transport
+# Copyright (c) 2024 Borui Zhang. All Rights Reserved.
+# Licensed under the MIT License [see LICENSE for details]
+# ------------------------------------------------------------------------------
+
+import time
+import datetime
+from typing import List
+import functools
+import os
+from PIL import Image
+from termcolor import colored
+import sys
+import logging
+from omegaconf import OmegaConf
+import json
+
+try:
+    from torch.utils.tensorboard import SummaryWriter
+    from torch import Tensor
+    import torch
+except:
+    raise ImportError("Please install torch to use this module!")
+
+"""
+NOTE: The `log` instance is a global variable, which should be imported by other modules as:
+    `import optvq.utils.logger as logger` 
+rather than 
+    `from optvq.utils.logger import log`.
+"""
+
+def setup_printer(file_log_dir: str, use_console: bool = True):
+    printer = logging.getLogger("LOG")
+    printer.setLevel(logging.DEBUG)
+    printer.propagate = False
+
+    # create formatter
+    fmt = '[%(asctime)s %(name)s] (%(filename)s %(lineno)d): %(levelname)s %(message)s'
+    color_fmt = colored('[%(asctime)s %(name)s]', 'green') + \
+                colored('(%(filename)s %(lineno)d)', 'yellow') + ': %(levelname)s %(message)s'
+    
+    # create the console handler
+    if use_console:
+        console_handler = logging.StreamHandler(sys.stdout)
+        console_handler.setLevel(logging.DEBUG)
+        console_handler.setFormatter(
+            logging.Formatter(fmt=color_fmt, datefmt="%Y-%m-%d %H:%M:%S")
+        )
+        printer.addHandler(console_handler)
+
+    # create the file handler
+    file_handler = logging.FileHandler(os.path.join(file_log_dir, "record.txt"), mode="a")
+    file_handler.setLevel(logging.DEBUG)
+    file_handler.setFormatter(
+        logging.Formatter(fmt=fmt, datefmt="%Y-%m-%d %H:%M:%S")
+    )
+    printer.addHandler(file_handler)
+
+    return printer
+
+@functools.lru_cache()
+def config_loggers(log_dir: str, local_rank: int = 0, master_rank: int = 0):
+    global log
+
+    if local_rank == master_rank:
+        log = LogManager(log_dir=log_dir, main_logger=True)
+    else:
+        log = LogManager(log_dir=log_dir, main_logger=False)
+
+class ProgressWithIndices:
+    def __init__(self, total: int, sep_char: str = "| ", 
+                 num_per_row: int = 4):
+        self.total = total
+        self.sep_char = sep_char
+        self.num_per_row = num_per_row
+
+        self.count = 0
+        self.start_time = time.time()
+        self.past_time = None
+        self.current_time = None
+        self.eta = None
+        self.speed = None
+        self.used_time = 0
+
+    def update(self):
+        self.count += 1
+        if self.count <= self.total:
+            self.past_time = self.current_time
+            self.current_time = time.time()
+            # compute eta
+            if self.past_time is not None:
+                self.eta = (self.total - self.count) * (self.current_time - self.past_time)
+                self.eta = str(datetime.timedelta(seconds=int(self.eta)))
+                self.speed = 1 / (self.current_time - self.past_time + 1e-8)
+            # compute used time
+            self.used_time = self.current_time - self.start_time
+            self.used_time = str(datetime.timedelta(seconds=int(self.used_time)))
+        else:
+            self.eta = 0
+            self.speed = 0
+            self.past_time = None
+            self.current_time = None
+
+    def print(self, prefix: str = "", content: str = "", ):
+        global log
+        prefix_str = f"{prefix}\t" + f"[{self.count}/{self.total} {self.used_time}/Eta:{self.eta}], Speed:{self.speed}iters/s\n"
+        content_list = content.split(self.sep_char)
+        content_list = [content.strip() for content in content_list]
+        content_list = [
+            "\t\t" + self.sep_char.join(content_list[i:i + self.num_per_row]) 
+            for i in range(0, len(content_list), self.num_per_row)
+        ]
+        content = prefix_str + "\n".join(content_list)
+        log.info(content)
+
+class LogManager:
+    """
+    This class encapsulates the tensorboard writer, the statistic meters, the console printer, and the progress counters.
+
+    Args:
+        log_dir (str): the parent directory to save all the logs
+        init_meters (List[str]): the initial meters to be shown
+        show_avg (bool): whether to show the average value of the meters
+    """
+    def __init__(self, log_dir: str, init_meters: List[str] = [], 
+                 show_avg: bool = True, main_logger: bool = False):
+        
+        # initiate all the directories
+        self.show_avg = show_avg
+        self.log_dir = log_dir
+        self.main_logger = main_logger
+        self.setup_dirs()
+
+        # initiate the statistic meters
+        self.meters = {meter: AverageMeter() for meter in init_meters}
+        
+        # initiate the progress counters
+        self.total_steps = 0
+        self.total_epochs = 0
+
+        if self.main_logger:
+            # initiate the tensorboard writer
+            self.board = SummaryWriter(log_dir=self.tb_log_dir)
+
+            # initiate the console printer
+            self.printer = setup_printer(self.file_log_dir, use_console=True)
+    
+    def state_dict(self):
+        return {
+            "total_steps": self.total_steps,
+            "total_epochs": self.total_epochs,
+            "meters": {
+                meter_name: meter.state_dict() for meter_name, meter in self.meters.items()
+            }
+        }
+    
+    def load_state_dict(self, state_dict: dict):
+        self.total_steps = state_dict["total_steps"]
+        self.total_epochs = state_dict["total_epochs"]
+        for meter_name, meter_state_dict in state_dict["meters"].items():
+            if meter_name not in self.meters:
+                self.meters[meter_name] = AverageMeter()    
+            self.meters[meter_name].load_state_dict(meter_state_dict)
+    
+    ### About directories
+    def setup_dirs(self):
+        """
+        The structure of the log directory:
+        - log_dir: [tb_log, txt_log, img_log, model_log]
+        """
+        self.tb_log_dir = os.path.join(self.log_dir, "tb_log")
+        # NOTE: For now, we save the txt records in the parent directory
+        # self.file_log_dir = os.path.join(self.log_dir, "txt_log")
+        self.file_log_dir = self.log_dir 
+        self.img_log_dir = os.path.join(self.log_dir, "img_log")
+
+        self.config_path = os.path.join(self.log_dir, "config.yaml")
+        self.checkpoint_path = os.path.join(self.log_dir, "checkpoint.pth")
+        self.backup_checkpoint_path = os.path.join(self.log_dir, "checkpoint.pth")
+        self.save_logger_path = os.path.join(self.log_dir, "logger.json")
+
+        if self.main_logger:
+            os.makedirs(self.tb_log_dir, exist_ok=True)
+            os.makedirs(self.file_log_dir, exist_ok=True)
+            os.makedirs(self.img_log_dir, exist_ok=True)
+
+    ### About printer
+
+    def info(self, msg, *args, **kwargs):
+        if self.main_logger:
+            self.printer.info(msg, *args, **kwargs)
+
+    def show(self, include_key: str = ""):
+        if isinstance(include_key, str):
+            include_key = [include_key]
+        if self.show_avg:
+            return "| ".join([f"{meter_name}: {meter.val:.4f}/{meter.avg:.4f}" for meter_name, meter in self.meters.items() if any([k in meter_name for k in include_key])])
+        else:
+            return "| ".join([f"{meter_name}: {meter.val:.4f}" for meter_name, meter in self.meters.items() if any([k in meter_name for k in include_key])])
+
+    ### About counter
+
+    def update_steps(self):
+        self.total_steps += 1
+        return self.total_steps
+    
+    def update_epochs(self):
+        self.total_epochs += 1
+        return self.total_epochs
+    
+    ### About tensorboard
+    def add_histogram(self, tag: str, values: Tensor, global_step: int = None):
+        if self.main_logger:
+            global_step = self.total_steps if global_step is None else global_step
+            self.board.add_histogram(tag, values, global_step)
+
+    def add_scalar(self, tag: str, scalar_value: float, global_step: int = None):
+        if isinstance(scalar_value, Tensor):
+            scalar_value = scalar_value.item()
+        if tag in self.meters:
+            cur_step = self.meters[tag].update(scalar_value)
+            cur_step = cur_step if global_step is None else global_step
+            if self.main_logger:
+                self.board.add_scalar(tag, scalar_value, cur_step)
+        else:
+            self.meters[tag] = AverageMeter()
+            cur_step = self.meters[tag].update(scalar_value)
+            cur_step = cur_step if global_step is None else global_step
+            if self.main_logger:
+                print(f"Create new meter: {tag}!")
+                self.board.add_scalar(tag, scalar_value, cur_step)
+    
+    def add_scalar_dict(self, scalar_dict: dict, global_step: int = None):
+        for tag, scalar_value in scalar_dict.items():
+            self.add_scalar(tag, scalar_value, global_step)
+        
+    def add_images(self, tag: str, images: Tensor, global_step: int = None):
+        if self.main_logger:
+            global_step = self.total_steps if global_step is None else global_step
+            self.board.add_images(tag, images, global_step, dataformats="NCHW")
+
+    ### About saving and resuming
+    def save_configs(self, config):
+        if self.main_logger:
+            # save config as yaml file
+            OmegaConf.save(config, self.config_path)
+            self.info(f"Save config to {self.config_path}.")
+            
+            # save logger
+            state_dict = self.state_dict()
+            with open(self.save_logger_path, "w") as f:
+                json.dump(state_dict, f)
+    
+    def load_configs(self):
+        # load config
+        assert os.path.exists(self.config_path), f"Config {self.config_path} does not exist!"
+        config = OmegaConf.load(self.config_path)
+
+        # load logger
+        assert os.path.exists(self.save_logger_path), f"Logger {self.save_logger_path} does not exist!"
+        state_dict = json.load(open(self.save_logger_path, "r"))
+        self.load_state_dict(state_dict)
+
+        return config
+
+    def save_checkpoint(self, model, optimizers, schedulers, scalers, suffix: str = ""):
+        """
+        checkpoint_dict: model, optimizer, scheduler, scalers
+        """
+        if self.main_logger:
+            
+            # save checkpoint_dict
+            checkpoint_dict = {
+                "model": model.state_dict(),
+                "epoch": self.total_epochs,
+                "step": self.total_steps
+            }
+            checkpoint_dict.update({k: v.state_dict() for k, v in optimizers.items()})
+            checkpoint_dict.update({k: v.state_dict() for k, v in schedulers.items() if v is not None})
+            checkpoint_dict.update({k: v.state_dict() for k, v in scalers.items()})
+
+            checkpoint_path = self.checkpoint_path + suffix
+            torch.save(checkpoint_dict, checkpoint_path)
+            if os.path.exists(self.backup_checkpoint_path):
+                os.remove(self.backup_checkpoint_path)
+            self.backup_checkpoint_path = checkpoint_path + f".epoch{self.total_epochs}"
+            torch.save(checkpoint_dict, self.backup_checkpoint_path)
+
+            self.info(f"### Epoch: {self.total_epochs}| Steps: {self.total_steps}| Save checkpoint to {checkpoint_path}.")
+    
+    def load_checkpoint(self, device, model, optimizers, schedulers, scalers, resume: str = None):
+        resume_path = self.checkpoint_path if resume is None else resume
+        assert os.path.exists(resume_path), f"Resume {resume_path} does not exist!"
+
+        # load checkpoint_dict
+        checkpoint_dict = torch.load(resume_path, map_location=device)
+        model.load_state_dict(checkpoint_dict["model"])
+        self.total_epochs = checkpoint_dict["epoch"]
+        self.total_steps = checkpoint_dict["step"]
+        for k, v in optimizers.items():
+            v.load_state_dict(checkpoint_dict[k])
+        for k, v in schedulers.items():
+            v.load_state_dict(checkpoint_dict[k])
+        for k, v in scalers.items():
+            v.load_state_dict(checkpoint_dict[k])
+        
+        self.info(f"### Epoch: {self.total_epochs}| Steps: {self.total_steps}| Resume checkpoint from {resume_path}.")
+
+        return self.total_epochs
+
+class EmptyManager:
+    def __init__(self):
+        for func_name in LogManager.__dict__.keys():
+            if not func_name.startswith("_"):
+                setattr(self, func_name, lambda *args, **kwargs: print(f"Empty Manager! {func_name} is not available!"))
+
+class AverageMeter:
+    def __init__(self):
+        self.reset()
+
+    def state_dict(self):
+        return {
+            "val": self.val,
+            "avg": self.avg,
+            "sum": self.sum,
+            "count": self.count,
+        }
+    
+    def load_state_dict(self, state_dict: dict):
+        self.val = state_dict["val"]
+        self.avg = state_dict["avg"]
+        self.sum = state_dict["sum"]
+        self.count = state_dict["count"]
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+        return 0
+
+    def update(self, val: float, n: int = 1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+        return self.count
+
+    def __str__(self):
+        return f"{self.avg:.4f}"
+
+def save_image(x: Tensor, save_path: str, scale_to_256: bool = True):
+    """
+    Args:
+        x (tensor): default data range is [0, 1]
+    """
+    if scale_to_256:
+        x = x.mul(255).clamp(0, 255)
+    x = x.permute(1, 2, 0).detach().cpu().numpy().astype("uint8")
+    img = Image.fromarray(x)
+    img.save(save_path)
+
+def save_images(images_list, ids_list, meta_path):
+    for i, (image, id) in enumerate(zip(images_list, ids_list)):
+        save_path = os.path.join(meta_path, f"{id}.png")
+        save_image(image, save_path)
+
+def save_images_multithread(images_list, ids_list, meta_path):
+    n_workers = 32
+    from concurrent.futures import ThreadPoolExecutor
+    with ThreadPoolExecutor(max_workers=n_workers) as executor:
+        for i in range(0, len(images_list), n_workers):
+            cur_images = images_list[i:(i + n_workers)]
+            cur_ids = ids_list[i:(i + n_workers)]
+            executor.submit(save_images, cur_images, cur_ids, meta_path)
+
+def add_prefix(log_dict: dict, prefix: str):
+    return {
+        f"{prefix}/{key}": val for key, val in log_dict.items()
+    }
+
+##################### GLOBAL VARIABLES #####################
+log = EmptyManager()
+GET_STATS: bool = (os.environ.get("ENABLE_STATS", "1") == "1")
+###########################################################

optvq/utils/metrics.py

@@ -0,0 +1,60 @@
+# ------------------------------------------------------------------------------
+# OptVQ: Preventing Local Pitfalls in Vector Quantization via Optimal Transport
+# Copyright (c) 2024 Borui Zhang. All Rights Reserved.
+# Licensed under the MIT License [see LICENSE for details]
+# ------------------------------------------------------------------------------
+
+import numpy as np
+
+import torch
+
+from pytorch_fid.inception import InceptionV3
+from pytorch_fid.fid_score import calculate_frechet_distance
+
+class FIDMetric:
+    def __init__(self, device, dims=2048):
+        self.device = device
+        self.num_workers = 32
+        
+        block_idx = InceptionV3.BLOCK_INDEX_BY_DIM[dims]
+        self.model = InceptionV3([block_idx]).to(device)
+        self.model.eval()
+
+        self.reset_metrics()
+    
+    def reset_metrics(self):
+        self.x_pred = []
+        self.x_rec_pred = []
+
+    @torch.no_grad()
+    def get_activates(self, x: torch.Tensor):
+        pred = self.model(x)[0]
+        # If model output is not scalar, apply global spatial average pooling.
+        # This happens if you choose a dimensionality not equal 2048.
+        if pred.size(2) != 1 or pred.size(3) != 1:
+            pred = torch.nn.functional.adaptive_avg_pool2d(pred, output_size=(1, 1))
+        return pred.squeeze().cpu().numpy()
+
+    def update(self, x: torch.Tensor, x_rec: torch.Tensor):
+        """
+        Args:
+            x (torch.Tensor): input tensor range from 0 to 1
+            x_rec (torch.Tensor): reconstructed tensor range from 0 to 1
+        """
+        self.x_pred.append(self.get_activates(x))
+        self.x_rec_pred.append(self.get_activates(x_rec))
+    
+    def result(self):
+        assert len(self.x_pred) != 0, "No data to compute FID"
+        x = np.concatenate(self.x_pred, axis=0)
+        x_rec = np.concatenate(self.x_rec_pred, axis=0)
+
+        x_mu = np.mean(x, axis=0)
+        x_sigma = np.cov(x, rowvar=False)
+        
+        x_rec_mu = np.mean(x_rec, axis=0)
+        x_rec_sigma = np.cov(x_rec, rowvar=False)
+
+        fid_score = calculate_frechet_distance(x_mu, x_sigma, x_rec_mu, x_rec_sigma)
+        self.reset_metrics()
+        return fid_score

requirements.txt

@@ -0,0 +1,25 @@
+torch>=2.3.0
+torchvision
+torchaudio
+numpy
+matplotlib
+scikit-learn
+opencv-python
+einops>=0.7.0
+omegaconf>=2.3.0
+lightning>=2.1.3
+transformers>=4.36.2
+pudb
+wandb
+tensorboard
+termcolor
+lpips==0.1.4
+ftfy
+imageio
+imageio-ffmpeg
+pyiqa
+clean-fid
+albumentations
+pytorch-fid
+torchinfo
+gradio