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"
      },
      "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"
      },
      "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"
      },
      "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 git+https://github.com/styler00dollar/pytorch-lightning.git@fc86f4ca817d5ba1702a210a898ac2729c870112\n",
        "!pip install git+https://github.com/vballoli/nfnets-pytorch\n",
        "!pip install git+https://github.com/styler00dollar/piq.git\n",
        "!pip install albumentations wget tfrecord x-transformers adamp efficientnet_pytorch tensorboardX vit-pytorch swin-transformer-pytorch madgrad timm pillow-avif-plugin kornia omegaconf"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Optional"
      ],
      "metadata": {
        "id": "-mUlJnLCj9t5"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "NMF7hHh0g92B"
      },
      "outputs": [],
      "source": [
        "#@title (optional) download precompiled mmcv (for GLEAN)\n",
        "%cd /content/\n",
        "!pip uninstall mmcv -y\n",
        "!pip uninstall mmcv-full -y\n",
        "!gdown --id 1--PoTPGKwAqGJsmaLYSEiqJy1yWMTx0G\n",
        "!pip install mmcv_full-1.3.5-cp37-cp37m-linux_x86_64.whl"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "fipjH_ZAg_tw"
      },
      "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"
      },
      "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.8.2/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"
      },
      "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"
      },
      "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": "THSivTgoj6PV"
      },
      "outputs": [],
      "source": [
        "#@title (optional) upgrade pytoch\n",
        "!pip3 install --pre torch torchvision torchaudio torchtext -f https://download.pytorch.org/whl/nightly/cu111/torch_nightly.html -U --force-reinstall"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "G7TlT3ioHdWK"
      },
      "outputs": [],
      "source": [
        "#@title (optional) Adam8bit optimizer\n",
        "!pip install bitsandbytes-cuda111"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "3y7tvaaxBv0z"
      },
      "outputs": [],
      "source": [
        "#@title (Optional) Install turboJPEG\n",
        "%cd /content/\n",
        "!sudo apt-get install -y libturbojpeg \n",
        "!pip install PyTurboJPEG"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Nv9qZfjphDMt"
      },
      "outputs": [],
      "source": [
        "#@title TPU\n",
        "!pip install cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.8.1-cp37-cp37m-linux_x86_64.whl"
      ]
    },
    {
      "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"
      },
      "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"
      },
      "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"
      },
      "outputs": [],
      "source": [
        "#@title config.yaml\n",
        "%%writefile /content/Colab-traiNNer/code/config.yaml\n",
        "name: template\n",
        "scale: 4\n",
        "gpus: 0 # amount of gpus, 0 = cpu\n",
        "distributed_backend: ddp # dp, ddp (for multi-gpu training)\n",
        "tpu_cores: 8 # 8 if you use a Google Colab TPU\n",
        "use_tpu: False\n",
        "use_amp: False\n",
        "use_swa: False\n",
        "progress_bar_refresh_rate: 20\n",
        "default_root_dir: '/content/'\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",
        "    # pip install --extra-index-url https://developer.download.nvidia.com/compute/redist --upgrade nvidia-dali-cuda110\n",
        "\n",
        "    mode: DS_lrhr # DS_video | DS_video_direct | DS_inpaint_TF | DS_inpaint  | DS_lrhr\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: 0 # 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: 2\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: 100 # 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.3, 0.4, 0.5, 0.7, 0.9, 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",
        "    RandomAnIsoBlur: 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_D: \n",
        "    checkpoint_path:\n",
        "    checkpoint_save_path: '/content/'\n",
        "    validation_output_path: '/content/validation'\n",
        "    log_path: '/content/logs'\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) | MRRDB_net (modified/\"new\" arch)\n",
        "    norm_type: null\n",
        "    mode: CNA\n",
        "    nf: 64 # of discrim filters in the first conv layer (default: 64, good: 32)\n",
        "    nb: 23 # (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",
        "    group: 1\n",
        "    convtype: Conv2D # Conv2D | PartialConv2D\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",
        "    nr: 3\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",
        "    # 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",
        "    #upscale: 4\n",
        "    #act_type: 'prelu'\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",
        "    # ----Interpolation Generators----\n",
        "    # cain (2020)\n",
        "    #netG: CAIN\n",
        "    #depth: 3\n",
        "    #conv: WSConv # 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.0\n",
        "    #netG: rife\n",
        "    #conv: conv2d # doconv | conv2d | gated | TBC | dynamic | MBConv | fft | WSConv\n",
        "    ## for dynamic\n",
        "    #nof_kernels: 4\n",
        "    #reduce: 4\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",
        "    # 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",
        "# 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",
        "    #norm_layer: nn.BatchNorm2d\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",
        "train: \n",
        "    # Adam8bit needs \"pip install bitsandbytes-cuda111\"\n",
        "    # SGD_AGC needs \"pip3 install git+https://github.com/vballoli/nfnets-pytorch\"\n",
        "    scheduler: AdamP # Adam, AdamP, Adam, SGDP, MADGRAD, Adam8bit, SGD_AGC, cosangulargrad [maybe broken], tanangulargrad [maybe broken]\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: 0.0001 # 0.00005\n",
        "    lr_d: 0.0001\n",
        "    \n",
        "    # AdamP, AGDP, MADGRAD, SGD_AGC, cosangulargrad, tanangulargrad\n",
        "    weight_decay: 0.01\n",
        "\n",
        "    # SGDP, MAGDRAD\n",
        "    momentum: 0.9\n",
        "\n",
        "    # AdamP, cosangulargrad, tanangulargrad\n",
        "    betas0: 0.9\n",
        "    betas1: 0.999\n",
        "    \n",
        "    # SGDP\n",
        "    nesterov: True\n",
        "\n",
        "    # MADGRAD, SGD_AGC, cosangulargrad, tanangulargrad\n",
        "    eps: 1e-6\n",
        "\n",
        "    ############################\n",
        "\n",
        "    # Losses:\n",
        "    L1Loss_weight: 1\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: 0\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",
        "\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 # not supported for cpu\n",
        "\n",
        "    ColorLoss_weight: 0 # converts rgb to yuv and calculates l1, expected input is rgb\n",
        "    FrobeniusNormLoss_weight: 0\n",
        "    GradientLoss_weight: 0\n",
        "    MultiscalePixelLoss_weight: 0\n",
        "    SPLoss_weight: 0\n",
        "    FFLoss_weight: 0\n",
        "\n",
        "    # only if the network outputs 2 images, will use l1\n",
        "    stage1_weight: 0 \n",
        "\n",
        "    Lap_weight: 0\n",
        "\n",
        "    # pytorch loss functions\n",
        "    MSE_weight: 0\n",
        "    BCE_weight: 0\n",
        "    Huber_weight: 0\n",
        "    SmoothL1_weight: 0\n",
        "\n",
        "    # piq loss\n",
        "    SSIMLoss_weight: 0\n",
        "    MultiScaleSSIMLoss_weight: 0\n",
        "    VIFLoss_weight: 0\n",
        "    FSIMLoss_weight: 0 # unstable\n",
        "    GMSDLoss_weight: 0 # unstable\n",
        "    MultiScaleGMSDLoss_weight: 0 # unstable\n",
        "    VSILoss_weight: 0 # unstable\n",
        "    HaarPSILoss_weight: 0\n",
        "    MDSILoss_weight: 0 # unstable\n",
        "    BRISQUELoss_weight: 0 # unstable\n",
        "    PieAPP_weight: 0\n",
        "    DISTS_weight: 0\n",
        "    # the losses below require a GPU\n",
        "    # may be changed in the future, to have custom feature extractors\n",
        "    IS_weight: 0 # unstable\n",
        "    FID_weight: 0\n",
        "    KID_weight: 0\n",
        "    PR_weight: 0 # only fp32\n",
        "    # todo(?): LPIPS\n",
        "    force_piq_fp16: True\n",
        "\n",
        "    # loss for CTSDG\n",
        "    CTSDG_edge_weight: 0 #0.01\n",
        "    CTSDG_projected_weight: 0 #0.1\n",
        "\n",
        "    # loss for rife\n",
        "    SOBEL_weight: 0 # not recommended, leave it at 0\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: \"diffaug\" # 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"
      },
      "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": "markdown",
      "metadata": {
        "id": "aZHf9AJuklKR"
      },
      "source": [
        "## Run Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "both",
        "id": "xWaG4UHJrTye"
      },
      "outputs": [],
      "source": [
        "%cd '/content/Colab-traiNNer/code'\n",
        "!python train.py"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Creating Kernels\n",
        "\n",
        "Be aware that this takes a long time"
      ],
      "metadata": {
        "id": "Sf9dbPXQCUnV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content\n",
        "!git clone https://github.com/victorca25/DLIP\n",
        "%cd /content/DLIP/kgan/"
      ],
      "metadata": {
        "id": "L7JiQQlhB8tC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# only works with GPU\n",
        "!python train.py --input-dir /content/hr --X4 --output-dir /content/kernels"
      ],
      "metadata": {
        "id": "JE8jluLICPps"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HuL4WYndJUhh"
      },
      "source": [
        "# Misc (Data)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "FdOMMzyxJTlB"
      },
      "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-"
      },
      "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"
      },
      "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": "1Vi_1JkUJbVb"
      },
      "outputs": [],
      "source": [
        "#@title copy pasting data to create artificatial dataset for debugging\n",
        "import shutil\n",
        "from random import random\n",
        "from tqdm import tqdm\n",
        "for i in tqdm(range(5000)):\n",
        "  shutil.copy(\"/content/4k/0.jpg\", \"/content/4k/\"+str(random())+\"jpg\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "aEZUr1DBJcxx"
      },
      "outputs": [],
      "source": [
        "#@title pip list with space\n",
        "!pip list | tail -n +3 | awk '{print $1}' | xargs pip show | grep -E 'Location:|Name:' | cut -d ' ' -f 2 | paste -d ' ' - - | awk '{print $2 \"/\" tolower($1)}' | xargs du -sh 2> /dev/null | sort -hr"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Gw3VGGhnJeF1"
      },
      "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"
      },
      "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)))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "V3h4-1YCJg9j"
      },
      "outputs": [],
      "source": [
        "#@title download DictionaryCenter512 for DFDNet\n",
        "!mkdir /content/DictionaryCenter512\n",
        "%cd /content/DictionaryCenter512\n",
        "!gdown --id 1sEB9j3s7Wj9aqPai1NF-MR7B-c0zfTin\n",
        "!gdown --id 1H4kByBiVmZuS9TbrWUR5uSNY770Goid6\n",
        "!gdown --id 10ctK3d9znZ9nGN3d1Z77xW3GGshbeKBb\n",
        "!gdown --id 1gcwmrIZjPFVu-cHjdQD6P4luohkPsil-\n",
        "!gdown --id 1rJ8cORPxbJsIVAiNrjBag0ihaY_Mvurn\n",
        "!gdown --id 1LkfJv2a3ud-mefAc1eZMJuINuNdSYgYO\n",
        "!gdown --id 1LH-nxD__icSJvTiAbXAXDch03oDtbpkZ\n",
        "!gdown --id 1JRTStLFsQ8dwaQjQ8qG5fNyrOvo6Tcvd\n",
        "!gdown --id 1Z4AkU1pOYTYpdbfljCgNMmPilhdEd0Kl\n",
        "!gdown --id 1Z4e1ltB3ACbYKzkoMBuVtzZ7a310G4xc\n",
        "!gdown --id 1fqWmi6-8ZQzUtZTp9UH4hyom7n4nl8aZ\n",
        "!gdown --id 1wfHtsExLvSgfH_EWtCPjTF5xsw3YyvjC\n",
        "!gdown --id 1Jr3Luf6tmcdKANcSLzvt0sjXr0QUIQ2g\n",
        "!gdown --id 1sPd4_IMYgqGLol0gqhHjBedKKxFAxswR\n",
        "!gdown --id 1eVFjXJRnBH4mx7ZbAmZRwVXZNUbgCQec\n",
        "!gdown --id 1w0GfO_KY775ZVF3KMk74ya6QL_bNU4cJ\n",
        "\n",
        "#!mkdir /content/DFDNet/weights/\n",
        "#%cd /content/DFDNet/weights/\n",
        "#!gdown --id 1SfKKZJduOGhDD27Xl01yDx0-YSEkL2Aa"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "E9OhrimFJiX4"
      },
      "outputs": [],
      "source": [
        "#@title getting ffhq test data\n",
        "%cd /content/\n",
        "!gdown --id 1VE5tnOKcfL6MoV839IVCCw5FhJxIgml5\n",
        "!7z x data.zip"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3NKIQXm7JsOI"
      },
      "source": [
        "# Download pretrain"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "jcRMhosvJjsW"
      },
      "outputs": [],
      "source": [
        "#@title getting DFDNet pretrain\n",
        "%cd /content\n",
        "!gdown --id 1UCo7YEbLLa1_87b0AoWmzhTGyrw-26nb"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ru8tCAevJnHz"
      },
      "outputs": [],
      "source": [
        "#@title downloading places2 dfnet\n",
        "%cd /content/\n",
        "!gdown --id 1SGJ_Z9kpchdnZ3Qwwf4HnN-Cq-AeK7vH # dfnet places2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-sn5xnPpJlgQ"
      },
      "outputs": [],
      "source": [
        "#@title download rfr paris model and fix state_dict\n",
        "# rfr paris\n",
        "%cd /content/\n",
        "!gdown --id 1jnUb-EvBw9DcwyWUQyWDdN9o42BPH7uT\n",
        "\n",
        "#https://discuss.pytorch.org/t/dataparallel-changes-parameter-names-issue-with-load-state-dict/60211\n",
        "import torch\n",
        "from collections import OrderedDict\n",
        "state_dict = torch.load(\"/content/checkpoint_paris.pth\", map_location='cpu')\n",
        "new_state_dict = OrderedDict()\n",
        "\n",
        "for k, v in state_dict['generator'].items():\n",
        "  if k == 'Pconv1.weight':\n",
        "      name = 'conv1.weight'\n",
        "  elif k == 'Pconv2.weight':\n",
        "      name = 'conv2.weight'\n",
        "  elif k == 'Pconv21.weight':\n",
        "      name = 'conv21.weight'\n",
        "  elif k == 'Pconv22.weight':\n",
        "      name = 'conv22.weight'\n",
        "  else:\n",
        "      name = k\n",
        "\n",
        "  new_state_dict[name] = v\n",
        "\n",
        "torch.save(new_state_dict, '/content/converted.pth')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "HTP4cZcIizwb"
      },
      "outputs": [],
      "source": [
        "#@title download and create fixed lama pretrain\n",
        "!pip install omegaconf\n",
        "\n",
        "%cd /content\n",
        "!pip3 install wldhx.yadisk-direct\n",
        "!curl -L $(yadisk-direct https://disk.yandex.ru/d/ouP6l8VJ0HpMZg) -o big-lama.zip\n",
        "!unzip big-lama.zip\n",
        "\n",
        "import torch\n",
        "from collections import OrderedDict\n",
        "state_dict = torch.load(\"/content/big-lama/models/best.ckpt\", map_location='cpu')\n",
        "new_state_dict = OrderedDict()\n",
        "\n",
        "for k, v in state_dict['state_dict'].items():\n",
        "  name = k.replace(\"generator.\", \"\")\n",
        "  new_state_dict[name] = v\n",
        "\n",
        "torch.save(new_state_dict, '/content/converted.pth')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t6iT28AsJrPo"
      },
      "source": [
        "# Misc"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cGg8I1zlJo9p"
      },
      "source": [
        "A summary of all interesting inpainting generators that are not trainable with my code.\n",
        "\n",
        "--------------------------------------------------\n",
        "\n",
        "``Broken generators:``\n",
        "\n",
        "Generators that are not included here since I can't seem to make them work properly:\n",
        "\n",
        "AOT-GAN (2021): [researchmm/AOT-GAN-for-Inpainting](https://github.com/researchmm/AOT-GAN-for-Inpainting)\n",
        "\n",
        "    Couldn't get generator working.\n",
        "\n",
        "PenNet [no AMP] (2019): [researchmm/PEN-Net-for-Inpainting](https://github.com/researchmm/PEN-Net-for-Inpainting/)\n",
        "\n",
        "    Always outputs white for some reason.\n",
        "\n",
        "CRA [no AMP] (2019): [wangyx240/High-Resolution-Image-Inpainting-GAN](https://github.com/wangyx240/High-Resolution-Image-Inpainting-GAN)\n",
        "\n",
        "    Likes to create the color pink.\n",
        "\n",
        "Global [no AMP] (2020): [SayedNadim/Global-and-Local-Attention-Based-Free-Form-Image-Inpainting](https://github.com/SayedNadim/Global-and-Local-Attention-Based-Free-Form-Image-Inpainting)\n",
        "\n",
        "    Always outputs white for some reason.\n",
        "\n",
        "crfill (2020): [zengxianyu/crfill](https://github.com/zengxianyu/crfill)\n",
        "\n",
        "    No clear instructions/code result in broken results. Unreleased training code makes a correct implementation harder.\n",
        "\n",
        "--------------------------------------------------\n",
        "\n",
        "``Non-Pytorch generators:``\n",
        "\n",
        "PSR-Net (2020): [sfwyly/PSR-Net](https://github.com/sfwyly/PSR-Net)\n",
        "\n",
        "    Uses Tensorflow 2\n",
        "\n",
        "co-mod-gan (2021): [zsyzzsoft/co-mod-gan](https://github.com/zsyzzsoft/co-mod-gan)\n",
        "\n",
        "    Has a web demo and (a broken link to a) docker. Relies on Tensorflow 1.15 / StyleGAN2 code. A Colab by me for this can be found inside https://github.com/styler00dollar/Colab-co-mod-gan.\n",
        "\n",
        "Diverse-Structure-Inpainting (2021): [USTC-JialunPeng/Diverse-Structure-Inpainting](https://github.com/USTC-JialunPeng/Diverse-Structure-Inpainting)\n",
        "\n",
        "    Tensorflow 1\n",
        "\n",
        "R-MNet (2021): [Jireh-Jam/R-MNet-Inpainting-keras](https://github.com/Jireh-Jam/R-MNet-Inpainting-keras)\n",
        "\n",
        "    Not sure if there is much new and interesting stuff.\n",
        "\n",
        "Hypergraphs (2021): [GouravWadhwa/Hypergraphs-Image-Inpainting](https://github.com/GouravWadhwa/Hypergraphs-Image-Inpainting)\n",
        "\n",
        "    Uses custom conv layer (that is implemented with tensorflow). It sounds interesting, but I got errors when I tried to port it to pytorch.\n",
        "\n",
        "PEPSI (2019): [Forty-lock/PEPSI-Fast_image_inpainting_with_parallel_decoding_network](https://github.com/Forty-lock/PEPSI-Fast_image_inpainting_with_parallel_decoding_network)\n",
        "\n",
        "    The net dcpV2 uses.\n",
        "\n",
        "Region (2019): [vickyFox/Region-wise-Inpainting](https://github.com/vickyFox/Region-wise-Inpainting)\n",
        "\n",
        "--------------------------------------------------\n",
        "\n",
        "``Pytorch generators that I never tested:``\n",
        "\n",
        "SPL (2021): [WendongZh/SPL](https://github.com/WendongZh/SPL)\n",
        "\n",
        "WTAM (2020): [ChenWang8750/WTAM](https://github.com/ChenWang8750/WTAM)\n",
        "\n",
        "MPI (2020): [ChenWang8750/MPI-model](https://github.com/ChenWang8750/MPI-model)\n",
        "\n",
        "Edge-LBAM (2021): [wds1998/Edge-LBAM](https://github.com/wds1998/Edge-LBAM)\n",
        "\n",
        "VCNET (2020): [birdortyedi/vcnet-blind-image-inpainting](https://github.com/birdortyedi/vcnet-blind-image-inpainting)\n",
        "\n",
        "    Blind image inpainting without masks.\n",
        "\n",
        "DFMA (2020): [mprzewie/dmfa_inpainting](https://github.com/mprzewie/dmfa_inpainting)\n",
        "\n",
        "GIN (2020): [rlct1/gin-sg](https://github.com/rlct1/gin-sg) and [rlct1/gin](https://github.com/rlct1/gin)\n",
        "\n",
        "StructureFlow (2019): [RenYurui/StructureFlow](https://github.com/RenYurui/StructureFlow)\n",
        "\n",
        "    Needs special files.\n",
        "\n",
        "GMCNN (2018): [shepnerd/inpainting_gmcnn](https://github.com/shepnerd/inpainting_gmcnn)\n",
        "\n",
        "    The net dcpV1 used iirc.\n",
        "\n",
        "ShiftNet (2018): [Zhaoyi-Yan/Shift-Net_pytorch](https://github.com/Zhaoyi-Yan/Shift-Net_pytorch)\n",
        "\n",
        "--------------------------------------------------\n",
        "``Soon:``\n",
        "\n",
        "ICT (2021): [raywzy/ICT](https://github.com/raywzy/ICT) (code released, waiting for pre-trained models)\n",
        "\n",
        "MuFA-Net (2021): [ChenWang8750/MuFA-Net](https://github.com/ChenWang8750/MuFA-Net)\n",
        "\n",
        "GCM-Net (2021): [ZhengHuanCS/GCM-Net](https://github.com/ZhengHuanCS/GCM-Net)\n",
        "\n",
        "--------------------------------------------------\n",
        "\n",
        "``No training code:``\n",
        "\n",
        "SC-FEGAN (2019): [run-youngjoo/SC-FEGAN](https://github.com/run-youngjoo/SC-FEGAN)\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "HuL4WYndJUhh",
        "3NKIQXm7JsOI",
        "t6iT28AsJrPo"
      ],
      "name": "Colab-traiNNer.ipynb",
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}