Convert_pkl_to_pt.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Convert-pkl-to-pt.ipynb",
      "provenance": [],
      "machine_shape": "hm",
      "authorship_tag": "ABX9TyOjZ0guu9ebnvn1anWNAM1u",
      "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/ai/blob/master/Convert_pkl_to_pt.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KLTTMlI6UJZT",
        "colab_type": "text"
      },
      "source": [
        "# Convert StyleGAN2 .pkl model to .pt\n",
        "\n",
        "Many StyleGAN tools (outside of the official Tensorflow-based NVIDIA library) use a PyTorch fork of StyleGAN. This requires we convert a .pkl model file to a .pt file.\n",
        "\n",
        "This notebook shows you the steps to do so. I also have a [video](https://youtu.be/ASn2DcYLUEE) describing this process."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KUoPW4rXUO9k",
        "colab_type": "text"
      },
      "source": [
        "## Setup\n",
        "\n",
        "First, we’ll download both the Tensorflow and PyTorch versions of StyleGAN2. We also need to make sure Colab is running Tensorflow 1.15."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WpHI1-QxF15O",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "f9fa977b-a994-49d2-b0f4-5b1236545eb7"
      },
      "source": [
        "%tensorflow_version 1.x"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "TensorFlow 1.x selected.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5wd3zk15G9gl",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 360
        },
        "outputId": "1c629eca-447f-41ce-8f1d-01f5765c4852"
      },
      "source": [
        "!git clone https://github.com/rosinality/stylegan2-pytorch\n",
        "!git clone https://github.com/dvschultz/stylegan2\n",
        "!pip install ninja #required dependency for stylegan2-pytorch"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Cloning into 'stylegan2-pytorch'...\n",
            "remote: Enumerating objects: 66, done.\u001b[K\n",
            "remote: Counting objects: 100% (66/66), done.\u001b[K\n",
            "remote: Compressing objects: 100% (48/48), done.\u001b[K\n",
            "remote: Total 297 (delta 31), reused 46 (delta 18), pack-reused 231\u001b[K\n",
            "Receiving objects: 100% (297/297), 119.33 MiB | 49.00 MiB/s, done.\n",
            "Resolving deltas: 100% (133/133), done.\n",
            "Cloning into 'stylegan2'...\n",
            "remote: Enumerating objects: 24, done.\u001b[K\n",
            "remote: Counting objects: 100% (24/24), done.\u001b[K\n",
            "remote: Compressing objects: 100% (24/24), done.\u001b[K\n",
            "remote: Total 447 (delta 13), reused 0 (delta 0), pack-reused 423\u001b[K\n",
            "Receiving objects: 100% (447/447), 15.33 MiB | 38.65 MiB/s, done.\n",
            "Resolving deltas: 100% (252/252), done.\n",
            "Collecting ninja\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/f6/27/41b8a846496451ea9f103bf8e1d0d1b010d1180c30d84a87433ec5d7f88e/ninja-1.10.0.post1-py3-none-manylinux1_x86_64.whl (102kB)\n",
            "\u001b[K     |████████████████████████████████| 112kB 4.8MB/s \n",
            "\u001b[?25hInstalling collected packages: ninja\n",
            "Successfully installed ninja-1.10.0.post1\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8HZ3iOOaURe3",
        "colab_type": "text"
      },
      "source": [
        "## Add .pkl\n",
        "\n",
        "Next we need to provide the .pkl file for converting. You can upload it directly to Colab or use gdown if you have the file on Google Drive."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bRGem5kiOaVa",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 85
        },
        "outputId": "fe155f3c-33ff-4658-d1bf-6c69d564ef57"
      },
      "source": [
        "!gdown --id 1UlDmJVLLnBD9SnLSMXeiZRO6g-OMQCA_"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Downloading...\n",
            "From: https://drive.google.com/uc?id=1UlDmJVLLnBD9SnLSMXeiZRO6g-OMQCA_\n",
            "To: /content/stylegan2-ffhq-config-f.pkl\n",
            "382MB [00:03, 124MB/s]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xzo4Rhy8Ua9l",
        "colab_type": "text"
      },
      "source": [
        "## Convert file\n",
        "Lastly we need to give the conversion script the path to the .pkl file and then run the convesion script."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8l_cBk5vOvG2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "pkl_file = \"/content/stylegan2-ffhq-config-f.pkl\" # copy and paste path from sidebar"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9xeaf1jxL8Z4",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 68
        },
        "outputId": "9e517f7f-341a-4012-b188-e68bd990b107"
      },
      "source": [
        "!python /content/stylegan2-pytorch/convert_weight.py --repo ./stylegan2 $pkl_file"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "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",
            "tensor(4.9233e-05, device='cuda:0')\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EbC9p8-fVu4z",
        "colab_type": "text"
      },
      "source": [
        "You should now see a .pt file with the same name as your .pkl file in the Files sideebar. There is also a .png that you can look at to make sure the generated images from the .pt modeel look ok.\n",
        "\n",
        "I recommend you download the .pt file or move it to Google Drive for re-use."
      ]
    }
  ]
}