vae_mnist_variablelatentsize.py

# -*- coding: utf-8 -*-
"""VAE_Mnist_VariableLatentSize_Main.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1M9R-see6buUaLwAm2bcJuihywgW14ebw

# FULL
"""

from google.colab import drive
drive.mount('/content/drive')

# %cd /content/drive/My Drive/AI/Projects/VAE/Mnist/

# Commented out IPython magic to ensure Python compatibility.
# %cd /content

pip install torchsummary

!pip install imageio

import os
import pandas as pd
import cv2
from tqdm import tqdm
import pickle
import glob
import torch
from torchsummary import summary

"""#### mnist_dataset.py"""

import os

# Paths for the dataset
train_images_path = '/content/data/train/images'
test_images_path = '/content/data/test/images'

# Create directories for training and testing images
os.makedirs(train_images_path, exist_ok=True)
os.makedirs(test_images_path, exist_ok=True)

# Create subdirectories for each digit class
for i in range(10):
    os.makedirs(os.path.join(train_images_path, str(i)), exist_ok=True)
    os.makedirs(os.path.join(test_images_path, str(i)), exist_ok=True)

print("Folders created successfully!")

# Function to download and setup Kaggle API
def setup_kaggle():
    # Install Kaggle library
    !pip install -q kaggle

    # Upload the kaggle.json file (ensure you've uploaded this to your Colab session)
    from google.colab import files
    if not os.path.exists("/root/.kaggle/kaggle.json"):
        uploaded = files.upload()  # Upload kaggle.json
        assert 'kaggle.json' in uploaded, "kaggle.json should be uploaded!"

        # Move kaggle.json into the folder where the API expects to find it
        !mkdir -p ~/.kaggle
        !cp kaggle.json ~/.kaggle/
        !chmod 600 ~/.kaggle/kaggle.json

# Function to download and extract the MNIST dataset
def download_mnist():
    # Setup Kaggle and download dataset
    setup_kaggle()

    # Download the MNIST CSV dataset
    !kaggle datasets download -d oddrationale/mnist-in-csv -p /content/data --unzip

# Function to extract images from the CSV files
def extract_images_from_csv(data_path, save_dir):
    data = pd.read_csv(data_path)
    labels = data.iloc[:, 0].values
    images = data.iloc[:, 1:].values

    for idx, (label, img) in enumerate(tqdm(zip(labels, images), total=len(labels))):
        folder = os.path.join(save_dir, str(label))
        os.makedirs(folder, exist_ok=True)
        image = img.reshape(28, 28).astype('uint8')
        cv2.imwrite(os.path.join(folder, f'{idx}.png'), image)

download_mnist()  # Download dataset
extract_images_from_csv('/content/data/mnist_train.csv', train_images_path)
extract_images_from_csv('/content/data/mnist_test.csv', test_images_path)

import glob
import os
import torchvision
from PIL import Image
from tqdm import tqdm
from torch.utils.data.dataset import Dataset

def load_latents(latent_path):
    r"""
    Simple utility to save latents to speed up ldm training
    :param latent_path:
    :return:
    """
    latent_maps = {}
    for fname in glob.glob(os.path.join(latent_path, '*.pkl')):
        s = pickle.load(open(fname, 'rb'))
        for k, v in s.items():
            latent_maps[k] = v[0]
    return latent_maps


class MnistDataset(Dataset):
    r"""
    Nothing special here. Just a simple dataset class for mnist images.
    Created a dataset class rather using torchvision to allow
    replacement with any other image dataset
    """

    def __init__(self, split, im_path, im_size, im_channels,
                 use_latents=False, latent_path=None, condition_config=None):
        r"""
        Init method for initializing the dataset properties
        :param split: train/test to locate the image files
        :param im_path: root folder of images
        :param im_ext: image extension. assumes all
        images would be this type.
        """
        self.split = split
        self.im_size = im_size
        self.im_channels = im_channels

        # Should we use latents or not
        self.latent_maps = None
        self.use_latents = False

        # Conditioning for the dataset
        self.condition_types = [] if condition_config is None else condition_config['condition_types']

        self.images, self.labels = self.load_images(im_path)

        # Whether to load images and call vae or to load latents
        if use_latents and latent_path is not None:
            latent_maps = load_latents(latent_path)
            if len(latent_maps) == len(self.images):
                self.use_latents = True
                self.latent_maps = latent_maps
                print('Found {} latents'.format(len(self.latent_maps)))
            else:
                print('Latents not found')

    def load_images(self, im_path):
        r"""
        Gets all images from the path specified
        and stacks them all up
        :param im_path:
        :return:
        """
        assert os.path.exists(im_path), "images path {} does not exist".format(im_path)
        ims = []
        labels = []
        for d_name in tqdm(os.listdir(im_path)):
            fnames = glob.glob(os.path.join(im_path, d_name, '*.{}'.format('png')))
            fnames += glob.glob(os.path.join(im_path, d_name, '*.{}'.format('jpg')))
            fnames += glob.glob(os.path.join(im_path, d_name, '*.{}'.format('jpeg')))
            for fname in fnames:
                ims.append(fname)
                if 'class' in self.condition_types:
                    labels.append(int(d_name))
        print('Found {} images for split {}'.format(len(ims), self.split))
        return ims, labels

    def __len__(self):
        return len(self.images)

    def __getitem__(self, index):
        ######## Set Conditioning Info ########
        cond_inputs = {}
        if 'class' in self.condition_types:
            cond_inputs['class'] = self.labels[index]
        #######################################

        if self.use_latents:
            latent = self.latent_maps[self.images[index]]
            if len(self.condition_types) == 0:
                return latent
            else:
                return latent, cond_inputs
        else:
            im = Image.open(self.images[index])
            im_tensor = torchvision.transforms.ToTensor()(im)

            # Convert input to -1 to 1 range.
            im_tensor = (2 * im_tensor) - 1
            if len(self.condition_types) == 0:
                return im_tensor
            else:
                return im_tensor, cond_inputs

"""#### blocks.py"""

import torch
import torch.nn as nn

class DownBlock(nn.Module):
    r"""
    Down conv block with attention.
    Sequence of following block
    1. Resnet block with time embedding
    2. Attention block
    3. Downsample
    """

    def __init__(self, in_channels, out_channels, t_emb_dim,
                 down_sample, num_heads, num_layers, attn, norm_channels, cross_attn=False, context_dim=None):
        super().__init__()
        self.num_layers = num_layers
        self.down_sample = down_sample
        self.attn = attn
        self.context_dim = context_dim
        self.cross_attn = cross_attn
        self.t_emb_dim = t_emb_dim
        self.resnet_conv_first = nn.ModuleList(
            [
                nn.Sequential(
                    nn.GroupNorm(norm_channels, in_channels if i == 0 else out_channels),
                    nn.SiLU(),
                    nn.Conv2d(in_channels if i == 0 else out_channels, out_channels,
                              kernel_size=3, stride=1, padding=1),
                )
                for i in range(num_layers)
            ]
        )
        if self.t_emb_dim is not None:
            self.t_emb_layers = nn.ModuleList([
                nn.Sequential(
                    nn.SiLU(),
                    nn.Linear(self.t_emb_dim, out_channels)
                )
                for _ in range(num_layers)
            ])
        self.resnet_conv_second = nn.ModuleList(
            [
                nn.Sequential(
                    nn.GroupNorm(norm_channels, out_channels),
                    nn.SiLU(),
                    nn.Conv2d(out_channels, out_channels,
                              kernel_size=3, stride=1, padding=1),
                )
                for _ in range(num_layers)
            ]
        )

        if self.attn:
            self.attention_norms = nn.ModuleList(
                [nn.GroupNorm(norm_channels, out_channels)
                 for _ in range(num_layers)]
            )

            self.attentions = nn.ModuleList(
                [nn.MultiheadAttention(out_channels, num_heads, batch_first=True)
                 for _ in range(num_layers)]
            )

        if self.cross_attn:
            assert context_dim is not None, "Context Dimension must be passed for cross attention"
            self.cross_attention_norms = nn.ModuleList(
                [nn.GroupNorm(norm_channels, out_channels)
                 for _ in range(num_layers)]
            )
            self.cross_attentions = nn.ModuleList(
                [nn.MultiheadAttention(out_channels, num_heads, batch_first=True)
                 for _ in range(num_layers)]
            )
            self.context_proj = nn.ModuleList(
                [nn.Linear(context_dim, out_channels)
                 for _ in range(num_layers)]
            )

        self.residual_input_conv = nn.ModuleList(
            [
                nn.Conv2d(in_channels if i == 0 else out_channels, out_channels, kernel_size=1)
                for i in range(num_layers)
            ]
        )
        self.down_sample_conv = nn.Conv2d(out_channels, out_channels,
                                          4, 2, 1) if self.down_sample else nn.Identity()

    def forward(self, x, t_emb=None, context=None):
        out = x
        for i in range(self.num_layers):
            # Resnet block of Unet
            resnet_input = out
            out = self.resnet_conv_first[i](out)
            if self.t_emb_dim is not None:
                out = out + self.t_emb_layers[i](t_emb)[:, :, None, None]
            out = self.resnet_conv_second[i](out)
            out = out + self.residual_input_conv[i](resnet_input)

            if self.attn:
                # Attention block of Unet
                batch_size, channels, h, w = out.shape
                in_attn = out.reshape(batch_size, channels, h * w)
                in_attn = self.attention_norms[i](in_attn)
                in_attn = in_attn.transpose(1, 2)
                out_attn, _ = self.attentions[i](in_attn, in_attn, in_attn)
                out_attn = out_attn.transpose(1, 2).reshape(batch_size, channels, h, w)
                out = out + out_attn

            if self.cross_attn:
                assert context is not None, "context cannot be None if cross attention layers are used"
                batch_size, channels, h, w = out.shape
                in_attn = out.reshape(batch_size, channels, h * w)
                in_attn = self.cross_attention_norms[i](in_attn)
                in_attn = in_attn.transpose(1, 2)
                assert context.shape[0] == x.shape[0] and context.shape[-1] == self.context_dim
                context_proj = self.context_proj[i](context)
                out_attn, _ = self.cross_attentions[i](in_attn, context_proj, context_proj)
                out_attn = out_attn.transpose(1, 2).reshape(batch_size, channels, h, w)
                out = out + out_attn

        # Downsample
        out = self.down_sample_conv(out)
        return out


