StyleGAN2_latent_discovery.ipynb

{
    "nbformat": 4,
    "nbformat_minor": 0,
    "metadata": {
        "colab": {
            "name": "StyleGAN2_latent_directions.ipynb",
            "provenance": [],
            "collapsed_sections": [
                "1bQeqOve7IJu",
                "2T2N0CgV7BHO",
                "JXq82s4lMHx1",
                "00WhMaqCMME-",
                "37fnSs3HlcYr",
                "j76sPQFl_-bm"
            ]
        },
        "kernelspec": {
            "name": "python3",
            "display_name": "Python 3"
        },
        "accelerator": "GPU"
    },
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "bSH1WWCJSzBR",
                "colab_type": "text"
            },
            "source": [
                "# StyleGAN2: unsupervised discovery of latent directions\n",
                "\n",
                "Reference: https://github.com/woctezuma/steam-stylegan2\n"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "annCW8EwTR7D",
                "colab_type": "text"
            },
            "source": [
                "## Papers\n",
                "\n",
                "[1] [Voynov, Andrey, et al. *Unsupervised Discovery of Interpretable Directions in Latent Space*. ICML 2020.][GanLatentDiscovery-paper] **([code][GanLatentDiscovery])**\n",
                "\n",
                "[GanLatentDiscovery-paper]: <https://arxiv.org/abs/2002.03754>\n",
                "[GanLatentDiscovery]: <https://github.com/anvoynov/GanLatentDiscovery>\n",
                "\n",
                "[2] [Härkönen, Erik, et al. *GANSpace: Discovering Interpretable GAN Controls*. arXiv 2020.][ganspace-paper] **([code][ganspace])**\n",
                "\n",
                "[ganspace-paper]: <https://arxiv.org/abs/2004.02546>\n",
                "[ganspace]: <https://github.com/harskish/ganspace>\n",
                "\n",
                "[3] [Shen, Yujun, et al. *Closed-Form Factorization of Latent Semantics in GANs*. arXiv 2020.][closed-form-paper] **([code][closed-form])**\n",
                "\n",
                "[closed-form-paper]: <https://arxiv.org/abs/2007.06600>\n",
                "[closed-form]: <https://github.com/rosinality/stylegan2-pytorch#closed-form-factorization-httpsarxivorgabs200706600>"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "rQRVqo2c2grH",
                "colab_type": "text"
            },
            "source": [
                "### Issues"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "IJHYe2MD8bpk",
                "colab_type": "text"
            },
            "source": [
                "NB: code [1] cannot be run in a Google Colab session, due to [a memory issue](https://github.com/anvoynov/GANLatentDiscovery/issues/13):\n",
                "```\n",
                "RuntimeError: CUDA out of memory. Tried to allocate 512.00 MiB (GPU 0; 14.73 \n",
                "GiB total capacity; 13.20 GiB already allocated; 499.88 MiB free; 13.36 GiB \n",
                "reserved in total by PyTorch)\n",
                "```"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "nuBxfOym1pzE",
                "colab_type": "text"
            },
            "source": [
                "NB: code [2] [can be run](https://github.com/harskish/ganspace/issues/4) with `config-e` of StyleGAN2, thanks to [my `steam` fork](https://github.com/woctezuma/ganspace/tree/steam) which includes the most obvious fixes.\n",
                "\n",
                "If you try to follow [this notebook](https://colab.research.google.com/drive/1g-ShMzkRWDMHPyjom_p-5kqkn2f-GwBi) with my fork, you will encounter this error, but only if you forget to specify a layer to explore:\n",
                "```\n",
                "Traceback (most recent call last):\n",
                "  File \"visualize.py\", line 152, in <module>\n",
                "    inst = get_instrumented_model(args.model, args.output_class, layer_key, device, use_w=args.use_w)\n",
                "  File \"/usr/lib/python3.6/functools.py\", line 807, in wrapper\n",
                "    return dispatch(args[0].__class__)(*args, **kw)\n",
                "  File \"/content/ganspace/models/wrappers.py\", line 726, in get_instrumented_model\n",
                "    latent_shape = model.get_latent_shape()\n",
                "  File \"/content/ganspace/netdissect/modelconfig.py\", line 98, in create_instrumented_model\n",
                "    model.retain_layers(args.layers)\n",
                "  File \"/content/ganspace/netdissect/nethook.py\", line 63, in retain_layers\n",
                "    self.add_hooks(layernames)\n",
                "  File \"/content/ganspace/netdissect/nethook.py\", line 164, in add_hooks\n",
                "    raise ValueError('Layer %s not found in model' % name)\n",
                "ValueError: Layer g_mapping not found in model\n",
                "```"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "g6stBQOMUFJq",
                "colab_type": "text"
            },
            "source": [
                "## Requirements"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "1bQeqOve7IJu",
                "colab_type": "text"
            },
            "source": [
                "### Check the GPU"
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "jY7Iz8-i7Jbm",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/"
                },
                "outputId": "e53d637e-a10e-4c1c-a4d6-563892fbc4bd"
            },
            "source": [
                "!nvidia-smi -L"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "GPU 0: Tesla T4 (UUID: GPU-eb898157-ff44-253a-adc1-105e488192ac)\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "2T2N0CgV7BHO",
                "colab_type": "text"
            },
            "source": [
                "### Switch to Tensorflow 1.x"
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "nmp8hmPn7DWd",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/"
                },
                "outputId": "f5e3162e-b5c3-4d49-9e7b-7b305c360dd5"
            },
            "source": [
                "%tensorflow_version 1.x"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "TensorFlow 1.x selected.\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "PAA45hjIX9LV",
                "colab_type": "text"
            },
            "source": [
                "### Mount Google Drive"
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "CvlF79b_XFTs",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 90
                },
                "outputId": "62cce766-0371-4325-864d-90c4c22c87af"
            },
            "source": [
                "%pip install Google-Colab-Transfer"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "Collecting Google-Colab-Transfer\n",
                        "  Downloading https://files.pythonhosted.org/packages/67/36/5ea5df79964e0a860594075ee68f6c3ea7d8f3a5826a0a01ee0f0a63a768/Google_Colab_Transfer-0.1.5-py3-none-any.whl\n",
                        "Installing collected packages: Google-Colab-Transfer\n",
                        "Successfully installed Google-Colab-Transfer-0.1.5\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "PjhexFnQYCwe",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 35
                },
                "outputId": "72d79b12-d47b-4799-8261-95a82ea738b3"
            },
            "source": [
                "import colab_transfer\n",
                "\n",
                "colab_transfer.mount_google_drive()"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "Mounted at /content/drive/\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "YW_s6uLiKZIG",
                "colab_type": "text"
            },
            "source": [
                "## Load the last checkpoint from Google Drive"
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "y9J-pL1ysraC",
                "colab_type": "code",
                "colab": {}
            },
            "source": [
                "debug_with_ffhq = True\n",
                "use_config_e = False"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "RvK3XEAxLged",
                "colab_type": "code",
                "colab": {}
            },
            "source": [
                "if debug_with_ffhq:\n",
                "\n",
                "  if use_config_e:\n",
                "    tensorflow_model_name = 'stylegan2-ffhq-config-e.pkl'\n",
                "  else:\n",
                "    tensorflow_model_name = 'stylegan2-ffhq-config-f.pkl'\n",
                "    \n",
                "else:\n",
                "  tensorflow_model_name = 'network-snapshot-005000.pkl'\n",
                "\n",
                "pytorch_model_name = tensorflow_model_name.replace('.pkl', '.pt')"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "HbPNBx5zYF8x",
                "colab_type": "text"
            },
            "source": [
                "### Import the Tensorflow checkpoint from Google Drive"
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "SksyFM6h0W4c",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 35
                },
                "outputId": "5e86f262-46e8-425d-982e-01581de1b4bc"
            },
            "source": [
                "%cd /content/"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "/content\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "qevAQB-tXUKb",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 219
                },
                "outputId": "f9af3d68-951c-4844-8ed3-084f3563d079"
            },
            "source": [
                "if debug_with_ffhq:\n",
                "  # Reference: https://github.com/NVlabs/stylegan2/blob/master/pretrained_networks.py  \n",
                "  !wget http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/$tensorflow_model_name\n",
                "else:\n",
                "  colab_transfer.copy_file(file_name=tensorflow_model_name,\n",
                "                          source=colab_transfer.get_path_to_home_of_google_drive() + 'checkpoints/',\n",
                "                          destination=colab_transfer.get_path_to_home_of_local_machine())"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "--2020-09-20 14:38:54--  http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/stylegan2-ffhq-config-e.pkl\n",
                        "Resolving d36zk2xti64re0.cloudfront.net (d36zk2xti64re0.cloudfront.net)... 13.224.10.24, 13.224.10.81, 13.224.10.190, ...\n",
                        "Connecting to d36zk2xti64re0.cloudfront.net (d36zk2xti64re0.cloudfront.net)|13.224.10.24|:80... connected.\n",
                        "HTTP request sent, awaiting response... 200 OK\n",
                        "Length: 318021648 (303M) [application/x-www-form-urlencoded]\n",
                        "Saving to: ‘stylegan2-ffhq-config-e.pkl.1’\n",
                        "\n",
                        "stylegan2-ffhq-conf 100%[===================>] 303.29M   300MB/s    in 1.0s    \n",
                        "\n",
                        "2020-09-20 14:38:56 (300 MB/s) - ‘stylegan2-ffhq-config-e.pkl.1’ saved [318021648/318021648]\n",
                        "\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "vwzHk8hyYjpx",
                "colab_type": "text"
            },
            "source": [
                "### Convert the checkpoint from TensorFlow to PyTorch\n",
                "\n",
                "Reference: https://github.com/rosinality/stylegan2-pytorch"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "JXq82s4lMHx1",
                "colab_type": "text"
            },
            "source": [
                "#### Clone StyleGAN2 (TensorFlow)"
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "ioMPEAKDcGKS",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/"
                },
                "outputId": "fd328d0f-ca91-47fd-a968-1b90b66260c4"
            },
            "source": [
                "%cd /content/\n",
                "!git clone https://github.com/NVlabs/stylegan2"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "/content\n",
                        "Cloning into 'stylegan2'...\n",
                        "remote: Enumerating objects: 93, done.\u001b[K\n",
                        "remote: Total 93 (delta 0), reused 0 (delta 0), pack-reused 93\u001b[K\n",
                        "Unpacking objects: 100% (93/93), done.\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "00WhMaqCMME-",
                "colab_type": "text"
            },
            "source": [
                "#### Clone StyleGAN2 (PyTorch)"
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "mBGg1uzJY2kp",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 109
                },
                "outputId": "151ffc17-3096-42a0-b27a-54ea0a4ee894"
            },
            "source": [
                "%pip install ninja"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "Collecting ninja\n",
                        "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/1d/de/393468f2a37fc2c1dc3a06afc37775e27fde2d16845424141d4da62c686d/ninja-1.10.0.post2-py3-none-manylinux1_x86_64.whl (107kB)\n",
                        "\r\u001b[K     |███                             | 10kB 27.2MB/s eta 0:00:01\r\u001b[K     |██████                          | 20kB 4.2MB/s eta 0:00:01\r\u001b[K     |█████████▏                      | 30kB 5.2MB/s eta 0:00:01\r\u001b[K     |████████████▏                   | 40kB 6.0MB/s eta 0:00:01\r\u001b[K     |███████████████▎                | 51kB 4.8MB/s eta 0:00:01\r\u001b[K     |██████████████████▎             | 61kB 5.3MB/s eta 0:00:01\r\u001b[K     |█████████████████████▍          | 71kB 5.8MB/s eta 0:00:01\r\u001b[K     |████████████████████████▍       | 81kB 6.4MB/s eta 0:00:01\r\u001b[K     |███████████████████████████▍    | 92kB 6.7MB/s eta 0:00:01\r\u001b[K     |██████████████████████████████▌ | 102kB 6.6MB/s eta 0:00:01\r\u001b[K     |████████████████████████████████| 112kB 6.6MB/s \n",
                        "\u001b[?25hInstalling collected packages: ninja\n",
                        "Successfully installed ninja-1.10.0.post2\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "BeDWLeL6YhiQ",
                "colab_type": "code",
                "colab": {}
            },
            "source": [
                "# %cd /content/\n",
                "# !git clone https://github.com/rosinality/stylegan2-pytorch.git"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "CfrQYtpG840t",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 164
                },
                "outputId": "fcabce8c-2b8f-4c85-9483-a2ccf72e7e31"
            },
            "source": [
                "%cd /content/\n",
                "!git clone https://github.com/woctezuma/stylegan2-pytorch.git"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "/content\n",
                        "Cloning into 'stylegan2-pytorch'...\n",
                        "remote: Enumerating objects: 67, done.\u001b[K\n",
                        "remote: Counting objects: 100% (67/67), done.\u001b[K\n",
                        "remote: Compressing objects: 100% (41/41), done.\u001b[K\n",
                        "remote: Total 379 (delta 40), reused 44 (delta 26), pack-reused 312\u001b[K\n",
                        "Receiving objects: 100% (379/379), 122.50 MiB | 42.94 MiB/s, done.\n",
                        "Resolving deltas: 100% (189/189), done.\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "glY9nZGc8PLY",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 90
                },
                "outputId": "fdf54074-9e6b-477b-bccb-6bc015013ae7"
            },
            "source": [
                "%cd /content/stylegan2-pytorch/\n",
                "!git pull\n",
                "!git checkout quality-of-life-merge"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "/content/stylegan2-pytorch\n",
                        "Already up to date.\n",
                        "Branch 'quality-of-life-merge' set up to track remote branch 'quality-of-life-merge' from 'origin'.\n",
                        "Switched to a new branch 'quality-of-life-merge'\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "reZe2SIPMRNv",
                "colab_type": "text"
            },
            "source": [
                "#### Convert"
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "mAq1R7RlPk1r",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 35
                },
                "outputId": "d07f8c40-d149-42d0-e2ab-0bb4152a81eb"
            },
            "source": [
                "%cd /content/stylegan2-pytorch/"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "/content/stylegan2-pytorch\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "Sxzy0t5ZhIfn",
                "colab_type": "text"
            },
            "source": [
                "For `config-e`, set `--channel_multiplier` to 1.\n",
                "\n",
                "Reference: https://github.com/rosinality/stylegan2-pytorch/issues/52"
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "BrMAvkyE6b0Z",
                "colab_type": "code",
                "colab": {}
            },
            "source": [
                "if 'config-f' in tensorflow_model_name:\n",
                "  channel_multiplier = 2\n",
                "else:\n",
                "  channel_multiplier = 1"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "LI5cbqxpZ1V3",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 90
                },
                "outputId": "4a4cdd0f-6536-4e27-a53a-ac372c70f90d"
            },
            "source": [
                "if debug_with_ffhq:\n",
                "  !python convert_weight.py \\\n",
                "  --repo /content/stylegan2/ \\\n",
                "  --channel_multiplier $channel_multiplier \\\n",
                "  /content/$tensorflow_model_name"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "Setting up TensorFlow plugin \"fused_bias_act.cu\": Preprocessing... Loading... Done.\n",
                        "Setting up TensorFlow plugin \"upfirdn_2d.cu\": Preprocessing... Loading... Done.\n",
                        "2020-09-20 14:39:30.427629: W tensorflow/core/framework/cpu_allocator_impl.cc:81] Allocation of 16777216 exceeds 10% of system memory.\n",
                        "tensor(6.5267e-05, device='cuda:0')\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "37fnSs3HlcYr",
                "colab_type": "text"
            },
            "source": [
                "#### Export the converted checkpoint to Google Drive"
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "anrT9W-vktoO",
                "colab_type": "code",
                "colab": {}
            },
            "source": [
                "# if not debug_with_ffhq:\n",
                "#   colab_transfer.copy_file(file_name=pytorch_model_name,\n",
                "#                           source=colab_transfer.get_path_to_home_of_local_machine() + 'stylegan2-pytorch/',\n",
                "#                           destination=colab_transfer.get_path_to_home_of_google_drive() + 'checkpoints/pytorch/')"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "fwshRAJIluDn",
                "colab_type": "text"
            },
            "source": [
                "#### Import the PyTorch checkpoint from Google Drive"
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "q-ixS0DGlvZN",
                "colab_type": "code",
                "colab": {}
            },
            "source": [
                "if not debug_with_ffhq:\n",
                "  colab_transfer.copy_file(file_name=pytorch_model_name,\n",
                "                          source=colab_transfer.get_path_to_home_of_google_drive() + 'checkpoints/pytorch/',\n",
                "                          destination=colab_transfer.get_path_to_home_of_local_machine() + 'stylegan2-pytorch/')"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "LeRVojARNP1A",
                "colab_type": "text"
            },
            "source": [
                "## Check the checkpoints"
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "VCZo0FT-kdiC",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 54
                },
                "outputId": "876b66f3-94cc-4339-d27f-d3491d6fc245"
            },
            "source": [
                "# TensorFlow\n",
                "!du -sh /content/$tensorflow_model_name\n",
                "\n",
                "# PyTorch\n",
                "!du -sh /content/stylegan2-pytorch/$pytorch_model_name"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "304M\t/content/stylegan2-ffhq-config-e.pkl\n",
                        "106M\t/content/stylegan2-pytorch/stylegan2-ffhq-config-e.pt\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "j76sPQFl_-bm",
                "colab_type": "text"
            },
            "source": [
                "## Show the average image\n",
                "\n",
                "Reference: https://gist.github.com/woctezuma/139cedb92a94c5ef2675cc9f06851b31"
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "F5EOiGM5Dg6a",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 35
                },
                "outputId": "f8bd1c73-0a23-44b7-bdf4-7ffca15a0e17"
            },
            "source": [
                "%cd /content/"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "/content\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "ByCNqT80DiGS",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 421
                },
                "outputId": "893e5280-3ab3-458a-de87-3cc91877e900"
            },
            "source": [
                "from pathlib import Path\n",
                "\n",
                "ffhq_model_names = [\n",
                "                  'stylegan2-ffhq-config-e.pkl',\n",
                "                  'stylegan2-ffhq-config-f.pkl',\n",
                "                  ]\n",
                "\n",
                "for model_name in ffhq_model_names:\n",
                "  if not Path('/content/' + model_name).exists():\n",
                "    url = 'http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/' + model_name\n",
                "    !wget $url"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "--2020-09-21 13:53:46--  http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/stylegan2-ffhq-config-e.pkl\n",
                        "Resolving d36zk2xti64re0.cloudfront.net (d36zk2xti64re0.cloudfront.net)... 13.224.10.190, 13.224.10.184, 13.224.10.81, ...\n",
                        "Connecting to d36zk2xti64re0.cloudfront.net (d36zk2xti64re0.cloudfront.net)|13.224.10.190|:80... connected.\n",
                        "HTTP request sent, awaiting response... 200 OK\n",
                        "Length: 318021648 (303M) [application/x-www-form-urlencoded]\n",
                        "Saving to: ‘stylegan2-ffhq-config-e.pkl’\n",
                        "\n",
                        "stylegan2-ffhq-conf 100%[===================>] 303.29M  16.2MB/s    in 20s     \n",
                        "\n",
                        "2020-09-21 13:54:07 (15.2 MB/s) - ‘stylegan2-ffhq-config-e.pkl’ saved [318021648/318021648]\n",
                        "\n",
                        "--2020-09-21 13:54:07--  http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/stylegan2-ffhq-config-f.pkl\n",
                        "Resolving d36zk2xti64re0.cloudfront.net (d36zk2xti64re0.cloudfront.net)... 13.224.10.81, 13.224.10.190, 13.224.10.24, ...\n",
                        "Connecting to d36zk2xti64re0.cloudfront.net (d36zk2xti64re0.cloudfront.net)|13.224.10.81|:80... connected.\n",
                        "HTTP request sent, awaiting response... 200 OK\n",
                        "Length: 381673535 (364M) [application/x-www-form-urlencoded]\n",
                        "Saving to: ‘stylegan2-ffhq-config-f.pkl’\n",
                        "\n",
                        "stylegan2-ffhq-conf 100%[===================>] 363.99M   279MB/s    in 1.3s    \n",
                        "\n",
                        "2020-09-21 13:54:10 (279 MB/s) - ‘stylegan2-ffhq-config-f.pkl’ saved [381673535/381673535]\n",
                        "\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "IwosbNJTDBIK",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 109
                },
                "outputId": "6e0e62c8-1415-4958-ce05-464bd4772c5e"
            },
            "source": [
                "model_names = [\n",
                "                'network-snapshot-001014.pkl',\n",
                "                'network-snapshot-002003.pkl',\n",
                "                'network-snapshot-003001.pkl',\n",
                "                'network-snapshot-004006.pkl',\n",
                "                'network-snapshot-005000.pkl',\n",
                "               ]\n",
                "\n",
                "for model_name in model_names:\n",
                "  colab_transfer.copy_file(file_name=model_name,\n",
                "                          source=colab_transfer.get_path_to_home_of_google_drive() + 'checkpoints/',\n",
                "                          destination=colab_transfer.get_path_to_home_of_local_machine())"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "Copying /content/drive/My Drive/checkpoints/network-snapshot-001014.pkl to /content/network-snapshot-001014.pkl\n",
                        "Copying /content/drive/My Drive/checkpoints/network-snapshot-002003.pkl to /content/network-snapshot-002003.pkl\n",
                        "Copying /content/drive/My Drive/checkpoints/network-snapshot-003001.pkl to /content/network-snapshot-003001.pkl\n",
                        "Copying /content/drive/My Drive/checkpoints/network-snapshot-004006.pkl to /content/network-snapshot-004006.pkl\n",
                        "Copying /content/drive/My Drive/checkpoints/network-snapshot-005000.pkl to /content/network-snapshot-005000.pkl\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "iwF9w3OS8mjB",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 109
                },
                "outputId": "36563a98-e0bb-4bd4-f03b-a944fa31ce44"
            },
            "source": [
                "network_pkl_list = []\n",
                "network_pkl_list += [\n",
                "                     '/content/' + model_name\n",
                "                     for model_name in sorted(model_names)\n",
                "                     ]\n",
                "\n",
                "print('\\n'.join(network_pkl_list))                     "
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "/content/network-snapshot-001014.pkl\n",
                        "/content/network-snapshot-002003.pkl\n",
                        "/content/network-snapshot-003001.pkl\n",
                        "/content/network-snapshot-004006.pkl\n",
                        "/content/network-snapshot-005000.pkl\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "foN0N1ZVBV8N",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 54
                },
                "outputId": "98ee4001-19f0-4cf7-aca4-728eeeddae8e"
            },
            "source": [
                "network_pkl_list = []\n",
                "network_pkl_list += [\n",
                "                     '/content/' + model_name\n",
                "                     for model_name in sorted(ffhq_model_names)\n",
                "                     ]\n",
                "\n",
                "print('\\n'.join(network_pkl_list))"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "/content/stylegan2-ffhq-config-e.pkl\n",
                        "/content/stylegan2-ffhq-config-f.pkl\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "2OaeCIVCA4Nl",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 35
                },
                "outputId": "4a819389-eb34-46f0-8e61-4fe437d7483e"
            },
            "source": [
                "%cd /content/stylegan2/"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "/content/stylegan2\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "dyWVjoSaACTj",
                "colab_type": "code",
                "colab": {}
            },
            "source": [
                "import dnnlib\n",
                "import dnnlib.tflib as tflib\n",
                "\n",
                "Gs_syn_kwargs = dnnlib.EasyDict()\n",
                "Gs_syn_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8,\n",
                "                                      nchw_to_nhwc=True)\n",
                "Gs_syn_kwargs.randomize_noise = False\n",
                "Gs_syn_kwargs.minibatch_size = 4"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "kAokfLdXIY7f",
                "colab_type": "code",
                "colab": {}
            },
            "source": [
                "output_folder_for_average_images = 'average_image/'\n",
                "%mkdir -p $output_folder_for_average_images"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "0CTlq6fkJSVD",
                "colab_type": "code",
                "colab": {}
            },
            "source": [
                "file_extension = '.jpg'"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "Rzm0BzemBZTm",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 90
                },
                "outputId": "f9e42d78-251f-4061-844d-a924b4100cef"
            },
            "source": [
                "import pretrained_networks\n",
                "import numpy as np\n",
                "import PIL.Image\n",
                "from pathlib import Path\n",
                "\n",
                "n = len(network_pkl_list)\n",
                "size = 256\n",
                "canvas = PIL.Image.new('RGB', (n * size, size))\n",
                "for (i, network_pkl)  in enumerate(network_pkl_list):  \n",
                "    print(network_pkl)\n",
                "\n",
                "    _, _, Gs = pretrained_networks.load_networks(network_pkl)\n",
                "    \n",
                "    w_avg = Gs.get_var('dlatent_avg')\n",
                "\n",
                "    # Reference: https://github.com/NVlabs/stylegan2/blob/cec605e0834de5404d5c7e5cead7053bdd0e4dde/projector.py#L95\n",
                "    w = np.tile(w_avg, [1, Gs.components.synthesis.input_shape[1], 1])\n",
                "\n",
                "    image = Gs.components.synthesis.run(w, **Gs_syn_kwargs)[0]\n",
                "    image = PIL.Image.fromarray(image)\n",
                "\n",
                "    output_root_file_name = Path(network_pkl).name.replace('.pkl', '')\n",
                "\n",
                "    image.save('{}{}{}'.format(output_folder_for_average_images,\n",
                "                             output_root_file_name,\n",
                "                             file_extension))\n",
                "\n",
                "    image = image.resize((size, size), PIL.Image.LANCZOS)\n",
                "    canvas.paste(image, (i * size, 0))\n",
                "\n",
                "canvas.save('{}average_image{}'.format(output_folder_for_average_images,\n",
                "                                       file_extension))"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "/content/stylegan2-ffhq-config-e.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",
                        "/content/stylegan2-ffhq-config-f.pkl\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "17kE2eBSJnBk",
                "colab_type": "code",
                "colab": {}
            },
            "source": [
                "!tar -cf average_image.tar.gz average_image/"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "vNQXk1b50soC",
                "colab_type": "text"
            },
            "source": [
                "## GANspace\n",
                "\n",
                "[2] [Härkönen, Erik, et al. *GANSpace: Discovering Interpretable GAN Controls*. arXiv 2020.][ganspace-paper] **([code][ganspace])**\n",
                "\n",
                "[ganspace-paper]: <https://arxiv.org/abs/2004.02546>\n",
                "[ganspace]: <https://github.com/harskish/ganspace>"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "mumta2Hf-Bxa",
                "colab_type": "text"
            },
            "source": [
                "My fork: https://github.com/woctezuma/ganspace/tree/steam\n",
                "\n",
                "The original GANSpace notebook: [Github][Ganspace_original-github] and [Colab][Ganspace_original-colab]\n",
                "\n",
                "A Colab notebook (very similar to the original GANSpace notebook above): [Colab][Ganspace_similar-colab]\n",
                "\n",
                "My version of the GANSpace notebook (adapted to use my fork): [Github][Ganspace_colab_for_steam-github] and [Colab][Ganspace_colab_for_steam-colab]\n",
                "\n",
                "[Ganspace_original-github]: <https://github.com/harskish/ganspace/blob/master/notebooks/Ganspace_colab.ipynb>\n",
                "[Ganspace_original-colab]: <https://colab.research.google.com/github/harskish/ganspace/blob/master/notebooks/Ganspace_colab.ipynb>\n",
                "[Ganspace_similar-colab]: <https://colab.research.google.com/drive/1g-ShMzkRWDMHPyjom_p-5kqkn2f-GwBi>\n",
                "[Ganspace_colab_for_steam-github]: <https://github.com/woctezuma/steam-stylegan2/blob/ganspace/Ganspace_colab_for_steam.ipynb>\n",
                "[Ganspace_colab_for_steam-colab]: <https://colab.research.google.com/github/woctezuma/steam-stylegan2/blob/ganspace/Ganspace_colab_for_steam.ipynb>"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "bPd0WQja9qNs",
                "colab_type": "text"
            },
            "source": [
                "## Closed-Form Factorization\n",
                "\n",
                "[3] [Shen, Yujun, et al. *Closed-Form Factorization of Latent Semantics in GANs*. arXiv 2020.][closed-form-paper] **([code][closed-form])**\n",
                "\n",
                "[closed-form-paper]: <https://arxiv.org/abs/2007.06600>\n",
                "[closed-form]: <https://github.com/rosinality/stylegan2-pytorch#closed-form-factorization-httpsarxivorgabs200706600>"
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "E8WO3ZQ0078y",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 35
                },
                "outputId": "502ff110-d7cf-4b63-edde-cd5d8b952d7e"
            },
            "source": [
                "%cd /content/stylegan2-pytorch/"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "/content/stylegan2-pytorch\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "CPU9VytW9mif",
                "colab_type": "text"
            },
            "source": [
                "### Extract eigenvectors of the weight matrix"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "__eDSZQjAWTx",
                "colab_type": "text"
            },
            "source": [
                "#### Run"
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "wsCLzp0p1goQ",
                "colab_type": "code",
                "colab": {}
            },
            "source": [
                "!python closed_form_factorization.py $pytorch_model_name"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "1k4e4mcBAYQ_",
                "colab_type": "text"
            },
            "source": [
                "#### Check"
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "QOyGcy-8AN_2",
                "colab_type": "code",
                "colab": {}
            },
            "source": [
                "import torch\n",
                "\n",
                "eigvec = torch.load('factor.pt')[\"eigvec\"].to('cuda')"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "RGtRYPD9AaAc",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 35
                },
                "outputId": "95b02db0-3384-413a-e957-66fa7fe8648e"
            },
            "source": [
                "eigvec.shape"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "execute_result",
                    "data": {
                        "text/plain": [
                            "torch.Size([512, 512])"
                        ]
                    },
                    "metadata": {
                        "tags": []
                    },
                    "execution_count": 87
                }
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "qIl4o5Vs9uPa",
                "colab_type": "text"
            },
            "source": [
                "### Test the meaning of extracted directions"
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "OVlwqaQshlOB",
                "colab_type": "code",
                "colab": {}
            },
            "source": [
                "degree = 5\n",
                "num_degrees = 5"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "xhYGUpPT7xvI",
                "colab_type": "code",
                "colab": {}
            },
            "source": [
                "truncation = 0.75\n",
                "\n",
                "assert(0 <= truncation <= 1)"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "5tVFsTnbffyc",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 35
                },
                "outputId": "b9d0fe2d-6010-4ad2-ba1c-526eab65c1e4"
            },
            "source": [
                "output_keyword = 'temp_degree_{:03d}_truncation_{}'.format(degree, \n",
                "                                                           truncation)\n",
                "\n",
                "print(output_keyword)"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "temp_degree_005_truncation_0.75\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "viuyn3eWiArh",
                "colab_type": "code",
                "colab": {}
            },
            "source": [
                "output_folder = output_keyword + '/'\n",
                "out_prefix = output_folder + 'factor'"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "QNbb2Y46hheK",
                "colab_type": "code",
                "colab": {}
            },
            "source": [
                "%rm -rf $output_folder\n",
                "%mkdir -p $output_folder"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "BS7Rh0rM7sp6",
                "colab_type": "text"
            },
            "source": [
                "##### Run"
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "S-Av6CUO4o7w",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 54
                },
                "outputId": "e9740d7a-f67a-4154-d8d7-a2783b83438d"
            },
            "source": [
                "print('PyTorch model name: {}'.format(pytorch_model_name))\n",
                "print('Files will be saved to {}'.format(output_folder))"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "PyTorch model name: stylegan2-ffhq-config-e.pt\n",
                        "Files will be saved to temp_degree_005_truncation_0.75/\n"
                    ],
                    "name": "stdout"
                }
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "ioyI9Izf1SuD",
                "colab_type": "code",
                "colab": {}
            },
            "source": [
                "if debug_with_ffhq:\n",
                "  image_size = 1024\n",
                "else:\n",
                "  image_size = 256"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "uJu1enLb1mA5",
                "colab_type": "code",
                "colab": {}
            },
            "source": [
                "num_eigenvectors = 100\n",
                "num_samples = 10"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "sFA5OAg-1pzx",
                "colab_type": "code",
                "colab": {}
            },
            "source": [
                "for eigenvector_index in range(num_eigenvectors):\n",
                "\n",
                "  print('Eigenvector n°{:02d}/{:02d}'.format(eigenvector_index+1, \n",
                "                                             num_eigenvectors))\n",
                "\n",
                "  !python apply_factor.py \\\n",
                "    --index $eigenvector_index \\\n",
                "    --degree $degree \\\n",
                "    --degree_num $num_degrees \\\n",
                "    --channel_multiplier 1 \\\n",
                "    --torch_seed 2021 \\\n",
                "    --ckpt $pytorch_model_name \\\n",
                "    --size $image_size \\\n",
                "    --n_sample $num_samples \\\n",
                "    --truncation $truncation \\\n",
                "    --out_prefix $out_prefix \\\n",
                "    factor.pt"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "id": "-uOO4caL-nRx",
                "colab_type": "text"
            },
            "source": [
                "#### Archive"
            ]
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "EThK_rV0f4oe",
                "colab_type": "code",
                "colab": {}
            },
            "source": [
                "archive_name = output_keyword + '.tar.gz'"
            ],
            "execution_count": null,
            "outputs": []
        },
        {
            "cell_type": "code",
            "metadata": {
                "id": "vwhqzFadKI0S",
                "colab_type": "code",
                "colab": {
                    "base_uri": "https://localhost:8080/",
                    "height": 35
                },
                "outputId": "70bd45cd-b5ec-405d-f5c7-b385c057f153"
            },
            "source": [
                "!tar -cf $archive_name $output_folder\n",
                "!du -sh $archive_name"
            ],
            "execution_count": null,
            "outputs": [
                {
                    "output_type": "stream",
                    "text": [
                        "527M\ttemp_degree_005_truncation_0.75.tar.gz\n"
                    ],
                    "name": "stdout"
                }
            ]
        }
    ]
}