util.py

from pathlib import Path
import numpy as np
import argparse

def parse_args_VQGAN():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--reconstruct', action='store_true', default=False, help='reconstruct model')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--batch_size', type=int, default=256, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--size', type=int, default=None, help='image size')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    argparser.add_argument('--warmup', type=int, default=10, help='warmup epochs')
    argparser.add_argument('--channels', type=int, default=64, help='Base channel count for encoder/decoder')
    argparser.add_argument('--z_channels', type=int, default=64, help='Number of latent channels')
    argparser.add_argument('--ch_mult', type=int, nargs='+', default=[1,2,2], help='Channel multiplier for each resolution')
    argparser.add_argument('--num_res_blocks', type=int, default=2, help='Number of residual blocks per resolution')
    argparser.add_argument('--attn_resolutions', type=int, nargs='+', default=[16], help='Resolutions with attention')
    argparser.add_argument('--dropout', type=float, default=0.0, help='Dropout rate')
    argparser.add_argument('--double_z', action='store_true', default=False, help='Use double z in encoder')
    argparser.add_argument('--disc_start', type=int, default=10000, help='Step to start discriminator loss')
    argparser.add_argument('--disc_weight', type=float, default=0.8, help='Discriminator loss weight')
    argparser.add_argument('--codebook_weight', type=float, default=1.0, help='Codebook loss weight')
    argparser.add_argument('--n_embed', type=int, default=128, help='Number of embeddings in codebook')
    argparser.add_argument('--embed_dim', type=int, default=64, help='Embedding dimension')
    argparser.add_argument('--remap', type=str, default=None, help='Remap indices for codebook')
    argparser.add_argument('--sane_index_shape', action='store_true', default=False, help='Use sane index shape for quantizer')
    argparser.add_argument('--checkpoint', type=str, default=None, help='Path to checkpoint')
    argparser.add_argument('--colorize_nlabels', type=int, default=None, help='Number of labels for colorization')
    argparser.add_argument('--lr', type=float, default=4.5e-6, help='Learning rate')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='Disable wandb logging')
    argparser.add_argument('--sample_and_save_freq', type=int, default=20, help='Sample and save frequency')
    argparser.add_argument('--disc_num_layers', type=int, default=3, help='Number of layers in discriminator')
    return argparser.parse_args()

def parse_args_GPT():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample from model')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--batch_size', type=int, default=256, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--size', type=int, default=None, help='image size')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    argparser.add_argument('--warmup', type=int, default=10, help='warmup epochs')
    argparser.add_argument('--channels', type=int, default=64, help='Base channel count for encoder/decoder')
    argparser.add_argument('--z_channels', type=int, default=64, help='Number of latent channels')
    argparser.add_argument('--ch_mult', type=int, nargs='+', default=[1,2,2], help='Channel multiplier for each resolution')
    argparser.add_argument('--num_res_blocks', type=int, default=2, help='Number of residual blocks per resolution')
    argparser.add_argument('--attn_resolutions', type=int, nargs='+', default=[16], help='Resolutions with attention')
    argparser.add_argument('--dropout', type=float, default=0.0, help='Dropout rate')
    argparser.add_argument('--double_z', action='store_true', default=False, help='Use double z in encoder')
    argparser.add_argument('--disc_start', type=int, default=10000, help='Step to start discriminator loss')
    argparser.add_argument('--disc_weight', type=float, default=0.8, help='Discriminator loss weight')
    argparser.add_argument('--codebook_weight', type=float, default=1.0, help='Codebook loss weight')
    argparser.add_argument('--n_embed', type=int, default=128, help='Number of embeddings in codebook')
    argparser.add_argument('--embed_dim', type=int, default=64, help='Embedding dimension for VQGAN')
    argparser.add_argument('--disc_num_layers', type=int, default=3, help='Number of layers in discriminator')
    argparser.add_argument('--embed_dim_t', type=int, default=64, help='Embedding dimension for transformer')
    argparser.add_argument('--remap', type=str, default=None, help='Remap indices for codebook')
    argparser.add_argument('--sane_index_shape', action='store_true', default=False, help='Use sane index shape for quantizer')
    argparser.add_argument('--checkpoint_vae', type=str, default=None, help='Path to checkpoint')
    argparser.add_argument('--colorize_nlabels', type=int, default=None, help='Number of labels for colorization')
    argparser.add_argument('--lr', type=float, default=1e-4, help='Learning rate')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='Disable wandb logging')
    argparser.add_argument('--sample_and_save_freq', type=int, default=20, help='Sample and save frequency')
    argparser.add_argument('--n_layer', type=int, default=6, help='Number of layers in transformer')
    argparser.add_argument('--n_head', type=int, default=8, help='Number of heads in transformer')
    argparser.add_argument('--block_size', type=int, default=64, help='Block size for transformer')
    argparser.add_argument('--bias', action='store_true', default=False, help='Use bias in transformer')
    argparser.add_argument('--dropout_t', type=float, default=0.1, help='Dropout rate in transformer')
    argparser.add_argument('--betas', type=float, nargs='+', default=[0.9, 0.95], help='Betas for Adam optimizer')
    argparser.add_argument('--weight_decay', type=float, default=0.0, help='Weight decay for Adam optimizer')
    argparser.add_argument('--num_samples', type=int, default=16, help='Number of samples to generate')
    argparser.add_argument('--checkpoint_gpt', type=str, default=None, help='Path to checkpoint for transformer')
    argparser.add_argument('--temperature', type=float, default=1.0, help='Temperature for sampling')
    argparser.add_argument('--top_k', type=int, default=None, help='Top k for sampling')
    argparser.add_argument('--outlier_detection', action='store_true', default=False, help='outlier detection')
    argparser.add_argument('--out_dataset', type=str, default='fashionmnist', help='outlier dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet'])
    argparser.add_argument('--ema_decay', type=float, default=0.999, help='EMA decay for model parameters')
    return argparser.parse_args()