class MidBlock(nn.Module):
    r"""
    Mid conv block with attention.
    Sequence of following blocks
    1. Resnet block with time embedding
    2. Attention block
    3. Resnet block with time embedding
    """

    def __init__(self, in_channels, out_channels, t_emb_dim, num_heads, num_layers, norm_channels, cross_attn=None, context_dim=None):
        super().__init__()
        self.num_layers = num_layers
        self.t_emb_dim = t_emb_dim
        self.context_dim = context_dim
        self.cross_attn = cross_attn
        self.resnet_conv_first = nn.ModuleList(
            [
                nn.Sequential(
                    nn.GroupNorm(norm_channels, in_channels if i == 0 else out_channels),
                    nn.SiLU(),
                    nn.Conv2d(in_channels if i == 0 else out_channels, out_channels, kernel_size=3, stride=1,
                              padding=1),
                )
                for i in range(num_layers + 1)
            ]
        )

        if self.t_emb_dim is not None:
            self.t_emb_layers = nn.ModuleList([
                nn.Sequential(
                    nn.SiLU(),
                    nn.Linear(t_emb_dim, out_channels)
                )
                for _ in range(num_layers + 1)
            ])
        self.resnet_conv_second = nn.ModuleList(
            [
                nn.Sequential(
                    nn.GroupNorm(norm_channels, out_channels),
                    nn.SiLU(),
                    nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),
                )
                for _ in range(num_layers + 1)
            ]
        )

        self.attention_norms = nn.ModuleList(
            [nn.GroupNorm(norm_channels, out_channels)
             for _ in range(num_layers)]
        )

        self.attentions = nn.ModuleList(
            [nn.MultiheadAttention(out_channels, num_heads, batch_first=True)
             for _ in range(num_layers)]
        )
        if self.cross_attn:
            assert context_dim is not None, "Context Dimension must be passed for cross attention"
            self.cross_attention_norms = nn.ModuleList(
                [nn.GroupNorm(norm_channels, out_channels)
                 for _ in range(num_layers)]
            )
            self.cross_attentions = nn.ModuleList(
                [nn.MultiheadAttention(out_channels, num_heads, batch_first=True)
                 for _ in range(num_layers)]
            )
            self.context_proj = nn.ModuleList(
                [nn.Linear(context_dim, out_channels)
                 for _ in range(num_layers)]
            )
        self.residual_input_conv = nn.ModuleList(
            [
                nn.Conv2d(in_channels if i == 0 else out_channels, out_channels, kernel_size=1)
                for i in range(num_layers + 1)
            ]
        )

    def forward(self, x, t_emb=None, context=None):
        out = x

        # First resnet block
        resnet_input = out
        out = self.resnet_conv_first[0](out)
        if self.t_emb_dim is not None:
            out = out + self.t_emb_layers[0](t_emb)[:, :, None, None]
        out = self.resnet_conv_second[0](out)
        out = out + self.residual_input_conv[0](resnet_input)

        for i in range(self.num_layers):
            # Attention Block
            batch_size, channels, h, w = out.shape
            in_attn = out.reshape(batch_size, channels, h * w)
            in_attn = self.attention_norms[i](in_attn)
            in_attn = in_attn.transpose(1, 2)
            out_attn, _ = self.attentions[i](in_attn, in_attn, in_attn)
            out_attn = out_attn.transpose(1, 2).reshape(batch_size, channels, h, w)
            out = out + out_attn

            if self.cross_attn:
                assert context is not None, "context cannot be None if cross attention layers are used"
                batch_size, channels, h, w = out.shape
                in_attn = out.reshape(batch_size, channels, h * w)
                in_attn = self.cross_attention_norms[i](in_attn)
                in_attn = in_attn.transpose(1, 2)
                assert context.shape[0] == x.shape[0] and context.shape[-1] == self.context_dim
                context_proj = self.context_proj[i](context)
                out_attn, _ = self.cross_attentions[i](in_attn, context_proj, context_proj)
                out_attn = out_attn.transpose(1, 2).reshape(batch_size, channels, h, w)
                out = out + out_attn


            # Resnet Block
            resnet_input = out
            out = self.resnet_conv_first[i + 1](out)
            if self.t_emb_dim is not None:
                out = out + self.t_emb_layers[i + 1](t_emb)[:, :, None, None]
            out = self.resnet_conv_second[i + 1](out)
            out = out + self.residual_input_conv[i + 1](resnet_input)

        return out


class UpBlock(nn.Module):
    r"""
    Up conv block with attention.
    Sequence of following blocks
    1. Upsample
    1. Concatenate Down block output
    2. Resnet block with time embedding
    3. Attention Block
    """

    def __init__(self, in_channels, out_channels, t_emb_dim,
                 up_sample, num_heads, num_layers, attn, norm_channels):
        super().__init__()
        self.num_layers = num_layers
        self.up_sample = up_sample
        self.t_emb_dim = t_emb_dim
        self.attn = attn
        self.resnet_conv_first = nn.ModuleList(
            [
                nn.Sequential(
                    nn.GroupNorm(norm_channels, in_channels if i == 0 else out_channels),
                    nn.SiLU(),
                    nn.Conv2d(in_channels if i == 0 else out_channels, out_channels, kernel_size=3, stride=1,
                              padding=1),
                )
                for i in range(num_layers)
            ]
        )

        if self.t_emb_dim is not None:
            self.t_emb_layers = nn.ModuleList([
                nn.Sequential(
                    nn.SiLU(),
                    nn.Linear(t_emb_dim, out_channels)
                )
                for _ in range(num_layers)
            ])

        self.resnet_conv_second = nn.ModuleList(
            [
                nn.Sequential(
                    nn.GroupNorm(norm_channels, out_channels),
                    nn.SiLU(),
                    nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),
                )
                for _ in range(num_layers)
            ]
        )
        if self.attn:
            self.attention_norms = nn.ModuleList(
                [
                    nn.GroupNorm(norm_channels, out_channels)
                    for _ in range(num_layers)
                ]
            )

            self.attentions = nn.ModuleList(
                [
                    nn.MultiheadAttention(out_channels, num_heads, batch_first=True)
                    for _ in range(num_layers)
                ]
            )

        self.residual_input_conv = nn.ModuleList(
            [
                nn.Conv2d(in_channels if i == 0 else out_channels, out_channels, kernel_size=1)
                for i in range(num_layers)
            ]
        )
        self.up_sample_conv = nn.ConvTranspose2d(in_channels, in_channels,
                                                 4, 2, 1) \
            if self.up_sample else nn.Identity()

    def forward(self, x, out_down=None, t_emb=None):
        # Upsample
        x = self.up_sample_conv(x)

        # Concat with Downblock output
        if out_down is not None:
            x = torch.cat([x, out_down], dim=1)

        out = x
        for i in range(self.num_layers):
            # Resnet Block
            resnet_input = out
            out = self.resnet_conv_first[i](out)
            if self.t_emb_dim is not None:
                out = out + self.t_emb_layers[i](t_emb)[:, :, None, None]
            out = self.resnet_conv_second[i](out)
            out = out + self.residual_input_conv[i](resnet_input)

            # Self Attention
            if self.attn:
                batch_size, channels, h, w = out.shape
                in_attn = out.reshape(batch_size, channels, h * w)
                in_attn = self.attention_norms[i](in_attn)
                in_attn = in_attn.transpose(1, 2)
                out_attn, _ = self.attentions[i](in_attn, in_attn, in_attn)
                out_attn = out_attn.transpose(1, 2).reshape(batch_size, channels, h, w)
                out = out + out_attn
        return out


