StyleGAN2 Network Interpolation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [ ]
    },
    {
     "name": "stdout",
     "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"
     ]
    }
   ],
   "source": [
    "import ipywidgets as widgets\n",
    "import pretrained_networks\n",
    "import PIL.Image\n",
    "import numpy as np\n",
    "src_model = '2020-01-11-skylion-stylegan2-animeportraits-networksnapshot-024664.pkl'\n",
    "dst_model = 'e621-network-snapshot-011573.pkl'\n",
    "_G, _D, Gs = pretrained_networks.load_networks(src_model)\n",
    "_Gd, _Dd, Gsd = pretrained_networks.load_networks(dst_model)\n",
    "bGs = Gs.clone()\n",
    "import dnnlib\n",
    "import dnnlib.tflib as tflib\n",
    "Gs_syn_kwargs = dnnlib.EasyDict()\n",
    "batch_size = 1\n",
    "Gs_syn_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)\n",
    "Gs_syn_kwargs.randomize_noise = True\n",
    "Gs_syn_kwargs.minibatch_size = batch_size\n",
    "\n",
    "noise_vars = [var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise')]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dnnlib.tflib import tfutil\n",
    "def weighted_average(src_net, dst_net, t):\n",
    "    names = []\n",
    "    for name in src_net.trainables.keys():\n",
    "        if name not in src_net.trainables:\n",
    "            print(\"Not restoring (not present):     {}\".format(name))\n",
    "        elif dst_net.trainables[name].shape != src_net.trainables[name].shape:\n",
    "            print(\"Not restoring (different shape): {}\".format(name))\n",
    "\n",
    "        if name in src_net.trainables and dst_net.trainables[name].shape == src_net.trainables[name].shape:\n",
    "            names.append(name)\n",
    "\n",
    "    tfutil.set_vars(tfutil.run({bGs.vars[name]: (t*dst_net.vars[name] + (1-t)*src_net.vars[name]) for name in names}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2919b7365b6740dca1f1d73865ec8df0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(HBox(children=(IntSlider(value=0, continuous_update=False, description='Seed: ', max=100000), F…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1b97876f1cf04d818b86bd98467feae5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "seed = widgets.IntSlider(min=0, max=100000, step=1, value=0, description='Seed: ', continuous_update=False)\n",
    "scale = widgets.FloatSlider(min=0, max=5, step=0.01, value=1, description='Scale: ', continuous_update=False)\n",
    "truncation = widgets.FloatSlider(min=-2, max=2, step=0.1, value=1, description='Truncation: ', continuous_update=False)\n",
    "blending = widgets.FloatSlider(min=0, max=1, step=0.01, value=0, description='Blending: ', continuous_update=False)\n",
    "\n",
    "bot_box = widgets.HBox([seed, scale, truncation, blending])\n",
    "ui = widgets.VBox([bot_box])\n",
    "\n",
    "def display_sample(seed, scale, truncation, blending):\n",
    "    weighted_average(Gs, Gsd, blending)\n",
    "    \n",
    "    Gs_kwargs = dnnlib.EasyDict()\n",
    "    Gs_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)\n",
    "    Gs_kwargs.randomize_noise = False\n",
    "    if truncation is not None:\n",
    "        Gs_kwargs.truncation_psi = truncation\n",
    "    rnd = np.random.RandomState(seed)\n",
    "    tflib.set_vars({var: rnd.randn(*var.shape.as_list()) for var in noise_vars}) # [height, width]\n",
    "    \n",
    "    batch_size = 1\n",
    "    all_seeds = [seed] * batch_size\n",
    "    all_z = np.stack([np.random.RandomState(seed).randn(*bGs.input_shape[1:]) for seed in all_seeds]) # [minibatch, component]\n",
    "    all_w = bGs.components.mapping.run(scale*all_z, None) # [minibatch, layer, component]\n",
    "    if truncation != 1:\n",
    "        w_avg = bGs.get_var('dlatent_avg')\n",
    "        all_w = w_avg + (all_w - w_avg) * truncation # [minibatch, layer, component]\n",
    "    all_images = bGs.components.synthesis.run(all_w, **Gs_syn_kwargs)\n",
    "    display(PIL.Image.fromarray(np.median(all_images, axis=0).astype(np.uint8)))\n",
    "\n",
    "out = widgets.interactive_output(display_sample, {'seed': seed, 'scale': scale, 'truncation': truncation, 'blending': blending})\n",
    "\n",
    "display(ui, out)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}