Add application files

Dockerfile

@@ -0,0 +1,10 @@
+FROM python:3.9.7
+
+WORKDIR /app
+COPY requirements.txt .
+RUN pip install -r --no-cache-dir requirements.txt
+# preload models
+
+COPY . .
+
+CMD ["python", "app.py"]

app.py

@@ -0,0 +1,143 @@
+import gradio as gr
+from PIL import Image
+from pathlib import Path
+from vq_gan_3d import load_VQGAN
+
+# Numerical computing
+import numpy as np
+
+# PyTorch
+import torch
+import torch.nn.functional as F
+
+# Utilities to calculate grids from SMILES and visualization
+from utils import get_grid_from_smiles, plot_voxel_grid, change_grid_size
+
+ckpt_path   = Path("./vq_gan_3d/weights/3DGrid-VQGAN_43.pt")
+folder      = str(ckpt_path.parent)      
+ckpt_file   = ckpt_path.name
+vqgan = load_VQGAN(folder=folder, ckpt_filename=ckpt_file).eval()
+
+
+def comparison(SMILES):
+    density_grids = get_grid_from_smiles([SMILES])
+
+    # 1) Prepare density grids → list of ready-to-use tensors
+    processed_tensors = []
+    for item in density_grids:
+        rho   = item["rho"]                   # raw NumPy array from cube generation
+        smi   = item["smiles"]
+        name  = item["name"]
+    
+        tensor = torch.from_numpy(rho).float() # convert grid to float32 tensor
+        tensor = torch.log1p(tensor)           # apply log(ρ + 1) normalization
+    
+        # enforce consistent 128×128×128 input size for VQGAN
+        if tensor.shape != torch.Size([128, 128, 128]):
+            tensor = tensor.unsqueeze(0).unsqueeze(0)  # add batch & channel dims
+            tensor = F.interpolate(
+                tensor,
+                size=(128, 128, 128),
+                mode="trilinear",
+                align_corners=False
+            )[0, 0]                                    # remove extra dims after resizing
+            print(f"[info] {smi} was interpolated to 128³")
+    
+        # store metadata alongside the processed tensor
+        processed_tensors.append({
+            "name":   name,
+            "smiles": smi,
+            "tensor": tensor
+        })
+    
+        # log shape and min/max to verify normalization and sizing
+        print(
+            f"{smi}: shape={tuple(tensor.shape)}, "
+            f"min={tensor.min():.4f}, max={tensor.max():.4f}"
+        )
+
+    # 2) Encode → Decode (inference with VQGAN)
+    reconstructions = []
+    for item in processed_tensors:
+        smi  = item["smiles"]       # original SMILES string
+        name = item["name"]         # unique grid name
+        vol  = item["tensor"]       # preprocessed [128³] tensor
+    
+        # add batch & channel dims and move to the selected device
+        x = vol.unsqueeze(0).unsqueeze(0)  # shape [1,1,128,128,128]
+    
+        with torch.no_grad():  # disable gradient computation for faster inference
+            indices = vqgan.encode(x)   # map input volume to discrete latent codes
+            recon   = vqgan.decode(indices)  # reconstruct volume from latent codes
+    
+        # convert reconstructed tensor and original tensor to NumPy arrays
+        recon_np = recon.cpu().numpy()[0, 0]
+        orig_np  = vol.cpu().numpy()
+    
+        # compute mean squared error between original and reconstruction
+        mse = np.mean((orig_np - recon_np) ** 2)
+        print(f"{smi} → reconstruction done | MSE={mse:.6f}")
+    
+        # collect results for later visualization
+        reconstructions.append({
+            "smiles":       smi,
+            "name":         name,
+            "original":     orig_np,
+            "reconstructed": recon_np
+        })
+
+    original_grid_plot = plot_voxel_grid(
+        change_grid_size(
+            torch.from_numpy(reconstructions[0]["original"]).unsqueeze(0).unsqueeze(0), 
+            size=(48, 48, 48)
+        ),
+        title=f"Original 3D Grid Plot from {SMILES}"
+    )
+    rec_grid_plot = plot_voxel_grid(
+        change_grid_size(
+            torch.from_numpy(reconstructions[0]["reconstructed"]).unsqueeze(0).unsqueeze(0), 
+            size=(48, 48, 48)
+        ),
+        title=f"Reconstructed 3D Grid Plot from {SMILES}"
+    )
+
+    np.save("original_grid.npy", reconstructions[0]["original"])
+    np.save("reconstructed_grid.npy", reconstructions[0]["reconstructed"])
+
+    original_grid_plot.savefig("original_grid_plot.png", format='png')
+    rec_grid_plot.savefig("reconstructed_grid_plot.png", format='png')
+    original_grid_plot = Image.open("original_grid_plot.png")
+    rec_grid_plot = Image.open("reconstructed_grid_plot.png")
+    
+    return [original_grid_plot, rec_grid_plot], mse, "original_grid.npy", "reconstructed_grid.npy" 
+
+
+with gr.Blocks() as demo:
+    gr.Markdown(
+        """
+        # 3DGrid-VQGAN SMILES to 3D Grid Reconstruction
+        **Single mode:** paste a SMILES string in the left box.  
+        **Batch mode:** upload a CSV file where each row has a SMILES in the first column.  
+        - **Maximum 1000 SMILES per batch.** Processing time increases with batch size due to Hugging Face environment limits.  
+        _This is just a demo environment; for heavy-duty usage, please visit:_  
+        https://github.com/IBM/materials/tree/main/models/smi_ted  
+        to download the model and run your own experiments.
+        - In both cases, an `embeddings.csv` file will be extracted for download, with the first column as SMILES and the embedding values in the following columns.
+        """
+    )
+    gr.Interface(
+        fn=comparison, 
+        inputs=[
+            gr.Dropdown(choices=["CCCO", "CC", "CCO"], label="Provide a SMILES or pre-select one", allow_custom_value=True)
+        ], 
+        outputs=[
+            gr.Gallery(label="3D Grid Reconstruction Comparison", columns=2),
+            gr.Number(label="MSE"),
+            gr.File(label="Original 3D Grid numpy file"),
+            gr.File(label="Reconstructed 3D Grid numpy file")
+        ]
+    )
+
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0")

requirements.txt

@@ -0,0 +1,18 @@
+torch>=2.1.0                
+numpy==1.26.4        
+pandas==1.4.0
+gradio>=4.33.1
+huggingface-hub
+pydantic==2.10.6
+rdkit>=2024.3.5
+imageio==2.34.1
+hydra-core==1.3.2
+omegaconf==2.3.0
+pdbpp==0.10.2
+torchvision==0.18.0
+tqdm==4.66.4
+requests>=2.32.0
+lpips==0.1.4
+pyscf>=2.10.0
+pyberny>=0.6.3
+git+https://github.com/pyscf/semiempirical