class UpBlockUnet(nn.Module):
    r"""
    Up conv block with attention.
    Sequence of following blocks
    1. Upsample
    1. Concatenate Down block output
    2. Resnet block with time embedding
    3. Attention Block
    """

    def __init__(self, in_channels, out_channels, t_emb_dim, up_sample,
                 num_heads, num_layers, norm_channels, cross_attn=False, context_dim=None):
        super().__init__()
        self.num_layers = num_layers
        self.up_sample = up_sample
        self.t_emb_dim = t_emb_dim
        self.cross_attn = cross_attn
        self.context_dim = context_dim
        self.resnet_conv_first = nn.ModuleList(
            [
                nn.Sequential(
                    nn.GroupNorm(norm_channels, in_channels if i == 0 else out_channels),
                    nn.SiLU(),
                    nn.Conv2d(in_channels if i == 0 else out_channels, out_channels, kernel_size=3, stride=1,
                              padding=1),
                )
                for i in range(num_layers)
            ]
        )

        if self.t_emb_dim is not None:
            self.t_emb_layers = nn.ModuleList([
                nn.Sequential(
                    nn.SiLU(),
                    nn.Linear(t_emb_dim, out_channels)
                )
                for _ in range(num_layers)
            ])

        self.resnet_conv_second = nn.ModuleList(
            [
                nn.Sequential(
                    nn.GroupNorm(norm_channels, out_channels),
                    nn.SiLU(),
                    nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),
                )
                for _ in range(num_layers)
            ]
        )

        self.attention_norms = nn.ModuleList(
            [
                nn.GroupNorm(norm_channels, out_channels)
                for _ in range(num_layers)
            ]
        )

        self.attentions = nn.ModuleList(
            [
                nn.MultiheadAttention(out_channels, num_heads, batch_first=True)
                for _ in range(num_layers)
            ]
        )

        if self.cross_attn:
            assert context_dim is not None, "Context Dimension must be passed for cross attention"
            self.cross_attention_norms = nn.ModuleList(
                [nn.GroupNorm(norm_channels, out_channels)
                 for _ in range(num_layers)]
            )
            self.cross_attentions = nn.ModuleList(
                [nn.MultiheadAttention(out_channels, num_heads, batch_first=True)
                 for _ in range(num_layers)]
            )
            self.context_proj = nn.ModuleList(
                [nn.Linear(context_dim, out_channels)
                 for _ in range(num_layers)]
            )
        self.residual_input_conv = nn.ModuleList(
            [
                nn.Conv2d(in_channels if i == 0 else out_channels, out_channels, kernel_size=1)
                for i in range(num_layers)
            ]
        )
        self.up_sample_conv = nn.ConvTranspose2d(in_channels // 2, in_channels // 2,
                                                 4, 2, 1) \
            if self.up_sample else nn.Identity()

    def forward(self, x, out_down=None, t_emb=None, context=None):
        x = self.up_sample_conv(x)
        if out_down is not None:
            x = torch.cat([x, out_down], dim=1)

        out = x
        for i in range(self.num_layers):
            # Resnet
            resnet_input = out
            out = self.resnet_conv_first[i](out)
            if self.t_emb_dim is not None:
                out = out + self.t_emb_layers[i](t_emb)[:, :, None, None]
            out = self.resnet_conv_second[i](out)
            out = out + self.residual_input_conv[i](resnet_input)
            # Self Attention
            batch_size, channels, h, w = out.shape
            in_attn = out.reshape(batch_size, channels, h * w)
            in_attn = self.attention_norms[i](in_attn)
            in_attn = in_attn.transpose(1, 2)
            out_attn, _ = self.attentions[i](in_attn, in_attn, in_attn)
            out_attn = out_attn.transpose(1, 2).reshape(batch_size, channels, h, w)
            out = out + out_attn
            # Cross Attention
            if self.cross_attn:
                assert context is not None, "context cannot be None if cross attention layers are used"
                batch_size, channels, h, w = out.shape
                in_attn = out.reshape(batch_size, channels, h * w)
                in_attn = self.cross_attention_norms[i](in_attn)
                in_attn = in_attn.transpose(1, 2)
                assert len(context.shape) == 3, \
                    "Context shape does not match B,_,CONTEXT_DIM"
                assert context.shape[0] == x.shape[0] and context.shape[-1] == self.context_dim,\
                    "Context shape does not match B,_,CONTEXT_DIM"
                context_proj = self.context_proj[i](context)
                out_attn, _ = self.cross_attentions[i](in_attn, context_proj, context_proj)
                out_attn = out_attn.transpose(1, 2).reshape(batch_size, channels, h, w)
                out = out + out_attn

        return out

"""#### discriminator.py"""

import torch
import torch.nn as nn


class Discriminator(nn.Module):
    r"""
    PatchGAN Discriminator.
    Rather than taking IMG_CHANNELSxIMG_HxIMG_W all the way to
    1 scalar value , we instead predict grid of values.
    Where each grid is prediction of how likely
    the discriminator thinks that the image patch corresponding
    to the grid cell is real
    """

    def __init__(self, im_channels=3,
                 conv_channels=[64, 128, 256],
                 kernels=[4,4,4,4],
                 strides=[2,2,2,1],
                 paddings=[1,1,1,1]):
        super().__init__()
        self.im_channels = im_channels
        activation = nn.LeakyReLU(0.2)
        layers_dim = [self.im_channels] + conv_channels + [1]
        self.layers = nn.ModuleList([
            nn.Sequential(
                nn.Conv2d(layers_dim[i], layers_dim[i + 1],
                          kernel_size=kernels[i],
                          stride=strides[i],
                          padding=paddings[i],
                          bias=False if i !=0 else True),
                nn.BatchNorm2d(layers_dim[i + 1]) if i != len(layers_dim) - 2 and i != 0 else nn.Identity(),
                activation if i != len(layers_dim) - 2 else nn.Identity()
            )
            for i in range(len(layers_dim) - 1)
        ])

    def forward(self, x):
        out = x
        for layer in self.layers:
            out = layer(out)
        return out


if __name__ == '__main__':
    x = torch.randn((2,3, 256, 256))
    prob = Discriminator(im_channels=3)(x)
    print(prob.shape)

"""#### lpips.py"""

from __future__ import absolute_import
from collections import namedtuple
import torch
import torch.nn as nn
import torch.nn.init as init
from torch.autograd import Variable
import numpy as np
import torch.nn
import torchvision

# Taken from https://github.com/richzhang/PerceptualSimilarity/blob/master/lpips/lpips.py

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


def spatial_average(in_tens, keepdim=True):
    return in_tens.mean([2, 3], keepdim=keepdim)


class vgg16(torch.nn.Module):
    def __init__(self, requires_grad=False, pretrained=True):
        super(vgg16, self).__init__()
        # Load pretrained vgg model from torchvision
        vgg_pretrained_features = torchvision.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])

        # Freeze vgg model
        if not requires_grad:
            for param in self.parameters():
                param.requires_grad = False

    def forward(self, X):
        # Return output of vgg features
        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


# Learned perceptual metric
class LPIPS(nn.Module):
    def __init__(self, net='vgg', version='0.1', use_dropout=True, weights_path="/content/"):
        super(LPIPS, self).__init__()
        self.version = version
        # Imagenet normalization
        self.scaling_layer = ScalingLayer()
        ########################

        # Instantiate vgg model
        self.chns = [64, 128, 256, 512, 512]
        self.L = len(self.chns)
        self.net = vgg16(pretrained=True, requires_grad=False)

        # Add 1x1 convolutional Layers
        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.lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4]
        self.lins = nn.ModuleList(self.lins)
        ########################

        model_path = weights_path  # Use the provided weights path
        print('Loading model from: %s' % model_path)
        self.load_state_dict(torch.load(model_path, map_location=device), strict=False)

        # Load the weights of trained LPIPS model
        # import inspect
        # model_path = os.path.abspath(
        #     os.path.join(inspect.getfile(self.__init__), '..', 'weights/v%s/%s.pth' % (version, net)))
        # print('Loading model from: %s' % model_path)
        # self.load_state_dict(torch.load(model_path, map_location=device), strict=False)
        ########################

        # Freeze all parameters
        self.eval()
        for param in self.parameters():
            param.requires_grad = False
        ########################

    def forward(self, in0, in1, normalize=False):
        # Scale the inputs to -1 to +1 range if needed
        if normalize:  # turn on this flag if input is [0,1] so it can be adjusted to [-1, +1]
            in0 = 2 * in0 - 1
            in1 = 2 * in1 - 1
        ########################

        # Normalize the inputs according to imagenet normalization
        in0_input, in1_input = self.scaling_layer(in0), self.scaling_layer(in1)
        ########################

        # Get VGG outputs for image0 and image1
        outs0, outs1 = self.net.forward(in0_input), self.net.forward(in1_input)
        feats0, feats1, diffs = {}, {}, {}
        ########################

        # Compute Square of Difference for each layer output
        for kk in range(self.L):
            feats0[kk], feats1[kk] = torch.nn.functional.normalize(outs0[kk], dim=1), torch.nn.functional.normalize(
                outs1[kk])
            diffs[kk] = (feats0[kk] - feats1[kk]) ** 2
        ########################

        # 1x1 convolution followed by spatial average on the square differences
        res = [spatial_average(self.lins[kk](diffs[kk]), keepdim=True) for kk in range(self.L)]
        val = 0

        # Aggregate the results of each layer
        for l in range(self.L):
            val += res[l]
        return val


