Colab-traiNNer.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZY_bDEi1J7bJ"
      },
      "source": [
        "# Colab-traiNNer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "My897iQFpTHH"
      },
      "source": [
        "victorca25's BasicSR fork: [victorca25/traiNNer](https://github.com/victorca25/traiNNer)\n",
        "\n",
        "My fork: [styler00dollar/Colab-traiNNer](https://github.com/styler00dollar/Colab-traiNNer)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "_MrEgrBSu4Qu",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "#@title Check GPU\n",
        "\n",
        "gpu = !nvidia-smi --query-gpu=gpu_name --format=csv\n",
        "print(\"GPU: \" + gpu[1])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zBbSxWWBOZEI",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IsK_NAi963qK"
      },
      "source": [
        "# Install"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "l32IresQs-oW",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "#@title Install dependencies\n",
        "%cd /content/\n",
        "\n",
        "# create empty folders\n",
        "!mkdir /content/hr\n",
        "!mkdir /content/lr\n",
        "!mkdir /content/val_hr\n",
        "!mkdir /content/val_lr\n",
        "\n",
        "!mkdir /content/masks\n",
        "!mkdir /content/validation\n",
        "!mkdir /content/data\n",
        "!mkdir /content/logs/\n",
        "\n",
        "!git clone https://github.com/styler00dollar/Colab-traiNNer\n",
        "!pip install lion-pytorch pytorch-lightning==2.0.4 \\\n",
        "    git+https://github.com/vballoli/nfnets-pytorch \\\n",
        "    git+https://github.com/styler00dollar/BasicSR albumentations \\\n",
        "    IPython scipy pandas opencv-python pillow wget \\\n",
        "    tfrecord x-transformers adamp efficientnet_pytorch \\\n",
        "    tensorboardX vit-pytorch swin-transformer-pytorch madgrad \\\n",
        "    git+https://github.com/huggingface/pytorch-image-models pillow-avif-plugin \\\n",
        "    kornia omegaconf git+https://github.com/styler00dollar/pytorch_optimizer \\\n",
        "    git+https://github.com/huggingface/transformers gdown PyTurboJPEG wavemix pyiqa\n",
        "\n",
        "!pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121 --force-reinstall"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-mUlJnLCj9t5"
      },
      "source": [
        "## Optional"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "fipjH_ZAg_tw",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "#@title (optional) compiling and installing mmcv (for GLEAN)\n",
        "!pip install torch torchvision torchaudio -U\n",
        "!pip uninstall mmcv -y\n",
        "!pip uninstall mmcv-full -y\n",
        "!pip install mmcv-full"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "riIB6Ev5hBKM",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "#@title (optional) ninja (for GFPGAN / GPEN / co-mod-gan)\n",
        "%cd /content\n",
        "!wget https://github.com/ninja-build/ninja/releases/download/v1.12.0/ninja-linux.zip\n",
        "!sudo unzip ninja-linux.zip -d /usr/local/bin/\n",
        "!sudo update-alternatives --install /usr/bin/ninja ninja /usr/local/bin/ninja 1 --force"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "5W23p3k65Yvd",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "#@title (optional) correlation package (for ABME)\n",
        "%cd /content/\n",
        "!sudo rm -rf ABME\n",
        "!git clone https://github.com/JunHeum/ABME\n",
        "%cd /content/ABME/correlation_package\n",
        "#!python setup.py install\n",
        "!python setup.py build install"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "oVlAejWyj3SP",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "#@title (optional) install cupy (for EDSC)\n",
        "!curl https://colab.chainer.org/install | sh -"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "G7TlT3ioHdWK",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "#@title (optional) Adam8bit optimizer\n",
        "!pip install bitsandbytes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "3y7tvaaxBv0z",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "#@title (Optional) Install turboJPEG\n",
        "%cd /content/\n",
        "!sudo apt-get install -y libturbojpeg\n",
        "!pip install PyTurboJPEG"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mu7HxxdGJdir"
      },
      "source": [
        "# Make Directories and Copy Data\n",
        "You need to upload the data and then extract it within colab. You can use Google Drive for that."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QGU28ZJIJlEC",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "%cd /content/\n",
        "!cp \"/content/drive/MyDrive/dataset.tar\" \"/content/data.tar\"\n",
        "!7z x /content/data.tar"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W7p3pQML4S_T",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "# if dataset is not big, copy it into ram instead\n",
        "%cd /content/\n",
        "!mkdir /content/ramdisk\n",
        "!sudo mount -t tmpfs none /content/ramdisk\n",
        "!sudo chmod 777 /content/ramdisk/\n",
        "\n",
        "%cd /content/ramdisk/\n",
        "!cp \"/content/drive/MyDrive/dataset.tar\" \"/content/ramdisk/data.tar\"\n",
        "!7z x \"/content/ramdisk/data.tar\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dG6BqJD5J41q"
      },
      "source": [
        "# Train"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "6_2VP2MV0eGv",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "#@title config.yaml\n",
        "%%writefile /content/Colab-traiNNer/code/config.yaml\n",
        "name: template\n",
        "scale: 2\n",
        "gpus: 1 # amount of gpus, 0 = cpu\n",
        "strategy: ddp # dp, ddp, ddp_spawn, ddp2, horovod, bagua\n",
        "tpu_cores: 8 # 8 if you use a Google Colab TPU\n",
        "use_tpu: False\n",
        "use_precision: fp32 # bf16 | fp16 | fp32\n",
        "use_swa: False\n",
        "progress_bar_refresh_rate: 20\n",
        "default_root_dir: '/content/'\n",
        "logging: True # colab easily crashes with logging, disable in colab\n",
        "\n",
        "# Dataset options:\n",
        "datasets:\n",
        "  train:\n",
        "    # DS_inpaint: hr is from dataroot_HR, loads masks\n",
        "    # DS_lrhr: loads lr from dataroot_LR and hr from dataroot_HR\n",
        "    # DS_video: video dataloader which has 3 frames as input (look into data/data_video.py for more details)\n",
        "    # DS_inpaint_TF: takes one tfrecord file as dataset input, but the validation is still just green masked images like in DS_inpaint\n",
        "    # DS_video_direct: direcly copy .npy files into GPU and avoiding CPU processing (upgrade to newest nvidia drivers and cuda, linux only)\n",
        "    # only works with n_workers = 0, use pipeline_threads instead\n",
        "    # DS_realesrgan: will use the realesrgan dataloader (only uses hr folder)\n",
        "    # pip install --extra-index-url https://developer.download.nvidia.com/compute/redist --upgrade nvidia-dali-cuda110\n",
        "\n",
        "    mode: DS_realesrgan # DS_video | DS_video_direct | DS_inpaint_TF | DS_inpaint  | DS_lrhr | DS_realesrgan\n",
        "    amount_files: 7 # tfrecord files do not store amount of images and are infinite, specify the images inside of it\n",
        "\n",
        "    tfrecord_path: \"/content/tfrecord/tfrecord-r09.tfrecords\"\n",
        "    dataroot_HR: '/content/hr' # Original, with a single directory. Inpainting will use this directory as source image.\n",
        "    dataroot_LR: '/content/lr' # Original, with a single directory\n",
        "    loading_backend: 'OpenCV' # 'PIL' | 'OpenCV' | 'turboJPEG' # install needed for turboJPEG, turboJPEG only for DS_video, 'PIL' for DS_inpaint_TF\n",
        "\n",
        "    n_workers: 16 # 0 to disable CPU multithreading, or an integrer representing CPU threads to use for dataloading\n",
        "    pipeline_threads: 16 # only for DS_video_direct\n",
        "    batch_size: 1 # 6\n",
        "\n",
        "    # does not apply to video dataloaders, look into python file instead\n",
        "    HR_size: 128 # The resolution the network will get. Random crop gets applied if that resolution does not match.\n",
        "    image_channels: 3 # number of channels to load images in\n",
        "\n",
        "    masks: '/content/masks/' # only for inpainting\n",
        "    mask_invert_ratio: 0.3 # 0.3 = 30% of masks will be inverted\n",
        "    max_epochs: 2000\n",
        "    save_step_frequency: 5000 # also validation frequency\n",
        "\n",
        "    # if edge data is required\n",
        "    canny_min: 100\n",
        "    canny_max: 150\n",
        "\n",
        "    # OTF downscaling\n",
        "    # This will downscale the HR image with a randomly chosen filter and ignore the LR folder.\n",
        "    # otf_filter_probs defines the cumulative (additive) weights used to select a downscaling filter\n",
        "    # KERNEL will randomly apply one of the kernels generated using\n",
        "    # https://github.com/victorca25/DLIP/tree/main/kgan\n",
        "    apply_otf_downscale: False\n",
        "    otf_filter_types: ['KERNEL', 'NEAREST', 'BILINEAR', 'AREA', 'BICUBIC', 'LANCZOS']\n",
        "    otf_filter_probs: [0.25, 0.4, 0.55, 0.70, 0.85, 1.0]\n",
        "    kernel_path: '/content/kernels'\n",
        "\n",
        "    # Image augmentations (only for lrhr dataloader). Set 'True' to use.\n",
        "    # To customize individual augmentations, edit aug_config.yaml.\n",
        "    # Augmentations will apply in random order.\n",
        "    ColorJitter: False\n",
        "    RandomGaussianNoise: False\n",
        "    RandomPoissonNoise: False\n",
        "    RandomSPNoise: False\n",
        "    RandomSpeckleNoise: False\n",
        "    RandomCompression: False\n",
        "    RandomAverageBlur: False\n",
        "    RandomBilateralBlur: False\n",
        "    RandomBoxBlur: False\n",
        "    RandomGaussianBlur: False\n",
        "    RandomMedianBlur: False\n",
        "    RandomMotionBlur: False\n",
        "    RandomComplexMotionBlur: False\n",
        "    RandomAnIsoBlur: False\n",
        "    RandomSincBlur: False\n",
        "    BayerDitherNoise: False\n",
        "    FSDitherNoise: False\n",
        "    FilterMaxRGB: False\n",
        "    FilterColorBalance: False\n",
        "    FilterUnsharp: False\n",
        "    FilterCanny: False\n",
        "    SimpleQuantize: False\n",
        "    KMeansQuantize: False\n",
        "    CLAHE: False\n",
        "    RandomGamma: False\n",
        "    Superpixels: False\n",
        "    RandomCameraNoise: False\n",
        "\n",
        "  val:\n",
        "    loading_backend: 'OpenCV' # 'OpenCV' | 'turboJPEG' # install needed for turboJPEG, currently only for DS_video\n",
        "    dataroot_HR: '/content/val_hr/'\n",
        "    dataroot_LR: '/content/val_lr/' # Inpainting will use this directory as input\n",
        "\n",
        "path:\n",
        "    pretrain_model_G: #\"/content/model.pth\"\n",
        "    pretrain_model_G_teacher:\n",
        "    pretrain_model_D:\n",
        "    checkpoint_path:\n",
        "    checkpoint_save_path: '/content/'\n",
        "    validation_output_path: '/content/validation'\n",
        "    log_path: '/content/logs'\n",
        "\n",
        "# using a teacher model to generate hr, has same configuration as normal netG\n",
        "# if teacher is used, same loss functions will be applied to both images\n",
        "# if you don't use a teacher, set netG to None\n",
        "network_G_teacher:\n",
        "    # CEM (for esrgan, not 1x)\n",
        "    CEM: False # uses hardcoded torch.cuda.FloatTensor\n",
        "\n",
        "    # comparing feature maps of teacher and student with l1\n",
        "    l1_feature_maps_weight: 1\n",
        "\n",
        "    netG:\n",
        "\n",
        "    #netG: MRRDBNet_FM # RRDB_net (original ESRGAN arch) | MRRDBNet_FM (modified/\"new\" arch) with feature map knowledge distillation\n",
        "    #norm_type: null\n",
        "    #mode: CNA\n",
        "    #nf: 64 # of discrim filters in the first conv layer (default: 64, good: 32)\n",
        "    #nb: 3 # (default: 23, good: 8)\n",
        "    #in_nc: 3 # of input image channels: 3 for RGB and 1 for grayscale\n",
        "    #out_nc: 3 # of output image channels: 3 for RGB and 1 for grayscale\n",
        "    #gc: 32\n",
        "    #convtype: Conv2D # Conv2D | PartialConv2D | doconv | gated | TBC | dynamic | MBConv | CondConv | fft | WSConv\n",
        "    #nr: 3\n",
        "    ## for dynamic\n",
        "    #nof_kernels: 4\n",
        "    #reduce: 4\n",
        "    ## RRDB_net\n",
        "    #net_act: leakyrelu # swish | leakyrelu\n",
        "    #gaussian: false # true | false # esrgan plus, does not work on TPU because of cuda()\n",
        "    #plus: false # true | false\n",
        "    #finalact: None #tanh # Test. Activation function to make outputs fit in [-1, 1] range. Default = None. Coordinate with znorm.\n",
        "    #upsample_mode: 'upconv'\n",
        "    #strided_conv: False\n",
        "    ##group: 1 # unused for now\n",
        "\n",
        "    # RealESRGAN Anime Compact model (2021)\n",
        "    #netG: SRVGGNetCompact\n",
        "    #num_in_ch: 3\n",
        "    #num_out_ch: 3\n",
        "    #num_feat: 64\n",
        "    #num_conv: 16\n",
        "    #act_type: 'prelu'\n",
        "    #conv_mode: 3 # 2 | 3\n",
        "    # only for rrdb\n",
        "    #rrdb: False\n",
        "    #rrdb_blocks: 2\n",
        "    #convtype: \"Conv2D\"\n",
        "\n",
        "# Generator options:\n",
        "network_G:\n",
        "    # CEM (for esrgan, not 1x)\n",
        "    CEM: False # uses hardcoded torch.cuda.FloatTensor\n",
        "    sigmoid_range_limit: False\n",
        "\n",
        "    finetune: False # Important for further rfr/dsnet training. Apply that after training for a while. https://github.com/jingyuanli001/RFR-Inpainting/issues/33\n",
        "\n",
        "    # ESRGAN:\n",
        "    #netG: RRDB_net # RRDB_net (original ESRGAN arch) | MRRDBNet_FM (modified/\"new\" arch) with feature map knowledge distillation\n",
        "    #norm_type: null\n",
        "    #mode: CNA\n",
        "    #nf: 64 # of discrim filters in the first conv layer (default: 64, good: 32)\n",
        "    #nb: 3 # (default: 23, good: 8)\n",
        "    #in_nc: 3 # of input image channels: 3 for RGB and 1 for grayscale\n",
        "    #out_nc: 3 # of output image channels: 3 for RGB and 1 for grayscale\n",
        "    #gc: 32\n",
        "    #convtype: Conv2D # Conv2D | PartialConv2D | doconv | gated | TBC | dynamic | MBConv | CondConv | fft | WSConv\n",
        "    #nr: 3\n",
        "    # for dynamic\n",
        "    #nof_kernels: 4\n",
        "    #reduce: 4\n",
        "    # RRDB_net\n",
        "    #net_act: leakyrelu # swish | leakyrelu\n",
        "    #gaussian: false # true | false # esrgan plus, does not work on TPU because of cuda()\n",
        "    #plus: false # true | false\n",
        "    #finalact: None #tanh # Test. Activation function to make outputs fit in [-1, 1] range. Default = None. Coordinate with znorm.\n",
        "    #upsample_mode: 'upconv'\n",
        "    #strided_conv: False\n",
        "    #group: 1 # unused for now\n",
        "\n",
        "    # CUGAN\n",
        "    # netG: cugan\n",
        "    #in_channels: 3\n",
        "    #out_channels: 3\n",
        "    #pro_mode: True\n",
        "\n",
        "    # OmniSR (2023)\n",
        "    #netG: OmniSR\n",
        "    #num_in_ch: 3\n",
        "    #num_out_ch: 3\n",
        "    #num_feat: 64\n",
        "    #window_size: 8\n",
        "    #res_num: 5\n",
        "    #bias: True\n",
        "    #block_num: 4\n",
        "    #pe: True\n",
        "    #ffn_bias: True\n",
        "\n",
        "    # EMT (2023)\n",
        "    #netG: EMT\n",
        "    #dim: 60\n",
        "    #n_blocks: 6\n",
        "    #n_layers: 6\n",
        "    #num_heads: 3\n",
        "    #mlp_ratio: 2\n",
        "    #n_GTLs: 2\n",
        "    #window_list: [[32, 8], [8, 32]]\n",
        "    #shift_list: [[16, 4], [4, 16]]\n",
        "    #task: \"lsr\"\n",
        "\n",
        "    # lkdn (2023)\n",
        "    #netG: lkdn\n",
        "    #num_in_ch: 3\n",
        "    #num_out_ch: 3\n",
        "    #num_feat: 56\n",
        "    #num_atten: 56\n",
        "    #num_block: 8\n",
        "    #upscale: 2\n",
        "    #num_in: 4\n",
        "    #conv: \"BSConvU\"\n",
        "    #upsampler: \"pixelshuffledirect\"\n",
        "\n",
        "    # wavemix (2023) (only 2x)\n",
        "    #netG: wavemix\n",
        "    #depth: 4\n",
        "    #mult: 1\n",
        "    #ff_channel: 144\n",
        "    #final_dim: 144\n",
        "    #dropout: 0.3\n",
        "\n",
        "    # DITN (2023)\n",
        "    #netG: DITN\n",
        "    #inp_channels: 3\n",
        "    #dim: 60\n",
        "    #ITL_blocks: 4\n",
        "    #SAL_blocks: 4\n",
        "    #UFONE_blocks: 1\n",
        "    #ffn_expansion_factor: 2\n",
        "    #bias:  False\n",
        "    #LayerNorm_type: 'WithBias'\n",
        "    #patch_size: 8\n",
        "\n",
        "    # dat (2023)\n",
        "    #netG: dat\n",
        "    #img_size: 64\n",
        "    #in_chans: 3\n",
        "    #embed_dim: 180\n",
        "    #split_size: [2,4]\n",
        "    #depth: [2,2,2,2]\n",
        "    #num_heads: [2,2,2,2]\n",
        "    #expansion_factor: 4.\n",
        "    #qkv_bias: True\n",
        "    #qk_scale:\n",
        "    #drop_rate: 0.\n",
        "    #attn_drop_rate: 0.\n",
        "    #drop_path_rate: 0.1\n",
        "    #use_chk: False\n",
        "    #upscale: 2,\n",
        "    #img_range: 1.\n",
        "    #resi_connection: '1conv'\n",
        "    #upsampler: 'pixelshuffle'\n",
        "\n",
        "    # grl (2023)\n",
        "    #netG: grl\n",
        "    #img_size: 64\n",
        "    #window_size: 16\n",
        "    #depths: [4, 4, 8, 8, 8, 4, 4]\n",
        "    #embed_dim: 180\n",
        "    #num_heads_window: [3, 3, 3, 3, 3, 3, 3]\n",
        "    #num_heads_stripe: [3, 3, 3, 3, 3, 3, 3]\n",
        "    #mlp_ratio: 2\n",
        "    #qkv_proj_type: \"linear\"\n",
        "    #anchor_proj_type: \"avgpool\"\n",
        "    #anchor_window_down_factor: 2\n",
        "    #out_proj_type: \"linear\"\n",
        "    #conv_type: \"1conv\"\n",
        "    #upsampler: \"pixelshuffle\"\n",
        "    #local_connection: True\n",
        "\n",
        "    # craft (2023)\n",
        "    #netG: craft\n",
        "    #in_chans: 3\n",
        "    #embed_dim: 96\n",
        "    #depths: [6, 6, 6, 6]\n",
        "    #num_heads: [6, 6, 6, 6]\n",
        "    #split_size_0: 4\n",
        "    #split_size_1: 16\n",
        "    #mlp_ratio: 2.0\n",
        "    #qkv_bias: True\n",
        "    #qk_scale:\n",
        "    #img_range: 1.0\n",
        "    #upsampler: \"\"\n",
        "    #resi_connection: \"1conv\"\n",
        "\n",
        "    # srformer (2023)\n",
        "    #netG: srformer\n",
        "    #img_size: 48\n",
        "    #patch_size: 1\n",
        "    #in_chans: 3\n",
        "    #embed_dim: 180\n",
        "    #depths: [6, 6, 6, 6, 6, 6]\n",
        "    #num_heads: [6, 6, 6, 6, 6, 6]\n",
        "    #window_size: 24\n",
        "    #mlp_ratio: 4.0\n",
        "    #qkv_bias: True\n",
        "    #qk_scale:\n",
        "    #drop_rate: 0.0\n",
        "    #attn_drop_rate: 0.0\n",
        "    #drop_path_rate: 0.1\n",
        "    #ape: False\n",
        "    #patch_norm: True\n",
        "    #use_checkpoint: False\n",
        "    #upscale: 2\n",
        "    #img_range: 1.0\n",
        "    #upsampler: \"\"\n",
        "    #resi_connection: \"1conv\"\n",
        "\n",
        "    # DCTLSA (2023)\n",
        "    #netG: DCTLSA\n",
        "    #in_nc: 3\n",
        "    #nf: 55\n",
        "    #num_modules: 6\n",
        "    #out_nc: 3\n",
        "    #num_head: 5\n",
        "\n",
        "    # SAFMN (2023)\n",
        "    #netG: SAFMN\n",
        "    #dim: 192\n",
        "    #n_blocks: 20\n",
        "    #ffn_scale: 2.0\n",
        "\n",
        "    # MFRAN (2023)\n",
        "    #netG: MFRAN\n",
        "    #n_feats: 64\n",
        "    #n_blocks: 16\n",
        "    #kernel_size_MFRAN: 3\n",
        "    #scale: 2\n",
        "    #n_blocks: 16\n",
        "    #div: 2\n",
        "    #rgb_range: 1\n",
        "    #n_colors: 3\n",
        "    #path: 4\n",
        "\n",
        "    # ASRGAN:\n",
        "    #which_model_G: asr_resnet # asr_resnet | asr_cnn\n",
        "    #nf: 64\n",
        "\n",
        "    # PPON:\n",
        "    #netG: ppon # | ppon\n",
        "    ##norm_type: null\n",
        "    #mode: CNA\n",
        "    #nf: 64\n",
        "    #nb: 24\n",
        "    #in_nc: 3\n",
        "    #out_nc: 3\n",
        "    ##gc: 32\n",
        "    #group: 1\n",
        "    ##convtype: Conv2D #Conv2D | PartialConv2D\n",
        "\n",
        "    # SRGAN:\n",
        "    #netG: sr_resnet # RRDB_net | sr_resnet\n",
        "    #norm_type: null\n",
        "    #mode: CNA\n",
        "    #nf: 64\n",
        "    #nb: 16\n",
        "    #in_nc: 3\n",
        "    #out_nc: 3\n",
        "\n",
        "    # SR:\n",
        "    #netG: RRDB_net # RRDB_net | sr_resnet\n",
        "    #norm_type: null\n",
        "    #mode: CNA\n",
        "    #nf: 64\n",
        "    #nb: 23\n",
        "    #in_nc: 3\n",
        "    #out_nc: 3\n",
        "    #gc: 32\n",
        "    #group: 1\n",
        "\n",
        "    # PAN:\n",
        "    # netG: pan_net\n",
        "    # in_nc: 3\n",
        "    # out_nc: 3\n",
        "    # nf: 40\n",
        "    # unf: 24\n",
        "    # nb: 16\n",
        "    # self_attention: true\n",
        "    # double_scpa: false\n",
        "\n",
        "    # edge-informed-sisr\n",
        "    #which_model_G: sisr\n",
        "    #use_spectral_norm: True\n",
        "\n",
        "    # USRNet\n",
        "    #netG: USRNet\n",
        "    #in_nc=4\n",
        "    #out_nc=3\n",
        "    #nc=[64, 128, 256, 512]\n",
        "    #nb=2\n",
        "    #act_mode='R'\n",
        "    #downsample_mode='strideconv'\n",
        "    #upsample_mode='convtranspose'\n",
        "\n",
        "    # GLEAN (2021)\n",
        "    # Warning: Does require \"pip install mmcv-full\"\n",
        "    #netG: GLEAN\n",
        "    #in_size: 512\n",
        "    #out_size: 512\n",
        "    #img_channels: 4\n",
        "    #img_channels_out: 3\n",
        "    #rrdb_channels: 16 # 64\n",
        "    #num_rrdbs: 8 # 23\n",
        "    #style_channels: 512 # 512\n",
        "    #num_mlps: 4 # 8\n",
        "    #channel_multiplier: 2\n",
        "    #blur_kernel: [1, 3, 3, 1]\n",
        "    #lr_mlp: 0.01\n",
        "    #default_style_mode: 'mix'\n",
        "    #eval_style_mode: 'single'\n",
        "    #mix_prob: 0.9\n",
        "    #pretrained: False # only works with official settings\n",
        "    #bgr2rgb: False\n",
        "\n",
        "    # srflow (upscaling factors: 4, 8, 16)\n",
        "    # Warning: Can be very unstable with batch_size 1, use higher batch_size\n",
        "    #netG: srflow\n",
        "    #in_nc: 3\n",
        "    #out_nc: 3\n",
        "    #nf: 64\n",
        "    #nb: 23\n",
        "    #train_RRDB: false\n",
        "    #train_RRDB_delay: 0.5\n",
        "    #flow:\n",
        "    #  K: 16\n",
        "    #  L: 3\n",
        "    #  noInitialInj: true\n",
        "    #  coupling: CondAffineSeparatedAndCond\n",
        "    #  additionalFlowNoAffine: 2\n",
        "    #  split:\n",
        "    #    enable: true\n",
        "    #  fea_up0: true\n",
        "    #  stackRRDB:\n",
        "    #    blocks: [ 1, 8, 15, 22 ]\n",
        "    #    concat: true\n",
        "    #nll_weight: 1\n",
        "    #freeze_iter: 100000\n",
        "\n",
        "    # DFDNet\n",
        "    # Warning: Expects \"DictionaryCenter512\" in the current folder, you can get the data here: https://drive.google.com/drive/folders/1bayYIUMCSGmoFPyd4Uu2Uwn347RW-vl5\n",
        "    # Also wants a folder called \"landmarks\", you can generate that data yourself. Example: https://github.com/styler00dollar/Colab-DFDNet/blob/local/Colab-DFDNet-lightning-train.ipynb\n",
        "    # Hardcoded resolution: 512px\n",
        "    #netG: DFDNet\n",
        "    #dictionary_path: \"/content/DictionaryCenter512\"\n",
        "    #landmarkpath: \"/content/landmarks\"\n",
        "    #val_landmarkpath: \"/content/landmarks\"\n",
        "\n",
        "    # GFPGAN (2021) [EXPERIMENTAL]\n",
        "    # does require ninja\n",
        "    # because it compiles files, the startup time is quite long\n",
        "    #netG: GFPGAN\n",
        "    #input_channels: 4\n",
        "    #output_channels: 3\n",
        "    #out_size: 512\n",
        "    #num_style_feat: 512\n",
        "    #channel_multiplier: 1\n",
        "    #resample_kernel: [1, 3, 3, 1]\n",
        "    #decoder_load_path: # None\n",
        "    #fix_decoder: True\n",
        "    #num_mlp: 8\n",
        "    #lr_mlp: 0.01\n",
        "    #input_is_latent: False\n",
        "    #different_w: False\n",
        "    #narrow: 1\n",
        "    #sft_half: False\n",
        "\n",
        "    # GPEN\n",
        "    # does require ninja\n",
        "    # because it compiles files, the startup time is quite long\n",
        "    # output_channels is hardcoded to 3\n",
        "    #netG: GPEN\n",
        "    #input_channels: 4\n",
        "    #size: 512\n",
        "    #style_dim: 512\n",
        "    #n_mlp: 8\n",
        "    #channel_multiplier: 2\n",
        "    #blur_kernel: [1, 3, 3, 1]\n",
        "    #lr_mlp: 0.01\n",
        "    #pooling: True # Experimental, to have any input size\n",
        "\n",
        "    # comodgan (2021)\n",
        "    # needs ninja\n",
        "    # because it compiles files, the startup time is quite long\n",
        "    #netG: comodgan\n",
        "    #dlatent_size: 512\n",
        "    #num_channels: 3 # amount of channels without mask\n",
        "    #resolution: 512\n",
        "    #fmap_base: 16384 # 16 << 10\n",
        "    #fmap_decay: 1.0\n",
        "    #fmap_min: 1\n",
        "    #fmap_max: 512\n",
        "    #randomize_noise: True\n",
        "    #architecture: 'skip'\n",
        "    #nonlinearity: 'lrelu'\n",
        "    #resample_kernel: [1,3,3,1]\n",
        "    #fused_modconv: True\n",
        "    #pix2pix: False\n",
        "    #dropout_rate: 0.5\n",
        "    #cond_mod: True\n",
        "    #style_mod: True\n",
        "    #noise_injection: True\n",
        "\n",
        "    # swinir (2021)\n",
        "    #netG: swinir\n",
        "    #img_size: 64\n",
        "    #patch_size: 1\n",
        "    #upscale: 2\n",
        "    #in_chans: 3\n",
        "    #window_size: 8\n",
        "    #img_range: 1.\n",
        "    #depths: [6, 6, 6, 6, 6, 6]\n",
        "    #embed_dim: 180\n",
        "    #num_heads: [6, 6, 6, 6, 6, 6]\n",
        "    #mlp_ratio: 2\n",
        "    #upsampler: 'pixelshuffle'\n",
        "    #resi_connection: '1conv'\n",
        "\n",
        "    # swinir2 (2022)\n",
        "    #netG: swinir2\n",
        "    #img_size: 56 # 56/60 # 48\n",
        "    #window_size: 8\n",
        "    #img_range: 1.\n",
        "    #depths: [3, 3, 3, 3] # [2, 2, 2, 2]\n",
        "    #embed_dim: 180 # 24\n",
        "    #num_heads: [6, 6, 6, 6]\n",
        "    #mlp_ratio: 2\n",
        "    #upsampler: 'pixelshuffledirect'\n",
        "    #use_deformable_block: False\n",
        "    #first_conv: fft # Conv2D | doconv | TBC | dynamic | fft\n",
        "    # for dynamic\n",
        "    #nof_kernels: 4\n",
        "    #reduce: 4\n",
        "\n",
        "    # ESRT (2021)\n",
        "    #netG: ESRT\n",
        "    #hiddenDim: 32\n",
        "    #mlpDim: 128\n",
        "\n",
        "    # RealESRGAN Anime Compact model (2021)\n",
        "    #netG: SRVGGNetCompact\n",
        "    #num_in_ch: 3\n",
        "    #num_out_ch: 3\n",
        "    #num_feat: 64\n",
        "    #num_conv: 16\n",
        "    #act_type: 'prelu'\n",
        "    #conv_mode: 3 # 2 | 3\n",
        "    ## only for rrdb\n",
        "    #rrdb: False\n",
        "    #rrdb_blocks: 2\n",
        "    #convtype: \"Conv2D\"\n",
        "\n",
        "    # ELAN (2022)\n",
        "    #netG: elan\n",
        "    #scale: 4\n",
        "    #colors: 3\n",
        "    #window_sizes: [4, 8, 16]\n",
        "    #m_elan: 24\n",
        "    #c_elan: 60\n",
        "    #n_share: 1\n",
        "    #r_expand: 2\n",
        "    #rgb_range: 255\n",
        "    #conv: fft # Conv2D | fft\n",
        "\n",
        "    # LFT (2022)\n",
        "    #netG: lft\n",
        "    #channels: 64\n",
        "    #angRes: 5\n",
        "    #layer_num: 4\n",
        "    #temperature: 10000\n",
        "    #num_heads: 8\n",
        "    #dropout: 0.\n",
        "\n",
        "    # swift (2021)\n",
        "    #netG: swift\n",
        "    #in_channels: 3\n",
        "    #num_channels: 64\n",
        "    #num_blocks: 16\n",
        "\n",
        "    # hat (2022)\n",
        "    #netG: hat\n",
        "    #img_size: 64\n",
        "    #patch_size: 1\n",
        "    #in_chans: 3\n",
        "    #embed_dim: 180\n",
        "    #depths: [6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6]\n",
        "    #num_heads: [6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6]\n",
        "    #window_size: 16\n",
        "    #compress_ratio: 3\n",
        "    #squeeze_factor: 30\n",
        "    #conv_scale: 0.01\n",
        "    #overlap_ratio: 0.5\n",
        "    #mlp_ratio: 2\n",
        "    #qkv_bias: True\n",
        "    #qk_scale:\n",
        "    #drop_rate: 0.\n",
        "    #attn_drop_rate: 0.\n",
        "    #drop_path_rate: 0.1\n",
        "    #ape: False\n",
        "    #patch_norm: True\n",
        "    #use_checkpoint: False\n",
        "    #img_range: 1.\n",
        "    #upsampler: 'pixelshuffle'\n",
        "    #resi_connection: '1conv'\n",
        "    #conv: fft # fft | Conv2D\n",
        "\n",
        "    # RLFN (2022)\n",
        "    #netG: RLFN\n",
        "    #in_nc: 3\n",
        "    #out_nc: 3\n",
        "    #nf: 46\n",
        "    #mf: 48\n",
        "\n",
        "    # SCET (2022)\n",
        "    #netG: SCET\n",
        "    #hiddenDim: 32\n",
        "    #mlpDim: 128\n",
        "    #conv: fft # Conv2D | fft\n",
        "\n",
        "    # span (2023) (Swift Parameter-free Attention Network)\n",
        "    #netG: span\n",
        "    #num_in_ch: 3\n",
        "    #num_out_ch: 3\n",
        "    #feature_channels: 48\n",
        "    #bias: True\n",
        "    #img_range: 255.\n",
        "    #rgb_mean: [0.5, 0.5, 0.5]\n",
        "\n",
        "    # rgt (2023)\n",
        "    #netG: rgt\n",
        "    #img_size: 64\n",
        "    #in_chans: 3\n",
        "    #embed_dim: 180\n",
        "    #depth: [6,6,6,6,6,6]\n",
        "    #num_heads: [6,6,6,6,6,6]\n",
        "    #mlp_ratio: 2\n",
        "    #qkv_bias: True\n",
        "    #qk_scale:\n",
        "    #drop_rate: 0.0\n",
        "    #attn_drop_rate: 0.0\n",
        "    #drop_path_rate: 0.1\n",
        "    #use_chk: False\n",
        "    #upscale: 2\n",
        "    #img_range: 1.0\n",
        "    #resi_connection: \"1conv\"\n",
        "    #split_size: [8, 32]\n",
        "    #c_ratio: 0.5\n",
        "\n",
        "    # SwinFIR (2023)\n",
        "    #netG: SwinFIR\n",
        "    #img_size: 64\n",
        "    #patch_size: 1\n",
        "    #in_chans: 3\n",
        "    #embed_dim: 96\n",
        "    #depths: [6, 6, 6, 6]\n",
        "    #num_heads: [6, 6, 6, 6]\n",
        "    #window_size: 8\n",
        "    #mlp_ratio: 4.0\n",
        "    #qkv_bias: True\n",
        "    #qk_scale:\n",
        "    #drop_rate: 0.0\n",
        "    #attn_drop_rate: 0.0\n",
        "    #drop_path_rate: 0.1\n",
        "    #ape: False\n",
        "    #patch_norm: True\n",
        "    #use_checkpoint: False\n",
        "    #img_range: 1.0\n",
        "    #upsampler: \"pixelshuffle\"\n",
        "    #resi_connection: \"SFB\"\n",
        "\n",
        "    # drct (2024)\n",
        "    #netG: drct\n",
        "    #img_size: 64\n",
        "    #patch_size: 1\n",
        "    #in_chans: 3\n",
        "    #embed_dim: 180\n",
        "    #depths: [6, 6, 6, 6, 6, 6]\n",
        "    #num_heads: [6, 6, 6, 6, 6, 6]\n",
        "    #window_size: 16\n",
        "    #compress_ratio: 3\n",
        "    #squeeze_factor: 30\n",
        "    #conv_scale: 0.01\n",
        "    #overlap_ratio: 0.5\n",
        "    #mlp_ratio: 2\n",
        "    #qkv_bias: True\n",
        "    #qk_scale:\n",
        "    #drop_rate: 0.0\n",
        "    #attn_drop_rate: 0.0\n",
        "    #drop_path_rate: 0.1\n",
        "    #ape: False\n",
        "    #patch_norm: True\n",
        "    #upscale: 4\n",
        "    #img_range: 1.0\n",
        "    #upsampler: \"pixelshuffle\"\n",
        "    #resi_connection: \"1conv\"\n",
        "    #gc: 32\n",
        "\n",
        "    # atd (2024)\n",
        "    #netG: atd\n",
        "    #img_size: 96\n",
        "    #patch_size: 1\n",
        "    #in_chans: 3\n",
        "    #embed_dim: 210\n",
        "    #depths: [6, 6, 6, 6, 6, 6]\n",
        "    #num_heads: [6, 6, 6, 6, 6, 6]\n",
        "    #window_size: 16\n",
        "    #category_size: 256\n",
        "    #num_tokens: 128\n",
        "    #reducted_dim: 20\n",
        "    #convffn_kernel_size: 5\n",
        "    #mlp_ratio: 2.0\n",
        "    #qkv_bias: True\n",
        "    #ape: False\n",
        "    #patch_norm: True\n",
        "    #img_range: 1.0\n",
        "    #upsampler: \"pixelshuffle\"\n",
        "    #resi_connection: \"1conv\"\n",
        "    #norm: False\n",
        "\n",
        "    # cfat (2024) (no bf16)\n",
        "    #netG: cfat\n",
        "    #img_size: 64\n",
        "    #patch_size: 1\n",
        "    #in_chans: 3\n",
        "    #embed_dim: 180\n",
        "    #depths: [8, 8, 8, 8, 8]\n",
        "    #num_heads: [6, 6, 6, 6, 6]\n",
        "    #window_size: 16\n",
        "    #shift_size: [0, 0, 8, 8, 16, 16, 24, 24]\n",
        "    #interval: [0, 2, 0, 2, 0]\n",
        "    #compress_ratio: 3\n",
        "    #squeeze_factor: 30\n",
        "    #conv_scale: 0.01\n",
        "    #overlap_ratio: 0.5\n",
        "    #mlp_ratio: 2\n",
        "    #qkv_bias: True\n",
        "    #qk_scale:\n",
        "    #drop_rate: 0.0\n",
        "    #attn_drop_rate: 0.0\n",
        "    #drop_path_rate: 0.1\n",
        "    #ape: False\n",
        "    #patch_norm: True\n",
        "    #use_checkpoint: False\n",
        "    #upscale: 4\n",
        "    #img_range: 1.0\n",
        "    #upsampler: \"pixelshuffle\"\n",
        "    #resi_connection: \"1conv\"\n",
        "\n",
        "    # ttst (2024)\n",
        "    #netG: ttst\n",
        "    #img_size: 64\n",
        "    #patch_size: 1\n",
        "    #in_chans: 3\n",
        "    #embed_dim: 180\n",
        "    #depths: [6, 6, 6, 6, 6, 6]\n",
        "    #num_heads: [6, 6, 6, 6, 6, 6]\n",
        "    #window_size: 8\n",
        "    #compress_ratio: 3\n",
        "    #squeeze_factor: 30\n",
        "    #conv_scale: 0.01\n",
        "    #overlap_ratio: 0.5\n",
        "    #mlp_ratio: 2.\n",
        "    #qkv_bias: True\n",
        "    #qk_scale:\n",
        "    #drop_rate: 0.\n",
        "    #attn_drop_rate: 0.\n",
        "    #drop_path_rate: 0.1\n",
        "    #ape: False\n",
        "    #patch_norm: True\n",
        "    #use_checkpoint: False\n",
        "    #img_range: 1\n",
        "    #upsampler: 'pixelshuffle'\n",
        "    #resi_connection: '1conv'\n",
        "\n",
        "    # agdn (2024)\n",
        "    netG: agdn\n",
        "    num_in_ch: 3\n",
        "    num_feat: 42\n",
        "    num_block: 7\n",
        "    num_out_ch: 3\n",
        "    light: False\n",
        "    rgb_mean: [0.5, 0.5, 0.5]\n",
        "\n",
        "    # ----Inpainting Generators----\n",
        "    # DFNet (batch_size: 2+, needs 2^x image input and validation) (2019)\n",
        "    #netG: DFNet\n",
        "    #c_img: 3\n",
        "    #c_mask: 1\n",
        "    #c_alpha: 3\n",
        "    #mode: nearest\n",
        "    #norm: batch\n",
        "    #act_en: relu\n",
        "    #act_de: leaky_relu\n",
        "    #en_ksize: [7, 5, 5, 3, 3, 3, 3, 3]\n",
        "    #de_ksize: [3, 3, 3, 3, 3, 3, 3, 3]\n",
        "    #blend_layers: [0, 1, 2, 3, 4, 5]\n",
        "    #conv_type: normal # partial | normal | deform\n",
        "\n",
        "\n",
        "    # EdgeConnect (2019)\n",
        "    #netG: EdgeConnect\n",
        "    #use_spectral_norm: True\n",
        "    #residual_blocks_edge: 8\n",
        "    #residual_blocks_inpaint: 8\n",
        "    #conv_type_edge: 'normal' # normal | partial | deform (has no spectral_norm)\n",
        "    #conv_type_inpaint: 'normal' # normal | partial | deform\n",
        "\n",
        "    # CSA (2019)\n",
        "    #netG: CSA\n",
        "    #c_img: 3\n",
        "    #norm: 'instance'\n",
        "    #act_en: 'leaky_relu'\n",
        "    #act_de: 'relu'\n",
        "\n",
        "    # RN (2020)\n",
        "    #netG: RN\n",
        "    #input_channels: 3\n",
        "    #residual_blocks: 8\n",
        "    #threshold: 0.8\n",
        "\n",
        "    # deepfillv1 (2018)\n",
        "    #netG:  deepfillv1\n",
        "\n",
        "    # deepfillv2 (2019)\n",
        "    #netG: deepfillv2\n",
        "    #in_channels:  4\n",
        "    #out_channels:  3\n",
        "    #latent_channels:  64\n",
        "    #pad_type:  'zero'\n",
        "    #activation:  'lrelu'\n",
        "    #norm: 'in'\n",
        "    #conv_type: partial # partial | normal\n",
        "\n",
        "    # Adaptive (2020)\n",
        "    #netG: Adaptive\n",
        "    #in_channels: 3\n",
        "    #residual_blocks: 1\n",
        "    #init_weights: True\n",
        "\n",
        "    # Global (2020)\n",
        "    #netG: Global\n",
        "    #input_dim: 5\n",
        "    #ngf: 32\n",
        "    #use_cuda: True\n",
        "    #device_ids: [0]\n",
        "\n",
        "    # Pluralistic (2019)\n",
        "    #netG: Pluralistic\n",
        "    #ngf_E: 32\n",
        "    #z_nc_E: 128\n",
        "    #img_f_E: 128\n",
        "    #layers_E: 5\n",
        "    #norm_E: 'none'\n",
        "    #activation_E: 'LeakyReLU'\n",
        "    #ngf_G: 32\n",
        "    #z_nc_G: 128\n",
        "    #img_f_G: 128\n",
        "    #L_G: 0\n",
        "    #output_scale_G: 1\n",
        "    #norm_G: 'instance'\n",
        "    #activation_G: 'LeakyReLU'\n",
        "\n",
        "    # crfill (2020)\n",
        "    #netG: crfill\n",
        "    #cnum: 48\n",
        "\n",
        "    # DeepDFNet (experimental)\n",
        "    #netG: DeepDFNet\n",
        "    #in_channels:  4\n",
        "    #out_channels:  3\n",
        "    #latent_channels:  64\n",
        "    #pad_type:  'zero'\n",
        "    #activation:  'lrelu'\n",
        "    #norm: 'in'\n",
        "\n",
        "    # partial (2018)\n",
        "    #netG: partial\n",
        "\n",
        "    # DMFN (2020)\n",
        "    #netG: DMFN\n",
        "    #in_nc: 4\n",
        "    #out_nc: 3\n",
        "    #nf: 64\n",
        "    #n_res: 8\n",
        "    #norm: 'in'\n",
        "    #activation: 'relu'\n",
        "\n",
        "    # pennet (2019)\n",
        "    #netG: pennet\n",
        "\n",
        "    # LBAM (2019)\n",
        "    #netG: LBAM\n",
        "    #inputChannels: 4\n",
        "    #outputChannels: 3\n",
        "\n",
        "    # RFR (use_swa: false, no TPU) (2020)\n",
        "    #netG: RFR\n",
        "    #conv_type: partial # partial | deform\n",
        "\n",
        "    # FRRN (2019)\n",
        "    #netG: FRRN\n",
        "\n",
        "    # PRVS (2019)\n",
        "    #netG: PRVS\n",
        "\n",
        "    # CRA (HR_size: 512) (2020)\n",
        "    #netG: CRA\n",
        "    #activation: 'elu'\n",
        "    #norm: 'none'\n",
        "\n",
        "    # atrous (2020)\n",
        "    #netG: atrous\n",
        "\n",
        "    # MEDFE (batch_size: 1) (2020)\n",
        "    #netG: MEDFE\n",
        "\n",
        "    # AdaFill (2021)\n",
        "    #netG: AdaFill\n",
        "\n",
        "    # lightweight_gan (2021)\n",
        "    #netG: lightweight_gan\n",
        "    #image_size: 512\n",
        "    #latent_dim: 256\n",
        "    #fmap_max: 512\n",
        "    #fmap_inverse_coef: 12\n",
        "    #transparent: False\n",
        "    #greyscale: False\n",
        "    #freq_chan_attn: False\n",
        "\n",
        "    # CTSDG (2021)\n",
        "    #netG: CTSDG\n",
        "\n",
        "    # lama (2022) (no AMP)\n",
        "    #netG: lama\n",
        "\n",
        "    # MST (2021)\n",
        "    #netG: MST\n",
        "\n",
        "    # mat (2022) (compiling)\n",
        "    #netG: mat\n",
        "    #z_dim: 512\n",
        "    #c_dim: 0\n",
        "    #w_dim: 512\n",
        "    #img_resolution: 512\n",
        "    #img_channels: 3\n",
        "    #noise_mode: const # const | random\n",
        "\n",
        "    # ----Interpolation Generators----\n",
        "    # cain (2020)\n",
        "    #netG: CAIN\n",
        "    #depth: 3\n",
        "    #conv: dynamic # doconv | conv2d | gated | TBC | dynamic | MBConv | Involution | CondConv | fft | WSConv\n",
        "    ## Warning: Configure OutlookAttention dimension according to resolution manually (160 for 720p)\n",
        "    #attention: CA # CA | OutlookAttention | A2Atttention | CBAM | CoTAttention | CoordAttention | ECAAttention | HaloAttention | ParNetAttention | TripletAttention | SKAttention | SGE | SEAttention | PolarizedSelfAttention\n",
        "    #RG: 2 # ResidualGroup amount\n",
        "    ## for dynamic\n",
        "    #nof_kernels: 4\n",
        "    #reduce: 4\n",
        "\n",
        "    # rife 4.x\n",
        "    #netG: rife\n",
        "    #arch_ver: 4.6 # any number between 4.0 and 4.6\n",
        "    #fastmode: False\n",
        "    #ensemble: True\n",
        "\n",
        "    # RRIN (2020)\n",
        "    #netG: RRIN\n",
        "\n",
        "    # ABME (2021)\n",
        "    #netG: ABME\n",
        "\n",
        "    # EDSC (2021) (2^x image size)\n",
        "    # pip install cupy\n",
        "    #netG: EDSC\n",
        "\n",
        "    # sepconv enhanced (2021) (needs cupy)\n",
        "    #netG: sepconv_enhanced\n",
        "\n",
        "    # sepconv realtime (enet) (2022) (needs cupy)\n",
        "    #netG: sepconv_rt\n",
        "    #real_time: True\n",
        "    #device: cuda\n",
        "    #in_channels: 64\n",
        "    #out_channels: 51\n",
        "\n",
        "    # AdaCoFNet (compressed) / CDFI (2021) (needs cupy)\n",
        "    #netG: CDFI\n",
        "\n",
        "    #-----------Misc---------------\n",
        "    # Restormer (2021) (1x model)\n",
        "    #netG: restormer\n",
        "    #inp_channels: 3\n",
        "    #out_channels: 3\n",
        "    #dim: 48\n",
        "    #num_blocks: [4,6,6,8]\n",
        "    #num_refinement_blocks: 4\n",
        "    #heads: [1,2,4,8]\n",
        "    #ffn_expansion_factor: 2.66\n",
        "    #bias: False\n",
        "    #LayerNorm_type: 'WithBias' # 'WithBias' | 'BiasFree'\n",
        "\n",
        "    # GMFSS_union\n",
        "    #netG: GMFSS_union\n",
        "\n",
        "# Discriminator options:\n",
        "network_D:\n",
        "    discriminator_criterion: MSE # MSE\n",
        "\n",
        "    d_loss_fool_weight: 1 # inside the generator loop, trying to fool the disciminator\n",
        "    d_loss_weight: 1 # inside own discriminator update\n",
        "\n",
        "    # needs \"pip3 install git+https://github.com/vballoli/nfnets-pytorch\"\n",
        "    WSConv_replace: True\n",
        "\n",
        "    netD: # in case there is no discriminator, leave it empty\n",
        "\n",
        "    # VGG\n",
        "    #netD: VGG\n",
        "    #size: 256\n",
        "    #in_nc: 3 #3\n",
        "    #base_nf: 64\n",
        "    #norm_type: 'batch'\n",
        "    #act_type: 'leakyrelu'\n",
        "    #mode: 'CNA'\n",
        "    #convtype: 'Conv2D'\n",
        "    #arch: 'ESRGAN'\n",
        "\n",
        "    # VGG fea\n",
        "    #netD: VGG_fea\n",
        "    #size: 256\n",
        "    #in_nc: 3\n",
        "    #base_nf: 64\n",
        "    #norm_type: 'batch'\n",
        "    #act_type: 'leakyrelu'\n",
        "    #mode: 'CNA'\n",
        "    #convtype: 'Conv2D'\n",
        "    #arch: 'ESRGAN'\n",
        "    #spectral_norm: False\n",
        "    #self_attention: False\n",
        "    #max_pool: False\n",
        "    #poolsize: 4\n",
        "\n",
        "\n",
        "    #netD: VGG_128_SN\n",
        "\n",
        "    # VGGFeatureExtractor\n",
        "    #netD: VGGFeatureExtractor\n",
        "    #feature_layer: 34\n",
        "    #use_bn: False\n",
        "    #use_input_norm: True\n",
        "    #device: 'cpu'\n",
        "    #z_norm: False\n",
        "\n",
        "    # PatchGAN\n",
        "    #netD: NLayerDiscriminator\n",
        "    #input_nc: 3\n",
        "    #ndf: 64\n",
        "    #n_layers: 3\n",
        "    #use_sigmoid: False\n",
        "    #getIntermFeat: False\n",
        "    #patch: True\n",
        "    #use_spectral_norm: False\n",
        "\n",
        "    # Multiscale\n",
        "    #netD: MultiscaleDiscriminator\n",
        "    #input_nc: 3\n",
        "    #ndf: 64\n",
        "    #n_layers: 3\n",
        "    #norm_layer: nn.BatchNorm2d\n",
        "    #use_sigmoid: False\n",
        "    #num_D: 3\n",
        "    #get_feats: False\n",
        "\n",
        "    #netD: ResNet101FeatureExtractor\n",
        "    #use_input_norm: True\n",
        "    #device: 'cpu'\n",
        "    #z_norm: False\n",
        "\n",
        "    # MINC\n",
        "    #netD: MINCNet\n",
        "\n",
        "    # Pixel\n",
        "    #netD: PixelDiscriminator\n",
        "    #input_nc: 3\n",
        "    #ndf: 64\n",
        "    #norm_layer: nn.BatchNorm2d\n",
        "\n",
        "    # EfficientNet (3-channel input)\n",
        "    #netD: EfficientNet\n",
        "    #EfficientNet_pretrain: 'efficientnet-b0'\n",
        "    #num_classes: 1 # should be 1\n",
        "\n",
        "    # ResNeSt (not working)\n",
        "    #netD: ResNeSt\n",
        "    #ResNeSt_pretrain: 'resnest50' # [\"resnest50\", \"resnest101\", \"resnest200\", \"resnest269\"]\n",
        "    #pretrained: False # cant be true currently\n",
        "    #num_classes: 1\n",
        "\n",
        "    # Transformer (not working)\n",
        "    #netD: TranformerDiscriminator\n",
        "    #img_size: 256\n",
        "    #patch_size: 1\n",
        "    #in_chans: 3\n",
        "    #num_classes: 1\n",
        "    #embed_dim: 64\n",
        "    #depth: 7\n",
        "    #num_heads: 4\n",
        "    #mlp_ratio: 4.\n",
        "    #qkv_bias: False\n",
        "    #qk_scale: None\n",
        "    #drop_rate: 0.\n",
        "    #attn_drop_rate: 0.\n",
        "    #drop_path_rate: 0.\n",
        "    #hybrid_backbone: None\n",
        "    #norm_layer:\n",
        "\n",
        "    # context_encoder (num_classes can't be set, broadcasting warning will be shown, training works, but I am not sure if it will work correctly)\n",
        "    #netD: context_encoder\n",
        "\n",
        "    # Transformer (doesn't do init)\n",
        "    #netD: ViT\n",
        "    #image_size: 256\n",
        "    #patch_size: 32\n",
        "    #num_classes: 1\n",
        "    #dim: 1024\n",
        "    #depth: 6\n",
        "    #heads: 16\n",
        "    #mlp_dim: 2048\n",
        "    #dropout: 0.1\n",
        "    #emb_dropout: 0.1\n",
        "\n",
        "    # Transformer (doesn't do init)\n",
        "    #netD: DeepViT\n",
        "    #image_size: 256\n",
        "    #patch_size: 32\n",
        "    #num_classes: 1\n",
        "    #dim: 1024\n",
        "    #depth: 6\n",
        "    #heads: 16\n",
        "    #mlp_dim: 2048\n",
        "    #dropout: 0.1\n",
        "    #emb_dropout: 0.1\n",
        "\n",
        "    # RepVGG\n",
        "    #netD: RepVGG\n",
        "    #RepVGG_arch: RepVGG-A0 # RepVGG-A0, RepVGG-A1, RepVGG-A2, RepVGG-B0, RepVGG-B1, RepVGG-B1g2, RepVGG-B1g4, , RepVGG-B2, RepVGG-B2g2, RepVGG-B2g4, RepVGG-B3, RepVGG-B3g2, RepVGG-B3g4\n",
        "    #num_classes: 1\n",
        "\n",
        "    # squeezenet\n",
        "    #netD: squeezenet\n",
        "    #version: \"1_1\" # 1_0, 1_1\n",
        "    #num_classes: 1\n",
        "\n",
        "    # SwinTransformer (doesn't do init)\n",
        "    #netD: SwinTransformer\n",
        "    #hidden_dim: 96\n",
        "    #layers: [2, 2, 6, 2]\n",
        "    #heads: [3, 6, 12, 24]\n",
        "    #channels: 3\n",
        "    #num_classes: 1\n",
        "    #head_dim: 32\n",
        "    #window_size: 8\n",
        "    #downscaling_factors: [4, 2, 2, 2]\n",
        "    #relative_pos_embedding: True\n",
        "\n",
        "    # mobilenetV3 (doesn't do init)\n",
        "    #netD: mobilenetV3\n",
        "    #mode: small # small, large\n",
        "    #n_class: 1\n",
        "    #input_size: 256\n",
        "\n",
        "    # resnet\n",
        "    #netD: resnet\n",
        "    #resnet_arch: resnet50 # resnet50, resnet101, resnet152\n",
        "    #num_classes: 1\n",
        "    #pretrain: True\n",
        "\n",
        "    # NFNet\n",
        "    #netD: NFNet\n",
        "    #num_classes: 1\n",
        "    #variant: 'F0'         # F0 - F7\n",
        "    #stochdepth_rate: 0.25 # 0-1, the probability that a layer is dropped during one step\n",
        "    #alpha: 0.2            # Scaling factor at the end of each block\n",
        "    #se_ratio: 0.5         # Squeeze-Excite expansion ratio\n",
        "    #activation: 'gelu'    # or 'relu'\n",
        "\n",
        "    # lvvit (2021)\n",
        "    # Warning: Needs 'pip install timm==0.4.5'\n",
        "    #netD: lvvit\n",
        "    #img_size: 224\n",
        "    #patch_size: 16\n",
        "    #in_chans: 3\n",
        "    #num_classes: 1\n",
        "    #embed_dim: 768\n",
        "    #depth: 12\n",
        "    #num_heads: 12\n",
        "    #mlp_ratio: 4.\n",
        "    #qkv_bias: False\n",
        "    #qk_scale: # None\n",
        "    #drop_rate: 0.\n",
        "    #attn_drop_rate: 0.\n",
        "    #drop_path_rate: 0.\n",
        "    #drop_path_decay: 'linear'\n",
        "    #hybrid_backbone: # None\n",
        "    ##norm_layer: nn.LayerNorm # Deafault: nn.LayerNorm / can't be configured\n",
        "    #p_emb: '4_2'\n",
        "    #head_dim: # None\n",
        "    #skip_lam: 1.0\n",
        "    #order: # None\n",
        "    #mix_token: False\n",
        "    #return_dense: False\n",
        "\n",
        "    # timm\n",
        "    # pip install timm\n",
        "    # you can loop up models here: https://rwightman.github.io/pytorch-image-models/\n",
        "    #netD: timm\n",
        "    #timm_model: \"tf_efficientnetv2_b0\"\n",
        "\n",
        "    #netD: resnet3d\n",
        "    #model_depth: 50 # [10, 18, 34, 50, 101, 152, 200]\n",
        "\n",
        "    # from lama (2021)\n",
        "    #netD: FFCNLayerDiscriminator\n",
        "    #FFCN_feature_weight: 1\n",
        "\n",
        "    #netD: effV2\n",
        "    #conv: fft # fft | conv2d\n",
        "    #size: s # s | m | l | xl\n",
        "\n",
        "    # x-transformers\n",
        "    # pip install x-transformers\n",
        "    #netD: x_transformers\n",
        "    #image_size: 512\n",
        "    #patch_size: 32\n",
        "    #dim: 512\n",
        "    #depth: 6\n",
        "    #heads: 8\n",
        "\n",
        "    #netD: mobilevit\n",
        "    #size: xxs # xxs | xs | s\n",
        "\n",
        "    # because of too many parameters, a seperate config file named \"hrt_config.yaml\" is available\n",
        "    #netD: hrt\n",
        "\n",
        "    # Attention U-Net (from asergan) (2021)\n",
        "    #netD: attention_unet\n",
        "    #num_in_ch: 3\n",
        "    #num_feat: 64\n",
        "    #skip_connection: True\n",
        "\n",
        "    # Multiscale Attention U-Net (from asergan) (2021)\n",
        "    #netD: multiscale_attention_unet\n",
        "    #num_in_ch: 3\n",
        "    #num_feat: 64\n",
        "    #num_D: 2\n",
        "\n",
        "    # U-Net (from RealESRGAN) (2021)\n",
        "    #netD: unet\n",
        "    #num_in_ch: 3\n",
        "    #num_feat: 64\n",
        "    #skip_connection: True\n",
        "\n",
        "    #netD: a2fpn\n",
        "    #band: 3\n",
        "    #class_num: 1\n",
        "    #encoder_channels: [512, 256, 128, 64]\n",
        "    #pyramid_channels: 64\n",
        "    #segmentation_channels: 64\n",
        "    #dropout: 0.2\n",
        "\n",
        "train:\n",
        "    # Adam8bit needs \"pip install bitsandbytes-cuda111\"\n",
        "    # SGD_AGC needs \"pip3 install git+https://github.com/vballoli/nfnets-pytorch\"\n",
        "    # every other optimizer uses pytorch-optimizers, for further information\n",
        "    # https://pytorch-optimizers.readthedocs.io/en/latest/optimizer_api.html\n",
        "    # to further customize, edit optimizer.py\n",
        "\n",
        "    # Examples are: AdamW, AdamP, Lamb, DAdaptAdan, Prodigy\n",
        "\n",
        "    scheduler: Prodigy\n",
        "    # ACG needs \"pip3 install git+https://github.com/vballoli/nfnets-pytorch\"\n",
        "    AGC: False # currently only for generator, only works with certain optimizers\n",
        "    lr_g: 1 # 0.0005\n",
        "    lr_d: 1 # 0.0005\n",
        "    gradient_clipping_G: true\n",
        "    gradient_clipping_G_value: 1\n",
        "    gradient_clipping_D: true\n",
        "    gradient_clipping_D_value: 1\n",
        "\n",
        "    ############################\n",
        "    ## Losses ##\n",
        "    ############################\n",
        "\n",
        "    L1Loss_weight: 0\n",
        "\n",
        "    # HFENLoss\n",
        "    HFEN_weight: 0\n",
        "    loss_f: L1CosineSim # L1Loss | L1CosineSim\n",
        "    kernel: 'log'\n",
        "    kernel_size: 15\n",
        "    sigma: 2.5\n",
        "    norm: False\n",
        "\n",
        "    # Elastic\n",
        "    Elastic_weight: 0\n",
        "    a: 0.2\n",
        "    reduction_elastic: 'mean'\n",
        "\n",
        "    # Relative L1\n",
        "    Relative_l1_weight: 0\n",
        "    l1_eps: .01\n",
        "    reduction_relative: 'mean'\n",
        "\n",
        "    # L1CosineSim (3-channel input)\n",
        "    L1CosineSim_weight: 1\n",
        "    loss_lambda: 5\n",
        "    reduction_L1CosineSim: 'mean'\n",
        "\n",
        "    # ClipL1\n",
        "    ClipL1_weight: 0\n",
        "    clip_min: 0.0\n",
        "    clip_max: 10.0\n",
        "\n",
        "    # FFTLoss\n",
        "    FFTLoss_weight: 0\n",
        "    loss_f_fft: L1Loss\n",
        "    reduction_fft: 'mean'\n",
        "\n",
        "    OFLoss_weight: 0\n",
        "\n",
        "    # GPLoss\n",
        "    GPLoss_weight: 0\n",
        "    gp_trace: False\n",
        "    gp_spl_denorm: False\n",
        "\n",
        "    # CPLoss\n",
        "    CPLoss_weight: 0\n",
        "    rgb: True\n",
        "    yuv: True\n",
        "    yuvgrad: True\n",
        "    cp_trace: False\n",
        "    cp_spl_denorm: False\n",
        "    yuv_denorm: False\n",
        "\n",
        "    # TVLoss\n",
        "    TVLoss_weight: 0\n",
        "    tv_type: 'tv'\n",
        "    p: 1\n",
        "\n",
        "    # Contextual_Loss (3-channel input)\n",
        "    Contextual_weight: 0\n",
        "    crop_quarter: False\n",
        "    max_1d_size: 100\n",
        "    distance_type: 'cosine' # [\"11\", \"l2\", \"consine\"]\n",
        "    b: 1.0\n",
        "    band_width: 0.5\n",
        "    # for vgg\n",
        "    use_vgg: False\n",
        "    net_contextual: 'vgg19'\n",
        "    layers_weights: {'conv_1_1': 1.0, 'conv_3_2': 1.0}\n",
        "    # for timm\n",
        "    use_timm: True\n",
        "    timm_model: \"tf_efficientnetv2_b0\"\n",
        "    # for both\n",
        "    calc_type: 'regular' # [\"bilateral\" | \"symetric\" | None]\n",
        "\n",
        "    # Style (3-channel input)\n",
        "    StyleLoss_weight: 0\n",
        "\n",
        "    # PerceptualLoss\n",
        "    perceptual_weight: 0\n",
        "    net: PNetLin # PNetLin, DSSIM (?)\n",
        "    pnet_type: 'vgg' # alex, squeeze, vgg\n",
        "    pnet_rand: False\n",
        "    pnet_tune: False\n",
        "    use_dropout: True\n",
        "    spatial: False\n",
        "    version: '0.1' # only version\n",
        "    lpips: True\n",
        "    force_fp16_perceptual: False # not supported for cpu\n",
        "    # you need to have tensorrt and torch_tensorrt, use docker or install it manually\n",
        "    # adjust optimum model compile values in CustomTrainClass\n",
        "    # fp16 will only work on fp16 compatible hardware, convert will only work with a GPU, FP16 untested\n",
        "    perceptual_tensorrt: False\n",
        "\n",
        "    # high receptive field (HRF) perceptual loss\n",
        "    # you can download it manually with this command, but the code will download it automatically if you don't have it\n",
        "    # wget -P /content/ http://sceneparsing.csail.mit.edu/model/pytorch/ade20k-resnet50dilated-ppm_deepsup/encoder_epoch_20.pth\n",
        "    hrf_perceptual_weight: 0\n",
        "    force_fp16_hrf: True\n",
        "\n",
        "    ldl_weight: 0\n",
        "    FrobeniusNormLoss_weight: 0\n",
        "    GradientLoss_weight: 0\n",
        "    MultiscalePixelLoss_weight: 0\n",
        "    SPLoss_weight: 0\n",
        "    FFLoss_weight: 0\n",
        "    Lap_weight: 0\n",
        "    SoftMargin_weight: 0\n",
        "\n",
        "    # pytorch loss functions\n",
        "    MSE_weight: 0\n",
        "    BCE_weight: 0\n",
        "    Huber_weight: 0\n",
        "    SmoothL1_weight: 0\n",
        "    SoftMargin_weight: 0\n",
        "\n",
        "    ############################\n",
        "    ## ARCH SPECIFIC LOSS ##\n",
        "    ############################\n",
        "    # only if the network outputs 2 images, will use l1\n",
        "    stage1_weight: 0\n",
        "\n",
        "    # loss for CTSDG\n",
        "    CTSDG_edge_weight: 0 #0.01\n",
        "    CTSDG_projected_weight: 0 #0.1\n",
        "    # loss for rife\n",
        "    SOBEL_weight: 0 # not recommended, leave it at 0\n",
        "\n",
        "    ############################\n",
        "    ## EXPERIMENTAL LOSS ##\n",
        "    ############################\n",
        "    YUVColorLoss_weight: 0\n",
        "    XYZColorLoss_weight: 0\n",
        "    KullbackHistogramLoss_weight: 0 # fp16 doesn't work\n",
        "    KullbackHistogramLossV2_weight: 0\n",
        "    SalientRegionLoss_weight: 0 # fp16 doesn't work, big values, use ~0.1 weight\n",
        "    glcmLoss_weight: 0 # big values, use ~0.001 weight\n",
        "    GradientDomainLoss_weight: 0\n",
        "    SobelLoss_weight: 0\n",
        "    SobelLossV2_weight: 0\n",
        "    ColorHarmonyLoss_weight: 0\n",
        "    VIT_FeatureLoss_weight: 0\n",
        "    VIT_MMD_FeatureLoss_weight: 0 # big values, use ~0.001 weight\n",
        "    LaplacianLoss_weight: 0\n",
        "    textured_loss_weight: 0\n",
        "\n",
        "    # Canny\n",
        "    Canny_weight: 0 # 0.01\n",
        "    canny_threshold: 5\n",
        "    canny_blurred_img_weight: 0\n",
        "    canny_grad_mag_weight: 0\n",
        "    canny_grad_orientation_weight: 0\n",
        "    canny_thin_edges_weight: 0\n",
        "    canny_thresholded_weight: 0\n",
        "    canny_early_threshold: 0.5\n",
        "\n",
        "    TIMM_FeatureLoss_weight: 0\n",
        "    # https://github.com/huggingface/pytorch-image-models/blob/main/results/results-imagenet.csv\n",
        "    TIMM_FeatureLoss_arch: tf_efficientnetv2_b0 # davit_small (224), davit_base (224), maxvit_xlarge_tf_512, tf_efficientnetv2_b0 (224), tf_efficientnet_l2 (800)\n",
        "    TIMM_FeatureLoss_resolution: 512\n",
        "    TIMM_FeatureLoss_fp16: True\n",
        "    TIMM_FeatureLoss_criterion: huber # huber, mse, l1, cross_entropy, kullback\n",
        "    TIMM_FeatureLoss_normalize: True\n",
        "    TIMM_FeatureLoss_last_feature: False # False means that all features will be used in a sum\n",
        "\n",
        "    # IQA-PyTorch Loss\n",
        "    iqa_weight: 0\n",
        "    # https://github.com/chaofengc/IQA-PyTorch/blob/main/pyiqa/default_model_configs.py\n",
        "    # only fr methods\n",
        "    ## method FR\n",
        "    # ahiq | ckdn | lpips | lpips-vgg | stlpips | stlpips-vgg | dists | ssim | ssimc |\n",
        "    # psnr | psnry | fsim | ms_ssim | vif | gmsd | nlpd | vsi | cw_ssim | mad | pieapp |\n",
        "    # topiq_fr | topiq_fr-pipal\n",
        "    ## method NR\n",
        "    # niqe | ilniqe | brisque | nrqm | pi | cnniqa | musiq | musiq-ava | musiq-koniq |\n",
        "    # musiq-paq2piq | musiq-spaq | nima | nima-vgg16-ava | pieapp | paq2piq | dbcnn |\n",
        "    # fid | maniqa | maniqa-koniq | maniqa-pipal | maniqa-kadid | clipiqa | clipiqa+ |\n",
        "    # clipiqa+_vitL14_512 | clipiqa+_rn50_512 | tres | tres-koniq | tres-flive | hyperiqa |\n",
        "    # uranker | clipscore | entropy | topiq_nr | topiq_nr-flive | topiq_nr-spaq | topiq_iaa\n",
        "    # topiq_iaa_res50 | laion_aes\n",
        "    iqa_metric: \"cw_ssim\"\n",
        "    iqa_method: FR\n",
        "    # inform yourself if value going to 1 or to 0 means higher quality\n",
        "    # if set to true, it will apply 1 - result instead of returning result\n",
        "    # higher value should return worse image quality\n",
        "    iqa_invert: True\n",
        "\n",
        "    ############################\n",
        "    ## AUG OPTIONS ##\n",
        "    ############################\n",
        "    # Three different augmentation options, diffaug and muaraugment apply to both training loops,\n",
        "    # while batch_aug only applies to the discriminator training loop\n",
        "    augmentation_method: batch_aug # diffaug, MuarAugment, batch_aug, None\n",
        "    # diffaug\n",
        "    policy: 'color,translation,cutout'\n",
        "    # MuarAugment\n",
        "    N_TFMS: 3 # Number of transformations\n",
        "    MAGN: 3 # Magnitude of augmentation applied\n",
        "    N_COMPS: 4 # Number of compositions placed on each image\n",
        "    N_SELECTED: 2 # Number of selected compositions for each image\n",
        "    # batch_aug\n",
        "    mixopts: [\"blend\", \"rgb\", \"mixup\", \"cutmix\", \"cutmixup\", \"cutout\"]\n",
        "    mixprob: [1.0, 1.0, 1.0, 1.0, 1.0, 1.0]  # , 1.0]\n",
        "    mixalpha: [0.6, 1.0, 1.2, 0.7, 0.7, 0.001]  # , 0.7]\n",
        "    aux_mixprob: 1.0\n",
        "    aux_mixalpha: 1.2\n",
        "\n",
        "    ## Metrics ##\n",
        "    metrics: [] # PSNR | SSIM | AE | MSE"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "vSs4mnExi1OA",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "#@title aug_config.yaml\n",
        "%%writefile /content/Colab-traiNNer/code/aug_config.yaml\n",
        "# Settings to configure OTF image augmentations for lrhr dataloader\n",
        "# p is always the probability of applying the augmentation\n",
        "# kernel_size must always be odd\n",
        "# sigma variables are standard deviation\n",
        "# Leave variable blank if setting to None\n",
        "  ColorJitter:\n",
        "    p: 0.5\n",
        "    brightness: 0\n",
        "    contrast: 0\n",
        "    saturation: 0\n",
        "    hue: 0\n",
        "  RandomGaussianNoise:\n",
        "    p: 0.5\n",
        "    mean: 0.0\n",
        "    var_limit: [10.0, 50.0]  # Variance range; 0-255 for sigma_calc 'sig', squared for 'var'\n",
        "    prob_color: 0.5  # Probability of color rather than grayscale noise\n",
        "    multi: True\n",
        "    mode: 'gauss'  # ['gauss' | 'speckle]\n",
        "    sigma_calc: 'sig'  # ['sig' | 'var']\n",
        "  RandomPoissonNoise:\n",
        "    p: 0.5\n",
        "    prob_color: 0.5  # Probability of color rather than grayscale noise\n",
        "    scale_range: [0.5, 1.0]  # Range for random selection noise scale\n",
        "  RandomSPNoise:\n",
        "    p: 0.5\n",
        "    prob: 0.1  # Threshold to control level of noise\n",
        "  RandomSpeckleNoise:  # Just Gaussian noise set to speckle\n",
        "    p: 0.5\n",
        "    mean: 0.0\n",
        "    var_limit: [0.04, 0.12]\n",
        "    prob_color: 0.5\n",
        "    sigma_calc: 'var'\n",
        "  RandomCompression:\n",
        "    p: 0.5\n",
        "    min_quality: 20\n",
        "    max_quality: 90\n",
        "    compression_type: '.jpg'  # ['.jpg' | '.webp']\n",
        "  RandomAverageBlur:\n",
        "    p: 0.5\n",
        "    kernel_size: 3  # Maximum kernel size (must be >=3)\n",
        "  RandomBilateralBlur:\n",
        "    p: 0.5\n",
        "    kernel_size: 3  # Maximum kernel size (must be >=3)\n",
        "    sigmaX: 0.5\n",
        "    sigmaY: 0.5\n",
        "  RandomBoxBlur:\n",
        "    p: 0.5\n",
        "    kernel_size: 3  # Maximum kernel size (must be >=3)\n",
        "  RandomGaussianBlur:\n",
        "    p: 0.5\n",
        "    kernel_size: 3  # Maximum kernel size (must be >=3)\n",
        "    sigmaX: 0.5\n",
        "    sigmaY: 0.5\n",
        "  RandomMedianBlur:\n",
        "    p: 0.5\n",
        "    kernel_size: 3  # Maximum kernel size (must be >=3)\n",
        "  RandomMotionBlur:\n",
        "    p: 0.5\n",
        "    kernel_size: 3  # Maximum kernel size (must be >=3)\n",
        "    per_channel: False  # Apply motion blur simultaneously or individually per channel\n",
        "  RandomComplexMotionBlur:\n",
        "    p: 0.5\n",
        "    size: [100, 100]\n",
        "    complexity: 0  # Modifies length and variance of motion blur path (may need to be range 0-1?)\n",
        "    eps: 0.1  # Small error for numerical stability\n",
        "  RandomAnIsoBlur:\n",
        "    p: 0.5\n",
        "    min_kernel_size: 1\n",
        "    kernel_size: 3  # Maximum kernel size (must be >=3)\n",
        "    sigmaX: 0.5\n",
        "    sigmaY: 0.5\n",
        "    angle:  # Rotation angle for anisotropic filters\n",
        "    noise:  # Multiplicative kernel noise\n",
        "    scale: 1  # To prevent filter misalignment (i.e. with nearest neighbor); scale 1 does not shift pixels\n",
        "  RandomSincBlur:\n",
        "    p: 0.5\n",
        "    min_kernel_size: 7\n",
        "    kernel_size: 21  # Maximum kernel size (must be >=3)\n",
        "    min_cutoff:  # Minimum omega cutoff frequency in radians (max: pi)\n",
        "  BayerDitherNoise:\n",
        "    p: 0.5\n",
        "  FSDitherNoise:\n",
        "    p: 0.5\n",
        "  FilterMaxRGB:\n",
        "    p: 0.5\n",
        "  FilterColorBalance:\n",
        "    p: 0.5\n",
        "    percent: 1  # Amount of balance to apply\n",
        "    random_params: False  # If true, randomizes percent from 0 to percent\n",
        "  FilterUnsharp:\n",
        "    p: 0.5\n",
        "    blur_algo: 'median'  # ['median' | None]; only used for 'laplacian'\n",
        "    kernel_size:  # Leave blank to select randomly from [1, 3, 5]\n",
        "    strength: 0.3  # Strength of filter applied (range 0-1)\n",
        "    unsharp_algo: 'laplacian'  # ['DoG' | 'laplacian]\n",
        "  FilterCanny:\n",
        "    p: 0.5\n",
        "    sigma: 0.33\n",
        "    bin_thresh: False  # Flag to apply binarize operation\n",
        "    threshold: 127  # Cutoff value for binarize (0-255)\n",
        "  SimpleQuantize:\n",
        "    p: 0.5\n",
        "    rgb_range: 40  # Higher values increase value range contained in each bin (1-255)\n",
        "  KMeansQuantize:\n",
        "    p: 0.5\n",
        "    n_colors: 128  # Number of colors in quantized image (1-255)\n",
        "  CLAHE:  # Contrast-Limited Adaptive Histogram Equalization\n",
        "    p: 0.5\n",
        "    clip_limit: 4.0  # Upper threshold value for contrast limiting (min 1)\n",
        "    tile_grid_size: [8, 8]\n",
        "  RandomGamma:\n",
        "    p: 0.5\n",
        "    gamma_range: [80, 120]  #  Range to randomly select gamma from\n",
        "    gain: 1  # Constant multiplier for gamma adjustment\n",
        "  Superpixels:\n",
        "    p: 0.5\n",
        "    p_replace: 0.1  # Probability for any segment of pixels within being replaced by their aggregate color\n",
        "    n_segments: 100  # Approximate target of number of superpixels to generate\n",
        "    cs:  # Colorspace conversion; ['lab' | 'hsv' | None]\n",
        "    algo: 'slic'  # Superpixels algorithm; ['seeds' | 'slic' | 'slico' | 'mslic' | 'sk_slic' | 'sk_felzenszwalb']\n",
        "    n_iters: 10  # Only applies for certain algorithms\n",
        "    kind: 'mix'  # How to aggregate colors; ['avg' | 'median' | 'mix']\n",
        "    reduction:  # Post-process segment to reduce colors for algos that produce more than n_segments (sk_felzenszwalb); ['selective' | 'cluster' | 'rag']\n",
        "    max_size: 128  # Max image size at which augmentation applied, will be downscaled if necessary [int or None]\n",
        "    interpolation: 'BILINEAR'  # ['NEAREST' | 'BILINEAR' | 'AREA' | 'BICUBIC' | 'LANCZOS']\n",
        "  RandomCameraNoise:\n",
        "    p: 0.5\n",
        "    demosaic_fn: 'malvar'  # ['malvar' | 'pixelshuffle' | 'menon' | 'bilinear']\n",
        "    xyz_arr: 'D50'  # Matrix to use for RGB to XYZ conversion; ['D50 | 'D65']\n",
        "    rg_range: [1.2, 2.4]  # Red gain range for white balance\n",
        "    bg_range: [1.2, 2.4]  # Blue gain range for white balance\n",
        "    random_params: False  # Initialize with random parameters if True"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "f0_srcEM1xYD",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "#@title realesrgan_aug_config.yaml\n",
        "%%writefile /content/Colab-traiNNer/code/realesrgan_aug_config.yaml\n",
        "# https://github.com/xinntao/Real-ESRGAN/blob/35ee6f781e9a5a80d5f2f1efb9102c9899a81ae1/options/train_realesrgan_x2plus.yml\n",
        "blur_kernel_size: 21\n",
        "kernel_list: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']\n",
        "kernel_prob: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]\n",
        "sinc_prob: 0.1\n",
        "blur_sigma: [0.2, 3]\n",
        "betag_range: [0.5, 4]\n",
        "betap_range: [1, 2]\n",
        "\n",
        "blur_kernel_size2: 21\n",
        "kernel_list2: ['iso', 'aniso', 'generalized_iso', 'generalized_aniso', 'plateau_iso', 'plateau_aniso']\n",
        "kernel_prob2: [0.45, 0.25, 0.12, 0.03, 0.12, 0.03]\n",
        "sinc_prob2: 0.1\n",
        "blur_sigma2: [0.2, 1.5]\n",
        "betag_range2: [0.5, 4]\n",
        "betap_range2: [1, 2]\n",
        "\n",
        "final_sinc_prob: 0.8\n",
        "\n",
        "use_hflip: True\n",
        "use_rot: False\n",
        "\n",
        "# the first degradation process\n",
        "resize_prob: [0.2, 0.7, 0.1]  # up, down, keep\n",
        "resize_range: [0.15, 1.5]\n",
        "gaussian_noise_prob: 0.5\n",
        "noise_range: [1, 30]\n",
        "poisson_scale_range: [0.05, 3]\n",
        "gray_noise_prob: 0.4\n",
        "jpeg_range: [30, 95]\n",
        "\n",
        "# the second degradation process\n",
        "second_blur_prob: 0.8\n",
        "resize_prob2: [0.3, 0.4, 0.3]  # up, down, keep\n",
        "resize_range2: [0.3, 1.2]\n",
        "gaussian_noise_prob2: 0.5\n",
        "noise_range2: [1, 25]\n",
        "poisson_scale_range2: [0.05, 2.5]\n",
        "gray_noise_prob2: 0.4\n",
        "jpeg_range2: [30, 95]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aZHf9AJuklKR"
      },
      "source": [
        "## Run Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "id": "xWaG4UHJrTye",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "%cd '/content/Colab-traiNNer/code'\n",
        "!python train.py"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Sf9dbPXQCUnV"
      },
      "source": [
        "# Creating Kernels\n",
        "\n",
        "Be aware that this takes a long time"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "L7JiQQlhB8tC",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "%cd /content\n",
        "!git clone https://github.com/victorca25/DLIP\n",
        "%cd /content/DLIP/kgan/"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JE8jluLICPps",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "# only works with GPU\n",
        "!python train.py --input-dir /content/hr --X4 --output-dir /content/kernels"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HuL4WYndJUhh"
      },
      "source": [
        "# Misc (Data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "FdOMMzyxJTlB",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "#@title Resize folder\n",
        "from tqdm import tqdm\n",
        "import os\n",
        "import glob\n",
        "import cv2\n",
        "import threading\n",
        "import shutil\n",
        "import hashlib\n",
        "import PIL\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "import hashlib\n",
        "\n",
        "rootdir = \"/content/val_lr/\"\n",
        "destination = \"/content/val_lr/\"\n",
        "broken_folder = \"/content/\"\n",
        "\n",
        "resize_method = 'PIL' #@param [\"OpenCV\", \"PIL\"] {allow-input: false}\n",
        "\n",
        "files = glob.glob(rootdir + '/**/*.png', recursive=True)\n",
        "files_jpg = glob.glob(rootdir + '/**/*.jpg', recursive=True)\n",
        "files_jpeg = glob.glob(rootdir + '/**/*.jpeg', recursive=True)\n",
        "files_webp = glob.glob(rootdir + '/**/*.webp', recursive=True)\n",
        "files.extend(files_jpg)\n",
        "files.extend(files_jpeg)\n",
        "files.extend(files_webp)\n",
        "err_files=[]\n",
        "\n",
        "image_size = 64\n",
        "\n",
        "for file in tqdm(files):\n",
        "    image = cv2.imread(file)\n",
        "    if image is not None:\n",
        "        #####################################\n",
        "        # resize with opencv\n",
        "        if resize_method == \"OpenCV\":\n",
        "          resized = cv2.resize(image, (image_size,image_size), interpolation=cv2.INTER_AREA)\n",
        "\n",
        "        # resize with PIL\n",
        "        elif resize_method == \"PIL\":\n",
        "          image = Image.fromarray(image)\n",
        "          image = image.resize((image_size,image_size))\n",
        "          resized = np.asarray(image)\n",
        "        #####################################\n",
        "\n",
        "        hash_md5 = hashlib.md5()\n",
        "        with open(file, \"rb\") as f:\n",
        "          for chunk in iter(lambda: f.read(4096), b\"\"):\n",
        "            hash_md5.update(chunk)\n",
        "\n",
        "        cv2.imwrite(file, resized)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "6PbYvgTnJX4-",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "#@title creating tiled images (image grids) (with skip)\n",
        "import cv2\n",
        "import numpy\n",
        "import glob\n",
        "import shutil\n",
        "import tqdm\n",
        "import os\n",
        "import PIL\n",
        "from PIL import Image\n",
        "import numpy as np\n",
        "\n",
        "resize_method = 'PIL' #@param [\"OpenCV\", \"PIL\"] {allow-input: false}\n",
        "grayscale = False #@param {type:\"boolean\"}\n",
        "\n",
        "rootdir = '/content/' #@param {type:\"string\"}\n",
        "destination_dir = \"/content/\" #@param {type:\"string\"}\n",
        "broken_dir = '/content/opencv_fail/' #@param {type:\"string\"}\n",
        "\n",
        "files = glob.glob(rootdir + '/**/*.png', recursive=True)\n",
        "files_jpg = glob.glob(rootdir + '/**/*.jpg', recursive=True)\n",
        "files_jpeg = glob.glob(rootdir + '/**/*.jpeg', recursive=True)\n",
        "files_webp = glob.glob(rootdir + '/**/*.webp', recursive=True)\n",
        "files.extend(files_jpg)\n",
        "files.extend(files_jpeg)\n",
        "files.extend(files_webp)\n",
        "err_files=[]\n",
        "\n",
        "amount_tiles = 8 #@param\n",
        "image_size = 256 #@param\n",
        "\n",
        "filepos = 0\n",
        "img_cnt = 0\n",
        "filename_cnt = 0\n",
        "\n",
        "if grayscale == True:\n",
        "  tmp_img = numpy.zeros((amount_tiles*image_size,amount_tiles*image_size))\n",
        "elif grayscale == False:\n",
        "  tmp_img = numpy.zeros((amount_tiles*image_size,amount_tiles*image_size, 3))\n",
        "\n",
        "with tqdm.tqdm(files) as pbar:\n",
        "  while True:\n",
        "      if grayscale == True:\n",
        "        image = cv2.imread(files[filepos], cv2.IMREAD_GRAYSCALE)\n",
        "      elif grayscale == False:\n",
        "        image = cv2.imread(files[filepos])\n",
        "\n",
        "      filepos += 1\n",
        "\n",
        "      if image is not None:\n",
        "\n",
        "        i = img_cnt % amount_tiles\n",
        "        j = img_cnt // amount_tiles\n",
        "\n",
        "        #####################################\n",
        "        # resize with opencv\n",
        "        if resize_method == \"OpenCV\":\n",
        "          image = cv2.resize(image, (image_size,image_size), interpolation=cv2.INTER_AREA)\n",
        "\n",
        "        # resize with PIL\n",
        "        elif resize_method == \"PIL\":\n",
        "          if grayscale == True:\n",
        "            image = Image.fromarray(image)\n",
        "            image = image.resize((image_size,image_size))\n",
        "            image = np.asarray(image)\n",
        "          if grayscale == False:\n",
        "            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
        "            image = Image.fromarray(image)\n",
        "            image = image.resize((image_size,image_size), resample=PIL.Image.LANCZOS)\n",
        "            image = np.asarray(image)\n",
        "            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)\n",
        "        #####################################\n",
        "\n",
        "        tmp_img[i*image_size:(i+1)*image_size, j*image_size:(j+1)*image_size] = image\n",
        "        img_cnt += 1\n",
        "      else:\n",
        "        print(files[filepos])\n",
        "        print(f'{broken_dir}/{os.path.basename(files[filepos])}')\n",
        "        shutil.move(files[filepos], f'{broken_dir}/{os.path.basename(files[filepos])}')\n",
        "\n",
        "      if img_cnt == (amount_tiles*amount_tiles):\n",
        "        #cv2.imwrite(destination_dir+str(filename_cnt)+\".jpg\", tmp_img, [int(cv2.IMWRITE_JPEG_QUALITY), 95])\n",
        "        cv2.imwrite(destination_dir+str(filename_cnt)+\".webp\", tmp_img)\n",
        "        filename_cnt += 1\n",
        "        img_cnt = 0\n",
        "      pbar.update(1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "YVaGjBR4JZpa",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "#@title convert to onnx\n",
        "#@markdown Make sure the input dimensions are correct. Maybe a runtime restart is needed if it complains about ``TypeError: forward() missing 1 required positional argument``. Make sure you only run the required cells.\n",
        "from torch.autograd import Variable\n",
        "model = CustomTrainClass()\n",
        "checkpoint_path = '/content/Checkpoint_0_0.ckpt' #@param\n",
        "output_path = '/content/output.onnx' #@param\n",
        "model = model.load_from_checkpoint(checkpoint_path) # start training from checkpoint, warning: apperantly global_step will be reset to zero and overwriting validation images, you could manually make an offset\n",
        "dummy_input = Variable(torch.randn(1, 1, 64, 64))\n",
        "\n",
        "model.to_onnx(output_path, input_sample=dummy_input)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Gw3VGGhnJeF1",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "#@title tiling script\n",
        "import cv2\n",
        "import numpy\n",
        "import glob\n",
        "import shutil\n",
        "import tqdm\n",
        "import os\n",
        "from multiprocessing.pool import ThreadPool as ThreadPool\n",
        "\n",
        "rootdir = '/content/' #@param {type:\"string\"}\n",
        "destination_dir = \"/content/\" #@param {type:\"string\"}\n",
        "broken_dir = '/content/' #@param {type:\"string\"}\n",
        "threads = 2 #@param\n",
        "tile_size = 256 #@param\n",
        "\n",
        "files = glob.glob(rootdir + '/**/*.png', recursive=True)\n",
        "files_jpg = glob.glob(rootdir + '/**/*.jpg', recursive=True)\n",
        "files_jpeg = glob.glob(rootdir + '/**/*.jpeg', recursive=True)\n",
        "files_webp = glob.glob(rootdir + '/**/*.webp', recursive=True)\n",
        "files.extend(files_jpg)\n",
        "files.extend(files_jpeg)\n",
        "files.extend(files_webp)\n",
        "err_files=[]\n",
        "\n",
        "pool = ThreadPool(threads)\n",
        "\n",
        "def tiling(f):\n",
        "  image = cv2.imread(f)\n",
        "  if image is not None:\n",
        "      counter = 0\n",
        "\n",
        "      x = image.shape[0]\n",
        "      y = image.shape[1]\n",
        "\n",
        "      x_amount = x // tile_size\n",
        "      y_amount = y // tile_size\n",
        "\n",
        "      for i in range(x_amount):\n",
        "        for j in range(y_amount):\n",
        "          crop = image[i*tile_size:(i+1)*tile_size, (j*tile_size):(j+1)*tile_size]\n",
        "          cv2.imwrite(os.path.join(destination_dir, os.path.splitext(os.path.basename(f))[0] + str(counter) + \".png\"), crop)\n",
        "          counter += 1\n",
        "\n",
        "    else:\n",
        "        print(f'Broken file: {os.path.basename(f)}')\n",
        "        shutil.move(f, f'{broken_dir}/{os.path.basename(f)}')\n",
        "\n",
        "\n",
        "pool.map(tiling, files)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "bJGFJf7MJfeE",
        "vscode": {
          "languageId": "python"
        }
      },
      "outputs": [],
      "source": [
        "#@title create landmarks (for DFDNet)\n",
        "!pip install face-alignment\n",
        "!pip install matplotlib --upgrade\n",
        "\n",
        "import face_alignment\n",
        "from skimage import io\n",
        "import numpy as np\n",
        "import glob\n",
        "from tqdm import tqdm\n",
        "import os\n",
        "import shutil\n",
        "\n",
        "fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, flip_input=False)\n",
        "\n",
        "unchecked_input_path = '/content/ffhq' #@param {type:\"string\"}\n",
        "checked_output_path = '/content/ffhq' #@param {type:\"string\"}\n",
        "failed_output_path = '/content/ffhq' #@param {type:\"string\"}\n",
        "landmark_output_path = '/content/landmarks' #@param {type:\"string\"}\n",
        "\n",
        "if not os.path.exists(unchecked_input_path):\n",
        "    os.makedirs(unchecked_input_path)\n",
        "if not os.path.exists(checked_output_path):\n",
        "    os.makedirs(checked_output_path)\n",
        "if not os.path.exists(failed_output_path):\n",
        "    os.makedirs(failed_output_path)\n",
        "if not os.path.exists(landmark_output_path):\n",
        "    os.makedirs(landmark_output_path)\n",
        "\n",
        "files = glob.glob(unchecked_input_path + '/**/*.png', recursive=True)\n",
        "files_jpg = glob.glob(unchecked_input_path + '/**/*.jpg', recursive=True)\n",
        "files.extend(files_jpg)\n",
        "err_files=[]\n",
        "\n",
        "for f in tqdm(files):\n",
        "  input = io.imread(f)\n",
        "  preds = fa.get_landmarks(input)\n",
        "  if preds is not None:\n",
        "    np.savetxt(os.path.join(landmark_output_path, os.path.basename(f)+\".txt\"), preds[0], delimiter=' ', fmt='%1.3f')   # X is an array\n",
        "    shutil.move(f, os.path.join(checked_output_path,os.path.basename(f)))\n",
        "  else:\n",
        "    shutil.move(f, os.path.join(failed_output_path,os.path.basename(f)))"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "IsK_NAi963qK",
        "-mUlJnLCj9t5",
        "Mu7HxxdGJdir",
        "dG6BqJD5J41q",
        "Sf9dbPXQCUnV",
        "HuL4WYndJUhh"
      ],
      "private_outputs": true,
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}