SG2_ADA_PyTorch.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "SG2-ADA-PyTorch.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true,
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/dvschultz/stylegan2-ada-pytorch/blob/main/SG2_ADA_PyTorch.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jG7ZEc_982io"
      },
      "source": [
        "# StyleGAN2-ADA-PyTorch\n",
        "\n",
        "**Notes**\n",
        "* Training section should be fairly stable. I’ll slowly add features but it should work for most mainstream use cases\n",
        "* Inference section is a work in progress. If you come across bug or have feature requests please post them in [Slack](https://ml-images.slack.com/archives/CLJGF384R) or on [Github](https://github.com/dvschultz/stylegan2-ada-pytorch/issues)\n",
        "\n",
        "---\n",
        "\n",
        "If you find this notebook useful, consider signing up for my [Patreon](https://www.patreon.com/bustbright) or [YouTube channel](https://www.youtube.com/channel/UCaZuPdmZ380SFUMKHVsv_AA/join). You can also send me a one-time payment on [Venmo](https://venmo.com/Derrick-Schultz)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Vj4PG4_i9Alt"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qGEXPcFJ9UTY"
      },
      "source": [
        "Let’s start by checking to see what GPU we’ve been assigned. Ideally we get a V100, but a P100 is fine too. Other GPUs may lead to issues."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7VVICTCvd4mc"
      },
      "source": [
        "!nvidia-smi -L"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rSV_HEoD9dxo"
      },
      "source": [
        "Next let’s connect our Google Drive account. This is optional but highly recommended."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IuVPuJmbigRs"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nTjVmfSK9CYa"
      },
      "source": [
        "## Install repo\n",
        "\n",
        "The next cell will install the StyleGAN repository in Google Drive. If you have already installed it it will just move into that folder. If you don’t have Google Drive connected it will just install the necessary code in Colab."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "B8ADVNpBh8Ox"
      },
      "source": [
        "import os\n",
        "if os.path.isdir(\"/content/drive/MyDrive/colab-sg2-ada-pytorch\"):\n",
        "    %cd \"/content/drive/MyDrive/colab-sg2-ada-pytorch/stylegan2-ada-pytorch\"\n",
        "elif os.path.isdir(\"/content/drive/\"):\n",
        "    #install script\n",
        "    %cd \"/content/drive/MyDrive/\"\n",
        "    !mkdir colab-sg2-ada-pytorch\n",
        "    %cd colab-sg2-ada-pytorch\n",
        "    !git clone https://github.com/dvschultz/stylegan2-ada-pytorch\n",
        "    %cd stylegan2-ada-pytorch\n",
        "    !mkdir downloads\n",
        "    !mkdir datasets\n",
        "    !mkdir pretrained\n",
        "    !gdown --id 1-5xZkD8ajXw1DdopTkH_rAoCsD72LhKU -O /content/drive/MyDrive/colab-sg2-ada-pytorch/stylegan2-ada-pytorch/pretrained/wikiart.pkl\n",
        "else:\n",
        "    !git clone https://github.com/dvschultz/stylegan2-ada-pytorch\n",
        "    %cd stylegan2-ada-pytorch\n",
        "    !mkdir downloads\n",
        "    !mkdir datasets\n",
        "    !mkdir pretrained\n",
        "    %cd pretrained\n",
        "    !gdown --id 1-5xZkD8ajXw1DdopTkH_rAoCsD72LhKU\n",
        "    %cd ../\n",
        "\n",
        "!pip install ninja opensimplex"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9jMmUpn4DWRe"
      },
      "source": [
        "You probably don’t need to run this, but this will update your repo to the latest and greatest."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uV9bdvzeDRPd"
      },
      "source": [
        "%cd \"/content/drive/My Drive/colab-sg2-ada-pytorch/stylegan2-ada-pytorch\"\n",
        "!git config --global user.name \"test\"\n",
        "!git config --global user.email \"test@test.com\"\n",
        "!git fetch origin\n",
        "!git pull\n",
        "!git checkout origin/main -- train.py generate.py"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cZkcJ58P97Ls"
      },
      "source": [
        "## Dataset Preparation\n",
        "\n",
        "Upload a .zip of square images to the `datasets` folder. Previously you had to convert your model to .tfrecords. That’s no longer needed :)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5B-h6FpB9FaK"
      },
      "source": [
        "## Train model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bNc-3wTO-MUd"
      },
      "source": [
        "Below are a series of variables you need to set to run the training. You probably won’t need to touch most of them.\n",
        "\n",
        "* `dataset_path`: this is the path to your .zip file\n",
        "* `resume_from`: if you’re starting a new dataset I recommend `'ffhq1024'` or `'./pretrained/wikiart.pkl'`\n",
        "* `mirror_x` and `mirror_y`: Allow the dataset to use horizontal or vertical mirroring."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JV0W6yxP-UIn"
      },
      "source": [
        "#required: definitely edit these!\n",
        "dataset_path = './datasets/chin-morris.zip'\n",
        "resume_from = './pretrained/wikiart.pkl'\n",
        "aug_strength = 0.0\n",
        "train_count = 0\n",
        "mirror_x = True\n",
        "\n",
        "#broken, don't use for now :(\n",
        "#mirror_y = False\n",
        "\n",
        "#optional: you might not need to edit these\n",
        "gamma_value = 50.0\n",
        "augs = 'bg'\n",
        "config = '11gb-gpu'\n",
        "snapshot_count = 4"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EL-M7WnnfMDI"
      },
      "source": [
        "!python train.py --gpus=1 --cfg=$config --metrics=None --outdir=./results --data=$dataset_path --snap=$snapshot_count --resume=$resume_from --augpipe=$augs --initstrength=$aug_strength --gamma=$gamma_value --mirror=$mirror_x --mirrory=False --nkimg=$train_count"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RgvSvfyi_R_-"
      },
      "source": [
        "### Resume Training\n",
        "\n",
        "Once Colab has shutdown, you’ll need to resume your training. Reset the variables above, particularly the `resume_from` and `aug_strength` settings.\n",
        "\n",
        "1. Point `resume_from` to the last .pkl you trained (you’ll find these in the `results` folder)\n",
        "2. Update `aug_strength` to match the augment value of the last pkl file. Often you’ll see this in the console, but you may need to look at the `log.txt`. Updating this makes sure training stays as stable as possible.\n",
        "3. You may want to update `train_count` to keep track of your training progress.\n",
        "\n",
        "Once all of this has been reset, run that variable cell and the training command cell after it."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VznRirOE5ENI"
      },
      "source": [
        "## Convert Legacy Model\n",
        "\n",
        "If you have an older version of a model (Tensorflow based StyleGAN, or Runway downloaded .pkl file) you’ll need to convert to the newest version. If you’ve trained in this notebook you do **not** need to use this cell.\n",
        "\n",
        "`--source`: path to model that you want to convert\n",
        "\n",
        "`--dest`: path and file name to convert to."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CzkP-Rww5Np9"
      },
      "source": [
        "!python legacy.py --source=/content/drive/MyDrive/runway.pkl --dest=/content/drive/MyDrive/colab-sg2-ada-pytorch/stylegan2-ada-pytorch/runway.pkl"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L6EtrPqL9ILk"
      },
      "source": [
        "## Testing/Inference\n",
        "\n",
        "Also known as \"Inference\", \"Evaluation\" or \"Testing\" the model. This is the process of usinng your trained model to generate new material, usually images or videos."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mYdyfH0O8In_"
      },
      "source": [
        "### Generate Single Images\n",
        "\n",
        "`--network`: Make sure the `--network` argument points to your .pkl file. (My preferred method is to right click on the file in the Files pane to your left and choose `Copy Path`, then paste that into the argument after the `=` sign).\n",
        "\n",
        "`--seeds`: This allows you to choose random seeds from the model. Remember that our input to StyleGAN is a 512-dimensional array. These seeds will generate those 512 values. Each seed will generate a different, random array. The same seed value will also always generate the same random array, so we can later use it for other purposes like interpolation.\n",
        "\n",
        "`--truncation`: Truncation, well, truncates the latent space. This can have a subtle or dramatic affect on your images depending on the value you use. The smaller the number the more realistic your images should appear, but this will also affect diversity. Most people choose between 0.5 and 1.0, but technically it's infinite. \n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VYRXenMoZSHf"
      },
      "source": [
        "!python generate.py --outdir=/content/out/images/ --trunc=1.0 --seeds=0-49 --network=/content/stylegan2-ada-pytorch/pretrained/wikiart.pkl"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VjOTCWVonoVL"
      },
      "source": [
        "### Truncation Traversal\n",
        "\n",
        "Below you can take one seed and look at the changes to it across any truncation amount. -1 to 1 will be pretty realistic images, but the further out you get the weirder it gets.\n",
        "\n",
        "#### Options \n",
        "`--network`: Again, this should be the path to your .pkl file.\n",
        "\n",
        "`--seeds`: Pass this only one seed. Pick a favorite from your generated images.\n",
        "\n",
        "`--start`: Starting truncation value.\n",
        "\n",
        "`--stop`: Stopping truncation value. This should be larger than the start value. (Will probably break if its not).\n",
        "\n",
        "`--increment`: How much each frame should increment the truncation value. Make this really small if you want a long, slow interpolation. (stop-start/increment=total frames)\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nyzdGr7OnrMG"
      },
      "source": [
        "!python generate.py --process=\"truncation\" --outdir=/content/out/trunc-trav-3/ --start=-0.8 --stop=2.8 --increment=0.02 --seeds=470 --network=/content/drive/MyDrive/stylegan2-transfer-models/mixed6k-network-snapshot-016470.pkl"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OSzj0igO8Lfu"
      },
      "source": [
        "### Interpolations\n",
        "\n",
        "Interpolation is the process of generating very small changes to a vector in order to make it appear animated from frame to frame.\n",
        "\n",
        "We’ll look at different examples of interpolation below.\n",
        "\n",
        "#### Options\n",
        "\n",
        "`--network`: path to your .pkl file\n",
        "\n",
        "`--interpolation`: Walk type defines the type of interpolation you want. In some cases it can also specify whether you want the z space or the w space.\n",
        "\n",
        "`--frames`: How many frames you want to produce. Use this to manage the length of your video.\n",
        "\n",
        "`--trunc`: truncation value"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OJSqafIzNwhx"
      },
      "source": [
        "#### Linear Interpolation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sqkiskly8S5_"
      },
      "source": [
        "!python generate.py --outdir=/content/out/video1-w-0.5/ --space=\"z\" --trunc=0.5 --process=\"interpolation\" --seeds=463,470 --network=/content/drive/MyDrive/stylegan2-transfer-models/mixed6k-network-snapshot-016470.pkl --frames=48"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DCUEV3aO8s_X"
      },
      "source": [
        "!python generate.py --outdir=out/video1-w/ --space=\"w\" --trunc=1 --process=\"interpolation\" --seeds=85,265,297,849 --network=/content/stylegan2-ada-pytorch/pretrained/wikiart.pkl"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pmKbwZDD8gjM"
      },
      "source": [
        "!zip -r vid1.zip /content/out/video1-w-0.5"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yi3d7xzpN2Uj"
      },
      "source": [
        "#### Slerp Interpolation\n",
        "\n",
        "This gets a little heady, but technically linear interpolations are not the best in high-dimensional GANs. [This github link](https://github.com/soumith/dcgan.torch/issues/14) is one of the more popular explanations ad discussions.\n",
        "\n",
        "In reality I do not find a huge difference between linear and spherical interpolations (the difference in z- and w-space is enough in many cases), but I’ve implemented slerp here for anyone interested.\n",
        "\n",
        "Note: Slerp in w space currently isn’t supported. I’m working on it."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "I0-cUd3fB_kJ"
      },
      "source": [
        "!python generate.py --outdir=out/video1/ --trunc=1 --process=\"interpolation\" --interpolation=\"slerp\" --seeds=85,265,297,849 --network=/content/stylegan2-ada-pytorch/pretrained/wikiart.pkl"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uP1HsU_CPcF5"
      },
      "source": [
        "#### Noise Loop\n",
        "\n",
        "If you want to just make a random but fun interpolation of your model the noise loop is the way to go. It creates a random path thru the z space to show you a diverse set of images.\n",
        "\n",
        "`--interpolation=\"noiseloop\"`: set this to use the noise loop funtion\n",
        "\n",
        "`--diameter`: This controls how \"wide\" the loop is. Make it smaller to show a less diverse range of samples. Make it larger to cover a lot of samples. This plus `--frames` can help determine how fast the video feels.\n",
        "\n",
        "`--random_seed`: this allows you to change your starting place in the z space. Note: this value has nothing to do with the seeds you use to generate images. It just allows you to randomize your start point (and if you want to return to it you can use the same seed multiple times).\n",
        "\n",
        "Noise loops currently only work in z space."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gfR6DhfvN8b_"
      },
      "source": [
        "!python generate.py --outdir=out/video-noiseloop-0.9d/ --trunc=0.8 --process=\"interpolation\" --interpolation=\"noiseloop\" --diameter=0.9 --random_seed=100 --network=/content/stylegan2-ada-pytorch/pretrained/wikiart.pkl"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PkKFb-4CedOq"
      },
      "source": [
        "#### Circular Loop\n",
        "\n",
        "The noise loop is, well, noisy. This circular loop will feel much more even, while still providing a random loop.\n",
        "\n",
        "I recommend using a higher `--diameter` value than you do with noise loops. Something between `50.0` and `500.0` alongside `--frames` can help control speed and diversity."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ao62za9_QfOF"
      },
      "source": [
        "!python generate.py --outdir=out/video-circularloop/ --trunc=1 --process=\"interpolation\" --interpolation=\"circularloop\" --diameter=800.00 --frames=720 --random_seed=90 --network=/content/stylegan2-ada-pytorch/pretrained/wikiart.pkl"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T8Ld_ozNJCOn"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qywlaS5pgzyH"
      },
      "source": [
        "## Combine NPZ files together"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "M2VooqrNfIpw"
      },
      "source": [
        "!python combine_npz.py --outdir=/content/npz --npzs='/content/pb-proj1-3-clip.npz,/content/pb-proj1-clip-nocenter.npz,/content/pb-proj2-clip-nocenter.npz,/content/projected_w (2).npz'"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}