Colab notebook

Author: rbbb

Finetune_VAE.ipynb

@@ -0,0 +1,176 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "gpuType": "T4"
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "source": [
+        "**Install conda, taming-transformers and its dependencies**"
+      ],
+      "metadata": {
+        "id": "EZx3a2Dsc_wq"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "!pip install -q condacolab\n",
+        "import condacolab\n",
+        "condacolab.install()\n",
+        "\n",
+        "!git clone https://github.com/rbbb/taming-transformers\n",
+        "%cd /content/taming-transformers\n",
+        "!git checkout stable-diff-vae-finetuning\n",
+        "\n",
+        "!conda env update -n taming-transformers -f environment.yaml\n",
+        "!conda run -n taming-transformers pip install -e ."
+      ],
+      "metadata": {
+        "id": "UYDCEq4iE_iF"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "**Put some images in /content/src_img_here/**"
+      ],
+      "metadata": {
+        "id": "uyDtJvL-dOWr"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "!mkdir /content/src_img_here/\n",
+        "#directory /0000/ from danbooru 2020 data containing 3.5k images\n",
+        "!unzip /content/drive/MyDrive/danbooru_tiny.zip -d /content/src_img_here/"
+      ],
+      "metadata": {
+        "id": "Zme4fnUSWwyo"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "**Find all images, put their name in train_img.txt and val_img.txt (10 validation images)**"
+      ],
+      "metadata": {
+        "id": "S55xG8OEdYjI"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import glob\n",
+        "img_filenames = list(glob.glob(\"/content/src_img_here/**/*.*\", recursive=True))\n",
+        "img_filenames = [i for i in img_filenames if (i.endswith(\".png\") or i.endswith(\".jpg\") or i.endswith(\".jpeg\"))]\n",
+        "with open(\"/content/taming-transformers/train_img.txt\",\"w\") as f:\n",
+        "  f.write(\"\\n\".join(img_filenames[10:]))\n",
+        "with open(\"/content/taming-transformers/val_img.txt\",\"w\") as f:\n",
+        "  f.write(\"\\n\".join(img_filenames[0:10]))\n"
+      ],
+      "metadata": {
+        "id": "XvInztU0Lx6y"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "**Get a VAE from civitai, convert from safetensors to ckpt**"
+      ],
+      "metadata": {
+        "id": "z71bh_0vdoap"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "!mkdir /content/src_vae_here/\n",
+        "!wget -O /content/src_vae_here/vae.safetensors https://civitai.com/api/download/models/88156 #Clear VAE\n",
+        "\n",
+        "!pip install safetensors torch\n",
+        "from safetensors import safe_open\n",
+        "import torch\n",
+        "vae = {}\n",
+        "with safe_open(\"/content/src_vae_here/vae.safetensors\", framework=\"pt\", device=0) as f:\n",
+        "    vae = {k: f.get_tensor(k) for k in f.keys()}\n",
+        "torch.save({'state_dict':vae}, \"/content/src_vae_here/vae.pt\")"
+      ],
+      "metadata": {
+        "id": "kPgbHeCpW-Et"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "**Change the parameters in the /content/taming-transformers/configs/finetune_vae.yaml (name of VAE file and batch size)**"
+      ],
+      "metadata": {
+        "id": "jg7BpQ4ad5xG"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import yaml\n",
+        "with open('/content/taming-transformers/configs/finetune_vae.yaml','r') as file:\n",
+        "    conf = yaml.load(file, Loader=yaml.FullLoader)\n",
+        "\n",
+        "conf['model']['params']['ckpt_path'] = \"\\\"/content/src_vae_here/vae.pt\\\"\"\n",
+        "conf['data']['params']['batch_size'] = \"5\"\n",
+        "\n",
+        "with open('/content/taming-transformers/configs/finetune_vae.yaml','w') as file:\n",
+        "    yaml.dump(conf, file)\n",
+        "\n"
+      ],
+      "metadata": {
+        "id": "0AX9Q_6zNu1C"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "**Run**"
+      ],
+      "metadata": {
+        "id": "-GXysMjAeLuC"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "%cd /content/taming-transformers\n",
+        "!conda run --no-capture-output -n taming-transformers python main.py --base configs/finetune_vae.yaml -t True --gpus 0, --accumulate_grad_batches 6"
+      ],
+      "metadata": {
+        "id": "GSOiQtWKYzeV"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}