def parse_args_MaskGiT():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample from model')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--batch_size', type=int, default=256, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--size', type=int, default=None, help='image size')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    argparser.add_argument('--warmup', type=int, default=10, help='warmup epochs')
    argparser.add_argument('--channels', type=int, default=64, help='Base channel count for encoder/decoder')
    argparser.add_argument('--z_channels', type=int, default=64, help='Number of latent channels')
    argparser.add_argument('--ch_mult', type=int, nargs='+', default=[1,2,2], help='Channel multiplier for each resolution')
    argparser.add_argument('--num_res_blocks', type=int, default=2, help='Number of residual blocks per resolution')
    argparser.add_argument('--attn_resolutions', type=int, nargs='+', default=[16], help='Resolutions with attention')
    argparser.add_argument('--dropout', type=float, default=0.0, help='Dropout rate')
    argparser.add_argument('--double_z', action='store_true', default=False, help='Use double z in encoder')
    argparser.add_argument('--disc_start', type=int, default=10000, help='Step to start discriminator loss')
    argparser.add_argument('--disc_weight', type=float, default=0.8, help='Discriminator loss weight')
    argparser.add_argument('--codebook_weight', type=float, default=1.0, help='Codebook loss weight')
    argparser.add_argument('--n_embed', type=int, default=128, help='Number of embeddings in codebook')
    argparser.add_argument('--embed_dim', type=int, default=64, help='Embedding dimension for VQGAN')
    argparser.add_argument('--remap', type=str, default=None, help='Remap indices for codebook')
    argparser.add_argument('--sane_index_shape', action='store_true', default=False, help='Use sane index shape for quantizer')
    argparser.add_argument('--checkpoint_vae', type=str, default=None, help='Path to checkpoint')
    argparser.add_argument('--colorize_nlabels', type=int, default=None, help='Number of labels for colorization')
    argparser.add_argument('--lr', type=float, default=1e-4, help='Learning rate')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='Disable wandb logging')
    argparser.add_argument("--cfg_w",        type=float, default=3,          help="classifier free guidance wight")
    argparser.add_argument("--r_temp",       type=float, default=4.5,        help="Gumbel noise temperature when sampling")
    argparser.add_argument("--sm_temp",      type=float, default=1.,         help="temperature before softmax when sampling")
    argparser.add_argument("--drop-label",   type=float, default=0.1,        help="drop rate for cfg")
    argparser.add_argument("--hidden_dim",    type=int,   default=128,        help="hidden dim for transformer")
    argparser.add_argument("--heads",    type=int,   default=8,          help="number of heads for transformer")
    argparser.add_argument("--depth",    type=int,   default=10,          help="number of layers for transformer")
    argparser.add_argument("--mlp_dim",    type=int,   default=384,       help="mlp dim for transformer")
    argparser.add_argument('--dropout_t', type=float, default=0.1, help='Dropout rate in transformer')
    argparser.add_argument('--sample_and_save_freq', type=int, default=20, help='Sample and save frequency')
    argparser.add_argument('--betas', type=float, nargs='+', default=[0.9, 0.95], help='Betas for Adam optimizer')
    argparser.add_argument('--weight_decay', type=float, default=0.1, help='Weight decay for Adam optimizer')
    argparser.add_argument("--step",         type=int,   default=8,          help="number of steps for sampling")
    argparser.add_argument("--sched_mode",   type=str,   default="arccos",   help="scheduler mode when sampling")
    argparser.add_argument("--mask-value",   type=int,   default=None,       help="mask value for sampling")
    argparser.add_argument("--n_classes",       type=int,   default=10,        help="number of classes for sampling")
    argparser.add_argument('--num_samples', type=int, default=16, help='Number of samples to generate')
    argparser.add_argument('--checkpoint_vit', type=str, default=None, help='Path to checkpoint for transformer')
    argparser.add_argument('--disc_num_layers', type=int, default=3, help='Number of layers in discriminator')
    return argparser.parse_args()

def parse_args_HierarchicalVAE():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample model')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--batch_size', type=int, default=256, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--lr', type=float, default=0.01, help='learning rate')
    argparser.add_argument('--latent_dim', type=int, default=512, help='latent dimension')
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path')
    argparser.add_argument('--sample_and_save_freq', type=int, default=5, help='sample and save frequency')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    return argparser.parse_args()