vq_gan_3d/__init__.py

@@ -0,0 +1,5 @@
+from vq_gan_3d.model import VQGAN, load_VQGAN
+from vq_gan_3d.model import Codebook
+from vq_gan_3d.model import LPIPS
+from vq_gan_3d.dataset import VQGANDataset
+from vq_gan_3d.utils import get_single_device

vq_gan_3d/dataset/__init__.py

@@ -0,0 +1,7 @@
+# from dataset.breast_uka import BreastUKA
+# from dataset.mrnet import MRNetDataset
+# from dataset.brats import BRATSDataset
+# from dataset.adni import ADNIDataset
+# from dataset.duke import DUKEDataset
+# from dataset.lidc import LIDCDataset
+from vq_gan_3d.dataset.default import VQGANDataset

vq_gan_3d/dataset/default.py

@@ -0,0 +1,34 @@
+import os
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+import torch.nn.functional as F
+import torch.multiprocessing as mp
+
+
+class VQGANDataset(Dataset):
+    def __init__(self, root_dir: str, file_paths: str, internal_resolution: int):
+        super().__init__()
+        self.root_dir = root_dir
+        self.file_paths = file_paths
+        self.internal_resolution = internal_resolution
+
+    def __len__(self):
+        return len(self.file_paths)
+
+    def __getitem__(self, idx: int):
+        filename = os.path.join(self.root_dir, self.file_paths[idx])
+        try:
+            numpy_file = np.load(filename)
+            torch_np = torch.from_numpy(numpy_file)
+            torch_np = torch_np.unsqueeze(0).unsqueeze(0).float()  # Convert to float and move to appropriate device
+            interpolated_data = F.interpolate(input=torch_np, size=(self.internal_resolution, self.internal_resolution, self.internal_resolution), mode='trilinear')
+
+            # Apply tanh and log operations
+            interpolated_data_tanh = torch.tanh(interpolated_data)
+            interpolated_data_log = torch.log(interpolated_data + 1).squeeze(0)  # Adding 1 to avoid log(0)
+
+            return interpolated_data_log
+        except Exception as e:
+            print(f"Error loading file '{filename}': {e}")
+            return None

vq_gan_3d/metrics.py

@@ -0,0 +1,159 @@
+import torch
+import torch.nn.functional as F
+
+import os
+import sys
+import math
+import lpips as lpips_metric
+import piq
+import numpy as np
+from numpy import cov
+from numpy import trace
+from numpy import iscomplexobj
+from numpy.random import random
+from scipy.linalg import sqrtm
+from skimage.metrics import structural_similarity as ssim
+from skimage.metrics import peak_signal_noise_ratio
+from tqdm import tqdm
+
+import warnings
+warnings.filterwarnings("ignore")
+
+
+def blockPrint():
+    sys.stdout = open(os.devnull, 'w')
+
+
+def enablePrint():
+    sys.stdout = sys.__stdout__
+
+
+def normalize_tensor(outmap):
+    flattened_outmap = outmap.view(outmap.shape[0], -1, 1, 1) # Use 1's to preserve the number of dimensions for broadcasting later, as explained
+    outmap_min, _ = torch.min(flattened_outmap, dim=1, keepdim=True)
+    outmap_max, _ = torch.max(flattened_outmap, dim=1, keepdim=True)
+    outmap = (outmap - outmap_min) / (outmap_max - outmap_min)
+    return outmap
+
+
+class ImageMetrics:
+
+    def __init__(self, grid_true, grid_pred, device='cpu'):
+        self.grid_true = grid_true.to(device)  # [N, H, W]
+        self.grid_pred = grid_pred.to(device) # [N, H, W]
+        self.num_sequence = grid_true.shape[0]
+        self.loss_fn_vgg = lpips_metric.LPIPS(net='vgg', verbose=False).to(device) # closer to "traditional" perceptual loss, when used for optimization
+        self.device = device
+
+    def ssim(self):
+        """Structured Similarity Index Metric"""
+        a = normalize_tensor(self.grid_true.unsqueeze(0))
+        b = normalize_tensor(self.grid_pred.unsqueeze(0))
+        ssim = piq.ssim(a, b, data_range=1., reduction='none').squeeze().item()
+        return ssim
+
+    def mssim(self):
+        """Mean Structured Similarity Index Metric"""
+        mssim = 0
+        for idx in range(self.num_sequence):
+            max_value = max([self.grid_true[idx].max(), self.grid_pred[idx].max()])
+            min_value = min([self.grid_true[idx].min(), self.grid_pred[idx].min()])
+            data_range = abs(max_value - min_value)
+
+            a = self.grid_true[idx].detach().cpu().numpy()
+            b = self.grid_pred[idx].detach().cpu().numpy()
+
+            mssim += ssim(a, b, data_range=data_range.item())
+        return mssim / self.num_sequence
+    
+    def multiscale_ssim(self):
+        """Multi-Scale SSIM"""
+        a = normalize_tensor(self.grid_true.unsqueeze(0))
+        b = normalize_tensor(self.grid_pred.unsqueeze(0))
+        ms_ssim_index = piq.multi_scale_ssim(a, b, data_range=1., kernel_size=7).item()
+        return ms_ssim_index
+
+    def psnr(self):
+        """Peak Signal-to-Noise Ratio"""
+        psnr = 0
+        for idx in range(self.num_sequence):
+            max_value = max([self.grid_true[idx].max(), self.grid_pred[idx].max()])
+            min_value = min([self.grid_true[idx].min(), self.grid_pred[idx].min()])
+            data_range = abs(max_value - min_value)
+
+            a = self.grid_true[idx].detach().cpu().numpy()
+            b = self.grid_pred[idx].detach().cpu().numpy()
+
+            psnr += peak_signal_noise_ratio(a, b, data_range=data_range.item())
+        return psnr / self.num_sequence
+
+    def _calculate_fid(self, img1, img2):
+        img1 = img1.detach().cpu().numpy()
+        img2 = img2.detach().cpu().numpy()
+
+        # calculate mean and covariance statistics
+        mu1, sigma1 = img1.mean(axis=0), cov(img1, rowvar=False)
+        mu2, sigma2 = img2.mean(axis=0), cov(img2, rowvar=False)
+        # calculate sum squared difference between means
+        ssdiff = np.sum((mu1 - mu2)**2.0)
+        # calculate sqrt of product between cov
+        covmean = sqrtm(sigma1.dot(sigma2))
+        # check and correct imaginary numbers from sqrt
+        if iscomplexobj(covmean):
+            covmean = covmean.real
+        # calculate score
+        fid = ssdiff + trace(sigma1 + sigma2 - 2.0 * covmean)
+        return fid
+
+    def fid(self):
+        """Frechet Inception Distance"""
+        fid = 0
+        for idx in range(self.num_sequence):
+            fid += self._calculate_fid(self.grid_true[idx], self.grid_pred[idx])
+        return fid / self.num_sequence
+
+    def _calculate_lpips(self, img1, img2):
+        img1 = (2 * img1 / img1.max() - 1)  # normalize between -1 to 1
+        img2 = (2 * img2 / img2.max() - 1)  # normalize between -1 to 1
+        perceptual_loss = self.loss_fn_vgg(img1, img2).squeeze()
+        return perceptual_loss.item()
+
+    def lpips(self):
+        """Learned Perceptual Image Patch Similarity"""
+        perceptual_loss = 0
+        for idx in range(self.num_sequence):
+            perceptual_loss += self._calculate_lpips(self.grid_true[idx], self.grid_pred[idx])
+        return perceptual_loss / self.num_sequence
+
+    def reconstruction(self):
+        return F.l1_loss(self.grid_true, self.grid_pred).item()
+
+    def IS(self):
+        """Inception Score"""
+        is_score = 0
+        for idx in range(self.num_sequence):
+            is_score += piq.IS(distance='l1')(self.grid_true[idx], self.grid_pred[idx])
+        return (is_score / self.num_sequence).item()
+
+    def kid(self):
+        """Kernel Inception Distance"""
+        kid_score = 0
+        for idx in range(self.num_sequence):
+            kid_score += piq.KID()(self.grid_true[idx], self.grid_pred[idx])
+        return (kid_score / self.num_sequence).item()
+
+    def get_metrics(self):
+        blockPrint()
+        metrics = dict(
+            SSIM=self.ssim(),
+            MSSIM=self.mssim(),
+            MS_SSIM=self.multiscale_ssim(),
+            PSNR=self.psnr(),
+            IS=self.IS(),
+            FID=self.fid(),
+            KID=self.kid(),
+            LPIPS=self.lpips(),
+            Reconstruction=self.reconstruction(),
+        )
+        enablePrint()
+        return metrics

