StyleGAN2_generate_feature_grids.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "StyleGAN2-generate-feature-grids.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "machine_shape": "hm",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/duskvirkus/colab-notebooks/blob/main/StyleGAN2_generate_feature_grids.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Fx0yRqdfsxOW"
      },
      "source": [
        "# Generate Feature Vector Grids\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j9oytDnqMur_"
      },
      "source": [
        "## Dependencies"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "MHTmbclksrJy"
      },
      "source": [
        "!git clone https://github.com/NVlabs/stylegan2-ada-pytorch\n",
        "%cd stylegan2-ada-pytorch\n",
        "\n",
        "!pip install ninja"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MYG__Wc7NCO7"
      },
      "source": [
        "## Load Model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IQ-SMHiDOdOW"
      },
      "source": [
        "load your model or connect google drive"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JTCz_oDnwamg"
      },
      "source": [
        "!wget http://d36zk2xti64re0.cloudfront.net/stylegan2/networks/stylegan2-cat-config-f.pkl"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "krlyw7x-QfOv"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CnZrdWI6OTAw"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Ef14ezRWv1PU",
        "cellView": "form"
      },
      "source": [
        "#@title imports (hidden code)\n",
        "import os\n",
        "import io\n",
        "\n",
        "import numpy as np\n",
        "import PIL.Image\n",
        "import cv2\n",
        "import IPython.display\n",
        "import dnnlib\n",
        "import torch\n",
        "\n",
        "import legacy"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uesuziPRwYcN"
      },
      "source": [
        "network_pkl = '/content/drive/MyDrive/art/stylegan/suicide-girls/pytorch-v1-cp1/suicide-girls-pytorch-v1-cp1.pkl'"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_RQbW3UYwebA"
      },
      "source": [
        "device = torch.device('cuda')\n",
        "with dnnlib.util.open_url(network_pkl) as f:\n",
        "    G = legacy.load_network_pkl(f)['G_ema'].to(device)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RcI3iREJWSpt"
      },
      "source": [
        "## Generate Feature Vectors"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "icQicrT2wJdv"
      },
      "source": [
        "features_save_path = '/content/drive/MyDrive/art/stylegan/suicide-girls/pytorch-v1-cp1/feature-vectors/feature-vectors.pt'\n",
        "\n",
        "modulate = {\n",
        "    k[0]: k[1]\n",
        "    for k in G.named_parameters()\n",
        "    if \"affine\" in k[0] and \"torgb\" not in k[0] and \"weight\" in k[0] or (\"torgb\" in k[0] and \"b4\" in k[0] and \"weight\" in k[0] and \"affine\" in k[0])\n",
        "}\n",
        "\n",
        "weight_mat = []\n",
        "for k, v in modulate.items():\n",
        "    weight_mat.append(v)\n",
        "\n",
        "W = torch.cat(weight_mat, 0)\n",
        "eigvec = torch.linalg.svd(W).V.to(\"cpu\")\n",
        "\n",
        "directory = os.path.dirname(features_save_path)\n",
        "if not os.path.exists(directory):\n",
        "    os.makedirs(directory)\n",
        "\n",
        "torch.save({\"ckpt\": network_pkl, \"eigvec\": eigvec}, features_save_path)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_MQ5-xGKrBbB"
      },
      "source": [
        "## Generating Grids"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "srIF_13bxCin",
        "cellView": "form"
      },
      "source": [
        "#@title functions (hidden code)\n",
        "# define functions\n",
        "\n",
        "def lerp(zs, steps):\n",
        "  out = []\n",
        "  for i in range(len(zs)-1):\n",
        "      for index in range(steps):\n",
        "          t = index/float(steps)\n",
        "          out.append(zs[i+1]*t + zs[i]*(1-t))\n",
        "  return out\n",
        "\n",
        "def concat_tile(im_list_2d):\n",
        "    return cv2.vconcat([cv2.hconcat(im_list_h) for im_list_h in im_list_2d])\n",
        "\n",
        "def save_grid(images, save_path, save_name):\n",
        "  # resize image\n",
        "  resized_images = []\n",
        "  for i in range(len(images)):\n",
        "    temp = []\n",
        "    for j in range(len(images[i])):\n",
        "      im = cv2.resize(np.float32(images[i][j]), (256, 256))\n",
        "      im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)\n",
        "      temp.append(im)\n",
        "    resized_images.append(temp)\n",
        "    \n",
        "\n",
        "  # save grid\n",
        "  img_grid = concat_tile(resized_images)\n",
        "\n",
        "  if not os.path.exists(save_path):\n",
        "    os.makedirs(save_path)\n",
        "\n",
        "  path = os.path.join(save_path, save_name)\n",
        "  cv2.imwrite(path, img_grid)\n",
        "\n",
        "def gen_feature_grid(rows, cols, feature_index, strength, truncation, save_location):\n",
        "\n",
        "  label = torch.zeros([1, G.c_dim], device=device)\n",
        "  truncation_psi = truncation\n",
        "  noise_mode = 'const'\n",
        "\n",
        "  zgrid = [[None for i in range(cols)] for j in range(rows)] # [row][col]\n",
        "\n",
        "  for row in range(rows):\n",
        "    z = np.random.RandomState(row).randn(1, G.z_dim)\n",
        "    z = torch.from_numpy(z)\n",
        "    z = z.cpu()\n",
        "\n",
        "    current_eigvec = eigvec[feature_index]\n",
        "    direction = strength * current_eigvec\n",
        "\n",
        "    col_zvals = lerp([z - direction, z + direction], cols)\n",
        "    for col in range(len(col_zvals)):\n",
        "      zgrid[row][col] = col_zvals[col]\n",
        "\n",
        "  all_images = []\n",
        "  for row in range(len(zgrid)):\n",
        "    zs = torch.cat(zgrid[row]).to(device)\n",
        "    generated = G(zs, label, truncation_psi=truncation_psi, noise_mode=noise_mode)\n",
        "    generated = (generated.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)\n",
        "    all_images.append(generated.cpu())\n",
        "\n",
        "  save_name = \"feature{feature_index}_seeds{seeds_start}-{seeds_end}.png\".format(feature_index=feature_index, seeds_start = 0, seeds_end = rows - 1)\n",
        "  save_grid(all_images, save_location, save_name)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZV5Dc0VUNhsW"
      },
      "source": [
        "# use it\n",
        "num_seeds = 20 # starts at seed 0\n",
        "num_rows = 7\n",
        "feature_index = 1\n",
        "strength = 20.0\n",
        "truncation = 0.8\n",
        "save_dir = '/content/output'\n",
        "\n",
        "gen_feature_grid(num_seeds, num_rows, feature_index, strength, truncation, save_dir)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "J7O2ufYePzSq"
      },
      "source": [
        "# in a loop\n",
        "\n",
        "feature_start = 0\n",
        "feature_end = 511\n",
        "\n",
        "num_seeds = 12 # starts at seed 0\n",
        "num_rows = 7\n",
        "strength = 20.0\n",
        "truncation = 0.8\n",
        "save_dir = '/content/drive/MyDrive/art/stylegan/suicide-girls/pytorch-v1-cp1/feature-vectors/grids/'\n",
        "\n",
        "for i in range(feature_start, feature_end):\n",
        "  gen_feature_grid(num_seeds, num_rows, i, strength, truncation, save_dir)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}