stylegan2_ada_pytorch_pokemon.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "stylegan2-ada-pytorch-pokemon.ipynb",
      "provenance": [],
      "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/freiheit2348/Flowise/blob/main/stylegan2_ada_pytorch_pokemon.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ba1NQzpO9iJ4"
      },
      "source": [
        "# Training Pokemon with StyleGAN\n",
        "\n",
        "In this tutorial, we only want to familiarize ourselves with running StyleGAN (in this case, the newest version) and seeing how training progresses."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oLqHshr2y-MO"
      },
      "source": [
        "## Verify Runtime is GPU\n",
        "\n",
        "In the menu, select Runtime -> Change Runtime Type and verify you are using the **GPU**.  Also select **High-RAM** if you are using Colab Pro.\n",
        "\n",
        "The `nvidia-smi` command below should **NOT** display *\"NVIDIA-SMI has failed because it couldn't communicate with the NVIDIA driver.\"*"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "R7JvPMLWy95f"
      },
      "source": [
        "!nvidia-smi -L"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WUl3Rt0ryX7m"
      },
      "source": [
        "## Mount your Google Drive\n",
        "\n",
        "You will be storing the training models and progress images on your Google Drive.  This is very convenient for viewing progress, and if your Colab notebook is disconnected you will not lose your models."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ag2Bb1pPzthT"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JIHuj2eEy13f"
      },
      "source": [
        "## Install Stylegan2-ada-pytorch Prerequisites"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A8r0Ca7Hpo5F"
      },
      "source": [
        "import torch"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2y195-nOz19D"
      },
      "source": [
        "Verify the next command results in \"1\".  If not, go back to the beginning and verify you have a GPU runtime."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KV07xnKTprut"
      },
      "source": [
        "torch.cuda.device_count()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6jY44x99pusN"
      },
      "source": [
        "import torchvision"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sZhoCdj4pREI"
      },
      "source": [
        "!pip install click requests tqdm pyspng ninja imageio-ffmpeg==0.4.3"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b94le1Kh_v1k"
      },
      "source": [
        "## Get the StyleGAN code"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "najHDP7gMeFy"
      },
      "source": [
        "!git clone https://github.com/NVlabs/stylegan2-ada-pytorch.git"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GCYmJ2Z60WSe"
      },
      "source": [
        "mkdir /content/stylegan2-ada-pytorch/datasets"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wpk9vGb50lee"
      },
      "source": [
        "## Get Pokemon dataset of images"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0OscsNQnMiqa"
      },
      "source": [
        "cd /content/stylegan2-ada-pytorch/datasets"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Cyv63QC7xc6l"
      },
      "source": [
        "!wget https://github.com/derekphilipau/machinelearningforartists/raw/main/docs/src/week2/pokemon256.zip"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TSNoTfiPcrNy"
      },
      "source": [
        "!unzip pokemon256.zip"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Anm9DkOccyMZ"
      },
      "source": [
        "!rm -rf __MACOSX"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CsBFXqK1NPGl"
      },
      "source": [
        "!rm pokemon256.zip"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WBg2Zy06dkqz"
      },
      "source": [
        "cd /content/stylegan2-ada-pytorch/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nQfprRHj_-gH"
      },
      "source": [
        "## Prepare Pokemon dataset for use by StyleGAN"
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "tYoumdb8V2Cx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_szXOaO0WI00"
      },
      "source": [
        "!python dataset_tool.py --source=./datasets/pokemon256 --dest=./datasets/pokemondataset.zip"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QTo4K3m8AEzh"
      },
      "source": [
        "## Create Folders on Your Google Drive\n",
        "\n",
        "If we accidentally close our browser or the Colab runtime disconnects, we will lose all of our training models and progress images.  Therefore we want to store the training data on our Google Drive.  The following cells will create a new folder on your Google Drive, **MachineLearningForArtists**.  The Pokemon training data will be stored in **MachineLearningForArtists/Pokemon**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MexWxMvLHP5O"
      },
      "source": [
        "mkdir /content/drive/MyDrive/MachineLearningForArtists"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Tja3AaYIHU9v"
      },
      "source": [
        "mkdir /content/drive/MyDrive/MachineLearningForArtists/Pokemon"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OS14PMAY6SUW"
      },
      "source": [
        "If you get an error \"cannot create directory\", it's probably because the folder already exists on your Google Drive and you can ignore the error."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HpR6A_qBAo8m"
      },
      "source": [
        "## Train from Scratch\n",
        "\n",
        "This cell block will train StyleGAN from scratch.  Training from scratch is much slower than using *transfer learning* on a previously trained model.  However, the purpose of this tutorial is to a) try training for the first time and b) notice how the progress images develop over time."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NuSkETl1V8pv"
      },
      "source": [
        "!python train.py --outdir=/content/drive/MyDrive/MachineLearningForArtists/Pokemon --data=./datasets/pokemondataset.zip --gpus=1 --augpipe=bg --gamma=10 --cfg=paper256 --mirror=1 --snap=10 --metrics=none\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NQVRH5rr4dki"
      },
      "source": [
        "## Wait an Hour or Two\n",
        "\n",
        "Normal training of a dataset can take days.  For this demonstration, we only want to train for a few hours in order to get familiar with the training process and see the initial results.\n",
        "\n",
        "Results are stored in your Google Drive under drive/MyDrive/MachineLearningForArtists/Pokemon\n",
        "\n",
        "Each training run is stored in a separate directory.  The initial training run is stored in `00000-pokemondataset...`.  If you run training twice, the second run will be stored in `00001-pokemondataset...`.  And so on.\n",
        "\n",
        "Inside the training run directory you will see various files.  `reals.png` shows a sample of the training dataset.  You should see various Pokemon in this image. `fakes000000.png` is a sample of generated images from the initial model.  `network-snapshot-XXXXXX.pkl` is the actual model which can be used later to generate \"fake\" images, videos, etc.\n",
        "\n",
        "As training progresses you will see more `fakes` & `network-snapshot` files created.\n",
        "\n",
        "Since we only want to test a training, once you get to around the 500th iteration, you can stop the above training cell.  Then take a look at the `fakes` images and familiarize yourself with how the training progresses.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RJ-xirKuHKEP"
      },
      "source": [
        "## Transfer Learning\n",
        "\n",
        "Now, let's train the same dataset but use *transfer learning*.  We will use a StyleGAN2 model already pre-trained on Pokemon to train a new model that is trained against a Tamagotchi dataset."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RM9mywxwPy8k"
      },
      "source": [
        "### Get the Tomagotchi dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Pm3cBBuTO4rC"
      },
      "source": [
        "%cd /content/stylegan2-ada-pytorch/datasets"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j68IvkgSOqwO"
      },
      "source": [
        "!wget https://github.com/derekphilipau/machinelearningforartists/raw/main/docs/src/week2/Tamagotchi256.zip"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aCQfcO9mO-3d"
      },
      "source": [
        "!unzip Tamagotchi256.zip"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4lNNF3pcPCwM"
      },
      "source": [
        "!rm -rf __MACOSX/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZrHIFInOPGi-"
      },
      "source": [
        "%cd /content/stylegan2-ada-pytorch"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8fa3gVmTPuXk"
      },
      "source": [
        "### Prepare Tomagotchi dataset for use by StyleGAN"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5GvnuNjSPLC5"
      },
      "source": [
        "!python dataset_tool.py --source=./datasets/Tamagotchi256 --dest=./datasets/tamagotchidataset.zip"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3XCXysDmOnqF"
      },
      "source": [
        "### Get the pre-trained Pokemon model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YN9c3nyQJWi-"
      },
      "source": [
        "!gdown --id 122_Nu7p6ESMqBarxuAHCXhQuAaNCsh_a"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bUbn23OUQIS7"
      },
      "source": [
        "### Train using the Pokemon model with the Tamagotchi dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nIqAPSgaHKEZ"
      },
      "source": [
        "!python train.py --resume=./stylegan2-ada-pytorch-pokemon-256px-11650.pkl --outdir=/content/drive/MyDrive/MachineLearningForArtists/Tamagotchi --data=./datasets/tamagotchidataset.zip --gpus=1 --augpipe=bg --gamma=10 --cfg=paper256 --mirror=1 --snap=10 --metrics=none\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}