vq_gan_3d/model/__init__.py

@@ -0,0 +1,3 @@
+from vq_gan_3d.model.vqgan import VQGAN, load_VQGAN
+from vq_gan_3d.model.codebook import Codebook
+from vq_gan_3d.model.lpips import LPIPS

vq_gan_3d/model/cache/vgg.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a78928a0af1e5f0fcb1f3b9e8f8c3a2a5a3de244d830ad5c1feddc79b8432868
+size 7289

vq_gan_3d/model/codebook.py

@@ -0,0 +1,107 @@
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.distributed as dist
+
+from vq_gan_3d.utils import shift_dim
+
+
+class Codebook(nn.Module):
+    def __init__(self, n_codes, embedding_dim, no_random_restart=False, restart_thres=1.0):
+        super().__init__()
+        self.register_buffer('embeddings', torch.randn(n_codes, embedding_dim))
+        self.register_buffer('N', torch.zeros(n_codes))
+        self.register_buffer('z_avg', self.embeddings.data.clone())
+
+        self.n_codes = n_codes
+        self.embedding_dim = embedding_dim
+        self._need_init = True
+        self.no_random_restart = no_random_restart
+        self.restart_thres = restart_thres
+
+    def _tile(self, x):
+        d, ew = x.shape
+        if d < self.n_codes:
+            n_repeats = (self.n_codes + d - 1) // d
+            std = 0.01 / np.sqrt(ew)
+            x = x.repeat(n_repeats, 1)
+            x = x + torch.randn_like(x) * std
+        return x
+
+    def _init_embeddings(self, z):
+        # z: [b, c, t, h, w]
+        self._need_init = False
+        flat_inputs = shift_dim(z, 1, -1).flatten(end_dim=-2)
+        y = self._tile(flat_inputs)
+
+        d = y.shape[0]
+        _k_rand = y[torch.randperm(y.shape[0])][:self.n_codes]
+        if dist.is_initialized():
+            dist.broadcast(_k_rand, 0)
+        self.embeddings.data.copy_(_k_rand)
+        self.z_avg.data.copy_(_k_rand)
+        self.N.data.copy_(torch.ones(self.n_codes))
+
+    def forward(self, z):
+        # z: [b, c, t, h, w]
+        if self._need_init and self.training:
+            self._init_embeddings(z)
+        flat_inputs = shift_dim(z, 1, -1).flatten(end_dim=-2)  # [bthw, c]
+        distances = (flat_inputs ** 2).sum(dim=1, keepdim=True) \
+            - 2 * flat_inputs @ self.embeddings.t() \
+            + (self.embeddings.t() ** 2).sum(dim=0, keepdim=True)  # [bthw, c]
+
+        encoding_indices = torch.argmin(distances, dim=1)
+        encode_onehot = F.one_hot(encoding_indices, self.n_codes).type_as(
+            flat_inputs)  # [bthw, ncode]
+        encoding_indices = encoding_indices.view(
+            z.shape[0], *z.shape[2:])  # [b, t, h, w, ncode]
+
+        embeddings = F.embedding(
+            encoding_indices, self.embeddings)  # [b, t, h, w, c]
+        embeddings = shift_dim(embeddings, -1, 1)  # [b, c, t, h, w]
+
+        commitment_loss = 0.25 * F.mse_loss(z, embeddings.detach())
+
+        # EMA codebook update
+        if self.training:
+            n_total = encode_onehot.sum(dim=0)
+            encode_sum = flat_inputs.t() @ encode_onehot
+            if dist.is_initialized():
+                dist.all_reduce(n_total)
+                dist.all_reduce(encode_sum)
+
+            self.N.data.mul_(0.99).add_(n_total, alpha=0.01)
+            self.z_avg.data.mul_(0.99).add_(encode_sum.t(), alpha=0.01)
+
+            n = self.N.sum()
+            weights = (self.N + 1e-7) / (n + self.n_codes * 1e-7) * n
+            encode_normalized = self.z_avg / weights.unsqueeze(1)
+            self.embeddings.data.copy_(encode_normalized)
+
+            y = self._tile(flat_inputs)
+            _k_rand = y[torch.randperm(y.shape[0])][:self.n_codes]
+            if dist.is_initialized():
+                dist.broadcast(_k_rand, 0)
+
+            if not self.no_random_restart:
+                usage = (self.N.view(self.n_codes, 1)
+                         >= self.restart_thres).float()
+                self.embeddings.data.mul_(usage).add_(_k_rand * (1 - usage))
+
+        embeddings_st = (embeddings - z).detach() + z
+        embeddings_st_exp = torch.exp(embeddings_st)
+
+
+        avg_probs = torch.mean(encode_onehot, dim=0)
+        perplexity = torch.exp(-torch.sum(avg_probs *
+                               torch.log(avg_probs + 1e-10)))
+
+        return dict(embeddings=embeddings_st, encodings=encoding_indices,
+                    commitment_loss=commitment_loss, perplexity=perplexity)
+
+    def dictionary_lookup(self, encodings):
+        embeddings = F.embedding(encodings, self.embeddings)
+        return embeddings

