{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.6.5"
    },
    "colab": {
      "name": "convert-frozen-sep.ipynb",
      "version": "0.3.2",
      "provenance": []
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "huNsqyrQjb_R",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 86
        },
        "outputId": "d780ebc3-e805-45d0-9437-8a9b64d4964c"
      },
      "source": [
        "!git clone https://github.com/NVlabs/stylegan"
      ],
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Cloning into 'stylegan'...\n",
            "remote: Enumerating objects: 77, done.\u001b[K\n",
            "remote: Total 77 (delta 0), reused 0 (delta 0), pack-reused 77\u001b[K\n",
            "Unpacking objects: 100% (77/77), done.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1BWYmZqDjxjt",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "os.chdir('stylegan')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MzaH2gk5jWwm",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 157
        },
        "outputId": "7eeee15d-e247-4893-eb8a-72ede28076b0"
      },
      "source": [
        "import time\n",
        "import re\n",
        "import bisect\n",
        "from collections import OrderedDict\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import scipy.ndimage\n",
        "import scipy.misc\n",
        "import pickle\n",
        "import dnnlib\n",
        "import dnnlib.tflib as tflib\n",
        "from PIL import Image\n",
        "from IPython.display import display"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "WARNING: Logging before flag parsing goes to stderr.\n",
            "W0828 23:46:56.914453 139852722513792 deprecation_wrapper.py:119] From /content/stylegan/dnnlib/tflib/tfutil.py:34: The name tf.Dimension is deprecated. Please use tf.compat.v1.Dimension instead.\n",
            "\n",
            "W0828 23:46:56.915941 139852722513792 deprecation_wrapper.py:119] From /content/stylegan/dnnlib/tflib/tfutil.py:74: The name tf.variable_scope is deprecated. Please use tf.compat.v1.variable_scope instead.\n",
            "\n",
            "W0828 23:46:56.916904 139852722513792 deprecation_wrapper.py:119] From /content/stylegan/dnnlib/tflib/tfutil.py:128: The name tf.Session is deprecated. Please use tf.compat.v1.Session instead.\n",
            "\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uY4h81mIjWw6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "sess = tf.InteractiveSession()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Fk699HyBj9S5",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 175
        },
        "outputId": "ae473fb2-ac22-4094-f5f7-5663fd477283"
      },
      "source": [
        "url = 'https://drive.google.com/uc?id=1MEGjdvVpUsu1jB4zrXZN7Y4kBBOzizDQ' # karras2019stylegan-ffhq-1024x1024.pkl\n",
        "with dnnlib.util.open_url(url, '.') as f:\n",
        "      _G, _D, Gs = pickle.load(f)"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Downloading https://drive.google.com/uc?id=1MEGjdvVpUsu1jB4zrXZN7Y4kBBOzizDQ .... done\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "W0828 23:47:01.073774 139852722513792 deprecation_wrapper.py:119] From /content/stylegan/dnnlib/tflib/tfutil.py:124: The name tf.get_default_session is deprecated. Please use tf.compat.v1.get_default_session instead.\n",
            "\n",
            "W0828 23:47:01.084098 139852722513792 deprecation_wrapper.py:119] From /content/stylegan/dnnlib/tflib/network.py:142: The name tf.get_default_graph is deprecated. Please use tf.compat.v1.get_default_graph instead.\n",
            "\n",
            "W0828 23:47:05.489010 139852722513792 deprecation.py:323] From <string>:364: add_dispatch_support.<locals>.wrapper (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.\n",
            "Instructions for updating:\n",
            "Use tf.where in 2.0, which has the same broadcast rule as np.where\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7iGxUNeQjWw-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "fmt = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O7FaWfSMjWxB",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "L = 5"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FVwHBbSOjWxD",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "src_latents = np.random.randn(L, *Gs.input_shape[1:])\n",
        "dst_latents = np.random.randn(L, *Gs.input_shape[1:])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Qm7NVRJnjWxF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "src_dlatents = Gs.components.mapping.run(src_latents, None) # [seed, layer, component]\n",
        "dst_dlatents = Gs.components.mapping.run(dst_latents, None) # [seed, layer, component]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0uVVisryjWxH",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 67
        },
        "outputId": "0e4d5314-112a-47a8-82ad-17f33c9fcca8"
      },
      "source": [
        "for k,v  in Gs.components.mapping._run_cache.items():\n",
        "    print(k)\n",
        "    print(v[0])\n",
        "    print(v[1])"
      ],
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[('assume_frozen', False), ('dynamic_kwargs', []), ('input_transform', None), ('num_gpus', 1), ('output_transform', None)]\n",
            "[<tf.Tensor 'G_mapping_1/_Run/latents_in:0' shape=<unknown> dtype=float32>, <tf.Tensor 'G_mapping_1/_Run/labels_in:0' shape=<unknown> dtype=float32>]\n",
            "[<tf.Tensor 'G_mapping_1/_Run/concat:0' shape=(?, 18, 512) dtype=float32>]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DMZZ9VB3jWxL",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "2c15c151-9ec2-4ff8-a8ac-12306f2eefe8"
      },
      "source": [
        "v[1][0].op.inputs[0]"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor 'G_mapping_1/_Run/G_mapping/dlatents_out:0' shape=(?, 18, 512) dtype=float32>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 11
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4xAFfKRajWxP",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "synthesis_kwargs = dict(output_transform=fmt)\n",
        "src_images = Gs.components.synthesis.run(src_dlatents, randomize_noise=False, **synthesis_kwargs)\n",
        "dst_images = Gs.components.synthesis.run(dst_dlatents, randomize_noise=False, **synthesis_kwargs)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "76WZLXlVozVY",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "Image.fromarray(src_images[0])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CGeqNpi4jWxR",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 87
        },
        "outputId": "994bee60-1224-4085-c0b4-c62309468c05"
      },
      "source": [
        "for k,v  in Gs.components.synthesis._run_cache.items():\n",
        "    print(k)\n",
        "    print(v[0])\n",
        "    print(v[1])"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "[('assume_frozen', False), ('dynamic_kwargs', [('randomize_noise', False)]), ('input_transform', None), ('num_gpus', 1), ('output_transform', [('func', 'dnnlib.tflib.tfutil.convert_images_to_uint8'), ('nchw_to_nhwc', True)])]\n",
            "[<tf.Tensor 'G_synthesis_1/_Run/dlatents_in:0' shape=<unknown> dtype=float32>]\n",
            "[<tf.Tensor 'G_synthesis_1/_Run/concat:0' shape=(?, 1024, 1024, 3) dtype=uint8>]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XlcUFsApjWxT",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "6c4413ab-3078-4dd8-e16c-38df15969b6e"
      },
      "source": [
        "v[1][0].op.inputs[0]"
      ],
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "<tf.Tensor 'G_synthesis_1/_Run/saturate_cast:0' shape=(?, 1024, 1024, 3) dtype=uint8>"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 14
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3hoZ9P40jWxV",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oOc7__DejWxX",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "57a1371a-0399-4691-fb89-9da8ccf8d814"
      },
      "source": [
        "src_dlatents.shape"
      ],
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "(5, 18, 512)"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 15
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7olCYpB4jWxZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "w,h=256,256\n",
        "style_ranges=[range(0,1), range(0,3), range(0,5), range(0,8), range(0,11)]\n",
        "canvas = Image.new('RGB', (w * (L + 1), h * (L + 1)) , 'white')\n",
        "for col, src_image in enumerate(list(src_images)):\n",
        "    canvas.paste(Image.fromarray(src_image, 'RGB').resize((w,h)), ((col + 1) * w, 0))\n",
        "for row, dst_image in enumerate(list(dst_images)):\n",
        "        canvas.paste(Image.fromarray(dst_image, 'RGB').resize((w,h)), (0, (row + 1) * h))\n",
        "        row_dlatents = np.stack([dst_dlatents[row]] * L)\n",
        "        row_dlatents[:, style_ranges[row]] = src_dlatents[:, style_ranges[row]]\n",
        "        row_images = Gs.components.synthesis.run(row_dlatents, randomize_noise=False, **synthesis_kwargs)\n",
        "        for col, image in enumerate(list(row_images)):\n",
        "            canvas.paste(Image.fromarray(image, 'RGB').resize((w,h)), ((col + 1) * w, (row + 1) * h))\n",
        "display(canvas)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PrERxHQhjWxk",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Uj247q3OjWxm",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lxI5ohkujWxq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}