class ScalingLayer(nn.Module):
    def __init__(self):
        super(ScalingLayer, self).__init__()
        # Imagnet normalization for (0-1)
        # mean = [0.485, 0.456, 0.406]
        # std = [0.229, 0.224, 0.225]
        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)

    def forward(self, x):
        out = self.model(x)
        return out

"""#### VAE"""

import torch
import torch.nn as nn

class VAE(nn.Module):
    def __init__(self, im_channels, model_config):
        super().__init__()
        self.down_channels = model_config['down_channels']
        self.mid_channels = model_config['mid_channels']
        self.down_sample = model_config['down_sample']
        self.num_down_layers = model_config['num_down_layers']
        self.num_mid_layers = model_config['num_mid_layers']
        self.num_up_layers = model_config['num_up_layers']

        # To disable attention in Downblock of Encoder and Upblock of Decoder
        self.attns = model_config['attn_down']

        # Latent Dimension
        self.z_channels = model_config['z_channels']
        self.norm_channels = model_config['norm_channels']
        self.num_heads = model_config['num_heads']

        # Assertion to validate the channel information
        assert self.mid_channels[0] == self.down_channels[-1]
        assert self.mid_channels[-1] == self.down_channels[-1]
        assert len(self.down_sample) == len(self.down_channels) - 1
        assert len(self.attns) == len(self.down_channels) - 1

        # Wherever we use downsampling in encoder correspondingly use
        # upsampling in decoder
        self.up_sample = list(reversed(self.down_sample))

        ##################### Encoder ######################
        self.encoder_conv_in = nn.Conv2d(im_channels, self.down_channels[0], kernel_size=3, padding=(1, 1))

        # Downblock + Midblock
        self.encoder_layers = nn.ModuleList([])
        for i in range(len(self.down_channels) - 1):
            self.encoder_layers.append(DownBlock(self.down_channels[i], self.down_channels[i + 1],
                                                 t_emb_dim=None, down_sample=self.down_sample[i],
                                                 num_heads=self.num_heads,
                                                 num_layers=self.num_down_layers,
                                                 attn=self.attns[i],
                                                 norm_channels=self.norm_channels))

        self.encoder_mids = nn.ModuleList([])
        for i in range(len(self.mid_channels) - 1):
            self.encoder_mids.append(MidBlock(self.mid_channels[i], self.mid_channels[i + 1],
                                              t_emb_dim=None,
                                              num_heads=self.num_heads,
                                              num_layers=self.num_mid_layers,
                                              norm_channels=self.norm_channels))

        self.encoder_norm_out = nn.GroupNorm(self.norm_channels, self.down_channels[-1])
        self.encoder_conv_out = nn.Conv2d(self.down_channels[-1], 2*self.z_channels, kernel_size=3, padding=1)

        # Latent Dimension is 2*Latent because we are predicting mean & variance
        self.pre_quant_conv = nn.Conv2d(2*self.z_channels, 2*self.z_channels, kernel_size=1)
        ####################################################


        ##################### Decoder ######################
        self.post_quant_conv = nn.Conv2d(self.z_channels, self.z_channels, kernel_size=1)
        self.decoder_conv_in = nn.Conv2d(self.z_channels, self.mid_channels[-1], kernel_size=3, padding=(1, 1))

        # Midblock + Upblock
        self.decoder_mids = nn.ModuleList([])
        for i in reversed(range(1, len(self.mid_channels))):
            self.decoder_mids.append(MidBlock(self.mid_channels[i], self.mid_channels[i - 1],
                                              t_emb_dim=None,
                                              num_heads=self.num_heads,
                                              num_layers=self.num_mid_layers,
                                              norm_channels=self.norm_channels))

        self.decoder_layers = nn.ModuleList([])
        for i in reversed(range(1, len(self.down_channels))):
            self.decoder_layers.append(UpBlock(self.down_channels[i], self.down_channels[i - 1],
                                               t_emb_dim=None, up_sample=self.down_sample[i - 1],
                                               num_heads=self.num_heads,
                                               num_layers=self.num_up_layers,
                                               attn=self.attns[i - 1],
                                               norm_channels=self.norm_channels))

        self.decoder_norm_out = nn.GroupNorm(self.norm_channels, self.down_channels[0])
        self.decoder_conv_out = nn.Conv2d(self.down_channels[0], im_channels, kernel_size=3, padding=1)

    def encode(self, x):
        out = self.encoder_conv_in(x)
        for idx, down in enumerate(self.encoder_layers):
            out = down(out)
        for mid in self.encoder_mids:
            out = mid(out)
        out = self.encoder_norm_out(out)
        out = nn.SiLU()(out)
        out = self.encoder_conv_out(out)
        out = self.pre_quant_conv(out)
        mean, logvar = torch.chunk(out, 2, dim=1)
        std = torch.exp(0.5 * logvar)
        sample = mean + std * torch.randn(mean.shape).to(device=x.device)
        return sample, out

    def decode(self, z):
        out = z
        out = self.post_quant_conv(out)
        out = self.decoder_conv_in(out)
        for mid in self.decoder_mids:
            out = mid(out)
        for idx, up in enumerate(self.decoder_layers):
            out = up(out)

        out = self.decoder_norm_out(out)
        out = nn.SiLU()(out)
        out = self.decoder_conv_out(out)
        return out

    def forward(self, x):
        z, encoder_output = self.encode(x)
        out = self.decode(z)
        return out, encoder_output, z

import requests
import os

# URL for the vgg.pth weights
url = 'https://github.com/richzhang/PerceptualSimilarity/raw/master/lpips/weights/v0.1/vgg.pth'

# Specify the path to save the weights in Colab's local environment
local_path = '/content/vgg.pth'

# Check if the weights file already exists to avoid re-downloading
if not os.path.exists(local_path):
    # Make the request to download the file
    response = requests.get(url, allow_redirects=True)

    # Save the content of the response
    with open(local_path, 'wb') as f:
        f.write(response.content)
    print(f"File downloaded and saved to {local_path}")
else:
    print(f"File already exists at {local_path}")

import numpy as np
import torch
import matplotlib.pyplot as plt
from skimage.metrics import peak_signal_noise_ratio as psnr, structural_similarity as ssim
from sklearn.manifold import TSNE

def select_specific_images(data_loader, class_labels):
    """Select one image per class label from the dataset."""
    images = []
    labels = []
    found_labels = set()
    for batch_images, batch_labels in data_loader:
        # Extract labels tensor from the dictionary
        batch_labels = batch_labels['class']  # Adjust this key based on your actual data structure
        for image, label in zip(batch_images, batch_labels):
            if label.item() in class_labels and label.item() not in found_labels:
                images.append(image.unsqueeze(0))  # Ensure image has batch dimension
                labels.append(label.item())
                found_labels.add(label.item())
            if len(found_labels) == len(class_labels):
                break
        if len(found_labels) == len(class_labels):
            break
    return images, labels

def compute_metrics(original, reconstructed):
    """Compute MSE, PSNR, and SSIM for a batch of images."""
    # Compute MSE for the entire batch at once
    mse = torch.nn.functional.mse_loss(reconstructed, original, reduction='mean').item()

    # Prepare to collect SSIM and PSNR for each image in the batch
    ssim_scores = []
    psnr_scores = []

    # Convert the tensors to numpy arrays and handle each image individually
    original_np = original.detach().cpu().numpy()
    reconstructed_np = reconstructed.detach().cpu().numpy()

    # Iterate over each image in the batch
    for i in range(original_np.shape[0]):
        # SSIM and PSNR expect single channel 2D arrays, hence [i, 0, :, :]
        ssim_score = ssim(original_np[i, 0], reconstructed_np[i, 0],
                          data_range=original_np[i, 0].max() - original_np[i, 0].min())
        psnr_score = psnr(original_np[i, 0], reconstructed_np[i, 0],
                          data_range=original_np[i, 0].max() - original_np[i, 0].min())

        ssim_scores.append(ssim_score)
        psnr_scores.append(psnr_score)

    # Compute average SSIM and PSNR over the batch
    avg_ssim = np.mean(ssim_scores)
    avg_psnr = np.mean(psnr_scores)

    return mse, avg_ssim, avg_psnr