vq_gan_3d/model/lpips.py

@@ -0,0 +1,181 @@
+"""Adapted from https://github.com/SongweiGe/TATS"""
+
+"""Stripped version of https://github.com/richzhang/PerceptualSimilarity/tree/master/models"""
+
+
+from collections import namedtuple
+from torchvision import models
+import torch.nn as nn
+import torch
+from tqdm import tqdm
+import requests
+import os
+import hashlib
+URL_MAP = {
+    "vgg_lpips": "https://heibox.uni-heidelberg.de/f/607503859c864bc1b30b/?dl=1"
+}
+
+CKPT_MAP = {
+    "vgg_lpips": "vgg.pth"
+}
+
+MD5_MAP = {
+    "vgg_lpips": "d507d7349b931f0638a25a48a722f98a"
+}
+
+
+def download(url, local_path, chunk_size=1024):
+    os.makedirs(os.path.split(local_path)[0], exist_ok=True)
+    with requests.get(url, stream=True) as r:
+        total_size = int(r.headers.get("content-length", 0))
+        with tqdm(total=total_size, unit="B", unit_scale=True) as pbar:
+            with open(local_path, "wb") as f:
+                for data in r.iter_content(chunk_size=chunk_size):
+                    if data:
+                        f.write(data)
+                        pbar.update(chunk_size)
+
+
+def md5_hash(path):
+    with open(path, "rb") as f:
+        content = f.read()
+    return hashlib.md5(content).hexdigest()
+
+
+def get_ckpt_path(name, root, check=False):
+    assert name in URL_MAP
+    path = os.path.join(root, CKPT_MAP[name])
+    if not os.path.exists(path) or (check and not md5_hash(path) == MD5_MAP[name]):
+        print("Downloading {} model from {} to {}".format(
+            name, URL_MAP[name], path))
+        download(URL_MAP[name], path)
+        md5 = md5_hash(path)
+        assert md5 == MD5_MAP[name], md5
+    return path
+
+
+class LPIPS(nn.Module):
+    # Learned perceptual metric
+    def __init__(self, use_dropout=True):
+        super().__init__()
+        self.scaling_layer = ScalingLayer()
+        self.chns = [64, 128, 256, 512, 512]  # vg16 features
+        self.net = vgg16(pretrained=False, requires_grad=True) # enabled grad
+        self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
+        self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
+        self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
+        self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
+        self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
+        self.load_from_pretrained()
+        # for param in self.parameters():
+        #     param.requires_grad = False
+
+    def load_from_pretrained(self, name="vgg_lpips"):
+        ckpt = get_ckpt_path(name, os.path.join(
+            os.path.dirname(os.path.abspath(__file__)), "cache"))
+        self.load_state_dict(torch.load(
+            ckpt, map_location=torch.device("cpu")), strict=False)
+        print("loaded pretrained LPIPS loss from {}".format(ckpt))
+
+    @classmethod
+    def from_pretrained(cls, name="vgg_lpips"):
+        if name != "vgg_lpips":
+            raise NotImplementedError
+        model = cls()
+        ckpt = get_ckpt_path(name, os.path.join(
+            os.path.dirname(os.path.abspath(__file__)), "cache"))
+        model.load_state_dict(torch.load(
+            ckpt, map_location=torch.device("cpu")), strict=False)
+        return model
+
+    def forward(self, input, target):
+        in0_input, in1_input = (self.scaling_layer(
+            input), self.scaling_layer(target))
+        outs0, outs1 = self.net(in0_input), self.net(in1_input)
+        feats0, feats1, diffs = {}, {}, {}
+        lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4]
+        for kk in range(len(self.chns)):
+            feats0[kk], feats1[kk] = normalize_tensor(
+                outs0[kk]), normalize_tensor(outs1[kk])
+            diffs[kk] = (feats0[kk] - feats1[kk]) ** 2
+
+        res = [spatial_average(lins[kk].model(diffs[kk]), keepdim=True)
+               for kk in range(len(self.chns))]
+        val = res[0]
+        for l in range(1, len(self.chns)):
+            val += res[l]
+        return val
+
+
+class ScalingLayer(nn.Module):
+    def __init__(self):
+        super(ScalingLayer, self).__init__()
+        self.register_buffer('shift', torch.Tensor(
+            [-.030, -.088, -.188])[None, :, None, None])
+        self.register_buffer('scale', torch.Tensor(
+            [.458, .448, .450])[None, :, None, None])
+
+    def forward(self, inp):
+        return (inp - self.shift) / self.scale
+
+
+class NetLinLayer(nn.Module):
+    """ A single linear layer which does a 1x1 conv """
+
+    def __init__(self, chn_in, chn_out=1, use_dropout=False):
+        super(NetLinLayer, self).__init__()
+        layers = [nn.Dropout(), ] if (use_dropout) else []
+        layers += [nn.Conv2d(chn_in, chn_out, 1, stride=1,
+                             padding=0, bias=False), ]
+        self.model = nn.Sequential(*layers)
+
+
+class vgg16(torch.nn.Module):
+    def __init__(self, requires_grad=False, pretrained=True):
+        super(vgg16, self).__init__()
+        vgg_pretrained_features = models.vgg16(pretrained=pretrained).features
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        self.N_slices = 5
+        for x in range(4):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(4, 9):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(9, 16):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(16, 23):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(23, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h = self.slice1(X)
+        h_relu1_2 = h
+        h = self.slice2(h)
+        h_relu2_2 = h
+        h = self.slice3(h)
+        h_relu3_3 = h
+        h = self.slice4(h)
+        h_relu4_3 = h
+        h = self.slice5(h)
+        h_relu5_3 = h
+        vgg_outputs = namedtuple(
+            "VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3', 'relu5_3'])
+        out = vgg_outputs(h_relu1_2, h_relu2_2,
+                          h_relu3_3, h_relu4_3, h_relu5_3)
+        return out
+
+
+def normalize_tensor(x, eps=1e-10):
+    norm_factor = torch.sqrt(torch.sum(x**2, dim=1, keepdim=True))
+    return x/(norm_factor+eps)
+
+
+def spatial_average(x, keepdim=True):
+    return x.mean([2, 3], keepdim=keepdim)

vq_gan_3d/model/vqgan.py

@@ -0,0 +1,565 @@
+"""Adapted from https://github.com/SongweiGe/TATS"""
+# Copyright (c) Meta Platforms, Inc. All Rights Reserved
+
+import math
+import argparse
+import numpy as np
+import pickle as pkl
+import random
+import gc
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.backends.cudnn as cudnn
+import torch.distributed as dist
+
+from vq_gan_3d.utils import shift_dim, adopt_weight, comp_getattr
+from vq_gan_3d.model.lpips import LPIPS
+from vq_gan_3d.model.codebook import Codebook
+
+
+def silu(x):
+    return x*torch.sigmoid(x)
+
+
+class SiLU(nn.Module):
+    def __init__(self):
+        super(SiLU, self).__init__()
+
+    def forward(self, x):
+        return silu(x)
+
+
+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 MeanPooling(nn.Module):
+    def __init__(self, kernel_size=16):
+        super(MeanPooling, self).__init__()
+        # Define a 3D average pooling layer
+        self.pool = nn.AvgPool3d(kernel_size=kernel_size)
+
+    def forward(self, x):
+        # Apply average pooling
+        x = self.pool(x)
+        # Flatten the tensor to a single dimension per batch element
+        x = x.view(x.size(0), -1)
+        return x
+
+
+class VQGAN(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+        self._set_seed(0)
+        self.embedding_dim = 256
+        self.n_codes = 16384
+
+        self.encoder = Encoder(16, [4,4,4], 1, 'group', 'replicate', 32)
+        self.decoder = Decoder(16, [4,4,4], 1, 'group', 32)
+        self.enc_out_ch = self.encoder.out_channels
+        self.pre_vq_conv = SamePadConv3d(self.enc_out_ch, 256, 1, padding_type='replicate')
+        self.post_vq_conv = SamePadConv3d(256, self.enc_out_ch, 1)
+
+        self.codebook = Codebook(16384, 256, no_random_restart=False, restart_thres=False)
+
+        self.pooling = MeanPooling(kernel_size=4)
+
+        self.gan_feat_weight = 4
+        # TODO: Changed batchnorm from sync to normal
+        self.image_discriminator = NLayerDiscriminator(1, 64, 3, norm_layer=nn.BatchNorm2d)
+
+        self.disc_loss = hinge_d_loss
+        self.perceptual_model = LPIPS()
+        self.image_gan_weight = 1
+        self.perceptual_weight = 4
+        self.l1_weight = 4
+
+    def encode(self, x, include_embeddings=False, quantize=True):
+        h = self.pre_vq_conv(self.encoder(x))
+        if quantize:
+            vq_output = self.codebook(h)
+            if include_embeddings:
+                return vq_output['embeddings'], vq_output['encodings']
+            else:
+                return vq_output['encodings']
+        return h
+
+    def decode(self, latent, quantize=False):
+        if quantize:
+            vq_output = self.codebook(latent)
+            latent = vq_output['encodings']
+        h = F.embedding(latent, self.codebook.embeddings)
+        h = self.post_vq_conv(shift_dim(h, -1, 1))
+        return self.decoder(h)
+
+    def feature_extraction(self, x):
+        """Extract embeddings given a grid."""
+        h = self.encode(x, include_embeddings=False, quantize=False)
+        return self.pooling(h.permute(0, 2, 3, 4, 1))
+
+    def forward(self, global_step, x, optimizer_idx=None, log_image=False, gpu_id=0):
+        B, C, T, H, W = x.shape
+
+        z = self.pre_vq_conv(self.encoder(x))
+        vq_output = self.codebook(z, gpu_id)
+
+        #vq_output['embeddings'] = torch.exp(vq_output['embeddings']) 
+        x_recon = self.decoder(self.post_vq_conv(vq_output['embeddings']))
+
+        recon_loss = (F.l1_loss(x_recon, x) * self.l1_weight)
+
+        # Selects one random 2D image from each 3D Image
+        frame_idx = torch.randint(0, T, [B]).to(gpu_id)
+        frame_idx_selected = frame_idx.reshape(-1,
+                                               1, 1, 1, 1).repeat(1, C, 1, H, W)
+        frames = torch.gather(x, 2, frame_idx_selected).squeeze(2)
+        frames_recon = torch.gather(x_recon, 2, frame_idx_selected).squeeze(2)
+
+        if log_image:
+            return frames, frames_recon, x, x_recon
+
+        if optimizer_idx == 0:
+            # Autoencoder - train the "generator"
+
+            # Perceptual loss
+            perceptual_loss = 0
+            if self.perceptual_weight > 0:
+                perceptual_loss = self.perceptual_model(
+                    frames, frames_recon).mean() * self.perceptual_weight
+                # perceptual_loss = .123
+
+            # Discriminator loss (turned on after a certain epoch)
+            logits_image_fake, pred_image_fake = self.image_discriminator(
+                frames_recon)
+            g_image_loss = -torch.mean(logits_image_fake)
+            g_loss = self.image_gan_weight*g_image_loss 
+            disc_factor = adopt_weight(
+                global_step, threshold=self.cfg.model.discriminator_iter_start)
+            aeloss = disc_factor * g_loss
+
+            # GAN feature matching loss - tune features such that we get the same prediction result on the discriminator
+            image_gan_feat_loss = 0
+            feat_weights = 4.0 / (3 + 1)
+            if self.image_gan_weight > 0:
+                logits_image_real, pred_image_real = self.image_discriminator(
+                    frames)
+                for i in range(len(pred_image_fake)-1):
+                    image_gan_feat_loss += feat_weights * \
+                        F.l1_loss(pred_image_fake[i], pred_image_real[i].detach(
+                        )) * (self.image_gan_weight > 0)
+
+            gan_feat_loss = disc_factor * self.gan_feat_weight * \
+                (image_gan_feat_loss)
+
+            return recon_loss, x_recon, vq_output, aeloss, perceptual_loss, gan_feat_loss, (g_image_loss, image_gan_feat_loss, vq_output['commitment_loss'], vq_output['perplexity'])
+
+        if optimizer_idx == 1:
+            # Train discriminator
+            logits_image_real, _ = self.image_discriminator(frames.detach())
+
+            logits_image_fake, _ = self.image_discriminator(
+                frames_recon.detach())
+
+            d_image_loss = self.disc_loss(logits_image_real, logits_image_fake)
+            disc_factor = adopt_weight(
+                global_step, threshold=self.cfg.model.discriminator_iter_start)
+            discloss = disc_factor * \
+                (self.image_gan_weight*d_image_loss)
+
+            return discloss
+
+        perceptual_loss = self.perceptual_model(
+            frames, frames_recon) * self.perceptual_weight
+        return recon_loss, x_recon, vq_output, perceptual_loss
+
+    def load_checkpoint(self, ckpt_path):
+        # load checkpoint file
+        ckpt_dict = torch.load(ckpt_path, map_location='cpu', weights_only=False)
+        
+        # load hyparameters
+        self.config = ckpt_dict['hparams']['_content']
+        self.embedding_dim = self.config['model']['embedding_dim']
+        self.n_codes = self.config['model']['n_codes']
+
+        # instantiate modules
+        self.encoder = Encoder(
+            self.config['model']['n_hiddens'],
+            self.config['model']['downsample'],
+            self.config['dataset']['image_channels'], 
+            self.config['model']['norm_type'], 
+            self.config['model']['padding_type'],
+            self.config['model']['num_groups'],
+        )
+        self.decoder = Decoder(
+            self.config['model']['n_hiddens'],
+            self.config['model']['downsample'], 
+            self.config['dataset']['image_channels'], 
+            self.config['model']['norm_type'], 
+            self.config['model']['num_groups']
+        )
+        self.enc_out_ch = self.encoder.out_channels
+        self.pre_vq_conv = SamePadConv3d(self.enc_out_ch, self.embedding_dim, 1, padding_type=self.config['model']['padding_type'])
+        self.post_vq_conv = SamePadConv3d(self.embedding_dim, self.enc_out_ch, 1)
+        self.codebook = Codebook(
+            self.n_codes, 
+            self.embedding_dim,
+            no_random_restart=self.config['model']['no_random_restart'], 
+            restart_thres=False
+        )
+        self.gan_feat_weight = self.config['model']['gan_feat_weight']
+        # TODO: Changed batchnorm from sync to normal
+        self.image_discriminator = NLayerDiscriminator(
+            self.config['dataset']['image_channels'], 
+            self.config['model']['disc_channels'],
+            self.config['model']['disc_layers'], 
+            norm_layer=nn.BatchNorm2d
+        )
+        self.disc_loss = hinge_d_loss
+        self.perceptual_model = LPIPS()
+        self.image_gan_weight = self.config['model']['gan_feat_weight']
+        self.perceptual_weight = self.config['model']['perceptual_weight'] 
+        self.l1_weight = self.config['model']['l1_weight']
+
+        # restore model weights
+        self.load_state_dict(ckpt_dict["MODEL_STATE"], strict=True)
+
+        # load RNG states each time the model and states are loaded from checkpoint
+        if 'rng' in self.config:
+            rng = self.config['rng']
+            for key, value in rng.items():
+                if key =='torch_state':
+                    torch.set_rng_state(value.cpu())
+                elif key =='cuda_state':
+                    torch.cuda.set_rng_state(value.cpu())
+                elif key =='numpy_state':
+                    np.random.set_state(value)
+                elif key =='python_state':
+                    random.setstate(value)
+                else:
+                    print('unrecognized state')
+
+    def log_images(self, batch, **kwargs):
+        log = dict()
+        x = batch['data']
+        x = x.to(self.device)
+        frames, frames_rec, _, _ = self(x, log_image=True)
+        log["inputs"] = frames
+        log["reconstructions"] = frames_rec
+        #log['mean_org'] = batch['mean_org']
+        #log['std_org'] = batch['std_org']
+        return log
+    
+    def _set_seed(self, value):
+        print('Random Seed:', value)
+        random.seed(value)
+        torch.manual_seed(value)
+        torch.cuda.manual_seed(value)
+        torch.cuda.manual_seed_all(value)
+        np.random.seed(value)
+        cudnn.deterministic = True
+        cudnn.benchmark = True
+        cudnn.enabled = True
+
+
+def Normalize(in_channels, norm_type='group', num_groups=32):
+    assert norm_type in ['group', 'batch']
+    if norm_type == 'group':
+        # TODO Changed num_groups from 32 to 8
+        return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
+    elif norm_type == 'batch':
+        return torch.nn.SyncBatchNorm(in_channels)
+
+
+class Encoder(nn.Module):
+    def __init__(self, n_hiddens = 16, downsample = [2,2,2] , image_channel=64, norm_type='group', padding_type='replicate', num_groups=32):
+        super().__init__()
+        n_times_downsample = np.array([int(math.log2(d)) for d in downsample])
+        self.conv_blocks = nn.ModuleList()
+        max_ds = n_times_downsample.max()
+
+        self.conv_first = SamePadConv3d(
+            image_channel, n_hiddens, kernel_size=3, padding_type=padding_type)
+
+        for i in range(max_ds):
+            block = nn.Module()
+            in_channels = n_hiddens * 2**i
+            out_channels = n_hiddens * 2**(i+1)
+            stride = tuple([2 if d > 0 else 1 for d in n_times_downsample])
+            block.down = SamePadConv3d(
+                in_channels, out_channels, 4, stride=stride, padding_type=padding_type)
+            block.res = ResBlock(
+                out_channels, out_channels, norm_type=norm_type, num_groups=num_groups)
+            self.conv_blocks.append(block)
+            n_times_downsample -= 1
+
+        self.final_block = nn.Sequential(
+            Normalize(out_channels, norm_type, num_groups=num_groups),
+            SiLU()
+        )
+
+        self.out_channels = out_channels
+
+    def forward(self, x):
+        h = self.conv_first(x)
+        for block in self.conv_blocks:
+            h = block.down(h)
+            h = block.res(h)
+        h = self.final_block(h)
+        return h
+
+
+class Decoder(nn.Module):
+    def __init__(self, n_hiddens = 16, upsample= [4,4,4], image_channel=1, norm_type='group', num_groups=1):
+        super().__init__()
+
+        n_times_upsample = np.array([int(math.log2(d)) for d in upsample])
+        print('n_times_upsample :', n_times_upsample)
+        max_us = n_times_upsample.max()
+        print('max_us :', max_us)
+
+
+        in_channels = n_hiddens*2**max_us
+        self.final_block = nn.Sequential(
+            Normalize(in_channels, norm_type, num_groups=num_groups),
+            SiLU()
+        )
+
+        self.conv_blocks = nn.ModuleList()
+        for i in range(max_us):
+            block = nn.Module()
+            in_channels = in_channels if i == 0 else n_hiddens*2**(max_us-i+1)
+            out_channels = n_hiddens*2**(max_us-i)
+            us = tuple([2 if d > 0 else 1 for d in n_times_upsample])
+            block.up = SamePadConvTranspose3d(
+                in_channels, out_channels, 4, stride=us)
+            block.res1 = ResBlock(
+                out_channels, out_channels, norm_type=norm_type, num_groups=num_groups)
+            block.res2 = ResBlock(
+                out_channels, out_channels, norm_type=norm_type, num_groups=num_groups)
+            self.conv_blocks.append(block)
+            n_times_upsample -= 1
+
+        self.conv_last = SamePadConv3d(
+            out_channels, image_channel, kernel_size=3)
+        
+
+    def forward(self, x):
+        h = self.final_block(x)
+        for i, block in enumerate(self.conv_blocks):
+            h = block.up(h)
+            h = block.res1(h)
+            h = block.res2(h)
+        h = self.conv_last(h)
+        return h
+
+
+class ResBlock(nn.Module):
+    def __init__(self, in_channels, out_channels=None, conv_shortcut=False, dropout=0.0, norm_type='group', padding_type='replicate', num_groups=32):
+        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, norm_type, num_groups=num_groups)
+        self.conv1 = SamePadConv3d(
+            in_channels, out_channels, kernel_size=3, padding_type=padding_type)
+        self.dropout = torch.nn.Dropout(dropout)
+        self.norm2 = Normalize(in_channels, norm_type, num_groups=num_groups)
+        self.conv2 = SamePadConv3d(
+            out_channels, out_channels, kernel_size=3, padding_type=padding_type)
+        if self.in_channels != self.out_channels:
+            self.conv_shortcut = SamePadConv3d(
+                in_channels, out_channels, kernel_size=3, padding_type=padding_type)
+
+    def forward(self, x):
+        h = x
+        h = self.norm1(h)
+        h = silu(h)
+        h = self.conv1(h)
+        h = self.norm2(h)
+        h = silu(h)
+        h = self.conv2(h)
+
+        if self.in_channels != self.out_channels:
+            x = self.conv_shortcut(x)
+
+        return x+h
+
+
+# Does not support dilation
+class SamePadConv3d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, bias=True, padding_type='replicate'):
+        super().__init__()
+        if isinstance(kernel_size, int):
+            kernel_size = (kernel_size,) * 3
+        if isinstance(stride, int):
+            stride = (stride,) * 3
+
+        # assumes that the input shape is divisible by stride
+        total_pad = tuple([k - s for k, s in zip(kernel_size, stride)])
+        pad_input = []
+        for p in total_pad[::-1]:  # reverse since F.pad starts from last dim
+            pad_input.append((p // 2 + p % 2, p // 2))
+        pad_input = sum(pad_input, tuple())
+        self.pad_input = pad_input
+        self.padding_type = padding_type
+
+        self.conv = nn.Conv3d(in_channels, out_channels, kernel_size,
+                              stride=stride, padding=0, bias=bias)
+
+    def forward(self, x):
+        return self.conv(F.pad(x, self.pad_input, mode=self.padding_type))
+
+
+class SamePadConvTranspose3d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, bias=True, padding_type='replicate'):
+        super().__init__()
+        if isinstance(kernel_size, int):
+            kernel_size = (kernel_size,) * 3
+        if isinstance(stride, int):
+            stride = (stride,) * 3
+
+        total_pad = tuple([k - s for k, s in zip(kernel_size, stride)])
+        pad_input = []
+        for p in total_pad[::-1]:  # reverse since F.pad starts from last dim
+            pad_input.append((p // 2 + p % 2, p // 2))
+        pad_input = sum(pad_input, tuple())
+        self.pad_input = pad_input
+        self.padding_type = padding_type
+
+        self.convt = nn.ConvTranspose3d(in_channels, out_channels, kernel_size,
+                                        stride=stride, bias=bias,
+                                        padding=tuple([k - 1 for k in kernel_size]))
+
+    def forward(self, x):
+        return self.convt(F.pad(x, self.pad_input, mode=self.padding_type))
+
+
+class NLayerDiscriminator(nn.Module):
+    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.SyncBatchNorm, use_sigmoid=False, getIntermFeat=True):
+        # def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, use_sigmoid=False, getIntermFeat=True):
+        super(NLayerDiscriminator, self).__init__()
+        self.getIntermFeat = getIntermFeat
+        self.n_layers = n_layers
+
+        kw = 4
+        padw = int(np.ceil((kw-1.0)/2))
+        sequence = [[nn.Conv2d(input_nc, ndf, kernel_size=kw,
+                               stride=2, padding=padw), nn.LeakyReLU(0.2, True)]]
+
+        nf = ndf
+        for n in range(1, n_layers):
+            nf_prev = nf
+            nf = min(nf * 2, 512)
+            sequence += [[
+                nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=2, padding=padw),
+                norm_layer(nf), nn.LeakyReLU(0.2, True)
+            ]]
+
+        nf_prev = nf
+        nf = min(nf * 2, 512)
+        sequence += [[
+            nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=1, padding=padw),
+            norm_layer(nf),
+            nn.LeakyReLU(0.2, True)
+        ]]
+
+        sequence += [[nn.Conv2d(nf, 1, kernel_size=kw,
+                                stride=1, padding=padw)]]
+
+        if use_sigmoid:
+            sequence += [[nn.Sigmoid()]]
+
+        if getIntermFeat:
+            for n in range(len(sequence)):
+                setattr(self, 'model'+str(n), nn.Sequential(*sequence[n]))
+        else:
+            sequence_stream = []
+            for n in range(len(sequence)):
+                sequence_stream += sequence[n]
+            self.model = nn.Sequential(*sequence_stream)
+
+    def forward(self, input):
+        if self.getIntermFeat:
+            res = [input]
+            for n in range(self.n_layers+2):
+                model = getattr(self, 'model'+str(n))
+                res.append(model(res[-1]))
+            return res[-1], res[1:]
+        else:
+            return self.model(input), _
+
+
+class NLayerDiscriminator3D(nn.Module):
+    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.SyncBatchNorm, use_sigmoid=False, getIntermFeat=True):
+        super(NLayerDiscriminator3D, self).__init__()
+        self.getIntermFeat = getIntermFeat
+        self.n_layers = n_layers
+
+        kw = 4
+        padw = int(np.ceil((kw-1.0)/2))
+        sequence = [[nn.Conv3d(input_nc, ndf, kernel_size=kw,
+                               stride=2, padding=padw), nn.LeakyReLU(0.2, True)]]
+
+        nf = ndf
+        for n in range(1, n_layers):
+            nf_prev = nf
+            nf = min(nf * 2, 512)
+            sequence += [[
+                nn.Conv3d(nf_prev, nf, kernel_size=kw, stride=2, padding=padw),
+                norm_layer(nf), nn.LeakyReLU(0.2, True)
+            ]]
+
+        nf_prev = nf
+        nf = min(nf * 2, 512)
+        sequence += [[
+            nn.Conv3d(nf_prev, nf, kernel_size=kw, stride=1, padding=padw),
+            norm_layer(nf),
+            nn.LeakyReLU(0.2, True)
+        ]]
+
+        sequence += [[nn.Conv3d(nf, 1, kernel_size=kw,
+                                stride=1, padding=padw)]]
+
+        if use_sigmoid:
+            sequence += [[nn.Sigmoid()]]
+
+        if getIntermFeat:
+            for n in range(len(sequence)):
+                setattr(self, 'model'+str(n), nn.Sequential(*sequence[n]))
+        else:
+            sequence_stream = []
+            for n in range(len(sequence)):
+                sequence_stream += sequence[n]
+            self.model = nn.Sequential(*sequence_stream)
+
+    def forward(self, input):
+        if self.getIntermFeat:
+            res = [input]
+            for n in range(self.n_layers+2):
+                model = getattr(self, 'model'+str(n))
+                res.append(model(res[-1]))
+            return res[-1], res[1:]
+        else:
+            return self.model(input), _
+
+
+def load_VQGAN(folder="../data/checkpoints/pretrained", ckpt_filename="VQGAN_43.pt"):
+    model = VQGAN()
+    model.load_checkpoint(os.path.join(folder, ckpt_filename))
+    model.eval()
+    return model