def parse_args_RealNVP():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample model')
    argparser.add_argument('--outlier_detection', action='store_true', default=False, help='outlier detection')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--out_dataset', type=str, default='fashionmnist', help='outlier dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet'])
    argparser.add_argument('--batch_size', type=int, default=128, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--lr', type=float, default=1e-3, help='learning rate')
    argparser.add_argument('--weight_decay', type=float, default=1e-5, help='weight decay')
    argparser.add_argument('--max_grad_norm', type=float, default=100.0, help='max grad norm')
    argparser.add_argument('--sample_and_save_freq', type=int, default=5, help='sample and save frequency')
    argparser.add_argument('--num_scales', type=int, default=2, help='number of scales')
    argparser.add_argument('--mid_channels', type=int, default=64, help='mid channels')
    argparser.add_argument('--num_blocks', type=int, default=8, help='number of blocks')
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    return argparser.parse_args()


def parse_args_PixelCNN():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample from model')
    argparser.add_argument('--outlier_detection', action='store_true', default=False, help='outlier detection')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--out_dataset', type=str, default='fashionmnist', help='outlier dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet'])
    argparser.add_argument('--batch_size', type=int, default=128, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--lr', type=float, default=1e-3, help='learning rate')
    argparser.add_argument('--gamma', type=float, default=0.99, help='gamma for the lr scheduler')
    argparser.add_argument('--sample_and_save_freq', type=int, default=5, help='sample and save frequency')
    argparser.add_argument('--hidden_channels', type=int, default=64, help='number of channels for the convolutional layers')
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    return argparser.parse_args()

def parse_args_VQGAN_Transformer():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample model')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--batch_size', type=int, default=128, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs for VQVAE')
    argparser.add_argument('--lr', type=float, default=0.0002, help='learning rate VQVAE')
    argparser.add_argument('--lr_d', type=float, default=0.0005, help='learning rate discriminator')
    argparser.add_argument('--adv_weight', type=float, default=0.01, help='adversarial weight')
    argparser.add_argument('--perceptual_weight', type=float, default=0.001, help='perceptual weight')
    argparser.add_argument('--lr_t', type=float, default=0.0005, help='learning rate transformer')
    argparser.add_argument('--n_epochs_t', type=int, default=100, help='number of epochs transformer')
    argparser.add_argument('--num_res_layers', type=int, default=2, help='number of residual layers')
    argparser.add_argument('--downsample_parameters', type=int, nargs='+', default=[2, 4, 1, 1], help='downsample parameters')
    argparser.add_argument('--upsample_parameters', type=int, nargs='+', default=[2, 4, 1, 1, 0], help='upsample parameters')
    argparser.add_argument('--num_channels', type=int, default=[256, 256], help='number of channels')
    argparser.add_argument('--num_res_channels', type=int, default=[256, 256], help='number of res channels')
    argparser.add_argument('--num_embeddings', type=int, default=256, help='number of embeddings')
    argparser.add_argument('--embedding_dim', type=int, default=32, help='embedding dimension')
    argparser.add_argument('--attn_layers_dim', type=int, default=96, help='attn layers dim')
    argparser.add_argument('--attn_layers_depth', type=int, default=12, help='attn layers depth')
    argparser.add_argument('--attn_layers_heads', type=int, default=8, help='attn layers heads')
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path to VQVAE')
    argparser.add_argument('--checkpoint_t', type=str, default=None, help='checkpoint path to Transformer')
    argparser.add_argument('--num_samples', type=int, default=16, help='number of samples')
    argparser.add_argument('--num_layers_d', type=int, default=3, help='number of layers in discriminator')
    argparser.add_argument('--num_channels_d', type=int, default=64, help='number of channels in discriminator')
    argparser.add_argument('--sample_and_save_freq', type=int, default=5, help='sample and save frequency')
    argparser.add_argument('--outlier_detection', action='store_true', default=False, help='outlier detection')
    argparser.add_argument('--out_dataset', type=str, default='fashionmnist', help='outlier dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet'])
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    args = argparser.parse_args()
    args.num_channels = tuple(args.num_channels)
    args.num_res_channels = tuple(args.num_res_channels)
    args.downsample_parameters = (tuple(args.downsample_parameters), tuple(args.downsample_parameters))
    args.upsample_parameters = (tuple(args.upsample_parameters), tuple(args.upsample_parameters))
    return args

def parse_args_FlowMatching():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample model')
    argparser.add_argument('--fid', action='store_true', default=False, help='calculate FID')
    argparser.add_argument('--batch_size', type=int, default=256, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--lr', type=float, default=1e-3, help='learning rate')
    argparser.add_argument('--model_channels', type=int, default = 64, help='number of features')
    argparser.add_argument('--num_res_blocks', type=int, default = 2, help='number of residual blocks per downsample')
    argparser.add_argument('--attention_resolutions', type=int, nargs='+', default = [4], help='downsample rates at which attention will take place')
    argparser.add_argument('--dropout', type=float, default = 0.0, help='dropout probability')
    argparser.add_argument('--channel_mult', type=int, nargs='+', default = [1, 2, 2], help='channel multiplier for each level of the UNet')
    argparser.add_argument('--conv_resample', type=bool, default = True, help='use learned convolutions for upsampling and downsampling')
    argparser.add_argument('--dims', type=int, default = 2, help='determines if the signal is 1D, 2D, or 3D')
    argparser.add_argument('--num_heads', type=int, default = 4, help='number of attention heads in each attention layer')
    argparser.add_argument('--num_head_channels', type=int, default = 32, help='use a fixed channel width per attention head')
    argparser.add_argument('--use_scale_shift_norm', type=bool, default = False, help='use a FiLM-like conditioning mechanism')
    argparser.add_argument('--resblock_updown', type=bool, default = False, help='use residual blocks for up/downsampling')
    argparser.add_argument('--use_new_attention_order', type=bool, default = False, help='use a different attention pattern for potentially increased efficiency')
    argparser.add_argument('--sample_and_save_freq', type=int, default=5, help='sample and save frequency')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path')
    argparser.add_argument('--num_samples', type=int, default=16, help='number of samples')
    argparser.add_argument('--out_dataset', type=str, default='fashionmnist', help='outlier dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet'])
    argparser.add_argument('--outlier_detection', action='store_true', default=False, help='outlier detection')
    argparser.add_argument('--interpolation', action='store_true', default=False, help='interpolation')
    argparser.add_argument('--solver_lib', type=str, default='none', help='solver library', choices=['torchdiffeq', 'zuko', 'none'])
    argparser.add_argument('--step_size', type=float, default=0.1, help='step size for ODE solver')
    argparser.add_argument('--solver', type=str, default='dopri5', help='solver for ODE', choices=['dopri5', 'rk4', 'dopri8', 'euler', 'bosh3', 'adaptive_heun', 'midpoint', 'explicit_adams', 'implicit_adams'])
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    argparser.add_argument('--warmup', type=int, default=10, help='warmup epochs')
    argparser.add_argument('--decay', type=float, default=1e-7, help='decay rate')
    argparser.add_argument('--latent', action='store_true', default=False, help='Use latent implementation')
    argparser.add_argument('--size', type=int, default=None, help='Size of the original image')
    argparser.add_argument('--ema_rate', type=float, default=0.999, help='ema rate')
    argparser.add_argument('--snapshots', type=int, default=10, help='how many snapshots during training')
    argparser.add_argument('--recon_factor', type=float, default=0.5, help='reconstruction factor')
    argparser.add_argument('--lpips', action='store_true', default=False, help='use lpips loss')
    args = argparser.parse_args()
    args.channel_mult = tuple(args.channel_mult)
    args.attention_resolutions = tuple(args.attention_resolutions)
    return args

def parse_args_CondFlowMatching():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample model')
    argparser.add_argument('--fid', action='store_true', default=False, help='calculate FID')
    argparser.add_argument('--batch_size', type=int, default=256, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--lr', type=float, default=1e-3, help='learning rate')
    argparser.add_argument('--model_channels', type=int, default = 64, help='number of features')
    argparser.add_argument('--num_res_blocks', type=int, default = 2, help='number of residual blocks per downsample')
    argparser.add_argument('--attention_resolutions', type=int, nargs='+', default = [4], help='downsample rates at which attention will take place')
    argparser.add_argument('--dropout', type=float, default = 0.0, help='dropout probability')
    argparser.add_argument('--channel_mult', type=int, nargs='+', default = [1, 2, 2], help='channel multiplier for each level of the UNet')
    argparser.add_argument('--conv_resample', type=bool, default = True, help='use learned convolutions for upsampling and downsampling')
    argparser.add_argument('--dims', type=int, default = 2, help='determines if the signal is 1D, 2D, or 3D')
    argparser.add_argument('--num_heads', type=int, default = 4, help='number of attention heads in each attention layer')
    argparser.add_argument('--num_head_channels', type=int, default = 32, help='use a fixed channel width per attention head')
    argparser.add_argument('--use_scale_shift_norm', type=bool, default = False, help='use a FiLM-like conditioning mechanism')
    argparser.add_argument('--resblock_updown', type=bool, default = False, help='use residual blocks for up/downsampling')
    argparser.add_argument('--use_new_attention_order', type=bool, default = False, help='use a different attention pattern for potentially increased efficiency')
    argparser.add_argument('--sample_and_save_freq', type=int, default=5, help='sample and save frequency')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path')
    argparser.add_argument('--num_samples', type=int, default=16, help='number of samples')
    argparser.add_argument('--solver_lib', type=str, default='none', help='solver library', choices=['torchdiffeq', 'zuko', 'none'])
    argparser.add_argument('--step_size', type=float, default=0.1, help='step size for ODE solver')
    argparser.add_argument('--solver', type=str, default='dopri5', help='solver for ODE', choices=['dopri5', 'rk4', 'dopri8', 'euler', 'bosh3', 'adaptive_heun', 'midpoint', 'explicit_adams', 'implicit_adams'])
    argparser.add_argument('--n_classes', type=int, default=10, help='number of classes')
    argparser.add_argument('--dropout_prob', type=float, default=0.2, help='probability of conditioning during training')
    argparser.add_argument('--cfg', type=float, default=2.0, help='guidance scale')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    argparser.add_argument('--warmup', type=int, default=10, help='warmup epochs')
    argparser.add_argument('--decay', type=float, default=1e-5, help='decay rate')
    argparser.add_argument('--latent', action='store_true', default=False, help='Use latent implementation')
    argparser.add_argument('--size', type=int, default=None, help='Size of the original image')
    argparser.add_argument('--ema_rate', type=float, default=0.999, help='ema rate')
    argparser.add_argument('--snapshots', type=int, default=10, help='how many snapshots during training')
    argparser.add_argument('--translation', action='store_true', default=False, help='translation')
    argparser.add_argument('--translation_factor', type=float, default=0.2, help='translation factor (as percentage of preserved original image)')
    argparser.add_argument('--classification', action='store_true', default=False, help='classification')
    argparser.add_argument('--recon_factor', type=float, default=0.5, help='reconstruction factor')
    argparser.add_argument('--lpips', action='store_true', default=False, help='use lpips loss')
    args = argparser.parse_args()
    args.channel_mult = tuple(args.channel_mult)
    args.attention_resolutions = tuple(args.attention_resolutions)
    return args

def parse_args_RectifiedFlows():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample model')
    argparser.add_argument('--outlier_detection', action='store_true', default=False, help='outlier detection')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--out_dataset', type=str, default='fashionmnist', help='outlier dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet'])
    argparser.add_argument('--batch_size', type=int, default=128, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--lr', type=float, default=5e-4, help='learning rate')
    argparser.add_argument('--patch_size', type=int, default=2, help='patch size')
    argparser.add_argument('--dim', type=int, default=64, help='dimension')
    argparser.add_argument('--n_layers', type=int, default=6, help='number of layers')
    argparser.add_argument('--n_heads', type=int, default=4, help='number of heads')
    argparser.add_argument('--multiple_of', type=int, default=256, help='multiple of')
    argparser.add_argument('--ffn_dim_multiplier', type=int, default=None, help='ffn dim multiplier')
    argparser.add_argument('--norm_eps', type=float, default=1e-5, help='norm eps')
    argparser.add_argument('--class_dropout_prob', type=float, default=0.1, help='class dropout probability')
    argparser.add_argument('--sample_and_save_freq', type=int, default=5, help='sample and save frequency')
    argparser.add_argument('--num_classes', type=int, default=10, help='number of classes')
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path')
    argparser.add_argument('--cfg', type=float, default=1.0, help='label guidance')
    argparser.add_argument('--sample_steps', type=int, default=50, help='number of steps for sampling')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    argparser.add_argument('--ema_rate', type=float, default=0.9999, help='ema rate')
    argparser.add_argument('--conditional', action='store_true', default=False, help='conditional')
    argparser.add_argument('--warmup', type=int, default=10, help='warmup epochs')
    argparser.add_argument('--decay', type=float, default=0.0, help='weight decay rate')
    argparser.add_argument('--latent', action='store_true', default=False, help='Use latent implementation')
    argparser.add_argument('--size', type=int, default=None, help='Size of the original image')
    argparser.add_argument('--snapshots', type=int, default=10, help='how many snapshots during training')
    argparser.add_argument('--solver_lib', type=str, default='torchdiffeq', help='solver library', choices=['torchdiffeq', 'none'])
    argparser.add_argument('--solver', type=str, default='euler', help='solver for ODE', choices=['dopri5', 'rk4', 'dopri8', 'euler', 'bosh3', 'adaptive_heun', 'midpoint', 'explicit_adams', 'implicit_adams', 'heun3'])
    argparser.add_argument('--fid', action='store_true', default=False, help='calculate FID')
    return argparser.parse_args()

def parse_args_VanillaFlow():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample model')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--batch_size', type=int, default=128, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--lr', type=float, default=1e-3, help='learning rate')
    argparser.add_argument('--c_hidden', type=int, default=16, help='Hidden units in the first coupling layer')
    argparser.add_argument('--multi_scale', action='store_true', default=False, help='use multi scale')
    argparser.add_argument('--vardeq', action='store_true', default=False, help='use variational dequantization')
    argparser.add_argument('--sample_and_save_freq', type=int, default=5, help='sample and save frequency')
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path')
    argparser.add_argument('--outlier_detection', action='store_true', default=False, help='outlier detection')
    argparser.add_argument('--out_dataset', type=str, default='fashionmnist', help='outlier dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet'])
    argparser.add_argument('--n_layers', type=int, default=8, help='number of layers')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    return argparser.parse_args()

def parse_args_VQVAE_Transformer():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample model')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--batch_size', type=int, default=128, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs for VQVAE')
    argparser.add_argument('--lr', type=float, default=0.0002, help='learning rate VQVAE')
    argparser.add_argument('--lr_t', type=float, default=0.0002, help='learning rate transformer')
    argparser.add_argument('--n_epochs_t', type=int, default=100, help='number of epochs transformer')
    argparser.add_argument('--num_res_layers', type=int, default=2, help='number of residual layers')
    argparser.add_argument('--downsample_parameters', type=int, nargs='+', default=[2, 4, 1, 1], help='downsample parameters')
    argparser.add_argument('--upsample_parameters', type=int, nargs='+', default=[2, 4, 1, 1, 0], help='upsample parameters')
    argparser.add_argument('--num_channels', type=int, default=[256, 256], help='number of channels')
    argparser.add_argument('--num_res_channels', type=int, default=[256, 256], help='number of res channels')
    argparser.add_argument('--num_embeddings', type=int, default=256, help='number of embeddings')
    argparser.add_argument('--embedding_dim', type=int, default=32, help='embedding dimension')
    argparser.add_argument('--attn_layers_dim', type=int, default=96, help='attn layers dim')
    argparser.add_argument('--attn_layers_depth', type=int, default=12, help='attn layers depth')
    argparser.add_argument('--attn_layers_heads', type=int, default=8, help='attn layers heads')
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path to VQVAE')
    argparser.add_argument('--checkpoint_t', type=str, default=None, help='checkpoint path to Transformer')
    argparser.add_argument('--num_samples', type=int, default=16, help='number of samples')
    argparser.add_argument('--sample_and_save_freq', type=int, default=5, help='sample and save frequency')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    args = argparser.parse_args()
    args.num_channels = tuple(args.num_channels)
    args.num_res_channels = tuple(args.num_res_channels)
    args.downsample_parameters = (tuple(args.downsample_parameters), tuple(args.downsample_parameters))
    args.upsample_parameters = (tuple(args.upsample_parameters), tuple(args.upsample_parameters))
    return args

def parse_args_FlowPP():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample from model')
    argparser.add_argument('--outlier_detection', action='store_true', default=False, help='outlier detection')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--out_dataset', type=str, default='fashionmnist', help='outlier dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet'])
    argparser.add_argument('--batch_size', type=int, default=8, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--lr', type=float, default=1e-3, help='learning rate')
    argparser.add_argument('--warm_up', type=int, default=200, help='warm up')
    argparser.add_argument('--grad_clip', type=float, default=1.0, help='gradient clip')
    argparser.add_argument('--drop_prob', type=float, default=0.2, help='dropout probability')
    argparser.add_argument('--num_blocks', type=int, default=10, help='number of blocks')
    argparser.add_argument('--num_components', default=32, type=int, help='Number of components in the mixture')
    argparser.add_argument('--num_dequant_blocks', default=2, type=int, help='Number of blocks in dequantization')
    argparser.add_argument('--num_channels', default=96, type=int, help='Number of channels in Flow++')
    argparser.add_argument('--use_attn', action='store_true', default=False, help='use attention')
    argparser.add_argument('--sample_and_save_freq', type=int, default=5, help='sample interval')
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    return argparser.parse_args()

def parse_args_VanillaVAE():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample model')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--out_dataset', type=str, default='fashionmnist', help='outlier dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet'])
    argparser.add_argument('--batch_size', type=int, default=128, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--lr', type=float, default=0.0002, help='learning rate')
    argparser.add_argument('--latent_dim', type=int, default=128, help='latent dimension')
    argparser.add_argument('--hidden_dims', type=int, nargs='+', default=None, help='hidden dimensions')
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path')
    argparser.add_argument('--num_samples', type=int, default=16, help='number of samples')
    argparser.add_argument('--sample_and_save_freq', type=int, default=5, help='sample and save frequency')
    argparser.add_argument('--loss_type', type=str, default='mse', help='loss type', choices=['mse', 'ssim'])
    argparser.add_argument('--kld_weight', type=float, default=1e-4, help='kl weight')
    argparser.add_argument('--outlier_detection', action='store_true', default=False, help='outlier detection')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    return argparser.parse_args()

def parse_args_ConditionalVAE():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample model')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--batch_size', type=int, default=128, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--lr', type=float, default=0.0002, help='learning rate')
    argparser.add_argument('--latent_dim', type=int, default=128, help='latent dimension')
    argparser.add_argument('--hidden_dims', type=int, nargs='+', default=None, help='hidden dimensions')
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path')
    argparser.add_argument('--num_samples', type=int, default=16, help='number of samples')
    argparser.add_argument('--num_classes', type=int, default=10, help='number of classes')
    argparser.add_argument('--kld_weight', type=float, default=1e-4, help='kl weight')
    argparser.add_argument('--loss_type', type=str, default='mse', help='loss type', choices=['mse', 'ssim'])
    argparser.add_argument('--sample_and_save_freq', type=int, default=5, help='sample and save frequency')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    return argparser.parse_args()

def parse_args_AdversarialVAE():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--test', action='store_true', default=False, help='test model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample model')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet', 'imagenet', 'imagenetpatch', 'tinyimagenetpatch'])
    argparser.add_argument('--batch_size', type=int, default=128, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--lr', type=float, default=0.0002, help='learning rate')
    argparser.add_argument('--latent_dim', type=int, default=128, help='latent dimension')
    argparser.add_argument('--hidden_dims', type=int, nargs='+', default=None, help='hidden dimensions')
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path')
    argparser.add_argument('--num_samples', type=int, default=16, help='number of samples')
    argparser.add_argument('--gen_weight', type=float, default=0.002, help='generator weight')
    argparser.add_argument('--recon_weight', type=float, default=0.002, help='reconstruction weight')
    argparser.add_argument('--sample_and_save_frequency', type=int, default=5, help='sample and save frequency')
    argparser.add_argument('--outlier_detection', action='store_true', default=False, help='outlier detection')
    argparser.add_argument('--discriminator_checkpoint', type=str, default=None, help='discriminator checkpoint path')
    argparser.add_argument('--out_dataset', type=str, default='fashionmnist', help='outlier dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet'])
    argparser.add_argument('--kld_weight', type=float, default=1e-4, help='kl weight')
    argparser.add_argument('--loss_type', type=str, default='mse', help='loss type', choices=['mse', 'ssim'])
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    argparser.add_argument('--size', type=int, default=None, help='size of image (leave None for default for each dataset)')
    argparser.add_argument('--patches', type=int, default=16, help='number of patches to divide image into')
    return argparser.parse_args()

