StyleGAN2_colab.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "StyleGAN2_colab.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "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.9"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xzA1-mt88AO_"
      },
      "source": [
        "# StyleGAN2 operations\n",
        "\n",
        "### Preparation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7IFfx8GQIAQm"
      },
      "source": [
        "\n",
        "**Run this cell after each session restart**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tkbcraCUaPEy"
      },
      "source": [
        "#@title General setup { display-mode: \"form\", run: \"auto\" }\n",
        "\n",
        "from IPython.display import HTML, Image, display\n",
        "from moviepy.editor import ImageSequenceClip, ipython_display\n",
        "import ipywidgets as widgets\n",
        "import os\n",
        "\n",
        "%tensorflow_version 1.x\n",
        "!apt-get -qq install ffmpeg\n",
        "from google.colab import drive\n",
        "drive.mount('/G', force_remount=True)\n",
        "gdir = !ls /G/\n",
        "gdir = '/G/%s/' % str(gdir[0])\n",
        "%cd $gdir\n",
        "\n",
        "#@markdown Copying StyleGAN2 to the directory below on your Google drive (creating it, if it doesn't exist):\n",
        "work_dir = 'sg2' #@param {type:\"string\"}\n",
        "#@markdown NB: All paths below are relative to this directory (except the archive with source images on the next step). \n",
        "\n",
        "#@markdown NB: Avoid connecting Google drive manually via the icon in Files section on the left. Doing so may break further operations.\n",
        "\n",
        "work_dir = gdir + work_dir + '/'\n",
        "if not os.path.isdir(work_dir):\n",
        "  !git clone git://github.com/eps696/stylegan2 $work_dir\n",
        "%cd $work_dir\n",
        "!pip install -r requirements.txt\n",
        "\n",
        "def makevid(seq_dir, size=None):\n",
        "  out_sequence = seq_dir + '/%06d.jpg'\n",
        "  out_video = seq_dir + '.mp4'\n",
        "  !ffmpeg -y -v warning -i $out_sequence $out_video\n",
        "  data_url = \"data:video/mp4;base64,\" + b64encode(open(out_video,'rb').read()).decode()\n",
        "  wh = '' if size is None else 'width=%d height=%d' % (size, size)\n",
        "  return \"\"\"<video %s controls><source src=\"%s\" type=\"video/mp4\"></video>\"\"\" % (wh, data_url)\n",
        "\n",
        "from src.util.utilgan import file_list, img_list, basename\n",
        "model = model_pkl = ''\n",
        "def model_select(work_dir):\n",
        "  models = file_list(work_dir, 'pkl', subdir=True)\n",
        "  global model, model_pkl\n",
        "  model_pkl = models[0]\n",
        "  model = model_pkl.replace('.pkl', '')\n",
        "  models = [m.replace(work_dir, '') for m in models if not basename(m) in ['submit_config', 'vgg16_zhang_perceptual']]\n",
        "  def on_change(change):\n",
        "    global model, model_pkl\n",
        "    if change['type'] == 'change' and change['name'] == 'value':\n",
        "      model = change['new']\n",
        "      model = os.path.splitext(model)[0] # filename without extension => load custom network\n",
        "      model_pkl = model + '.pkl' # filename with extension => load original network\n",
        "  model_select = widgets.Dropdown(options=models, description='Found models:', style={'description_width': 'initial'}, layout={'width': 'max-content'})\n",
        "  display(model_select)\n",
        "  model_select.observe(on_change)\n",
        "# model_select(work_dir)\n",
        "\n",
        "# Hardware check\n",
        "!ln -sf /opt/bin/nvidia-smi /usr/bin/nvidia-smi\n",
        "!pip install gputil \n",
        "import GPUtil as GPU\n",
        "gpu = GPU.getGPUs()[0]\n",
        "!nvidia-smi -L\n",
        "print(\"GPU RAM {0:.0f}MB | Free {1:.0f}MB)\".format(gpu.memoryTotal, gpu.memoryFree))\n",
        "print('\\nDone!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yWpWFeyO8APF"
      },
      "source": [
        "## Training"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wxiSThiS4r5Q"
      },
      "source": [
        ""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R_LGhTfV8APG"
      },
      "source": [
        "First, let's prepare the dataset. Ensure all your images have the same color channels (monochrome, RGB or RGBA). If you edit the images yourself (e.g. for non-square aspect ratios, or to keep the compositions), ensure their correct size.  \n",
        "For conditional model the images should be split by subfolders (mydata/1, mydata/2, ..).\n",
        "\n",
        "Upload zip-archive with images onto Google drive and type its path here (relative to G-drive root). \n",
        "\n",
        "If you work with patterns or shapes (rather than compostions), turn on `multicrop` to crop square fragments from bigger images (effectively increasing amount of data). They will be cut into `size`px fragments, overlapped with `step` shift by X and Y. If the image is smaller, it will be upscaled to the `size`. Edit these values according to your dataset. *Omit this, if you have edited the images yourself, e.g. to keep the compositions, or to work with non-square aspect ratios.*"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tEmNzQmm0t_o",
        "cellView": "form"
      },
      "source": [
        "#@title Data setup \n",
        "dataset = 'test' #@param {type:\"string\"}\n",
        "source = 'source.zip' #@param {type:\"string\"}\n",
        "data_dir = os.path.join('/content', dataset)\n",
        "\n",
        "#@markdown Multicrop\n",
        "multicrop = False  #@param {type:\"boolean\"}\n",
        "size = \"256\" #@param [256,512,1024]\n",
        "overlap = 128 #@param {type:\"integer\"}\n",
        "size = int(size)\n",
        "\n",
        "# cleanup previous attempts\n",
        "!rm -rf /content/tmp \n",
        "!rm -rf /content/*.tfr\n",
        "!rm -rf $data_dir\n",
        "\n",
        "!mkdir /content/tmp\n",
        "%cd /content/tmp\n",
        "fpath = os.path.join(gdir, source)\n",
        "!unzip -o -q $fpath\n",
        "%cd $work_dir\n",
        "\n",
        "if multicrop:\n",
        "  if os.path.isdir('/content/tmp'):\n",
        "    %run src/util/multicrop.py --in_dir /content/tmp --out_dir $data_dir --size $size --step $overlap\n",
        "  else:\n",
        "    print('/content/tmp not found!!')\n",
        "else:\n",
        "  !mv /content/tmp $data_dir\n",
        "\n",
        "%run src/training/dataset_tool.py --data $data_dir\n",
        "\n",
        "data_dir = os.path.join(work_dir, 'data')\n",
        "!mv /content/*.tfr $data_dir\n",
        "!ls $data_dir/*.tfr"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZfczVF0W8APH"
      },
      "source": [
        "Now, we can train StyleGAN2 on the prepared dataset."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YKmktOPb8APH",
        "cellView": "form"
      },
      "source": [
        "#@title Train\n",
        "%cd $work_dir\n",
        "dataset = 'test' #@param {type:\"string\"}\n",
        "kimg = 5000 #@param {type:\"integer\"}\n",
        "data_dir = os.path.join(work_dir, 'data', dataset)\n",
        "\n",
        "%run src/train.py --data $data_dir --kimg $kimg --d_aug"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wYlI1fXe8APH"
      },
      "source": [
        "> This will run training process, according to the options in `src/train.py`. Results (models and samples) are saved under `train` directory, similar to original Nvidia approach. There are two types of models saved: compact (containing only Gs network for inference) as `<dataset>-...pkl` (e.g. `test-256-0360.pkl`), and full (containing G/D/Gs networks for further training) as `snapshot-...pkl`. Check sample images there to follow the progress.\n",
        "\n",
        "> The length of the training is defined by the `kimg` argument (training duration in thousands of images). Reasonable duration for full training from scratch is 5000-8000, while for finetuning in `--d_aug` mode 1000-2000 may be sufficient.  \n",
        "\n",
        "> *By default, the most powerful SG2 config (F) is used; if you face OOM issue, you may set smaller `--batch_size`, or resort to `--config E`, requiring less memory. [Differential Augmentation](https://github.com/mit-han-lab/data-efficient-gans) mode `--d_aug` is turned on for more effective training (absolutely required for small datasets with ~100x images). \n",
        "These settings can be added to the command line inside the cell above.*\n",
        "\n",
        "Other training options:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sAqpEFzukpfs"
      },
      "source": [
        "%run src/train.py --help"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NpeUmhiH8API"
      },
      "source": [
        "If the training process was interrupted, resume it from the last saved model."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wr73IUPG8API",
        "cellView": "form"
      },
      "source": [
        "#@title Resume training\n",
        "%cd $work_dir\n",
        "dataset = 'test' #@param {type:\"string\"}\n",
        "resume_dir = 'train/000-test-256-f-daug' #@param {type:\"string\"}\n",
        "kimg = 5000 #@param {type:\"integer\"}\n",
        "data_dir = os.path.join(work_dir, 'data', dataset)\n",
        "\n",
        "%run src/train.py --data $data_dir --resume $resume_dir --d_aug --kimg $kimg "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kVxyOjz48API"
      },
      "source": [
        "NB: In most cases it's much easier to use a \"transfer learning\" trick, rather than perform full training from the scratch. For that, we use existing well-trained model as a starter, and \"finetune\" (uptrain) it with our data. This works pretty well, even if our dataset is very different from the original model. \n",
        "\n",
        "So here is a faster way to train our GAN (presuming we have full trained model `ffhq-256.pkl` in `train` directory already):"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z3kwDOh88APJ",
        "cellView": "form"
      },
      "source": [
        "#@title Finetune training\n",
        "%cd $work_dir\n",
        "dataset = 'test' #@param {type:\"string\"}\n",
        "resume_pkl = 'train/ffhq-256.pkl' #@param {type:\"string\"}\n",
        "kimg = 1000 #@param {type:\"integer\"}\n",
        "data_dir = os.path.join(work_dir, 'data', dataset)\n",
        "\n",
        "%run src/train.py --data $data_dir --resume $resume_pkl --d_aug --kimg $kimg"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1Njelbgu8APJ"
      },
      "source": [
        "## Generation\n",
        "\n",
        "Let's produce some imagery from the original cat model (get it [here](https://nvlabs-fi-cdn.nvidia.com/stylegan2/networks/stylegan2-cat-config-f.pkl) and put onto Google drive within our working directory).  \n",
        "Generated results are saved as sequences and videos (by default, under `_out` directory).  \n",
        "More cool models can be found [here](https://github.com/justinpinkney/awesome-pretrained-stylegan2)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9S73oknl8wLU",
        "cellView": "form"
      },
      "source": [
        "#@title ### Generator setup\n",
        "!cd $work_dir\n",
        "\n",
        "output = '_out/cats' #@param {type:\"string\"}\n",
        "\n",
        "frames = 50 #@param {type:\"integer\"}\n",
        "frame_step = 10 #@param {type:\"integer\"}\n",
        "timeframe = '%d-%d' % (frames, frame_step)\n",
        "\n",
        "cubic_smooth = True #@param {type:\"boolean\"}\n",
        "gauss_smooth = False #@param {type:\"boolean\"}\n",
        "cubic_smooth = '--cubic ' if cubic_smooth else ''\n",
        "gauss_smooth = '--gauss ' if gauss_smooth else ''\n",
        "smooth = cubic_smooth + gauss_smooth\n",
        "\n",
        "save_lat = False #@param {type:\"boolean\"}\n",
        "save_lat = '--save_lat ' if save_lat else ''\n",
        "\n",
        "seed = 0 #@param {type:\"integer\"}\n",
        "\n",
        "#@markdown Select model from the dropdown below:\n",
        "model_select(work_dir)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vlzqD3DEGzIH"
      },
      "source": [
        "> This means loading the model, and producing 50 frames, interpolating between random latent `z` space keypoints, with a step of 10 frames between keypoints. \n",
        "`save_lat` option = save all traversed dlatent points (in `w` space) as Numpy array in `*.npy` file (useful for further curating). `cubic` option changes linear interpolation to cubic for smoother animation; `gauss` provides additional smoothing. Set `seed` value to produce repeatable results (0 = random)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oBHD0n0i8APJ",
        "cellView": "form"
      },
      "source": [
        "#@title ### Native generation\n",
        "%cd $work_dir\n",
        "%run src/_genSGAN2.py --model $model_pkl --out_dir $output --frames $timeframe $smooth $save_lat --seed $seed\n",
        "ipython_display(ImageSequenceClip(img_list(output), fps=25), center=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GwoOBOcR8APK",
        "cellView": "form"
      },
      "source": [
        "#@title ### Custom size generation\n",
        "\n",
        "sizeX = 400 #@param {type:\"integer\"} \n",
        "sizeY = 300 #@param {type:\"integer\"}\n",
        "size = '%d-%d' % (sizeX, sizeY)\n",
        "scaling = 'pad' #@param ['pad', 'padside', 'symm', 'symmside']\n",
        "\n",
        "%cd $work_dir\n",
        "%run src/_genSGAN2.py --model $model --out_dir $output --frames $timeframe --size $size --scale_type $scaling $smooth $save_lat --seed $seed\n",
        "ipython_display(ImageSequenceClip(img_list(output), fps=25), center=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cOGVIJt18APL",
        "cellView": "form"
      },
      "source": [
        "#@title ### Multi-latent  generation\n",
        "\n",
        "sizeX = 768 #@param {type:\"integer\"} \n",
        "sizeY = 256 #@param {type:\"integer\"}\n",
        "size = '%d-%d' % (sizeX, sizeY)\n",
        "scaling = 'pad' #@param ['pad', 'padside', 'symm', 'symmside']\n",
        "\n",
        "split_X =  3#@param {type:\"integer\"} \n",
        "split_Y =  1#@param {type:\"integer\"}\n",
        "split = '%d-%d' % (split_X, split_Y)\n",
        "splitfine = 0. #@param {type:\"number\"}\n",
        "\n",
        "%cd $work_dir\n",
        "%run src/_genSGAN2.py --model $model --out_dir $output --frames $timeframe --size $size --scale_type $scaling -n $split --splitfine $splitfine $smooth --seed $seed\n",
        "ipython_display(ImageSequenceClip(img_list(output), fps=25), center=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VZ3zGyFy8APL"
      },
      "source": [
        "> Here we get animated composition of 3 independent frames, blended together horizontally.  \n",
        "`splitfine` controls boundary fineness (0 = smoothest/default, higher => thinner).  \n",
        "\n",
        "Instead of frame splitting, we can load external mask from b/w image file (or folder with image sequence):"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zB0-VZ798APL",
        "cellView": "form"
      },
      "source": [
        "#@title ### Masked  generation\n",
        "\n",
        "sizeX = 400 #@param {type:\"integer\"} \n",
        "sizeY = 300 #@param {type:\"integer\"}\n",
        "size = '%d-%d' % (sizeX, sizeY)\n",
        "\n",
        "lat_mask = '_in/mask.jpg' #@param {type:\"string\"} \n",
        "\n",
        "%cd $work_dir\n",
        "%run src/_genSGAN2.py --model $model --out_dir $output --frames $timeframe --size $size --latmask $lat_mask  $smooth --seed $seed\n",
        "ipython_display(ImageSequenceClip(img_list(output), fps=25), center=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sWUJDcGW8APM"
      },
      "source": [
        "`digress` adds some funky displacements by tweaking initial constant layer.  \n",
        "`truncation` controls variety, as usual  (0 = boring, >1 = weird). "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tITi-ILU8APM",
        "cellView": "form"
      },
      "source": [
        "#@title ### Other tricks\n",
        "\n",
        "sizeX = 400 #@param {type:\"integer\"} \n",
        "sizeY = 300 #@param {type:\"integer\"}\n",
        "size = '%d-%d' % (sizeX, sizeY)\n",
        "\n",
        "digress =  2#@param {type:\"number\"} \n",
        "truncation =  1.5#@param {type:\"number\"} \n",
        "\n",
        "%cd $work_dir\n",
        "%run src/_genSGAN2.py --model $model --out_dir $output --frames $timeframe --size $size --digress $digress --trunc $truncation $smooth --seed $seed\n",
        "ipython_display(ImageSequenceClip(img_list(output), fps=25), center=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OLBZAOSi8APM"
      },
      "source": [
        "## Latent space exploration\n",
        "\n",
        "For these experiments download [FFHQ model](https://nvlabs-fi-cdn.nvidia.com/stylegan2/networks/stylegan2-ffhq-config-f.pkl) and save to `models` directory."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VB7ujwkD8APM",
        "cellView": "form"
      },
      "source": [
        "#@title ### Generator setup\n",
        "!cd $work_dir\n",
        "\n",
        "output = '_out/ffhq' #@param {type:\"string\"}\n",
        "\n",
        "frames = 50 #@param {type:\"integer\"}\n",
        "frame_step = 10 #@param {type:\"integer\"}\n",
        "timeframe = '%d-%d' % (frames, frame_step)\n",
        "\n",
        "cubic_smooth = True #@param {type:\"boolean\"}\n",
        "gauss_smooth = False #@param {type:\"boolean\"}\n",
        "cubic_smooth = '--cubic ' if cubic_smooth else ''\n",
        "gauss_smooth = '--gauss ' if gauss_smooth else ''\n",
        "smooth = cubic_smooth + gauss_smooth\n",
        "\n",
        "save_lat = False #@param {type:\"boolean\"}\n",
        "save_lat = '--save_lat ' if save_lat else ''\n",
        "\n",
        "#@markdown Select model from the dropdown list below:\n",
        "model_select(work_dir)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dIbPEUyI8APN"
      },
      "source": [
        "Project external images onto the model dlatent `w` space, saving points as Numpy arrays in `*.npy` files.  \n",
        "Hint: download [VGG model](https://drive.google.com/uc?id=1N2-m9qszOeVC9Tq77WxsLnuWwOedQiD2) and save it as `models/vgg/vgg16_zhang_perceptual.pkl` to avoid auto-downloading errors."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BJx9ws7n8APN",
        "cellView": "form"
      },
      "source": [
        "#@title ### Images projection\n",
        "\n",
        "image_dir = '_in/photo' #@param {type:\"string\"} \n",
        "out_dir = '_out/proj' #@param {type:\"string\"} \n",
        "steps =  1000#@param {type:\"integer\"} \n",
        "\n",
        "%cd $work_dir\n",
        "%run src/project_latent.py --model $model_pkl --in_dir _in/photo --out_dir _out/proj --steps $steps\n",
        "\n",
        "from src.util.utilgan import img_list\n",
        "images = img_list(out_dir, subdir=True)\n",
        "images = [f for f in images if '-%04d.jpg' % steps in f]\n",
        "Image(images[0], width=512, height=512)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X6l471qJWsX2"
      },
      "source": [
        "Produce looped animation, interpolating between saved dlatent points (from a file or directory with `*.npy` or `*.npz` files). To select only few frames from a sequence `xxx.npy`, create text file with comma-delimited frame numbers and save it as `xxx.txt` in the same directory (check examples in `_in/dlats`). "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SIJJUfAN8APN",
        "cellView": "form"
      },
      "source": [
        "#@title ### \"Walk\" between saved dlatent points\n",
        "\n",
        "dlatents = '_in/dlats' #@param {type:\"string\"} \n",
        "out_dir = '_out/ffhq-dlats' #@param {type:\"string\"} \n",
        "\n",
        "%cd $work_dir\n",
        "%run src/_play_dlatents.py --model $model --dlatents $dlatents --out_dir $out_dir --fstep $frame_step $smooth\n",
        "ipython_display(ImageSequenceClip(img_list(out_dir), fps=25), center=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3NBBc2CV8APN"
      },
      "source": [
        "\n",
        "\n",
        "We can also load another dlatent point and apply it to higher network layers, effectively \"stylizing\" images with its' fine features.  \n",
        "`digress` and `truncation` are also applicable here."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "J84NdtP38APO",
        "cellView": "form"
      },
      "source": [
        "#@title ### \"Walk\" between saved points w/tricks\n",
        "\n",
        "dlatents = '_in/dlats' #@param {type:\"string\"} \n",
        "out_dir = '_out/ffhq-dlats' #@param {type:\"string\"} \n",
        "style_dlat = '_in/blonde458.npy' #@param {type:\"string\"} \n",
        "digress =  2#@param {type:\"number\"} \n",
        "truncation =  1.5#@param {type:\"number\"} \n",
        "\n",
        "%cd $work_dir\n",
        "%run src/_play_dlatents.py --model $model --dlatents $dlatents --out_dir $out_dir --fstep $frame_step $smooth --style_dlat $style_dlat --digress $digress --trunc $truncation\n",
        "ipython_display(ImageSequenceClip(img_list(out_dir), fps=25), center=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Iykr7uKL8APO"
      },
      "source": [
        "Generate animation by walking from saved dlatent point along feature direction vectors (aging/smiling/etc) one by one."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UOToqrX68APO",
        "scrolled": true,
        "cellView": "form"
      },
      "source": [
        "#@title ### \"Walk\" along feature directions\n",
        "\n",
        "base_lat = '_in/blonde458.npy' #@param {type:\"string\"} \n",
        "vectors = '_in/vectors_ffhq' #@param {type:\"string\"} \n",
        "out_dir = '_out/ffhq_looks' #@param {type:\"string\"} \n",
        "\n",
        "%cd $work_dir\n",
        "%run src/_play_vectors.py --model $model_pkl --base_lat $base_lat --vector_dir $vectors --out_dir $out_dir --fstep $frame_step\n",
        "ipython_display(ImageSequenceClip(img_list(out_dir), fps=25), center=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L4grqz6UHkcA"
      },
      "source": [
        "> Try discovering more such vectors:\n",
        "> * https://github.com/genforce/sefa\n",
        "> * https://github.com/harskish/ganspace"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aiItkMSM8APO"
      },
      "source": [
        "## Tweaking models\n",
        "\n",
        "NB: No real examples here! The commands are for reference, try with your own files.  \n",
        "All paths are relative to StyleGAN2 copy on G drive. Resulting models will be saved there as well."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4tldFcuR8APP"
      },
      "source": [
        "Strip G/D networks from a full model, leaving only Gs for inference. Resulting file is saved with `-Gs` suffix.  \n",
        "It's recommended to reconstruct the network, saving necessary arguments with it. Useful for foreign downloaded models."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YJwhP_IY8APP",
        "cellView": "form"
      },
      "source": [
        "model = 'snapshot-256.pkl' #@param {type:\"string\"} \n",
        "reconstruct = True #@param {type:\"boolean\"} \n",
        "reconstruct = '-r ' if reconstruct else ''\n",
        "%cd $work_dir\n",
        "%run src/model_convert.py --source $model $reconstruct"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iOS24WCS8APP"
      },
      "source": [
        "Add or remove layers (from a trained model) to adjust its resolution for further finetuning. For given example, this would produce new model with 512px resolution, populating weights on the layers up to 256px from the source snapshot; the rest will be initialized randomly. It also can decrease resolution (say, make 512 from 1024). Note that this effectively changes the number of layers in the model.   \n",
        "You can also add alpha channel to work with RGBA data.  \n",
        "This option works with complete (G/D/Gs) models only, since its' purpose is transfer-learning (the resulting model will contain either partially random weights, or wrong `ToRGB` params). "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9gqba1BY8APP",
        "cellView": "form"
      },
      "source": [
        "model = 'snapshot-256.pkl' #@param {type:\"string\"}\n",
        "resolution = \"1024\" #@param [256, 512, 1024]\n",
        "resolution = int(resolution)\n",
        "add_alpha = False #@param {type:\"boolean\"}\n",
        "\n",
        "%cd $work_dir\n",
        "if add_alpha:\n",
        "  %run src/model_convert.py --source $model --res $resolution --alpha\n",
        "else:\n",
        "  %run src/model_convert.py --source $model --res $resolution"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PUtk27qe8APP"
      },
      "source": [
        "Change aspect ratio of a trained model by cropping or padding layers (keeping their count). Originally from @Aydao. This is experimental function with some voluntary logic, so use with care.  \n",
        "This produces working non-square model. In case of basic aspect conversion (like 4x4 => 5x3), full models (G/D/Gs) will be trainable for further finetuning."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SKguO2Iw8APP",
        "cellView": "form"
      },
      "source": [
        "model = 'snapshot-1024.pkl' #@param {type:\"string\"} \n",
        "resX = \"1280\" #@param [256, 384, 512, 640, 768, 1024, 1280, 1536]\n",
        "resX = int(resX)\n",
        "resY = \"768\" #@param [256, 384, 512, 640, 768, 1024, 1280, 1536]\n",
        "resY = int(resY)\n",
        "res = '%d-%d' % (resX, resY)\n",
        "\n",
        "%cd $work_dir\n",
        "%run src/model_convert.py --source $model --res $res"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xJsXdF5P8APP"
      },
      "source": [
        "Combine lower layers (pose/shape) from one model with higher layers (style/texture) from another. The models are switched at the layer with size `resolution` (usually 32/64); `level` is 0 or 1.  \n",
        "For inference (generation) this method works properly only for models from one \"family\", i.e. uptrained (finetuned) from the same original model. For training may be useful in other cases too (not tested yet)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AJLnOqFQ8APQ",
        "cellView": "form"
      },
      "source": [
        "model1 = 'model1.pkl' #@param {type:\"string\"}\n",
        "model2 = 'model2.pkl' #@param {type:\"string\"}\n",
        "resolution = \"64\" #@param [16, 32, 64, 128]\n",
        "resolution = int(resolution)\n",
        "level = 0 #@param [0, 1]\n",
        "level = int(level)\n",
        "\n",
        "%cd $work_dir\n",
        "%run src/models_blend.py --pkl1 $model1 --pkl2 $model2 --res $resolution --level $level"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tRTS_bpI8APQ"
      },
      "source": [
        "Mix few models by stochastic averaging all weights.  \n",
        "This method also works properly only for models from one \"family\", i.e. uptrained (finetuned) from the same original model. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "PRatWRXC8APQ",
        "cellView": "form"
      },
      "source": [
        "models_dir = 'models' #@param {type:\"string\"}\n",
        "\n",
        "%cd $work_dir\n",
        "%run src/models_swa.py --in_dir $models_dir"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}