def evaluate_on_test_set(model, test_data_loader, specific_images, specific_labels, config, epoch_num):
    """Evaluate the model on the test set and visualize results."""
    model.eval()
    total_mse, total_ssim, total_psnr, count = 0, 0, 0, 0
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    all_latents = []
    all_labels = []

    # Compute metrics over the whole test dataset
    with torch.no_grad():
        for images, data in test_data_loader:
            # print(data)
            labels = data['class']  # Adjust 'labels' based on actual key in your dataset
            images, labels = images.to(device), labels.to(device)
            outputs, _, latents = model(images)  # Unpack the tuple and only use the reconstructed images

            latents = latents.view(latents.size(0), -1)  # Reshape to [batch_size, product_of_other_dims]
            all_latents.append(latents.cpu().numpy())
            all_labels.append(labels.cpu().numpy())

            for idx in range(images.size(0)):
                reconstructed = outputs[idx:idx+1]
                original = images[idx:idx+1]
                mse, ssim_score, psnr_score = compute_metrics(original, reconstructed)
                total_mse += mse
                total_ssim += ssim_score
                total_psnr += psnr_score
                count += 1

    # Average metrics
    avg_mse = total_mse / count
    avg_ssim = total_ssim / count
    avg_psnr = total_psnr / count

    # Visualization of specific images
    fig, axs = plt.subplots(2, len(specific_images), figsize=(15, 6))
    for idx, (image, label) in enumerate(zip(specific_images, specific_labels)):
        image = image.to(device)
        output, _, _ = model(image)  # Unpack the tuple and use only the reconstructed image
        image_mse, image_ssim, image_psnr = compute_metrics(image, output)

        axs[0, idx].imshow(image.squeeze().cpu().numpy(), cmap='gray')
        axs[0, idx].set_title(f"Original - Class: {label}")
        axs[0, idx].axis('off')

        axs[1, idx].imshow(output.squeeze().detach().cpu().numpy(), cmap='gray')
        axs[1, idx].set_title(f"Eeconstruction \n MSE: {image_mse:.2f}, \n SSIM: {image_ssim:.2f}, \n PSNR: {image_psnr:.2f}")
        axs[1, idx].axis('off')

    plt.tight_layout(rect=[0, 0, 1, 0.95])  # Leave space for the suptitle
    fig.suptitle(f"Image Reconstructions - Epoch {epoch_num:.2f}", fontsize=16, y=0.98)  # Adjust y to place it above
    plot_save_path = config['plot_save_path']
    plt.savefig(f"{plot_save_path}/evaluation_plot_epoch_{epoch_num:.2f}.png")
    plt.show()


    # Concatenate all batches of latents and labels
    all_latents = np.concatenate(all_latents, axis=0)
    all_labels = np.concatenate(all_labels, axis=0)

    # t-SNE for dimensionality reduction
    tsne = TSNE(n_components=2, random_state=42)
    latents_2d = tsne.fit_transform(all_latents)

    # Visualization of latent space
    plt.figure(figsize=(12, 10))
    scatter = plt.scatter(latents_2d[:, 0], latents_2d[:, 1], c=all_labels, cmap='viridis', alpha=0.6, edgecolors='w')
    plt.colorbar(scatter, label='Class Labels')
    plt.title(f"Epoch {epoch_num:.2f}: Latent Space Visualization on Entire Test Set", fontsize=16)
    plt.xlabel("Dimension 1", fontsize=12)
    plt.ylabel("Dimension 2", fontsize=12)
    plt.grid(True, linestyle="--", alpha=0.5)
    plt.savefig(f"{plot_save_path}/latent_visualization_epoch_{epoch_num:.2f}.png")
    plt.show()

    print(f"Test Metrics - Overall MSE: {avg_mse:.4f}, SSIM: {avg_ssim:.4f}, PSNR: {avg_psnr:.4f}")
    return avg_mse, avg_ssim, avg_psnr

"""#### Config"""

class Config:
    def __init__(self):
        self.config = {
            "dataset_params": {
                "im_path": "data/train/images",
                "im_channels": 1,
                "im_size": 28,
                "name": "mnist"
            },
            "autoencoder_params": {
                "z_channels": 16,
                "down_channels": [32, 64, 128],
                "mid_channels": [128, 128],
                "down_sample": [True, True],
                "attn_down": [False, False],
                "norm_channels": 32,
                "num_heads": 16,
                "num_down_layers": 1,
                "num_mid_layers": 1,
                "num_up_layers": 1
            },
            "train_params": {
                "seed": 1111,
                "task_name": "/content/drive/My Drive/AI/Projects/VAE/Mnist/variable_latent_sizes/latent_16/models",
                "ldm_batch_size": 64,
                "autoencoder_batch_size": 64,
                "disc_start": 1000,
                "disc_weight": 0.5,
                "codebook_weight": 1,
                "commitment_beta": 0.2,
                "perceptual_weight": 1,
                "kl_weight": 0.000005,
                "ldm_epochs": 100,
                "autoencoder_epochs": 20,
                "num_samples": 64,
                "num_grid_rows": 5,
                "ldm_lr": 0.00001,
                "autoencoder_lr": 0.0001,
                "autoencoder_acc_steps": 1,
                "autoencoder_img_save_steps": 8,
                "save_latents": True,
                "save_latent_plot": True,
                "vae_latent_dir_name": "vae_latents",
                "vae_latent_dir_name": "vae_latents",
                "ldm_ckpt_name": "ddpm_ckpt.pth",
                "vae_autoencoder_ckpt_name": "vae_autoencoder_ckpt.pth",
                "vae_autoencoder_ckpt_name": "vae_autoencoder_ckpt.pth",
                "lpips_weights_path" : "/content/vgg.pth",
                "lpips_folder_download_path" : "/content",
                "vae_discriminator_ckpt_name": "vae_discriminator_ckpt.pth",
                "vae_discriminator_ckpt_name": "vae_discriminator_ckpt.pth",
                "plot_save_path": "/content/drive/My Drive/AI/Projects/VAE/Mnist/variable_latent_sizes/latent_16/plots",

            }
        }

############

"""#### train VAE

### latent 16
"""

############

## Do not run again #####

############

import yaml
import argparse
import torch
import random
import torchvision
import os
import numpy as np
from tqdm import tqdm
from torch.optim import Adam
from torchvision.utils import make_grid
from torch.utils.data.dataloader import DataLoader
import json

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

