InterFaceGAN++.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EEqf7sm-RiMq"
      },
      "source": [
        "Authors: Apavou Clément & Belkada Younes\n",
        "\n",
        "Inspired by https://colab.research.google.com/github/genforce/interfacegan/blob/master/docs/InterFaceGAN.ipynb"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qJDJLE3v0HNr"
      },
      "source": [
        "# Fetch Codebase and Models"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qrMxUAyeBPVz"
      },
      "outputs": [],
      "source": [
        "!pip -qq install ninja"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JqiWKjpFa0ov"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "os.chdir('/content')\n",
        "CODE_DIR = 'interfacegan'\n",
        "!git clone https://github.com/younesbelkada/interfacegan.git $CODE_DIR\n",
        "os.chdir(f'./{CODE_DIR}')\n",
        "!wget https://www.dropbox.com/s/qyv37eaobnow7fu/stylegan_ffhq.pth?dl=1 -O models/pretrain/stylegan_ffhq.pth --quiet\n",
        "!wget https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-ffhq-1024x1024.pkl -O models/pretrain/stylegan2-ffhq-1024x1024.pkl --quiet\n",
        "!wget https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-t-ffhq-1024x1024.pkl -O models/pretrain/stylegan3-t-ffhq-1024x1024.pkl --quiet"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hQ_IXBZr8YcJ"
      },
      "source": [
        "# Define Utility Functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ijKTlG5GeTd3"
      },
      "outputs": [],
      "source": [
        "import os.path\n",
        "import io\n",
        "import IPython.display\n",
        "import numpy as np\n",
        "import cv2\n",
        "import PIL.Image\n",
        "import ipywidgets as widgets\n",
        "\n",
        "import torch\n",
        "\n",
        "from models.model_settings import MODEL_POOL\n",
        "from models.stylegan_generator import StyleGANGenerator\n",
        "from models.stylegan2_generator import StyleGAN2Generator\n",
        "from models.stylegan3_generator import StyleGAN3Generator\n",
        "from utils.manipulator import linear_interpolate\n",
        "\n",
        "\n",
        "def build_generator(model_name):\n",
        "  \"\"\"Builds the generator by model name.\"\"\"\n",
        "  gan_type = MODEL_POOL[model_name]['gan_type']\n",
        "  if gan_type == 'stylegan':\n",
        "    generator = StyleGANGenerator(model_name)\n",
        "  if gan_type == 'stylegan2':\n",
        "    generator = StyleGAN2Generator(model_name)\n",
        "  elif gan_type == 'stylegan3':\n",
        "    generator = StyleGAN3Generator(model_name)\n",
        "  return generator\n",
        "\n",
        "\n",
        "def sample_codes(generator, num, latent_space_type='Z', seed=0):\n",
        "  \"\"\"Samples latent codes randomly.\"\"\"\n",
        "  np.random.seed(seed)\n",
        "  codes = generator.easy_sample(num)\n",
        "  if generator.gan_type == 'stylegan' and latent_space_type == 'W':\n",
        "    codes = torch.from_numpy(codes).type(torch.FloatTensor).to(generator.run_device)\n",
        "    codes = generator.get_value(generator.model.mapping(codes))\n",
        "  return codes\n",
        "\n",
        "\n",
        "def imshow(images, col, viz_size=256):\n",
        "  \"\"\"Shows images in one figure.\"\"\"\n",
        "  num, height, width, channels = images.shape\n",
        "  assert num % col == 0\n",
        "  row = num // col\n",
        "\n",
        "  fused_image = np.zeros((viz_size * row, viz_size * col, channels), dtype=np.uint8)\n",
        "\n",
        "  for idx, image in enumerate(images):\n",
        "    i, j = divmod(idx, col)\n",
        "    y = i * viz_size\n",
        "    x = j * viz_size\n",
        "    if height != viz_size or width != viz_size:\n",
        "      image = cv2.resize(image, (viz_size, viz_size))\n",
        "    fused_image[y:y + viz_size, x:x + viz_size] = image\n",
        "\n",
        "  fused_image = np.asarray(fused_image, dtype=np.uint8)\n",
        "  data = io.BytesIO()\n",
        "  PIL.Image.fromarray(fused_image).save(data, 'jpeg')\n",
        "  im_data = data.getvalue()\n",
        "  disp = IPython.display.display(IPython.display.Image(im_data))\n",
        "  return disp"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q7gkmrVW8eR1"
      },
      "source": [
        "# Select a Model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GdDaWpEl3zPm"
      },
      "outputs": [],
      "source": [
        "import glob\n",
        "import os.path as osp\n",
        "attr_path = glob.glob(osp.join('boundaries','stylegan_ffhq_z', '*.npy'))\n",
        "ATTRS = [path.split('.')[0].split('/')[-1].replace('boundary_', '') for path in attr_path]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NoWI4fPQ6Gnf"
      },
      "outputs": [],
      "source": [
        "#@title { display-mode: \"form\", run: \"auto\" }\n",
        "model_name = \"stylegan3_ffhq\" #@param ['stylegan_ffhq', 'stylegan2_ffhq', 'stylegan3_ffhq']\n",
        "\n",
        "latent_space_type = 'Z' #@param ['Z', 'W']\n",
        "\n",
        "if latent_space_type == 'W' and ( '2' in model_name or '3' in model_name):\n",
        "  raise ValueError('Latent space is not available for StyleGAN 2 and 3')\n",
        "\n",
        "generator = build_generator(model_name)\n",
        "\n",
        "boundaries = {}\n",
        "for i, attr_name in enumerate(ATTRS):\n",
        "  boundaries[attr_name] = np.load(f'boundaries/{model_name}_{latent_space_type.lower()}/boundary_{attr_name}.npy')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zDStH1O5t1KC"
      },
      "source": [
        "# Sample latent codes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qlRGKZbJt9hA"
      },
      "outputs": [],
      "source": [
        "#@title { display-mode: \"form\", run: \"auto\" }\n",
        "\n",
        "num_samples = 5 #@param {type:\"slider\", min:1, max:8, step:1}\n",
        "noise_seed = 678 #@param {type:\"slider\", min:0, max:1000, step:1}\n",
        "\n",
        "latent_codes = sample_codes(generator, num_samples, latent_space_type, noise_seed)\n",
        "if generator.gan_type == 'stylegan' and latent_space_type == 'W':\n",
        "  synthesis_kwargs = {'latent_space_type': 'W'}\n",
        "else:\n",
        "  synthesis_kwargs = {}\n",
        "\n",
        "images = generator.easy_synthesize(latent_codes, **synthesis_kwargs)['image']\n",
        "imshow(images, col=num_samples)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MmRPN3xz8jCH"
      },
      "source": [
        "# Edit facial attributes"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ccONBF60mVir"
      },
      "outputs": [],
      "source": [
        "#@title { display-mode: \"form\", run: \"auto\" }\n",
        "\n",
        "Wearing_Hat = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Straight_Hair = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Big_Lips = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Receding_Hairline = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Heavy_Makeup = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Bags_Under_Eyes = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Big_Nose = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Bald = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Eyeglasses = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Bangs = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Black_Hair = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Oval_Face = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1} \n",
        "Bushy_Eyebrows = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Chubby = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Brown_Hair = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Wearing_Lipstick = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Pale_Skin = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Gender = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Wavy_Hair = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Wearing_Earrings = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Young = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Goatee = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Pointy_Nose = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Wearing_Necktie = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Attractive = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Double_Chin = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Narrow_Eyes = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Blond_Hair = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Wearing_Necklace = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Mouth_Slightly_Open = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "No_Beard = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Arched_Eyebrows = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Rosy_Cheeks = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Mustache = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Five_o_Clock_Shadow = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Sideburns = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Gray_Hair = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "Smiling = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "High_Cheekbones = 0 #@param {type:\"slider\", min:-3.0, max:3.0, step:0.1}\n",
        "\n",
        "new_codes = latent_codes.copy()\n",
        "for i, attr_name in enumerate(ATTRS):\n",
        "  if attr_name=='5_o_Clock_Shadow': \n",
        "    new_codes += boundaries[attr_name] * eval('Five_o_Clock_Shadow')\n",
        "  else:\n",
        "    new_codes += boundaries[attr_name] * eval(attr_name)\n",
        "\n",
        "new_images = generator.easy_synthesize(new_codes, **synthesis_kwargs)['image']\n",
        "imshow(new_images, col=num_samples)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "InterFaceGAN++",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}