def parse_args_SGM():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--batch_size', type=int, default=128, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample model')
    argparser.add_argument('--model_channels', type=int, default = 64, help='number of features')
    argparser.add_argument('--num_res_blocks', type=int, default = 2, help='number of residual blocks per downsample')
    argparser.add_argument('--attention_resolutions', type=int, nargs='+', default = [4], help='downsample rates at which attention will take place')
    argparser.add_argument('--dropout', type=float, default = 0.0, help='dropout probability of the U-Net')
    argparser.add_argument('--channel_mult', type=int, nargs='+', default = [1, 2, 2], help='channel multiplier for each level of the UNet')
    argparser.add_argument('--conv_resample', type=bool, default = True, help='use learned convolutions for upsampling and downsampling')
    argparser.add_argument('--dims', type=int, default = 2, help='determines if the signal is 1D, 2D, or 3D')
    argparser.add_argument('--num_heads', type=int, default = 4, help='number of attention heads in each attention layer')
    argparser.add_argument('--num_head_channels', type=int, default = 32, help='use a fixed channel width per attention head')
    argparser.add_argument('--use_scale_shift_norm', type=bool, default = False, help='use a FiLM-like conditioning mechanism')
    argparser.add_argument('--resblock_updown', type=bool, default = False, help='use residual blocks for up/downsampling')
    argparser.add_argument('--use_new_attention_order', type=bool, default = False, help='use a different attention pattern for potentially increased efficiency')
    argparser.add_argument('--sampler', type=str, default='ode', help='sampler for SGM', choices=['em', 'pc', 'ode'])
    argparser.add_argument('--snr', type=float, default=0.16, help='signal to noise ratio')
    argparser.add_argument('--solver', type=str, default='euler', help='solver for ODE', choices=['dopri5', 'rk4', 'dopri8', 'euler', 'bosh3', 'adaptive_heun', 'midpoint', 'explicit_adams', 'implicit_adams'])
    argparser.add_argument('--outlier_detection', action='store_true', default=False, help='outlier detection')
    argparser.add_argument('--atol', type=float, default=1e-6, help='absolute tolerance')
    argparser.add_argument('--rtol', type=float, default=1e-6, help='relative tolerance')
    argparser.add_argument('--eps', type=float, default=1e-3, help='smallest timestep for numeric stability')
    argparser.add_argument('--lr', type=float, default=5e-4, help='learning rate')
    argparser.add_argument('--sigma', type=float, default=25.0, help='sigma')
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path')
    argparser.add_argument('--num_samples', type=int, default=16, help='number of samples')
    argparser.add_argument('--num_steps', type=int, default=100, help='number of steps')
    argparser.add_argument('--sample_and_save_freq', type=int, default=10, help='sample and save frequency')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    argparser.add_argument('--ema_rate', type=float, default=0.999, help='ema rate')
    argparser.add_argument('--conditional', action='store_true', default=False, help='conditional')
    argparser.add_argument('--warmup', type=int, default=10, help='warmup epochs')
    argparser.add_argument('--decay', type=float, default=0.0, help='weight decay rate')
    argparser.add_argument('--latent', action='store_true', default=False, help='Use latent implementation')
    argparser.add_argument('--size', type=int, default=None, help='Size of the original image')
    argparser.add_argument('--n_classes', type=int, default=10, help='number of classes')
    argparser.add_argument('--cfg', type=float, default=1.0, help='label guidance')
    argparser.add_argument('--drop_prob', type=float, default=0.1, help='dropout probability')
    args = argparser.parse_args()
    args.channel_mult = tuple(args.channel_mult)
    args.attention_resolutions = tuple(args.attention_resolutions)
    return args