def train_VAE(custom_config=None):
    # Use custom_config if provided, otherwise use the default Config
    if custom_config:
        config_instance = custom_config
    else:
        config_instance = Config()
    config = config_instance.config  # Access the configuration dictionary
    print(config)

    dataset_config = config['dataset_params']
    autoencoder_config = config['autoencoder_params']
    train_config = config['train_params']

    # Set the desired seed value #
    seed = train_config['seed']
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    if device == 'cuda':
        torch.cuda.manual_seed_all(seed)
    print(device)
    #############################

    # Create the model and dataset #
    model = VAE(im_channels=dataset_config['im_channels'],
                  model_config=autoencoder_config).to(device)

    summary(model, input_size=(dataset_config['im_channels'], dataset_config['im_size'], dataset_config['im_size']))
    # Create the dataset
    im_dataset_cls = {
        'mnist': MnistDataset,
    }.get(dataset_config['name'])

    train_dataset = im_dataset_cls(split='train',
                                   im_path=dataset_config['im_path'],
                                   im_size=dataset_config['im_size'],
                                   im_channels=dataset_config['im_channels'])

    # Define condition configuration for the test set to always return labels
    test_condition_config = {
        'condition_types': ['class']
    }

    # Initialize test dataset with conditions
    test_dataset = MnistDataset(
        split='test',
        im_path=dataset_config['im_path'].replace('train', 'test'),
        im_size=dataset_config['im_size'],
        im_channels=dataset_config['im_channels'],
        condition_config=test_condition_config
    )

    im_dataset = im_dataset_cls(split='train',
                                im_path=dataset_config['im_path'],
                                im_size=dataset_config['im_size'],
                                im_channels=dataset_config['im_channels'])

    train_data_loader = DataLoader(train_dataset,
                                   batch_size=train_config['autoencoder_batch_size'],
                                   shuffle=True)

    test_data_loader = DataLoader(test_dataset,
                                  batch_size=train_config['autoencoder_batch_size'],
                                  shuffle=True)

    # Select specific images for consistent visualization
    specific_images, specific_labels = select_specific_images(test_data_loader, [1, 2, 3, 4, 5])

    # Create output directories
    if not os.path.exists(train_config['task_name']):
        os.mkdir(train_config['task_name'])

    num_epochs = train_config['autoencoder_epochs']

    # L1/L2 loss for Reconstruction
    recon_criterion = torch.nn.MSELoss()
    # Disc Loss can even be BCEWithLogits
    disc_criterion = torch.nn.MSELoss()


    # No need to freeze lpips as lpips.py takes care of that
    lpips_model = LPIPS(weights_path = train_config["lpips_weights_path"]).eval().to(device)
    discriminator = Discriminator(im_channels=dataset_config['im_channels']).to(device)

    optimizer_d = Adam(discriminator.parameters(), lr=train_config['autoencoder_lr'], betas=(0.5, 0.999))
    optimizer_g = Adam(model.parameters(), lr=train_config['autoencoder_lr'], betas=(0.5, 0.999))

    disc_step_start = train_config['disc_start']
    step_count = 0

    # This is for accumulating gradients incase the images are huge
    # And one cant afford higher batch sizes
    acc_steps = train_config['autoencoder_acc_steps']
    image_save_steps = train_config['autoencoder_img_save_steps']
    img_save_count = 0

    # Initialize metrics tracking dictionaries
    train_metrics = {
        "recon_losses": [],
        "perceptual_losses": [],
        "gen_losses": [],
        "disc_losses": []
    }
    test_metrics = {
        "mse": [],
        "ssim": [],
        "psnr": []
    }

    # Total steps in one epoch for first-epoch evaluation intervals
    total_steps = len(train_data_loader)
    interval = total_steps // 5  # Perform 5 evaluations during the first epoch

    for epoch_idx in range(num_epochs):
        recon_losses = []
        perceptual_losses = []
        disc_losses = []
        gen_losses = []
        losses = []

        optimizer_g.zero_grad()
        optimizer_d.zero_grad()

        for step_idx, im in enumerate(tqdm(train_data_loader)):
            step_count += 1
            im = im.float().to(device)

            # Fetch autoencoders output(reconstructions)
            model_output = model(im)
            output, enc_, z_  = model_output

            # Image Saving Logic
            if step_count % image_save_steps == 0 or step_count == 1:
                sample_size = min(8, im.shape[0])
                save_output = torch.clamp(output[:sample_size], -1., 1.).detach().cpu()
                save_output = ((save_output + 1) / 2)
                save_input = ((im[:sample_size] + 1) / 2).detach().cpu()

                grid = make_grid(torch.cat([save_input, save_output], dim=0), nrow=sample_size)
                img = torchvision.transforms.ToPILImage()(grid)
                if not os.path.exists(os.path.join(train_config['task_name'],'vae_autoencoder_samples')):
                    os.mkdir(os.path.join(train_config['task_name'], 'vae_autoencoder_samples'))
                img.save(os.path.join(train_config['task_name'],'vae_autoencoder_samples',
                                      'current_autoencoder_sample_{}.png'.format(img_save_count)))
                img_save_count += 1
                img.close()

            ######### Optimize Generator ##########
            # L2 Loss
            recon_loss = recon_criterion(output, im)
            recon_losses.append(recon_loss.item())
            recon_loss = recon_loss / acc_steps
            # g_loss = (recon_loss +
            #           (train_config['codebook_weight'] * quantize_losses['codebook_loss'] / acc_steps) +
            #           (train_config['commitment_beta'] * quantize_losses['commitment_loss'] / acc_steps))
            # codebook_losses.append(train_config['codebook_weight'] * quantize_losses['codebook_loss'].item())
            g_loss = recon_loss
            # Adversarial loss only if disc_step_start steps passed
            if step_count >= disc_step_start:
                if step_count == disc_step_start:
                    print("Adversarial Loss started")
                disc_fake_pred = discriminator(model_output[0])
                disc_fake_loss = disc_criterion(disc_fake_pred,
                                                torch.ones(disc_fake_pred.shape,
                                                           device=disc_fake_pred.device))
                gen_losses.append(train_config['disc_weight'] * disc_fake_loss.item())
                g_loss += train_config['disc_weight'] * disc_fake_loss / acc_steps
            lpips_loss = torch.mean(lpips_model(output, im))
            perceptual_losses.append(train_config['perceptual_weight'] * lpips_loss.item())
            g_loss += train_config['perceptual_weight']*lpips_loss / acc_steps
            losses.append(g_loss.item())
            g_loss.backward()
            #####################################

            ######### Optimize Discriminator #######
            if step_count > disc_step_start:
                fake = output
                disc_fake_pred = discriminator(fake.detach())
                disc_real_pred = discriminator(im)
                disc_fake_loss = disc_criterion(disc_fake_pred,
                                                torch.zeros(disc_fake_pred.shape,
                                                            device=disc_fake_pred.device))
                disc_real_loss = disc_criterion(disc_real_pred,
                                                torch.ones(disc_real_pred.shape,
                                                           device=disc_real_pred.device))
                disc_loss = train_config['disc_weight'] * (disc_fake_loss + disc_real_loss) / 2
                disc_losses.append(disc_loss.item())
                disc_loss = disc_loss / acc_steps
                disc_loss.backward()
                if step_count % acc_steps == 0:
                    optimizer_d.step()
                    optimizer_d.zero_grad()
            #####################################
            if step_count % acc_steps == 0:
                optimizer_g.step()
                optimizer_g.zero_grad()

            # Perform additional evaluations during the first epoch
            if epoch_idx == 0 and step_idx % interval == 0:
                fraction = (step_idx + 1) / total_steps
                print(f"\nFirst epoch intermediate evaluation at step {step_idx + 1} ({fraction:.2f} epoch)")
                avg_mse, avg_ssim, avg_psnr = evaluate_on_test_set(model, test_data_loader, specific_images, specific_labels, train_config, fraction)
                test_metrics["mse"].append(avg_mse)
                test_metrics["ssim"].append(avg_ssim)
                test_metrics["psnr"].append(avg_psnr)
                print(f"Test Evaluation: MSE={avg_mse:.4f}, SSIM={avg_ssim:.4f}, PSNR={avg_psnr:.4f}")

        optimizer_d.step()
        optimizer_d.zero_grad()
        optimizer_g.step()
        optimizer_g.zero_grad()
        if len(disc_losses) > 0:
            print(
                '\nFinished epoch: {} | Recon Loss : {:.4f} | Perceptual Loss : {:.4f} | '
                'G Loss : {:.4f} | D Loss {:.4f}'.
                format(epoch_idx + 1,
                       np.mean(recon_losses),
                       np.mean(perceptual_losses),
                       np.mean(gen_losses),
                       np.mean(disc_losses)))
        else:
            print('Finished epoch: {} | Recon Loss : {:.4f} | Perceptual Loss : {:.4f}'.
                  format(epoch_idx + 1,
                         np.mean(recon_losses),
                         np.mean(perceptual_losses)
                         ))

        # Record epoch metrics
        train_metrics["recon_losses"].append(np.mean(recon_losses))
        train_metrics["perceptual_losses"].append(np.mean(perceptual_losses))
        train_metrics["gen_losses"].append(np.mean(gen_losses))
        train_metrics["disc_losses"].append(np.mean(disc_losses) if disc_losses else 0)

        torch.save(model.state_dict(), os.path.join(train_config['task_name'],
                                                    train_config['vae_autoencoder_ckpt_name']))
        torch.save(discriminator.state_dict(), os.path.join(train_config['task_name'],
                                                            train_config['vae_discriminator_ckpt_name']))

        avg_mse, avg_ssim, avg_psnr = evaluate_on_test_set(model, test_data_loader, specific_images, specific_labels, train_config, epoch_idx+1)
        test_metrics["mse"].append(avg_mse)
        test_metrics["ssim"].append(avg_ssim)
        test_metrics["psnr"].append(avg_psnr)
        print(f"Epoch {epoch_idx+1} - Test Evaluation: MSE={avg_mse:.4f}, SSIM={avg_ssim:.4f}, PSNR={avg_psnr:.4f}")

    print('Done Training...')

"""#### Latent 16"""

import yaml
import argparse
import torch
import random
import torchvision
import os
import numpy as np
from tqdm import tqdm
from torch.optim import Adam
from torchvision.utils import make_grid
from torch.utils.data.dataloader import DataLoader
import json

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

