StyleGAN2a_colab.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "StyleGAN2a_colab.ipynb",
      "provenance": []
    },
    "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"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xzA1-mt88AO_"
      },
      "source": [
        "# StyleGAN2-ada 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",
        "import locale\n",
        "locale.getpreferredencoding = lambda: 'UTF-8'\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\", category=UserWarning)\n",
        "!pip install gputil\n",
        "\n",
        "from IPython.display import HTML, Image, display, clear_output\n",
        "from moviepy.editor import ImageSequenceClip, ipython_display\n",
        "import ipywidgets as widgets\n",
        "import os\n",
        "import numpy as np\n",
        "\n",
        "!apt-get -qq install ffmpeg\n",
        "!pip install ninja\n",
        "from google.colab import drive\n",
        "drive.mount('/G', force_remount=True)\n",
        "gdir = '/G/MyDrive/'\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 = 'sg2a_eps' #@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 https://github.com/eps696/stylegan2ada $work_dir\n",
        "%cd $work_dir\n",
        "!pip install -r requirements.txt\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",
        "clear_output()\n",
        "\n",
        "# Hardware check\n",
        "!ln -sf /opt/bin/nvidia-smi /usr/bin/nvidia-smi\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": "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), you may want to enable `multicrop` to crop square fragments from bigger images, effectively increasing amount of data (not suitable for conditional data, as it would break folder structure!). The images 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.   \n",
        "*Restart session (Runtime) after `multicrop` run*."
      ]
    },
    {
      "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",
        "content_data = os.path.join('/content', dataset)\n",
        "dataset_dir = os.path.join('/content/data/', 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_data\n",
        "!rm -rf $dataset_dir\n",
        "os.makedirs(dataset_dir, exist_ok=True)\n",
        "\n",
        "!mkdir /content/tmp\n",
        "%cd /content/tmp\n",
        "fpath = os.path.join(gdir, source)\n",
        "!unzip -o -q $fpath -d $dataset\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 $content_data --size $size --step $overlap\n",
        "    !mv $content_data $dataset_dir\n",
        "  else:\n",
        "    print('/content/tmp not found!!')\n",
        "else:\n",
        "  !mv /content/tmp/* $dataset_dir\n"
      ],
      "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 = 1000 #@param {type:\"integer\"}\n",
        "data_dir = os.path.join('/content/data', dataset)\n",
        "\n",
        "%run src/train.py --data $data_dir --kimg $kimg"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wYlI1fXe8APH"
      },
      "source": [
        "> This will run training process, according to the options and settings in `src/train.py` (check and explore those!!). 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 `--kimg X` argument (training duration in thousands of images). Reasonable `kimg` value for full training from scratch is 5000-8000, while for finetuning 1000-2000 may be sufficient.  \n",
        "\n",
        "> *NB: If you get `RuntimeError: context has already been set` (e.g. after training interruption or `multicrop` run) - restart session, run General setup and Train again.*\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-auto-apa-gf_bnc' #@param {type:\"string\"}\n",
        "kimg = 1000 #@param {type:\"integer\"}\n",
        "data_dir = os.path.join('/content/data', dataset)\n",
        "\n",
        "%run src/train.py --data $data_dir --resume $resume_dir --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 `train/ffhq-256.pkl` 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('/content/data', dataset)\n",
        "\n",
        "%run src/train.py --data $data_dir --resume $resume_pkl --kimg $kimg\n",
        "\n",
        "#@markdown NB: add `--fmaps_fix` argument if you finetune models, trained in the original Tensorflow-based version."
      ],
      "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",
        "extension = 'jpg' #@param ['jpg', 'png']\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 (`w`) points 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 --ext $extension\n",
        "\n",
        "ipython_display(ImageSequenceClip(img_list(output), fps=15), rd_kwargs={'logger':None}, 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 --ext $extension\n",
        "ipython_display(ImageSequenceClip(img_list(output), fps=15), rd_kwargs={'logger':None}, 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 --ext $extension\n",
        "ipython_display(ImageSequenceClip(img_list(output), fps=15), rd_kwargs={'logger':None}, 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 --ext $extension\n",
        "ipython_display(ImageSequenceClip(img_list(output), fps=15), rd_kwargs={'logger':None}, 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 --ext $extension\n",
        "ipython_display(ImageSequenceClip(img_list(output), fps=15), rd_kwargs={'logger':None}, 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",
        "extension = 'jpg' #@param ['jpg', 'png']\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."
      ]
    },
    {
      "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/projector.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 'target.jpg' not 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 --ext $extension\n",
        "ipython_display(ImageSequenceClip(img_list(out_dir), fps=15), rd_kwargs={'logger':None}, 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 --ext $extension\n",
        "ipython_display(ImageSequenceClip(img_list(out_dir), fps=15), rd_kwargs={'logger':None}, 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 --ext $extension\n",
        "ipython_display(ImageSequenceClip(img_list(out_dir), fps=25), rd_kwargs={'logger':None}, center=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XN_82CUzLWdL"
      },
      "source": [
        "> Try discovering more such vectors:\n",
        "> * https://github.com/genforce/sefa\n",
        "> * https://github.com/harskish/ganspace"
      ]
    },
    {
      "cell_type": "markdown",
      "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."
      ],
      "metadata": {
        "id": "7p3jmPGfwg5F"
      }
    },
    {
      "cell_type": "markdown",
      "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."
      ],
      "metadata": {
        "id": "ABv4-b2UwyQg"
      }
    },
    {
      "cell_type": "code",
      "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"
      ],
      "metadata": {
        "cellView": "form",
        "id": "ydJfoJ9iw2Wx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "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)."
      ],
      "metadata": {
        "id": "OSnTq23hw55r"
      }
    },
    {
      "cell_type": "code",
      "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"
      ],
      "metadata": {
        "cellView": "form",
        "id": "WAKdrONZw8zl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "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."
      ],
      "metadata": {
        "id": "HFwhCRf6xCPd"
      }
    },
    {
      "cell_type": "code",
      "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"
      ],
      "metadata": {
        "cellView": "form",
        "id": "gCERMDEzxGHZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Add fake labels to a trained model for further finetuning as conditional"
      ],
      "metadata": {
        "id": "x4UO0O5udkuw"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model = 'snapshot-1024.pkl' #@param {type:\"string\"}\n",
        "labels = 0 #@param {type:\"integer\"}\n",
        "\n",
        "%cd $work_dir\n",
        "%run src/model_convert.py --source $model --labels $labels"
      ],
      "metadata": {
        "cellView": "form",
        "id": "3f1FHjxmd62F"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Convert `src_pt_model`, created in [Rosinality](https://github.com/rosinality/stylegan2-pytorch/) or in [StyleGAN2-NADA](https://github.com/rinongal/StyleGAN-nada) repos to SG2-ada-pytorch PKL format. For the moment it requires also a `base_pkl_model` of the same resolution, to copy the weights onto it (if you don't have it - grab one from Nvidia and tweak it with the tools above to fit your configuration).\n",
        "\n",
        "Resulting SG2-ada-pytorch model is saved with the same name as `src_pt_model`, but with `pkl` extension."
      ],
      "metadata": {
        "id": "oJdTEmG4LH-k"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "src_model_pt = 'sg2-1024.pt' #@param {type:\"string\"}\n",
        "base_pkl_model = 'sg2-ada-pt-1024.pkl' #@param {type:\"string\"}\n",
        "\n",
        "%cd $work_dir\n",
        "%run src/model_pt2pkl.py --model_pt $src_model_pt --model_pkl $base_pkl_model"
      ],
      "metadata": {
        "cellView": "form",
        "id": "qkko3OClMbQY"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}