def parse_args_DDPM():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample model')
    argparser.add_argument('--outlier_detection', action='store_true', default=False, help='outlier detection')
    argparser.add_argument('--inpaint', action='store_true', default=False, help='inpainting')
    argparser.add_argument('--batch_size', type=int, default=128, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--lr', type=float, default=1e-3, help='learning rate')
    argparser.add_argument('--timesteps', type=int, default=1000, help='number of timesteps')
    argparser.add_argument('--model_channels', type=int, default = 64, help='number of features')
    argparser.add_argument('--num_res_blocks', type=int, default = 2, help='number of residual blocks per downsample')
    argparser.add_argument('--attention_resolutions', type=int, nargs='+', default = [4], help='downsample rates at which attention will take place')
    argparser.add_argument('--dropout', type=float, default = 0.0, help='dropout probability')
    argparser.add_argument('--channel_mult', type=int, nargs='+', default = [1, 2, 2], help='channel multiplier for each level of the UNet')
    argparser.add_argument('--conv_resample', type=bool, default = True, help='use learned convolutions for upsampling and downsampling')
    argparser.add_argument('--dims', type=int, default = 2, help='determines if the signal is 1D, 2D, or 3D')
    argparser.add_argument('--num_heads', type=int, default = 4, help='number of attention heads in each attention layer')
    argparser.add_argument('--num_head_channels', type=int, default = 32, help='use a fixed channel width per attention head')
    argparser.add_argument('--use_scale_shift_norm', type=bool, default = False, help='use a FiLM-like conditioning mechanism')
    argparser.add_argument('--resblock_updown', type=bool, default = False, help='use residual blocks for up/downsampling')
    argparser.add_argument('--use_new_attention_order', type=bool, default = False, help='use a different attention pattern for potentially increased efficiency')
    argparser.add_argument('--beta_start', type=float, default=0.0001, help='beta start')
    argparser.add_argument('--beta_end', type=float, default=0.02, help='beta end')
    argparser.add_argument('--sample_and_save_freq', type=int, default=5, help='sample and save frequency')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--ddpm', type=float, default=1.0, help='ddim sampling is 0.0, pure ddpm is 1.0')
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path')
    argparser.add_argument('--num_samples', type=int, default=16, help='number of samples')
    argparser.add_argument('--out_dataset', type=str, default='fashionmnist', help='outlier dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet'])
    argparser.add_argument('--loss_type', type=str, default='l2', help='loss type', choices=['huber','l2', 'l1'])
    argparser.add_argument('--sample_timesteps', type=int, default=1000, help='number of timesteps')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    argparser.add_argument('--recon_factor', type=float, default=0.5, help='reconstruction factor')
    argparser.add_argument('--lpips', action='store_true', default=False, help='use lpips loss')
    argparser.add_argument('--warmup', type=int, default=10, help='warmup epochs')
    argparser.add_argument('--decay', type=float, default=1e-7, help='decay rate')
    argparser.add_argument('--latent', action='store_true', default=False, help='Use latent implementation')
    argparser.add_argument('--size', type=int, default=None, help='Size of the original image')
    argparser.add_argument('--ema_rate', type=float, default=0.999, help='ema rate')
    argparser.add_argument('--fid', action='store_true', default=False, help='Sample for FID calculation')
    argparser.add_argument('--snapshots', type=int, default=10, help='how many snapshots during training')
    args = argparser.parse_args()
    args.channel_mult = tuple(args.channel_mult)
    args.attention_resolutions = tuple(args.attention_resolutions)

    return args