def train_VAE(custom_config=None):
    # Use custom_config if provided, otherwise use the default Config
    if custom_config:
        config_instance = custom_config
    else:
        config_instance = Config()
    config = config_instance.config  # Access the configuration dictionary
    print(config)

    dataset_config = config['dataset_params']
    autoencoder_config = config['autoencoder_params']
    train_config = config['train_params']

    # Set the desired seed value #
    seed = train_config['seed']
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    if device == 'cuda':
        torch.cuda.manual_seed_all(seed)
    print(device)
    #############################

    # Create the model and dataset #
    model = VAE(im_channels=dataset_config['im_channels'],
                  model_config=autoencoder_config).to(device)

    summary(model, input_size=(dataset_config['im_channels'], dataset_config['im_size'], dataset_config['im_size']))
    # Create the dataset
    im_dataset_cls = {
        'mnist': MnistDataset,
    }.get(dataset_config['name'])

    train_dataset = im_dataset_cls(split='train',
                                   im_path=dataset_config['im_path'],
                                   im_size=dataset_config['im_size'],
                                   im_channels=dataset_config['im_channels'])

    # Define condition configuration for the test set to always return labels
    test_condition_config = {
        'condition_types': ['class']
    }

    # Initialize test dataset with conditions
    test_dataset = MnistDataset(
        split='test',
        im_path=dataset_config['im_path'].replace('train', 'test'),
        im_size=dataset_config['im_size'],
        im_channels=dataset_config['im_channels'],
        condition_config=test_condition_config
    )

    im_dataset = im_dataset_cls(split='train',
                                im_path=dataset_config['im_path'],
                                im_size=dataset_config['im_size'],
                                im_channels=dataset_config['im_channels'])

    train_data_loader = DataLoader(train_dataset,
                                   batch_size=train_config['autoencoder_batch_size'],
                                   shuffle=True)

    test_data_loader = DataLoader(test_dataset,
                                  batch_size=train_config['autoencoder_batch_size'],
                                  shuffle=True)

    # Select specific images for consistent visualization
    specific_images, specific_labels = select_specific_images(test_data_loader, [1, 2, 3, 4, 5])

    # Create output directories
    if not os.path.exists(train_config['task_name']):
        os.mkdir(train_config['task_name'])

    num_epochs = train_config['autoencoder_epochs']

    # L1/L2 loss for Reconstruction
    recon_criterion = torch.nn.MSELoss()
    # Disc Loss can even be BCEWithLogits
    disc_criterion = torch.nn.MSELoss()


    # No need to freeze lpips as lpips.py takes care of that
    lpips_model = LPIPS(weights_path = train_config["lpips_weights_path"]).eval().to(device)
    discriminator = Discriminator(im_channels=dataset_config['im_channels']).to(device)

    optimizer_d = Adam(discriminator.parameters(), lr=train_config['autoencoder_lr'], betas=(0.5, 0.999))
    optimizer_g = Adam(model.parameters(), lr=train_config['autoencoder_lr'], betas=(0.5, 0.999))

    disc_step_start = train_config['disc_start']
    step_count = 0

    # This is for accumulating gradients incase the images are huge
    # And one cant afford higher batch sizes
    acc_steps = train_config['autoencoder_acc_steps']
    image_save_steps = train_config['autoencoder_img_save_steps']
    img_save_count = 0

    # Initialize metrics tracking dictionaries
    train_metrics = {
        "recon_losses": [],
        "perceptual_losses": [],
        "gen_losses": [],
        "disc_losses": []
    }
    test_metrics = {
        "mse": [],
        "ssim": [],
        "psnr": []
    }

    # Save metrics to JSON file
    metrics_save_path = os.path.join(train_config['task_name'], 'metrics.json')
    metrics = {
        "train_metrics": train_metrics,
        "test_metrics": test_metrics
    }


    # Total steps in one epoch for first-epoch evaluation intervals
    total_steps = len(train_data_loader)
    interval = total_steps // 5  # Perform 5 evaluations during the first epoch

    for epoch_idx in range(num_epochs):
        recon_losses = []
        perceptual_losses = []
        disc_losses = []
        gen_losses = []
        losses = []

        optimizer_g.zero_grad()
        optimizer_d.zero_grad()

        for step_idx, im in enumerate(tqdm(train_data_loader)):
            step_count += 1
            im = im.float().to(device)

            # Fetch autoencoders output(reconstructions)
            model_output = model(im)
            output, enc_, z_  = model_output

            # # Image Saving Logic
            # if step_count % image_save_steps == 0 or step_count == 1:
            #     sample_size = min(8, im.shape[0])
            #     save_output = torch.clamp(output[:sample_size], -1., 1.).detach().cpu()
            #     save_output = ((save_output + 1) / 2)
            #     save_input = ((im[:sample_size] + 1) / 2).detach().cpu()

            #     grid = make_grid(torch.cat([save_input, save_output], dim=0), nrow=sample_size)
            #     img = torchvision.transforms.ToPILImage()(grid)
            #     if not os.path.exists(os.path.join(train_config['task_name'],'vae_autoencoder_samples')):
            #         os.mkdir(os.path.join(train_config['task_name'], 'vae_autoencoder_samples'))
            #     img.save(os.path.join(train_config['task_name'],'vae_autoencoder_samples',
            #                           'current_autoencoder_sample_{}.png'.format(img_save_count)))
            #     img_save_count += 1
            #     img.close()

            ######### Optimize Generator ##########
            # L2 Loss
            recon_loss = recon_criterion(output, im)
            recon_losses.append(recon_loss.item())
            recon_loss = recon_loss / acc_steps
            # g_loss = (recon_loss +
            #           (train_config['codebook_weight'] * quantize_losses['codebook_loss'] / acc_steps) +
            #           (train_config['commitment_beta'] * quantize_losses['commitment_loss'] / acc_steps))
            # codebook_losses.append(train_config['codebook_weight'] * quantize_losses['codebook_loss'].item())
            g_loss = recon_loss
            # Adversarial loss only if disc_step_start steps passed
            if step_count >= disc_step_start:
                if step_count == disc_step_start:
                    print("Adversarial Loss started")
                disc_fake_pred = discriminator(model_output[0])
                disc_fake_loss = disc_criterion(disc_fake_pred,
                                                torch.ones(disc_fake_pred.shape,
                                                           device=disc_fake_pred.device))
                gen_losses.append(train_config['disc_weight'] * disc_fake_loss.item())
                g_loss += train_config['disc_weight'] * disc_fake_loss / acc_steps
            lpips_loss = torch.mean(lpips_model(output, im))
            perceptual_losses.append(train_config['perceptual_weight'] * lpips_loss.item())
            g_loss += train_config['perceptual_weight']*lpips_loss / acc_steps
            losses.append(g_loss.item())
            g_loss.backward()
            #####################################

            ######### Optimize Discriminator #######
            if step_count > disc_step_start:
                fake = output
                disc_fake_pred = discriminator(fake.detach())
                disc_real_pred = discriminator(im)
                disc_fake_loss = disc_criterion(disc_fake_pred,
                                                torch.zeros(disc_fake_pred.shape,
                                                            device=disc_fake_pred.device))
                disc_real_loss = disc_criterion(disc_real_pred,
                                                torch.ones(disc_real_pred.shape,
                                                           device=disc_real_pred.device))
                disc_loss = train_config['disc_weight'] * (disc_fake_loss + disc_real_loss) / 2
                disc_losses.append(disc_loss.item())
                disc_loss = disc_loss / acc_steps
                disc_loss.backward()
                if step_count % acc_steps == 0:
                    optimizer_d.step()
                    optimizer_d.zero_grad()
            #####################################
            if step_count % acc_steps == 0:
                optimizer_g.step()
                optimizer_g.zero_grad()

            # Perform additional evaluations during the first epoch
            if epoch_idx == 0 and step_idx % interval == 0:
                fraction = (step_idx + 1) / total_steps
                print(f"\nFirst epoch intermediate evaluation at step {step_idx + 1} ({fraction:.2f} epoch)")
                avg_mse, avg_ssim, avg_psnr = evaluate_on_test_set(model, test_data_loader, specific_images, specific_labels, train_config, fraction)
                test_metrics["mse"].append(avg_mse)
                test_metrics["ssim"].append(avg_ssim)
                test_metrics["psnr"].append(avg_psnr)
                print(f"Test Evaluation: MSE={avg_mse:.4f}, SSIM={avg_ssim:.4f}, PSNR={avg_psnr:.4f}")

        optimizer_d.step()
        optimizer_d.zero_grad()
        optimizer_g.step()
        optimizer_g.zero_grad()
        if len(disc_losses) > 0:
            print(
                '\nFinished epoch: {} | Recon Loss : {:.4f} | Perceptual Loss : {:.4f} | '
                'G Loss : {:.4f} | D Loss {:.4f}'.
                format(epoch_idx + 1,
                       np.mean(recon_losses),
                       np.mean(perceptual_losses),
                       np.mean(gen_losses),
                       np.mean(disc_losses)))
        else:
            print('Finished epoch: {} | Recon Loss : {:.4f} | Perceptual Loss : {:.4f}'.
                  format(epoch_idx + 1,
                         np.mean(recon_losses),
                         np.mean(perceptual_losses)
                         ))

        # Record epoch metrics
        train_metrics["recon_losses"].append(np.mean(recon_losses))
        train_metrics["perceptual_losses"].append(np.mean(perceptual_losses))
        train_metrics["gen_losses"].append(np.mean(gen_losses))
        train_metrics["disc_losses"].append(np.mean(disc_losses) if disc_losses else 0)

        torch.save(model.state_dict(), os.path.join(train_config['task_name'],
                                                    train_config['vae_autoencoder_ckpt_name']))
        torch.save(discriminator.state_dict(), os.path.join(train_config['task_name'],
                                                            train_config['vae_discriminator_ckpt_name']))

        avg_mse, avg_ssim, avg_psnr = evaluate_on_test_set(model, test_data_loader, specific_images, specific_labels, train_config, epoch_idx+1)
        test_metrics["mse"].append(avg_mse)
        test_metrics["ssim"].append(avg_ssim)
        test_metrics["psnr"].append(avg_psnr)
        print(f"Epoch {epoch_idx+1} - Test Evaluation: MSE={avg_mse:.4f}, SSIM={avg_ssim:.4f}, PSNR={avg_psnr:.4f}")

    with open(metrics_save_path, 'w') as f:
        json.dump(metrics, f, indent=4)
    print(f"Metrics saved to {metrics_save_path}")
    print('Done Training...')

# Define your custom configuration
custom_config = Config()
custom_config.config["autoencoder_params"]["z_channels"] = 16
custom_config.config["train_params"]["task_name"] = "/content/drive/My Drive/AI/Projects/VAE/Mnist/variable_latent_sizes/latent_16/V1/models"
custom_config.config["train_params"]["plot_save_path"] = "/content/drive/My Drive/AI/Projects/VAE/Mnist/variable_latent_sizes/latent_16/V1/plots"

# Call the train_VAE function with the custom configuration
train_VAE(custom_config=custom_config)

"""### Latent 4"""