vq_gan_3d/utils.py

@@ -0,0 +1,165 @@
+""" Adapted from https://github.com/SongweiGe/TATS"""
+# Copyright (c) Meta Platforms, Inc. All Rights Reserved
+
+import warnings
+import torch
+import imageio
+
+import math
+import numpy as np
+
+import sys
+import pdb as pdb_original
+# import SimpleITK as sitk
+import logging
+
+import imageio.core.util
+logging.getLogger("imageio_ffmpeg").setLevel(logging.ERROR)
+
+
+def get_single_device(cpu=True):
+    if cpu:
+        return torch.device('cpu')
+    elif torch.cuda.is_available():
+        return torch.device('cuda')
+    elif torch.xpu.is_available():
+        return torch.device('xpu')
+    elif torch.mps.is_available():
+        return torch.device('mps')
+    return None
+
+
+class ForkedPdb(pdb_original.Pdb):
+    """A Pdb subclass that may be used
+    from a forked multiprocessing child
+
+    """
+
+    def interaction(self, *args, **kwargs):
+        _stdin = sys.stdin
+        try:
+            sys.stdin = open('/dev/stdin')
+            pdb_original.Pdb.interaction(self, *args, **kwargs)
+        finally:
+            sys.stdin = _stdin
+
+
+# Shifts src_tf dim to dest dim
+# i.e. shift_dim(x, 1, -1) would be (b, c, t, h, w) -> (b, t, h, w, c)
+def shift_dim(x, src_dim=-1, dest_dim=-1, make_contiguous=True):
+    n_dims = len(x.shape)
+    if src_dim < 0:
+        src_dim = n_dims + src_dim
+    if dest_dim < 0:
+        dest_dim = n_dims + dest_dim
+
+    assert 0 <= src_dim < n_dims and 0 <= dest_dim < n_dims
+
+    dims = list(range(n_dims))
+    del dims[src_dim]
+
+    permutation = []
+    ctr = 0
+    for i in range(n_dims):
+        if i == dest_dim:
+            permutation.append(src_dim)
+        else:
+            permutation.append(dims[ctr])
+            ctr += 1
+    x = x.permute(permutation)
+    if make_contiguous:
+        x = x.contiguous()
+    return x
+
+
+# reshapes tensor start from dim i (inclusive)
+# to dim j (exclusive) to the desired shape
+# e.g. if x.shape = (b, thw, c) then
+# view_range(x, 1, 2, (t, h, w)) returns
+# x of shape (b, t, h, w, c)
+def view_range(x, i, j, shape):
+    shape = tuple(shape)
+
+    n_dims = len(x.shape)
+    if i < 0:
+        i = n_dims + i
+
+    if j is None:
+        j = n_dims
+    elif j < 0:
+        j = n_dims + j
+
+    assert 0 <= i < j <= n_dims
+
+    x_shape = x.shape
+    target_shape = x_shape[:i] + shape + x_shape[j:]
+    return x.view(target_shape)
+
+
+def accuracy(output, target, topk=(1,)):
+    """Computes the accuracy over the k top predictions for the specified values of k"""
+    with torch.no_grad():
+        maxk = max(topk)
+        batch_size = target.size(0)
+
+        _, pred = output.topk(maxk, 1, True, True)
+        pred = pred.t()
+        correct = pred.eq(target.reshape(1, -1).expand_as(pred))
+
+        res = []
+        for k in topk:
+            correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
+            res.append(correct_k.mul_(100.0 / batch_size))
+        return res
+
+
+def tensor_slice(x, begin, size):
+    assert all([b >= 0 for b in begin])
+    size = [l - b if s == -1 else s
+            for s, b, l in zip(size, begin, x.shape)]
+    assert all([s >= 0 for s in size])
+
+    slices = [slice(b, b + s) for b, s in zip(begin, size)]
+    return x[slices]
+
+
+def adopt_weight(global_step, threshold=0, value=0.):
+    weight = 1
+    if global_step < threshold:
+        weight = value
+    return weight
+
+def comp_getattr(args, attr_name, default=None):
+    if hasattr(args, attr_name):
+        return getattr(args, attr_name)
+    else:
+        return default
+
+
+def visualize_tensors(t, name=None, nest=0):
+    if name is not None:
+        print(name, "current nest: ", nest)
+    print("type: ", type(t))
+    if 'dict' in str(type(t)):
+        print(t.keys())
+        for k in t.keys():
+            if t[k] is None:
+                print(k, "None")
+            else:
+                if 'Tensor' in str(type(t[k])):
+                    print(k, t[k].shape)
+                elif 'dict' in str(type(t[k])):
+                    print(k, 'dict')
+                    visualize_tensors(t[k], name, nest + 1)
+                elif 'list' in str(type(t[k])):
+                    print(k, len(t[k]))
+                    visualize_tensors(t[k], name, nest + 1)
+    elif 'list' in str(type(t)):
+        print("list length: ", len(t))
+        for t2 in t:
+            visualize_tensors(t2, name, nest + 1)
+    elif 'Tensor' in str(type(t)):
+        print(t.shape)
+    else:
+        print(t)
+    return ""

vq_gan_3d/weights/3DGrid-VQGAN_43.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:22dbe6ccab1ee629c51e6b98bcc109dc898e7f033feb5c1f5a271d1ee3d0ab83
+size 260643354