StyleGAN2_Tazik_25GB_RAM.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "StyleGAN2_Tazik_25GB_RAM.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "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/ZKTKZ/thdne/blob/master/StyleGAN2_Tazik_25GB_RAM.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I4XmmpwHtwXq",
        "colab_type": "text"
      },
      "source": [
        "#Introduction\n",
        "\n",
        "This document will give you step-by-step instructions on training a GAN to make infinite images of an anime girl of your choice.\n",
        "\n",
        "This would not be possible without the work of many before me -- most notably Gwern, whose pre-trained StyleGAN 2 model is the basis for our transfer learning, and who has also written an in-depth guide on his site; random chinese user on CSDN, whose Colab-specific experiences and code samples were helpful; and nagadomi, for his anime face cropper.\n",
        "\n",
        "My original contribution is a color distance computer to filter undesirable Pixiv data. For characters with sufficient Danbooru images, this is not necessary; but for others, being able to draw on the Pixiv dataset is essential. In my case, Pixiv yielded 1500+ images, of which *hundreds* (25-50%) were not relevant; and the color distance script helped filter it down.\n",
        "\n",
        "This project was my induction into deep learning. I've learnt to parse papers and debug Tensorflow. Of course, this is only the beginning -- to gain a proper, first-principles understanding of the field, I have begun to re-implement important DL papers."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WDshosOloaLF",
        "colab_type": "text"
      },
      "source": [
        "# Scraping\n",
        "Two sources:\n",
        "\n",
        "1) Danbooru (https://github.com/Bionus/imgbrd-grabber)\n",
        "\n",
        "2) Pixiv (https://github.com/Redcxx/Pikax)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ptbi4aQgnEiV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from pikax import Pikax, settings\n",
        "\n",
        "pixiv = Pikax(settings.username, settings.password)\n",
        "\n",
        "results = pixiv.search(keyword='早坂愛')  # search\n",
        "\n",
        "pixiv.download(results)  # download\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CVTc7xE0odRT",
        "colab_type": "text"
      },
      "source": [
        "See https://github.com/Redcxx/Pikax for instructions on setting up your `username` & `password`.\n",
        "Next, run Dupeguru (https://github.com/arsenetar/dupeguru/) on your downloaded images. \n",
        "\n",
        "Danbooru consists primarily of high tier images for Pixiv, and this step prevents duplication.\n",
        "Now that we have our data, on to processing!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eRXDhX-Tv0Ru",
        "colab_type": "text"
      },
      "source": [
        "# Cropping"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "REoeE_gyn3i-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# https://github.com/nagadomi/lbpcascade_animeface/blob/master/examples/detect.py\n",
        "\n",
        "import imutils\n",
        "import cv2\n",
        "import sys\n",
        "import os.path\n",
        "\n",
        "def detect(abs_filename, cascade_file = \"../lbpcascade_animeface.xml\"):#, mode=\"display\"):\n",
        "    if not os.path.isfile(cascade_file):\n",
        "        raise RuntimeError(\"%s: not found\" % cascade_file)\n",
        "\n",
        "    cascade = cv2.CascadeClassifier(cascade_file)\n",
        "    image = cv2.imread(abs_filename, cv2.IMREAD_COLOR)\n",
        "    #height, width, channels = image.shape\n",
        "    #image = image[0: int(h/2), 0: w]\n",
        "    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)\n",
        "    gray = cv2.equalizeHist(gray)\n",
        "    \n",
        "    faces = cascade.detectMultiScale(gray,\n",
        "                                     # detector options\n",
        "                                     scaleFactor = 1.05,\n",
        "                                     minNeighbors = 5,\n",
        "                                     minSize = (250, 250)\n",
        "                                     #,maxSize = (int(0.4*w), int(0.4*h))\n",
        "                                     )\n",
        "    tag = 0\n",
        "    filename = os.path.basename(abs_filename)\n",
        "    for (x, y, w, h) in faces:\n",
        "\n",
        "        #cv2.rectangle(image, (int(x*0.85), int(y*0.1)), (x + int(w*1.5), y + h), (255, 0, 0), 50)\n",
        "        cropped = image[int(y*0.2): y + int(h*0.825), int(x*0.95): x + int(w*1.225)]\n",
        "        #cv2.imshow(\"AnimeFaceDetect\", imutils.resize(cropped, width=1080, height=1366))        cv2.waitKey(0)\n",
        "        cv2.imwrite(str(filename[0:-4] + '_' + str(tag) + filename[-4:]), cropped)\n",
        "        tag += 1\n",
        "\n",
        "if len(sys.argv) != 2:\n",
        "    sys.stderr.write(\"usage: detect.py <abs_filename>\\n\")\n",
        "    sys.exit(-1)\n",
        "detect(sys.argv[1])\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ilHhY9FZoYJc",
        "colab_type": "text"
      },
      "source": [
        "You may want to modify the parameters. I crop the images rather selectively, to minimize background noise. The `scaleFactor` determines how many scales of the image the classification is run on. A lower value means more results, but also more false positives. The other two parameters are self-descriptive."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l3hT1PWODvcv",
        "colab_type": "text"
      },
      "source": [
        "## Upscaling"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2M76qFwNeo4y",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!cat /usr/local/cuda/version.txt\n",
        "!wget http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1604/x86_64/cuda-repo-ubuntu1604_9.0.176-1_amd64.deb\n",
        "!sudo dpkg -i cuda-repo-ubuntu1604_9.0.176-1_amd64.deb\n",
        "!apt-get install libvulkan-dev\n",
        "!apt-get update\n",
        "\n",
        "!%cd /content/\n",
        "!git clone https://github.com/nihui/waifu2x-ncnn-vulkan.git\n",
        "!cd waifu2x-ncnn-vulkan/\n",
        "!git submodule update --init --recursive\n",
        "!wget https://github.com/nihui/waifu2x-ncnn-vulkan/releases/download/20200606/waifu2x-ncnn-vulkan-20200606-linux.zip\n",
        "!unzip waifu2x-ncnn-vulkan-20200606-linux.zip\n",
        "%cd waifu2x-ncnn-vulkan-20200606-linux\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ONHKknp7D1XY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# upload all files first ->\n",
        "%mkdir 2x\n",
        "!for img in *.??g; do ./waifu2x-ncnn-vulkan -i $img -o 2x/${img%.*}_2x.png; done\n",
        "\n",
        "# copying in gdrive\n",
        "#!gsutil -m cp -r hand_tuned_larger/ '/content/drive/My Drive/twist_moe/hand_tuned_larger_2x/'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nOvrhlKDpnPO",
        "colab_type": "text"
      },
      "source": [
        "# Cleaning\n",
        "\n",
        "After obtaining the cropped images, I run the shell scripts in this Git repository, courtesy of Gwern. The one change I made is to preserve the JPGs at 100% quality, as I have a small dataset.\n",
        "\n",
        "After changing the directory parameter, run them in the following order:\n",
        "\n",
        "delete -> convert -> resize -> final\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xCM1wNepqHf5",
        "colab_type": "text"
      },
      "source": [
        "I downloaded *all* images of Hayasaka from Pixiv. Unlike Danbooru, Pixiv does not have proper image tags. So, to separate images of Hayasaka from images we don't want (black & white, other characters), I devised the following script, which calculates distance of the dominant color in thte image from RGB yellow (255, 255, 0)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xAx-JARVqiWG",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from PIL import Image\n",
        "import os \n",
        "from time import sleep\n",
        "from colorthief import ColorThief\n",
        "\n",
        "for root, dirs, files in os.walk('/home/tazik/Nextcloud/code/lbpcascade_animeface/examples/datasets/hand_tuned_larger_2x/'):\n",
        "    #print(root, dirs, files)\n",
        "    for name in files:\n",
        "        print(name)\n",
        "        color_thief = ColorThief(os.path.join(root, name))\n",
        "        palette = color_thief.get_palette(color_count=2, quality=1)\n",
        "        rgb = palette[0]\n",
        "        delta_E = pow(rgb[0]-255, 2) + pow(rgb[1]-255,2) + pow(rgb[2], 2)\n",
        "        print(delta_E)\n",
        "        new_name = os.path.join(root, str(delta_E) + \".png\")\n",
        "        os.rename(os.path.join(root, name), new_name)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DqTxN576rBzt",
        "colab_type": "text"
      },
      "source": [
        "The script renames images according to \"yellowness\", making it easy to eliminate non-matching images.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yk6ezENIrQvO",
        "colab_type": "text"
      },
      "source": [
        "Now, it's time to use Colab. We use the pre-trained anime face StyleGan2 model to rank the our pre-processed images. This helps with filtering, as higher ranked images tend to be lower quality. This trick is also courtesy of Gwern."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fimqSTw_CSNX",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#CSDN blog\n",
        "#https://translate.googleusercontent.com/translate_c?depth=1&pto=aue&rurl=translate.google.com&sl=auto&sp=nmt4&tl=en&u=https://blog.csdn.net/DLW__/article/details/104222546&usg=ALkJrhjWEtjIz8Yklx8uSjuFQuv7O9bPnA\n",
        "\n",
        "import os\n",
        "import pickle\n",
        "import numpy as np\n",
        "import PIL.Image\n",
        "import dnnlib\n",
        "import dnnlib.tflib as tflib\n",
        "#import config\n",
        "import sys\n",
        "\n",
        "!pip install googledrivedownloader\n",
        "from google_drive_downloader import GoogleDriveDownloader as gdd\n",
        "\n",
        "import pretrained_networks\n",
        "\n",
        "# StyleGAN2 Danbooru Portrait\n",
        "url = 'https://drive.google.com/open?id=1WNQELgHnaqMTq3TlrnDaVkyrAH8Zrjez'\n",
        "#'https://drive.google.com/open?id=1BHeqOZ58WZ-vACR2MJkh1ZVbJK2B-Kle'\n",
        "model_id = url.replace('https://drive.google.com/open?id=', '')\n",
        "\n",
        "network_pkl = '/content/models/model_%s.pkl' % model_id#(hashlib.md5(model_id.encode()).hexdigest())\n",
        "gdd.download_file_from_google_drive(file_id=model_id,\n",
        "                                    dest_path=network_pkl)\n",
        "\n",
        "\n",
        "# If downloads fails, due to 'Google Drive download quota exceeded' you can try downloading manually from your own Google Drive account\n",
        "# network_pkl = \"/content/drive/My Drive/GAN/stylegan2-ffhq-config-f.pkl\"\n",
        "\n",
        "\n",
        "def main(origin_dir):\n",
        "    image_names = [files for root, dirs, files in os.walk(origin_dir)][0]\n",
        "    print('find %s files in %s' % (len(image_names), origin_dir))\n",
        "\n",
        "    tflib.init_tf()\n",
        "    print('Loading networks from \"%s\"...' % network_pkl)\n",
        "    _G, _D, Gs = pretrained_networks.load_networks(network_pkl)\n",
        "    noise_vars = [var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise')]\n",
        "\n",
        "    for index, image_name in enumerate(image_names):\n",
        "        image_path = os.path.join(origin_dir, image_name)\n",
        "        img = np.asarray(PIL.Image.open(image_path))\n",
        "        img = img.reshape(1, 3, 512, 512)\n",
        "        score = _D.run(img, None)\n",
        "        os.rename(image_path, os.path.join(origin_dir, '%s_%s.png' % (score[0][0], index)))\n",
        "        print(image_name, score[0][0])\n",
        "\n",
        "    print('Done!')\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    main('/content/drive/My Drive/twist_moe/hand_tuned/')\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VRbH4y4DtLM7",
        "colab_type": "text"
      },
      "source": [
        "Finally, I go through the images to look for potential outliers. I spent a bit of time on this step, as there were many low quality images of Hayasaka that I did not want the model to be learning from. This step could potentially be made redundant if one appropriately filters Pixiv images by art type. But the art categories are not immediately apparent to non-Pixiv users."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VR8n0yeImPCx",
        "colab_type": "text"
      },
      "source": [
        "# Training\n",
        "\n",
        "## Colab Hacks\n",
        "\n",
        "```\n",
        "i = []\n",
        "while True:\n",
        "  i.append(i)\n",
        "```\n",
        "\n",
        "The above is used to induce Google to offer you more RAM. Do note that this does not work on newly initialized notebooks, after a patch by Google; instead, you have to use an older notebook as your base (e.g. make a copy of this NB).\n",
        "\n",
        "Keep Colab from disconnecting after 1.5hrs.\n",
        "\n",
        "```\n",
        "\n",
        "function KeepClicking(){\n",
        "   console.log(\"Clicking\");\n",
        "   document.querySelector(\"colab-toolbar-button#connect\").click()\n",
        "}setInterval(KeepClicking,60000)\n",
        "```\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pE5GqZHcHHBd",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "\n",
        "%tensorflow_version 1.x\n",
        "import tensorflow as tf\n",
        "\n",
        "# Download the code\n",
        "!git clone https://github.com/ZKTKZ/stylegan2.git\n",
        "%cd stylegan2\n",
        "!nvcc test_nvcc.cu -o test_nvcc -run\n",
        "\n",
        "print('Tensorflow version: {}'.format(tf.__version__) )\n",
        "!nvidia-smi -L\n",
        "print('GPU Identified at: {}'.format(tf.test.gpu_device_name()))\n",
        "\n",
        "!pip install tensorboard\n",
        "\n",
        "url = 'https://drive.google.com/open?id=1WNQELgHnaqMTq3TlrnDaVkyrAH8Zrjez'\n",
        "#'https://drive.google.com/open?id=1BHeqOZ58WZ-vACR2MJkh1ZVbJK2B-Kle'\n",
        "model_id = url.replace('https://drive.google.com/open?id=', '')\n",
        "\n",
        "!pip install googledrivedownloader\n",
        "from google_drive_downloader import GoogleDriveDownloader as gdd\n",
        "\n",
        "network_pkl = './models/model_%s.pkl' % model_id#(hashlib.md5(model_id.encode()).hexdigest())\n",
        "gdd.download_file_from_google_drive(file_id=model_id,\n",
        "                                    dest_path=network_pkl)\n",
        "\n",
        "#!python dataset_tool.py create_from_images ./dataset/hayasaka /content/drive/'My Drive'/twist_moe/hand_tuned_larger_2x/"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vmeWpLdZdhq2",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "f4d02ffc-2f95-492a-fa58-fa395b9da517"
      },
      "source": [
        "!python run_training.py --spatial-augmentations=true --lr=0.0005 --num-gpus=1 --data-dir=./dataset --config=config-f --dataset=hayasaka --mirror-augment=true --metric=none --total-kimg=10000 --min-h=4 --min-w=4 --res-log2=7 --result-dir=\"/content/drive/My Drive/twist_moe/results/\" --resume-pkl='./models/model_1WNQELgHnaqMTq3TlrnDaVkyrAH8Zrjez.pkl'\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Local submit - run_dir: /content/drive/My Drive/twist_moe/results/00006-stylegan2-hayasaka-1gpu-config-f\n",
            "dnnlib: Running training.training_loop.training_loop() on localhost...\n",
            "Streaming data using training.dataset.TFRecordDataset...\n",
            "Dataset shape = [3, 512, 512]\n",
            "Dynamic range = [0, 255]\n",
            "Label size    = 0\n",
            "Loading networks from \"./models/model_1WNQELgHnaqMTq3TlrnDaVkyrAH8Zrjez.pkl\"...\n",
            "Setting up TensorFlow plugin \"fused_bias_act.cu\": Preprocessing... Compiling... Loading... Done.\n",
            "Setting up TensorFlow plugin \"upfirdn_2d.cu\": Preprocessing... Compiling... Loading... Done.\n",
            "\n",
            "G                             Params    OutputShape         WeightShape     \n",
            "---                           ---       ---                 ---             \n",
            "latents_in                    -         (?, 512)            -               \n",
            "labels_in                     -         (?, 0)              -               \n",
            "lod                           -         ()                  -               \n",
            "dlatent_avg                   -         (512,)              -               \n",
            "G_mapping/latents_in          -         (?, 512)            -               \n",
            "G_mapping/labels_in           -         (?, 0)              -               \n",
            "G_mapping/Normalize           -         (?, 512)            -               \n",
            "G_mapping/Dense0              262656    (?, 512)            (512, 512)      \n",
            "G_mapping/Dense1              262656    (?, 512)            (512, 512)      \n",
            "G_mapping/Dense2              262656    (?, 512)            (512, 512)      \n",
            "G_mapping/Dense3              262656    (?, 512)            (512, 512)      \n",
            "G_mapping/Dense4              262656    (?, 512)            (512, 512)      \n",
            "G_mapping/Dense5              262656    (?, 512)            (512, 512)      \n",
            "G_mapping/Dense6              262656    (?, 512)            (512, 512)      \n",
            "G_mapping/Dense7              262656    (?, 512)            (512, 512)      \n",
            "G_mapping/Broadcast           -         (?, 16, 512)        -               \n",
            "G_mapping/dlatents_out        -         (?, 16, 512)        -               \n",
            "Truncation/Lerp               -         (?, 16, 512)        -               \n",
            "G_synthesis/dlatents_in       -         (?, 16, 512)        -               \n",
            "G_synthesis/4x4/Const         8192      (?, 512, 4, 4)      (1, 512, 4, 4)  \n",
            "G_synthesis/4x4/Conv          2622465   (?, 512, 4, 4)      (3, 3, 512, 512)\n",
            "G_synthesis/4x4/ToRGB         264195    (?, 3, 4, 4)        (1, 1, 512, 3)  \n",
            "G_synthesis/8x8/Conv0_up      2622465   (?, 512, 8, 8)      (3, 3, 512, 512)\n",
            "G_synthesis/8x8/Conv1         2622465   (?, 512, 8, 8)      (3, 3, 512, 512)\n",
            "G_synthesis/8x8/Upsample      -         (?, 3, 8, 8)        -               \n",
            "G_synthesis/8x8/ToRGB         264195    (?, 3, 8, 8)        (1, 1, 512, 3)  \n",
            "G_synthesis/16x16/Conv0_up    2622465   (?, 512, 16, 16)    (3, 3, 512, 512)\n",
            "G_synthesis/16x16/Conv1       2622465   (?, 512, 16, 16)    (3, 3, 512, 512)\n",
            "G_synthesis/16x16/Upsample    -         (?, 3, 16, 16)      -               \n",
            "G_synthesis/16x16/ToRGB       264195    (?, 3, 16, 16)      (1, 1, 512, 3)  \n",
            "G_synthesis/32x32/Conv0_up    2622465   (?, 512, 32, 32)    (3, 3, 512, 512)\n",
            "G_synthesis/32x32/Conv1       2622465   (?, 512, 32, 32)    (3, 3, 512, 512)\n",
            "G_synthesis/32x32/Upsample    -         (?, 3, 32, 32)      -               \n",
            "G_synthesis/32x32/ToRGB       264195    (?, 3, 32, 32)      (1, 1, 512, 3)  \n",
            "G_synthesis/64x64/Conv0_up    2622465   (?, 512, 64, 64)    (3, 3, 512, 512)\n",
            "G_synthesis/64x64/Conv1       2622465   (?, 512, 64, 64)    (3, 3, 512, 512)\n",
            "G_synthesis/64x64/Upsample    -         (?, 3, 64, 64)      -               \n",
            "G_synthesis/64x64/ToRGB       264195    (?, 3, 64, 64)      (1, 1, 512, 3)  \n",
            "G_synthesis/128x128/Conv0_up  1442561   (?, 256, 128, 128)  (3, 3, 512, 256)\n",
            "G_synthesis/128x128/Conv1     721409    (?, 256, 128, 128)  (3, 3, 256, 256)\n",
            "G_synthesis/128x128/Upsample  -         (?, 3, 128, 128)    -               \n",
            "G_synthesis/128x128/ToRGB     132099    (?, 3, 128, 128)    (1, 1, 256, 3)  \n",
            "G_synthesis/256x256/Conv0_up  426369    (?, 128, 256, 256)  (3, 3, 256, 128)\n",
            "G_synthesis/256x256/Conv1     213249    (?, 128, 256, 256)  (3, 3, 128, 128)\n",
            "G_synthesis/256x256/Upsample  -         (?, 3, 256, 256)    -               \n",
            "G_synthesis/256x256/ToRGB     66051     (?, 3, 256, 256)    (1, 1, 128, 3)  \n",
            "G_synthesis/512x512/Conv0_up  139457    (?, 64, 512, 512)   (3, 3, 128, 64) \n",
            "G_synthesis/512x512/Conv1     69761     (?, 64, 512, 512)   (3, 3, 64, 64)  \n",
            "G_synthesis/512x512/Upsample  -         (?, 3, 512, 512)    -               \n",
            "G_synthesis/512x512/ToRGB     33027     (?, 3, 512, 512)    (1, 1, 64, 3)   \n",
            "G_synthesis/images_out        -         (?, 3, 512, 512)    -               \n",
            "G_synthesis/noise0            -         (1, 1, 4, 4)        -               \n",
            "G_synthesis/noise1            -         (1, 1, 8, 8)        -               \n",
            "G_synthesis/noise2            -         (1, 1, 8, 8)        -               \n",
            "G_synthesis/noise3            -         (1, 1, 16, 16)      -               \n",
            "G_synthesis/noise4            -         (1, 1, 16, 16)      -               \n",
            "G_synthesis/noise5            -         (1, 1, 32, 32)      -               \n",
            "G_synthesis/noise6            -         (1, 1, 32, 32)      -               \n",
            "G_synthesis/noise7            -         (1, 1, 64, 64)      -               \n",
            "G_synthesis/noise8            -         (1, 1, 64, 64)      -               \n",
            "G_synthesis/noise9            -         (1, 1, 128, 128)    -               \n",
            "G_synthesis/noise10           -         (1, 1, 128, 128)    -               \n",
            "G_synthesis/noise11           -         (1, 1, 256, 256)    -               \n",
            "G_synthesis/noise12           -         (1, 1, 256, 256)    -               \n",
            "G_synthesis/noise13           -         (1, 1, 512, 512)    -               \n",
            "G_synthesis/noise14           -         (1, 1, 512, 512)    -               \n",
            "images_out                    -         (?, 3, 512, 512)    -               \n",
            "---                           ---       ---                 ---             \n",
            "Total                         30276583                                      \n",
            "\n",
            "\n",
            "D                    Params    OutputShape         WeightShape     \n",
            "---                  ---       ---                 ---             \n",
            "images_in            -         (?, 3, 512, 512)    -               \n",
            "labels_in            -         (?, 0)              -               \n",
            "512x512/FromRGB      256       (?, 64, 512, 512)   (1, 1, 3, 64)   \n",
            "512x512/Conv0        36928     (?, 64, 512, 512)   (3, 3, 64, 64)  \n",
            "512x512/Conv1_down   73856     (?, 128, 256, 256)  (3, 3, 64, 128) \n",
            "512x512/Skip         8192      (?, 128, 256, 256)  (1, 1, 64, 128) \n",
            "256x256/Conv0        147584    (?, 128, 256, 256)  (3, 3, 128, 128)\n",
            "256x256/Conv1_down   295168    (?, 256, 128, 128)  (3, 3, 128, 256)\n",
            "256x256/Skip         32768     (?, 256, 128, 128)  (1, 1, 128, 256)\n",
            "128x128/Conv0        590080    (?, 256, 128, 128)  (3, 3, 256, 256)\n",
            "128x128/Conv1_down   1180160   (?, 512, 64, 64)    (3, 3, 256, 512)\n",
            "128x128/Skip         131072    (?, 512, 64, 64)    (1, 1, 256, 512)\n",
            "64x64/Conv0          2359808   (?, 512, 64, 64)    (3, 3, 512, 512)\n",
            "64x64/Conv1_down     2359808   (?, 512, 32, 32)    (3, 3, 512, 512)\n",
            "64x64/Skip           262144    (?, 512, 32, 32)    (1, 1, 512, 512)\n",
            "32x32/Conv0          2359808   (?, 512, 32, 32)    (3, 3, 512, 512)\n",
            "32x32/Conv1_down     2359808   (?, 512, 16, 16)    (3, 3, 512, 512)\n",
            "32x32/Skip           262144    (?, 512, 16, 16)    (1, 1, 512, 512)\n",
            "16x16/Conv0          2359808   (?, 512, 16, 16)    (3, 3, 512, 512)\n",
            "16x16/Conv1_down     2359808   (?, 512, 8, 8)      (3, 3, 512, 512)\n",
            "16x16/Skip           262144    (?, 512, 8, 8)      (1, 1, 512, 512)\n",
            "8x8/Conv0            2359808   (?, 512, 8, 8)      (3, 3, 512, 512)\n",
            "8x8/Conv1_down       2359808   (?, 512, 4, 4)      (3, 3, 512, 512)\n",
            "8x8/Skip             262144    (?, 512, 4, 4)      (1, 1, 512, 512)\n",
            "4x4/MinibatchStddev  -         (?, 513, 4, 4)      -               \n",
            "4x4/Conv             2364416   (?, 512, 4, 4)      (3, 3, 513, 512)\n",
            "4x4/Dense0           4194816   (?, 512)            (8192, 512)     \n",
            "Output               513       (?, 1)              (512, 1)        \n",
            "scores_out           -         (?, 1)              -               \n",
            "---                  ---       ---                 ---             \n",
            "Total                28982849                                      \n",
            "\n",
            "Building TensorFlow graph...\n",
            "Initializing logs...\n",
            "Augmenting fakes and reals\n",
            "Augmentation alpha at default setting of 0.1 - change by setting SPATIAL_AUGS_ALPHA environment variable\n",
            "Training for 10000 kimg...\n",
            "\n",
            "tick 0     kimg 0.0      lod 0.00  minibatch 4    time 18s          sec/tick 17.7    sec/kimg 1106.15 maintenance 0.0    gpumem 6.3\n",
            "tick 1     kimg 6.0      lod 0.00  minibatch 4    time 31m 46s      sec/tick 1869.0  sec/kimg 311.49  maintenance 18.9   gpumem 6.3\n",
            "tick 2     kimg 12.0     lod 0.00  minibatch 4    time 1h 03m 01s   sec/tick 1868.1  sec/kimg 311.36  maintenance 6.8    gpumem 6.3\n",
            "tick 3     kimg 18.0     lod 0.00  minibatch 4    time 1h 34m 15s   sec/tick 1867.5  sec/kimg 311.26  maintenance 6.7    gpumem 6.3\n",
            "tick 4     kimg 24.0     lod 0.00  minibatch 4    time 2h 05m 28s   sec/tick 1867.0  sec/kimg 311.17  maintenance 6.7    gpumem 6.3\n",
            "tick 5     kimg 30.0     lod 0.00  minibatch 4    time 2h 36m 43s   sec/tick 1867.5  sec/kimg 311.25  maintenance 6.7    gpumem 6.3\n",
            "tick 6     kimg 36.0     lod 0.00  minibatch 4    time 3h 07m 56s   sec/tick 1867.1  sec/kimg 311.18  maintenance 6.7    gpumem 6.3\n",
            "tick 7     kimg 42.0     lod 0.00  minibatch 4    time 3h 39m 11s   sec/tick 1867.8  sec/kimg 311.30  maintenance 6.6    gpumem 6.3\n",
            "tick 8     kimg 48.0     lod 0.00  minibatch 4    time 4h 10m 25s   sec/tick 1867.8  sec/kimg 311.30  maintenance 6.6    gpumem 6.3\n",
            "tick 9     kimg 54.0     lod 0.00  minibatch 4    time 4h 41m 39s   sec/tick 1867.0  sec/kimg 311.16  maintenance 6.7    gpumem 6.3\n",
            "tick 10    kimg 60.0     lod 0.00  minibatch 4    time 5h 12m 53s   sec/tick 1867.6  sec/kimg 311.27  maintenance 6.6    gpumem 6.3\n",
            "tick 11    kimg 66.0     lod 0.00  minibatch 4    time 5h 44m 07s   sec/tick 1867.7  sec/kimg 311.28  maintenance 6.7    gpumem 6.3\n",
            "tick 12    kimg 72.0     lod 0.00  minibatch 4    time 6h 15m 22s   sec/tick 1867.6  sec/kimg 311.27  maintenance 6.9    gpumem 6.3\n",
            "tick 13    kimg 78.0     lod 0.00  minibatch 4    time 6h 46m 34s   sec/tick 1864.8  sec/kimg 310.79  maintenance 6.9    gpumem 6.3\n",
            "tick 14    kimg 84.0     lod 0.00  minibatch 4    time 7h 17m 51s   sec/tick 1870.4  sec/kimg 311.73  maintenance 6.9    gpumem 6.3\n",
            "tick 15    kimg 90.0     lod 0.00  minibatch 4    time 7h 49m 04s   sec/tick 1866.1  sec/kimg 311.01  maintenance 7.0    gpumem 6.3\n",
            "tick 16    kimg 96.0     lod 0.00  minibatch 4    time 8h 20m 19s   sec/tick 1867.8  sec/kimg 311.30  maintenance 6.8    gpumem 6.3\n",
            "tick 17    kimg 102.0    lod 0.00  minibatch 4    time 8h 51m 33s   sec/tick 1867.2  sec/kimg 311.19  maintenance 6.8    gpumem 6.3\n",
            "tick 18    kimg 108.0    lod 0.00  minibatch 4    time 9h 22m 44s   sec/tick 1864.2  sec/kimg 310.71  maintenance 6.8    gpumem 6.3\n",
            "tick 19    kimg 114.0    lod 0.00  minibatch 4    time 9h 53m 54s   sec/tick 1863.2  sec/kimg 310.54  maintenance 6.7    gpumem 6.3\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J2Ln57WdwKcu",
        "colab_type": "text"
      },
      "source": [
        "# Generation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0mLChQRkmgPP",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "!python run_generator.py generate-images --seeds=0-100 --truncation-psi=0.7 --network='/content/drive/My Drive/twist_moe/results/00006-stylegan2-hayasaka-1gpu-config-f/network-snapshot-000114.pkl'\n",
        "#%cp -av /content/stylegan2/results/00000-generate-images /content/drive/'My Drive'/twist_moe/seeds-1.0\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gXzkpzBqnGBK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import pickle\n",
        "import numpy as np\n",
        "import PIL.Image\n",
        "import dnnlib\n",
        "import dnnlib.tflib as tflib\n",
        "import scipy\n",
        "import math\n",
        "import moviepy.editor\n",
        "from numpy import linalg\n",
        "\n",
        "\n",
        "def main():\n",
        "    tflib.init_tf()\n",
        "    _G, _D, Gs = pickle.load(open(\"/content/drive/My Drive/twist_moe/results/00006-stylegan2-hayasaka-1gpu-config-f/network-snapshot-000114.pkl\", \"rb\"))\n",
        "\n",
        "    rnd = np.random\n",
        "    latents_a = rnd.randn(1, Gs.input_shape[1])\n",
        "    latents_b = rnd.randn(1, Gs.input_shape[1])\n",
        "    latents_c = rnd.randn(1, Gs.input_shape[1])\n",
        "\n",
        "    def circ_generator(latents_interpolate):\n",
        "        radius = 40.0\n",
        "\n",
        "        latents_axis_x = (latents_a - latents_b).flatten() / linalg.norm(latents_a - latents_b)\n",
        "        latents_axis_y = (latents_a - latents_c).flatten() / linalg.norm(latents_a - latents_c)\n",
        "\n",
        "        latents_x = math.sin(math.pi * 2.0 * latents_interpolate) * radius\n",
        "        latents_y = math.cos(math.pi * 2.0 * latents_interpolate) * radius\n",
        "\n",
        "        latents = latents_a + latents_x * latents_axis_x + latents_y * latents_axis_y\n",
        "        return latents\n",
        "\n",
        "    def mse(x, y):\n",
        "        return (np.square(x - y)).mean()\n",
        "\n",
        "    def generate_from_generator_adaptive(gen_func):\n",
        "        max_step = 1.0\n",
        "        current_pos = 0.0\n",
        "\n",
        "        change_min = 10.0\n",
        "        change_max = 11.0\n",
        "\n",
        "        fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)\n",
        "\n",
        "        current_latent = gen_func(current_pos)\n",
        "        current_image = Gs.run(current_latent, None, truncation_psi=0.7, randomize_noise=False, output_transform=fmt)[0]\n",
        "        array_list = []\n",
        "\n",
        "        video_length = 1.0\n",
        "        while(current_pos < video_length):\n",
        "            array_list.append(current_image)\n",
        "\n",
        "            lower = current_pos\n",
        "            upper = current_pos + max_step\n",
        "            current_pos = (upper + lower) / 2.0\n",
        "\n",
        "            current_latent = gen_func(current_pos)\n",
        "            current_image = images = Gs.run(current_latent, None, truncation_psi=0.7, randomize_noise=False, output_transform=fmt)[0]\n",
        "            current_mse = mse(array_list[-1], current_image)\n",
        "\n",
        "            while current_mse < change_min or current_mse > change_max:\n",
        "                if current_mse < change_min:\n",
        "                    lower = current_pos\n",
        "                    current_pos = (upper + lower) / 2.0\n",
        "\n",
        "                if current_mse > change_max:\n",
        "                    upper = current_pos\n",
        "                    current_pos = (upper + lower) / 2.0\n",
        "\n",
        "\n",
        "                current_latent = gen_func(current_pos)\n",
        "                current_image = images = Gs.run(current_latent, None, truncation_psi=0.7, randomize_noise=False, output_transform=fmt)[0]\n",
        "                current_mse = mse(array_list[-1], current_image)\n",
        "            print(current_pos, current_mse)\n",
        "        return array_list\n",
        "\n",
        "    frames = generate_from_generator_adaptive(circ_generator)\n",
        "    frames = moviepy.editor.ImageSequenceClip(frames, fps=30)\n",
        "\n",
        "    # Generate video.\n",
        "    mp4_file = 'circular.mp4'\n",
        "    mp4_codec = 'libx264'\n",
        "    mp4_bitrate = '3M'\n",
        "    mp4_fps = 20\n",
        "\n",
        "    frames.write_videofile(mp4_file, fps=mp4_fps, codec=mp4_codec, bitrate=mp4_bitrate)\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    main()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3UMbcw83Z9VF",
        "colab_type": "text"
      },
      "source": [
        "## Video"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "f-HTjbmIDseN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#https://stackoverflow.com/a/57378660/8773953\n",
        "!pip install -U kora\n",
        "from kora.drive import upload_public\n",
        "url = upload_public('/content/drive/My Drive/twist_moe/videos/circular.mp4')\n",
        "\n",
        "from IPython.display import HTML\n",
        "HTML(f\"\"\"<video src={url} width=500 controls/>\"\"\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "974x9801ZisW",
        "colab_type": "text"
      },
      "source": [
        "## Comments\n",
        "\n",
        "The above is after training for about ~9 hours on Colab, and the interpolation is already pretty good!"
      ]
    }
  ]
}