# Define your custom configuration
custom_config = Config()
custom_config.config["autoencoder_params"]["z_channels"] = 4
custom_config.config["train_params"]["task_name"] = "/content/drive/My Drive/AI/Projects/VAE/Mnist/variable_latent_sizes/latent_4/V1/models"
custom_config.config["train_params"]["plot_save_path"] = "/content/drive/My Drive/AI/Projects/VAE/Mnist/variable_latent_sizes/latent_4/V1/plots"

# Call the train_VAE function with the custom configuration
train_VAE(custom_config=custom_config)

"""### Latent 2"""

# Define your custom configuration
custom_config = Config()
custom_config.config["autoencoder_params"]["z_channels"] = 2
custom_config.config["train_params"]["task_name"] = "/content/drive/My Drive/AI/Projects/VAE/Mnist/variable_latent_sizes/latent_2/V1/models"
custom_config.config["train_params"]["plot_save_path"] = "/content/drive/My Drive/AI/Projects/VAE/Mnist/variable_latent_sizes/latent_2/V1/plots"

# Call the train_VAE function with the custom configuration
train_VAE(custom_config=custom_config)

"""### Latent 64"""

# Define your custom configuration
custom_config = Config()
custom_config.config["autoencoder_params"]["z_channels"] = 64
custom_config.config["train_params"]["task_name"] = "/content/drive/My Drive/AI/Projects/VAE/Mnist/variable_latent_sizes/latent_64/models"
custom_config.config["train_params"]["plot_save_path"] = "/content/drive/My Drive/AI/Projects/VAE/Mnist/variable_latent_sizes/latent_64/plots"

# Call the train_VAE function with the custom configuration
train_VAE(custom_config=custom_config)

"""## Gif creation"""

import os
import json
import matplotlib.pyplot as plt
from PIL import Image
import glob
import numpy as np
from natsort import natsorted
from IPython.display import Image as IPImage, display

class VisualizationConfig:
    """Configuration class for visualization paths."""
    def __init__(self, metrics_path, visualization_images_path, gif_save_path, latent_size):
        self.metrics_path = metrics_path  # Path to the metrics JSON file
        self.visualization_images_path = visualization_images_path  # Path to visualization images
        self.gif_save_path = gif_save_path  # Path to save the generated GIFs
        self.latent_size = latent_size  # Latent size to include in titles


def plot_metrics_with_nan(ax, epochs, data, label, color, title, xlabel, ylabel):
    """
    Helper function to plot metrics, handling NaN values with gaps in the line.
    """
    nan_mask = np.isnan(data)
    ax.plot(
        epochs, np.where(nan_mask, None, data), label=label, color=color, linestyle="-", marker="o"
    )
    ax.scatter(
        np.array(epochs)[nan_mask], [0] * sum(nan_mask), color="red", marker="x", label="NaN"
    )
    ax.set_title(title)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.legend()
    ax.grid(alpha=0.3)
    ax.set_xlim(1, len(epochs) + 1)


def plot_metrics(config):
    """
    Plot train and test metrics (MSE, SSIM, PSNR, Gen Loss, Disc Loss) over epochs.
    Generates two figures: one for train and one for test metrics.

    Args:
        config (VisualizationConfig): Configuration object with paths.
    """
    with open(config.metrics_path, 'r') as f:
        metrics = json.load(f)

    train_metrics = metrics.get("train_metrics", {})
    test_metrics = metrics.get("test_metrics", {})

    train_epochs = range(1, len(train_metrics.get("recon_losses", [])) + 1)
    test_epochs = range(1, len(test_metrics.get("mse", [])) + 1)

    # Train Metrics
    fig_train, axs_train = plt.subplots(1, 4, figsize=(20, 5))
    fig_train.suptitle(f"Train Metrics for VAE with Latent Size {config.latent_size}", fontsize=22)
    for ax, (metric, color, label) in zip(
        axs_train,
        [
            (train_metrics["recon_losses"], "blue", "Recon Loss"),
            (train_metrics["perceptual_losses"], "green", "Perceptual Loss"),
            (train_metrics["gen_losses"], "orange", "Generator Loss"),
            (train_metrics["disc_losses"], "red", "Discriminator Loss"),
        ],
    ):
        plot_metrics_with_nan(ax, train_epochs, metric, label, color, label, "Epoch", "Loss")
    plt.tight_layout(rect=[0, 0, 1, 0.95])
    plt.savefig(os.path.join(config.gif_save_path, "train_metrics_summary.png"))
    plt.show()

    # Test Metrics
    fig_test, axs_test = plt.subplots(1, 3, figsize=(15, 5))
    fig_test.suptitle(f"Test Metrics for VAE with Latent Size {config.latent_size}", fontsize=22)
    for ax, (metric, color, label) in zip(
        axs_test,
        [
            (test_metrics["mse"], "blue", "MSE"),
            (test_metrics["ssim"], "green", "SSIM"),
            (test_metrics["psnr"], "orange", "PSNR"),
        ],
    ):
        plot_metrics_with_nan(ax, test_epochs, metric, label, color, label, "Epoch", label)
    plt.tight_layout(rect=[0, 0, 1, 0.95])
    plt.savefig(os.path.join(config.gif_save_path, "test_metrics_summary.png"))
    plt.show()


def categorize_images(file_paths):
    """
    Categorize image files based on their names.

    Args:
        file_paths (list): List of file paths.

    Returns:
        dict: Dictionary with categorized file paths.
    """
    categories = {
        "evaluation_plots": [],
        "latent_visualizations": []
    }

    for file_path in file_paths:
        if "evaluation_plot_epoch" in file_path:
            categories["evaluation_plots"].append(file_path)
        elif "latent_visualization_epoch" in file_path:
            categories["latent_visualizations"].append(file_path)

    return {key: natsorted(value) for key, value in categories.items()}


def create_gif(image_paths, gif_save_path, gif_name, duration=300):
    """
    Create a GIF from a list of images and display it.

    Args:
        image_paths (list): List of image paths.
        gif_save_path (str): Path to save the generated GIF.
        gif_name (str): Name of the GIF file.
        duration (int): Duration between frames in milliseconds.
    """
    if not image_paths:
        print(f"No images found for {gif_name}")
        return

    frames = [Image.open(img) for img in image_paths]
    gif_path = os.path.join(gif_save_path, gif_name)
    frames[0].save(gif_path, save_all=True, append_images=frames[1:], duration=duration, loop=0)
    print(f"GIF saved to {gif_path}")
    display(IPImage(filename=gif_path))


def process_and_create_gifs(image_path, gif_save_path):
    """
    Process images and create GIFs for evaluation plots and latent visualizations.

    Args:
        image_path (str): Path containing all the images.
        gif_save_path (str): Path to save the generated GIFs.
    """
    all_image_files = glob.glob(os.path.join(image_path, "*.png"))
    categorized_images = categorize_images(all_image_files)

    create_gif(
        image_paths=categorized_images["evaluation_plots"],
        gif_save_path=gif_save_path,
        gif_name="evaluation_plots.gif"
    )
    create_gif(
        image_paths=categorized_images["latent_visualizations"],
        gif_save_path=gif_save_path,
        gif_name="latent_visualizations.gif"
    )

"""### Latent 16"""

latent_size = 16
base_path = "/content/drive/My Drive/AI/Projects/VAE/Mnist/variable_latent_sizes"

config = VisualizationConfig(
    metrics_path=os.path.join(base_path, f"latent_{latent_size}/V1/models/metrics.json"),
    visualization_images_path=os.path.join(base_path, f"latent_{latent_size}/V1/plots"),
    gif_save_path=os.path.join(base_path, f"latent_{latent_size}/V1/plots"),
    latent_size=latent_size
)

plot_metrics(config)

image_path = config.visualization_images_path
gif_save_path = config.gif_save_path
process_and_create_gifs(image_path=image_path, gif_save_path=gif_save_path)

"""### Latent 4"""

latent_size = 4
base_path = "/content/drive/My Drive/AI/Projects/VAE/Mnist/variable_latent_sizes"

config = VisualizationConfig(
    metrics_path=os.path.join(base_path, f"latent_{latent_size}/V1/models/metrics.json"),
    visualization_images_path=os.path.join(base_path, f"latent_{latent_size}/V1/plots"),
    gif_save_path=os.path.join(base_path, f"latent_{latent_size}/V1/plots"),
    latent_size=latent_size
)

plot_metrics(config)

image_path = config.visualization_images_path
gif_save_path = config.gif_save_path
process_and_create_gifs(image_path=image_path, gif_save_path=gif_save_path)

"""### Latent 2"""

latent_size = 2
base_path = "/content/drive/My Drive/AI/Projects/VAE/Mnist/variable_latent_sizes"

config = VisualizationConfig(
    metrics_path=os.path.join(base_path, f"latent_{latent_size}/V1/models/metrics.json"),
    visualization_images_path=os.path.join(base_path, f"latent_{latent_size}/V1/plots"),
    gif_save_path=os.path.join(base_path, f"latent_{latent_size}/V1/plots"),
    latent_size=latent_size
)

plot_metrics(config)

image_path = config.visualization_images_path
gif_save_path = config.gif_save_path
process_and_create_gifs(image_path=image_path, gif_save_path=gif_save_path)