def parse_args_CDDPM():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample model')
    argparser.add_argument('--batch_size', type=int, default=128, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--lr', type=float, default=1e-3, help='learning rate')
    argparser.add_argument('--timesteps', type=int, default=500, help='number of timesteps')
    argparser.add_argument('--beta_start', type=float, default=0.0001, help='beta start')
    argparser.add_argument('--beta_end', type=float, default=0.02, help='beta end')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--ddpm', type=float, default=1.0, help='ddim sampling is 0.0, pure ddpm is 1.0')
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path')
    argparser.add_argument('--sample_timesteps', type=int, default=500, help='number of timesteps')
    argparser.add_argument('--model_channels', type=int, default = 64, help='number of features')
    argparser.add_argument('--num_res_blocks', type=int, default = 2, help='number of residual blocks per downsample')
    argparser.add_argument('--attention_resolutions', type=int, nargs='+', default = [4], help='downsample rates at which attention will take place')
    argparser.add_argument('--dropout', type=float, default = 0.0, help='dropout probability of the U-Net')
    argparser.add_argument('--channel_mult', type=int, nargs='+', default = [1, 2, 2], help='channel multiplier for each level of the UNet')
    argparser.add_argument('--conv_resample', type=bool, default = True, help='use learned convolutions for upsampling and downsampling')
    argparser.add_argument('--dims', type=int, default = 2, help='determines if the signal is 1D, 2D, or 3D')
    argparser.add_argument('--num_heads', type=int, default = 4, help='number of attention heads in each attention layer')
    argparser.add_argument('--num_head_channels', type=int, default = 32, help='use a fixed channel width per attention head')
    argparser.add_argument('--use_scale_shift_norm', type=bool, default = False, help='use a FiLM-like conditioning mechanism')
    argparser.add_argument('--resblock_updown', type=bool, default = False, help='use residual blocks for up/downsampling')
    argparser.add_argument('--use_new_attention_order', type=bool, default = False, help='use a different attention pattern for potentially increased efficiency')
    argparser.add_argument('--n_classes', type=int, default=10, help='number of classes')
    argparser.add_argument('--sample_and_save_freq', type=int, default=10, help='sample and save frequency')
    argparser.add_argument('--drop_prob', type=float, default=0.1, help='dropout probability')
    argparser.add_argument('--cfg', type=float, default=1.0, help='guide weight')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    argparser.add_argument('--warmup', type=int, default=10, help='warmup epochs')
    argparser.add_argument('--decay', type=float, default=1e-7, help='decay rate')
    argparser.add_argument('--latent', action='store_true', default=False, help='Use latent implementation')
    argparser.add_argument('--size', type=int, default=None, help='Size of the original image')
    argparser.add_argument('--ema_rate', type=float, default=0.999, help='ema rate')
    argparser.add_argument('--fid', action='store_true', default=False, help='Sample for FID calculation')
    argparser.add_argument('--snapshots', type=int, default=10, help='how many snapshots during training')
    args = argparser.parse_args()
    args.channel_mult = tuple(args.channel_mult)
    args.attention_resolutions = tuple(args.attention_resolutions)
    return args

def parse_args_DiffAE():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--manipulate', action='store_true', default=False, help='manipulate latents')
    argparser.add_argument('--batch_size', type=int, default=16, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--lr', type=float, default=1e-3, help='learning rate')
    argparser.add_argument('--timesteps', type=int, default=1000, help='number of timesteps')
    argparser.add_argument('--sample_timesteps', type=int, default=100, help='number of timesteps')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path')
    argparser.add_argument('--embedding_dim', type=int, default=512, help='embedding dimension')
    argparser.add_argument('--model_channels', type=int, nargs='+', default=[64, 128, 256], help='model channels')
    argparser.add_argument('--attention_levels', type=bool, nargs='+', default=[False, True, True], help='attention levels (must match len of model_channels)')
    argparser.add_argument('--num_res_blocks', type=int, default=1, help='number of res blocks')
    argparser.add_argument('--sample_and_save_freq', type=int, default=10, help='sample and save frequency')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    args = argparser.parse_args()
    args.model_channels = tuple(args.model_channels)
    args.attention_levels = tuple(args.attention_levels)
    return args

def parse_args_CycleGAN():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--test', action='store_true', default=False, help='test model')
    argparser.add_argument('--batch_size', type=int, default=1, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=200, help='number of epochs')
    argparser.add_argument('--lr', type=float, default=0.0002, help='learning rate')
    argparser.add_argument('--decay', type=float, default=100, help='epoch to start linearly decaying the learning rate to 0')
    argparser.add_argument('--sample_and_save_freq', type=int, default=5, help='sample and save frequency')
    argparser.add_argument('--dataset', type=str, default='horse2zebra', help='dataset name', choices=['horse2zebra'])
    argparser.add_argument('--checkpoint_A', type=str, default=None, help='checkpoint A path')
    argparser.add_argument('--checkpoint_B', type=str, default=None, help='checkpoint B path')
    argparser.add_argument('--input_size', type=int, default=128, help='input size')
    argparser.add_argument('--in_channels', type=int, default=3, help='in channels')
    argparser.add_argument('--out_channels', type=int, default=3, help='out channels')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    return argparser.parse_args()

def parse_args_CondGAN():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample from model')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--batch_size', type=int, default=128, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--lr', type=float, default=0.0002, help='learning rate')
    argparser.add_argument('--beta1', type=float, default=0.5, help='beta1')
    argparser.add_argument('--beta2', type=float, default=0.999, help='beta2')
    argparser.add_argument('--latent_dim', type=int, default=100, help='latent dimension')
    argparser.add_argument('--n_classes', type=int, default=10, help='number of classes')
    argparser.add_argument('--img_size', type=int, default=32, help='image size')
    argparser.add_argument('--channels', type=int, default=1, help='channels')
    argparser.add_argument('--sample_and_save_freq', type=int, default=5, help='sample interval')
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path')
    argparser.add_argument('--n_samples', type=int, default=9, help='number of samples')
    argparser.add_argument('--d', type=int, default=128, help='d')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    return argparser.parse_args()

def parse_args_DCGAN():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample from model')
    argparser.add_argument('--outlier_detection', action='store_true', default=False, help='outlier detection')
    argparser.add_argument('--batch_size', type=int, default=128, help='batch size')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--out_dataset', type=str, default='fashionmnist', help='outlier dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet'])
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--lrg', type=float, default=0.0002, help='learning rate generator')
    argparser.add_argument('--lrd', type=float, default=0.0002, help='learning rate discriminator')
    argparser.add_argument('--beta1', type=float, default=0.5, help='beta1')
    argparser.add_argument('--beta2', type=float, default=0.999, help='beta2')
    argparser.add_argument('--latent_dim', type=int, default=100, help='latent dimension')
    argparser.add_argument('--img_size', type=int, default=32, help='image size')
    argparser.add_argument('--channels', type=int, default=1, help='channels')
    argparser.add_argument('--sample_and_save_freq', type=int, default=5, help='sample interval')
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path')
    argparser.add_argument('--discriminator_checkpoint', type=str, default=None, help='discriminator checkpoint path')
    argparser.add_argument('--n_samples', type=int, default=9, help='number of samples')
    argparser.add_argument('--d', type=int, default=128, help='d')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    return argparser.parse_args()

def parse_args_WassersteinGAN():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample from model')
    argparser.add_argument('--outlier_detection', action='store_true', default=False, help='outlier detection')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--out_dataset', type=str, default='fashionmnist', help='outlier dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet'])
    argparser.add_argument('--batch_size', type=int, default=256, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--latent_dim', type=int, default=100, help='latent dimension')
    argparser.add_argument('--d', type=int, default=64, help='d')
    argparser.add_argument('--lrg', type=float, default=0.0002, help='learning rate generator')
    argparser.add_argument('--lrd', type=float, default=0.0002, help='learning rate discriminator')
    argparser.add_argument('--beta1', type=float, default=0.5, help='beta1')
    argparser.add_argument('--beta2', type=float, default=0.999, help='beta2')
    argparser.add_argument('--sample_and_save_freq', type=int, default=5, help='sample interval')
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path')
    argparser.add_argument('--discriminator_checkpoint', type=str, default=None, help='discriminator checkpoint path')
    argparser.add_argument('--gp_weight', type=float, default=10.0, help='gradient penalty weight')
    argparser.add_argument('--n_critic', type=int, default=5, help='number of critic updates per generator update')
    argparser.add_argument('--n_samples', type=int, default=9, help='number of samples')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')

    return argparser.parse_args()

def parse_args_PresGAN():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample from model')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    ###### Model arguments
    argparser.add_argument('--nz', type=int, default=100, help='size of the latent z vector')
    argparser.add_argument('--ngf', type=int, default=64)
    argparser.add_argument('--ndf', type=int, default=64)

    ###### Optimization arguments
    argparser.add_argument('--batch_size', type=int, default=64, help='input batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs to train for')
    argparser.add_argument('--lrD', type=float, default=0.0002, help='learning rate, default=0.0002')
    argparser.add_argument('--lrG', type=float, default=0.0002, help='learning rate, default=0.0002')
    argparser.add_argument('--lrE', type=float, default=0.0002, help='learning rate, default=0.0002')
    argparser.add_argument('--beta1', type=float, default=0.5, help='beta1 for adam')

    ###### Checkpointing and Logging arguments
    argparser.add_argument('--checkpoint', type=str, default=None, help='a given checkpoint file for generator')
    argparser.add_argument('--discriminator_checkpoint', type=str, default=None, help='a given checkpoint file for discriminator')
    argparser.add_argument('--sigma_checkpoint', type=str, default=None, help='a given file for logsigma for the generator')
    argparser.add_argument('--num_gen_images', type=int, default=16, help='number of images to generate for inspection')

    ###### PresGAN-specific arguments
    argparser.add_argument('--sigma_lr', type=float, default=0.0002, help='generator variance')
    argparser.add_argument('--lambda_', type=float, default=0.01, help='entropy coefficient')
    argparser.add_argument('--sigma_min', type=float, default=0.01, help='min value for sigma')
    argparser.add_argument('--sigma_max', type=float, default=0.3, help='max value for sigma')
    argparser.add_argument('--logsigma_init', type=float, default=-1.0, help='initial value for log_sigma_sian')
    argparser.add_argument('--num_samples_posterior', type=int, default=2, help='number of samples from posterior')
    argparser.add_argument('--burn_in', type=int, default=2, help='hmc burn in')
    argparser.add_argument('--leapfrog_steps', type=int, default=5, help='number of leap frog steps for hmc')
    argparser.add_argument('--flag_adapt', type=int, default=1, help='0 or 1')
    argparser.add_argument('--delta', type=float, default=1.0, help='delta for hmc')
    argparser.add_argument('--hmc_learning_rate', type=float, default=0.02, help='lr for hmc')
    argparser.add_argument('--hmc_opt_accept', type=float, default=0.67, help='hmc optimal acceptance rate')
    argparser.add_argument('--stepsize_num', type=float, default=1.0, help='initial value for hmc stepsize')
    argparser.add_argument('--restrict_sigma', type=int, default=0, help='whether to restrict sigma or not')
    argparser.add_argument('--sample_and_save_freq', type=int, default=5, help='sample and save frequency')
    argparser.add_argument('--outlier_detection', action='store_true', default=False, help='outlier detection')
    argparser.add_argument('--out_dataset', type=str, default='fashionmnist', help='outlier dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet'])
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    return argparser.parse_args()

def parse_args_Glow():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample from model')
    argparser.add_argument('--outlier_detection', action='store_true', default=False, help='outlier detection')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--out_dataset', type=str, default='fashionmnist', help='outlier dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet'])
    argparser.add_argument('--batch_size', type=int, default=128, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--lr', type=float, default=0.0002, help='learning rate')
    argparser.add_argument('--hidden_channels', type=int, default=64, help='hidden channels')
    argparser.add_argument('--K', type=int, default=8, help='Number of layers per block')
    argparser.add_argument('--L', type=int, default=3, help='number of blocks')
    argparser.add_argument('--actnorm_scale', type=float, default=1.0, help='act norm scale')
    argparser.add_argument('--flow_permutation', type=str, default='invconv', help='flow permutation', choices=['invconv', 'shuffle', 'reverse'])
    argparser.add_argument('--flow_coupling', type=str, default='affine', help='flow coupling, affine ', choices=['additive', 'affine'])
    argparser.add_argument('--LU_decomposed', action='store_true', default=False, help='Train with LU decomposed 1x1 convs')
    argparser.add_argument('--learn_top', action='store_true', default=False, help='learn top layer (prior)')
    argparser.add_argument('--y_condition', action='store_true', default=False, help='Class Conditioned Glow')
    argparser.add_argument('--y_weight', type=float, default=0.01, help='weight of class condition')
    argparser.add_argument('--num_classes', type=int, default=10, help='number of classes')
    argparser.add_argument('--sample_and_save_freq', type=int, default=5, help='sample and save frequency')
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path')
    argparser.add_argument('--n_bits', type=int, default=8, help='number of bits')
    argparser.add_argument('--max_grad_clip', type=float, default=0.0, help='max grad clip')
    argparser.add_argument('--max_grad_norm', type=float, default=0.0, help='max grad norm')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    argparser.add_argument('--warmup', type=int, default=10, help='warmup epochs')
    argparser.add_argument('--decay', type=float, default=1e-5, help='decay rate')
    return argparser.parse_args()

def parse_args_NCSNv2():
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--train', action='store_true', default=False, help='train model')
    argparser.add_argument('--sample', action='store_true', default=False, help='sample from model')
    argparser.add_argument('--dataset', type=str, default='mnist', help='dataset name', choices=['mnist', 'cifar10', 'cifar100', 'places365', 'dtd', 'fashionmnist', 'chestmnist', 'bloodmnist', 'dermamnist', 'dermamnist', 'octmnist', 'tissuemnist', 'pneumoniamnist', 'retinamnist', 'svhn', 'tinyimagenet','imagenet', 'celeba'])
    argparser.add_argument('--batch_size', type=int, default=128, help='batch size')
    argparser.add_argument('--n_epochs', type=int, default=100, help='number of epochs')
    argparser.add_argument('--lr', type=float, default=0.0002, help='learning rate')
    argparser.add_argument('--nf', type=int, default=128, help='number of filters')
    argparser.add_argument('--act', type=str, default='elu', help='activation', choices=['relu', 'elu', 'swish'])
    argparser.add_argument('--centered', action='store_true', default=False, help='centered')
    argparser.add_argument('--sigma_min', type=float, default=0.01, help='min value for sigma')
    argparser.add_argument('--sigma_max', type=float, default=50, help='max value for sigma')
    argparser.add_argument('--num_scales', type=int, default=232, help='number of scales')
    argparser.add_argument('--normalization', type=str, default='InstanceNorm++', help='Normalization', choices=['InstanceNorm', 'GroupNorm', 'VarianceNorm', 'InstanceNorm++'])
    argparser.add_argument('--num_classes', type=int, default=10, help='number of classes')
    argparser.add_argument('--ema_decay', type=float, default=0.999, help='ema decay')
    argparser.add_argument('--continuous', action='store_true', default=False, help='continuous')
    argparser.add_argument('--reduce_mean', action='store_true', default=False, help='reduce mean')
    argparser.add_argument('--likelihood_weighting', action='store_true', default=False, help='likelihood weighting')
    argparser.add_argument('--beta1', type=float, default=0.9, help='beta1')
    argparser.add_argument('--beta2', type=float, default=0.999, help='beta2')
    argparser.add_argument('--weight_decay', type=float, default=0.0, help='weight decay')
    argparser.add_argument('--warmup', type=int, default=0, help='warmup')
    argparser.add_argument('--grad_clip', type=float, default=-1.0, help='grad clip')
    argparser.add_argument('--sample_and_save_freq', type=int, default=5, help='sample and save frequency')
    argparser.add_argument('--sampler', type=str, default='pc', help='sampler name', choices=['pc', 'ode'])
    argparser.add_argument('--predictor', type=str, default='none', help='predictor', choices=['none', 'em', 'rd', 'as'])
    argparser.add_argument('--corrector', type=str, default='ald', help='corrector', choices=['none', 'l', 'ald'])
    argparser.add_argument('--snr', type=float, default=0.176, help='signal to noise ratio')
    argparser.add_argument('--n_steps', type=int, default=5, help='number of steps')
    argparser.add_argument('--probability_flow', action='store_true', default=False, help='probability flow')
    argparser.add_argument('--noise_removal', action='store_true', default=False, help='noise removal')
    argparser.add_argument('--checkpoint', type=str, default=None, help='checkpoint path')
    argparser.add_argument('--no_wandb', action='store_true', default=False, help='disable wandb logging')
    argparser.add_argument('--num_workers', type=int, default=0, help='number of workers for dataloader')
    return argparser.parse_args()

# EOF