StyleGAN2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# StyleGAN2 operations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# modules for notebook logging\n",
    "!pip install IPython\n",
    "!pip install ipywidgets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, let's prepare the data. Ensure all your images have the same color channels (monochrome, RGB or RGBA).  \n",
    "\n",
    "If you work with patterns or shapes (rather than compostions), you can crop square fragments from bigger images (effectively multiplying their amount). For that, edit source and target paths below; `size` is fragment resolution, `step` is shift between the fragments. This will cut every source image into 512x512px fragments, overlapped with 256px shift by X and Y. Edit `size` and `overlap`  according to your dataset. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "src_dir = 'data/src'\n",
    "data_dir = 'data/mydata'\n",
    "size = 512\n",
    "step = 256\n",
    "\n",
    "%run src/util/multicrop.py --in_dir $src_dir --out_dir $data_dir --size $size --step $step"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you edit the images yourself (e.g. for non-square aspect ratios), ensure their correct size and put to the `data_dir`. For conditional model split the data by subfolders (`mydata/1`, `mydata/2`, ..) and add `--cond` option to the training command below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = 'data/mydata'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Convert directory with images to TFRecords dataset (`mydata-512x512.tfr` file in `data` directory):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run src/training/dataset_tool.py --data $data_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we can train StyleGAN2 on the prepared dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run src/train.py --data $data_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> This will run training process, according to the options in `src/train.py` (check and explore those!!). If there was no TFRecords file from the previous step, it will be created at this point. 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. `mydata-512-0360.pkl`), and full (containing G/D/Gs networks for further training) as `snapshot-...pkl`. \n",
    "\n",
    "> By default, the most powerful SG2 config (F) is used; if you face OOM issue, you may resort to `--config E`, requiring less memory (with poorer results, of course). For small datasets (100x images instead of 10000x) one should add `--d_aug` option to use [Differential Augmentation](https://github.com/mit-han-lab/data-efficient-gans) for more effective training. \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 in `--d_aug` mode 1000-2000 may be sufficient."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the training process was interrupted, we can resume it from the last saved model as following:  \n",
    "*(replace `000-mydata-512-f` with existing training directory)*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run src/train.py --data $data_dir --resume train/000-mydata-512-f"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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 `ffhq-512.pkl` model already):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run src/train.py --data $data_dir --resume train/ffhq-512.pkl --d_aug --kimg 1000 --finetune"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generation\n",
    "\n",
    "Let's produce some imagery from the original cat model (download it from [here](https://nvlabs-fi-cdn.nvidia.com/stylegan2/networks/stylegan2-cat-config-f.pkl) and put to `models` directory).  \n",
    "More cool models can be found [here](https://github.com/justinpinkney/awesome-pretrained-stylegan2)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image, Video\n",
    "\n",
    "model = 'models/stylegan2-cat-config-f' # without \".pkl\" extension\n",
    "model_pkl = model + '.pkl' # with \".pkl\" extension\n",
    "output = '_out/cats'\n",
    "frames = '50-10'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate some animation to test the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run src/_genSGAN2.py --model $model_pkl --out_dir $output --frames $frames\n",
    "\n",
    "Image(filename = output + '/000000.jpg') # show first frame"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vlzqD3DEGzIH"
   },
   "source": [
    "> Here we loaded the model 'as is', and produced 50 frames in its natural resolution, interpolating between random latent `z` space keypoints, with a step of 10 frames between keypoints.\n",
    "\n",
    "If you have `ffmpeg` installed, you can convert it into video:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "out_sequence = output + '/%06d.jpg'\n",
    "out_video = output + '.mp4'\n",
    "\n",
    "!ffmpeg -y -v warning -i $out_sequence $out_video\n",
    "\n",
    "Video(out_video)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> \n",
    "Now let's generate custom animation. For that we omit model extension, so the script would load custom network, effectively enabling special features, e.g. arbitrary resolution (set by `--size` argument in `X-Y` format).  \n",
    "`--cubic` option changes linear interpolation to cubic for smoother animation (there is also `--gauss` option for additional smoothing)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run src/_genSGAN2.py --model $model --out_dir $output --frames $frames --size 400-300 --cubic\n",
    "Image(output+'/000000.jpg')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Run `ffmpeg` command above after each generation, if you want to check results in motion.**  \n",
    "Adding `--save_lat` option will save all traversed dlatent points (in `w` space) as Numpy array in `*.npy` file (useful for further curating). Set `--seed X` value to produce repeatable results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate more various imagery:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run src/_genSGAN2.py --model $model --out_dir $output --frames $frames --size 768-256 -n 3-1\n",
    "Image(output+'/000000.jpg')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Here we get animated composition of 3 independent frames, blended together horizontally (like the image in the repo header). Argument `--splitfine X` controls boundary fineness (0 = smoothest/default, higher => thinner).  \n",
    "\n",
    "Instead of frame splitting, we can load external mask from b/w image file (it also could be folder with file sequence):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run src/_genSGAN2.py --model $model --out_dir $output --frames $frames --size 400-300 --latmask _in/mask.jpg\n",
    "Image(output+'/000000.jpg')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> \n",
    "`--digress X` adds some funky displacements with X strength (by tweaking initial constant layer).  \n",
    "`--trunc X` controls truncation psi parameter (0 = boring, 1+ = weird). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run src/_genSGAN2.py --model $model --out_dir $output --frames $frames --digress 2 --trunc 0.5\n",
    "Image(output+'/000000.jpg')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Don't forget to check other options of `_genSGAN2.py` by `--help` argument."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image, Video\n",
    "\n",
    "model = 'models/stylegan2-ffhq-config-f' # without \".pkl\" extension\n",
    "model_pkl = model + '.pkl' # with \".pkl\" extension"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> \n",
    "Project external images (aligned face portraits) from `_in/photo` onto StyleGAN2 model dlatent `w` space. \n",
    "Results (found dlatent points as Numpy arrays in `*.npy` files, and video/still previews) are saved to `_out/proj` directory.  \n",
    "NB: first download [VGG model](https://drive.google.com/uc?id=1N2-m9qszOeVC9Tq77WxsLnuWwOedQiD2) and save it as `models/vgg/vgg16_zhang_perceptual.pkl`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run src/project_latent.py --model $model_pkl --in_dir _in/photo --out_dir _out/proj "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> \n",
    "Generate animation between saved dlatent points:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dlat = 'dlats'\n",
    "path_in = '_in/' + dlat\n",
    "path_out = '_out/ffhq-' + dlat\n",
    "out_sequence = path_out + '/%06d.jpg'\n",
    "out_video = path_out + '.mp4'\n",
    "\n",
    "%run src/_play_dlatents.py --model $model --dlatents $path_in --out_dir $path_out --fstep 10\n",
    "Image(path_out+'/000000.jpg', width=512, height=512)\n",
    "\n",
    "!ffmpeg -y -v warning -i $out_sequence $out_video\n",
    "Video(out_video, width=512, height=512)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> This loads saved dlatent points from `_in/dlats` and produces smooth looped animation between them (with interpolation step of 50 frames). `dlats` may be a file or a directory with `*.npy` or `*.npz` files. To select only few frames from a sequence `somename.npy`, create text file with comma-delimited frame numbers and save it as `somename.txt` in the same directory (check given examples for FFHQ model).\n",
    "\n",
    "Style-blending argument `--style_dlat blonde458.npy` would also load dlatent from `blonde458.npy` and apply it to higher network layers. `--cubic` smoothing and `--digress X` displacements are also applicable here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run src/_play_dlatents.py --model $model --dlatents $path_in --out_dir $path_out --fstep 10 --style_dlat _in/blonde458.npy --digress 2 --cubic\n",
    "!ffmpeg -y -v warning -i $out_sequence $out_video\n",
    "Video(out_video, width=512, height=512)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> \n",
    "Generate animation by moving saved dlatent point `_in/blonde458.npy` along feature direction vectors from `_in/vectors_ffhq` (aging/smiling/etc) one by one: (check preview window!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%run src/_play_vectors.py --model $model_pkl --base_lat _in/blonde458.npy --vector_dir _in/vectors_ffhq --out_dir _out/ffhq_looks\n",
    "\n",
    "!ffmpeg -y -v warning -i _out/ffhq_looks/%06d.jpg _out/ffhq-vectors.mp4\n",
    "Video('_out/ffhq-vectors.mp4', width=512, height=512)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Such vectors can be discovered, for example, with the following methods:\n",
    "> * https://github.com/genforce/sefa\n",
    "> * https://github.com/harskish/ganspace\n",
    "> * https://github.com/LoreGoetschalckx/GANalyze"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tweaking models\n",
    "\n",
    "NB: No real examples here! Just reference commands, try with your own files."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Strip G/D networks from a full model, leaving only Gs for inference. Resulting file is saved with `-Gs` suffix. It's recommended to add `-r` option to reconstruct the network, saving necessary arguments with it. Useful for foreign downloaded models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run src/model_convert.py --source snapshot-1024.pkl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add or remove layers (from a trained model) to adjust its resolution for further finetuning. This will 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",
    "This option works with complete (G/D/Gs) models only, since it's purposed for transfer-learning (the resulting model will contain either partially random weights, or wrong `ToRGB` params). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run src/model_convert.py --source snapshot-256.pkl --res 1024"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "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. This produces working non-square model. In case of basic aspect conversion (like 4x4 => 5x3), complete models (G/D/Gs) will be trainable for further finetuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run src/model_convert.py --source snapshot-1024.pkl --res 1280-768"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Add alpha channel to a trained model for further finetuning:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run src/model_convert.py --source snapshot-1024.pkl --alpha"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "All above (adding/cropping/padding layers + alpha channel) can be done in one shot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run src/model_convert.py --source snapshot-256.pkl --res 1280-768 --alpha"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Combine lower layers from one model with higher layers from another. `<res>` is resolution, at which the models are switched (usually 16/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",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run src/models_blend.py --pkl1 model1.pkl --pkl2 model2.pkl --res <res> --level <level>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Mix few models by stochastic averaging all weights. This would work properly only for models from one \"family\", i.e. uptrained (finetuned) from the same original model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run src/models_swa.py --in_dir <models_dir>"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "StyleGAN2.ipynb",
   "private_outputs": true,
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}