kohya-dreambooth.ipynb

{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "slgjeYgd6pWp"
      },
      "source": [
        "![visitors](https://visitor-badge.glitch.me/badge?page_id=linaqruf.kohya-dreambooth) [![](https://dcbadge.vercel.app/api/shield/850007095775723532?style=flat)](https://lookup.guru/850007095775723532) [![ko-fi](https://img.shields.io/badge/Support%20me%20on%20Ko--fi-F16061?logo=ko-fi&logoColor=white&style=flat)](https://ko-fi.com/linaqruf) <a href=\"https://saweria.co/linaqruf\"><img alt=\"Saweria\" src=\"https://img.shields.io/badge/Saweria-7B3F00?style=flat&logo=ko-fi&logoColor=white\"/></a>\n",
        "\n",
        "# **Kohya Dreambooth**\n",
        "A Colab Notebook For Dreambooth Training\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "OQqyYzkDFA0a"
      },
      "source": [
        "| Notebook Name | Description | Link | V14 |\n",
        "| --- | --- | --- | --- |\n",
        "| [Kohya LoRA Dreambooth](https://github.com/Linaqruf/kohya-trainer/blob/main/kohya-LoRA-dreambooth.ipynb) | LoRA Training (Dreambooth method) | [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/kohya-LoRA-dreambooth.ipynb) | [![](https://img.shields.io/static/v1?message=Older%20Version&logo=googlecolab&labelColor=5c5c5c&color=e74c3c&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/ff701379c65380c967cd956e4e9e8f6349563878/kohya-LoRA-dreambooth.ipynb) | \n",
        "| [Kohya LoRA Fine-Tuning](https://github.com/Linaqruf/kohya-trainer/blob/main/kohya-LoRA-finetuner.ipynb) | LoRA Training (Fine-tune method) | [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/kohya-LoRA-finetuner.ipynb) | [![](https://img.shields.io/static/v1?message=Older%20Version&logo=googlecolab&labelColor=5c5c5c&color=e74c3c&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/ff701379c65380c967cd956e4e9e8f6349563878/kohya-LoRA-finetuner.ipynb) | \n",
        "| [Kohya Trainer](https://github.com/Linaqruf/kohya-trainer/blob/main/kohya-trainer.ipynb) | Native Training | [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/kohya-trainer.ipynb) | [![](https://img.shields.io/static/v1?message=Older%20Version&logo=googlecolab&labelColor=5c5c5c&color=e74c3c&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/ff701379c65380c967cd956e4e9e8f6349563878/kohya-trainer.ipynb) | \n",
        "| [Kohya Dreambooth](https://github.com/Linaqruf/kohya-trainer/blob/main/kohya-dreambooth.ipynb) | Dreambooth Training | [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/kohya-dreambooth.ipynb) | [![](https://img.shields.io/static/v1?message=Older%20Version&logo=googlecolab&labelColor=5c5c5c&color=e74c3c&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/ff701379c65380c967cd956e4e9e8f6349563878/kohya-dreambooth.ipynb) | \n",
        "| [Cagliostro Colab UI](https://github.com/Linaqruf/sd-notebook-collection/blob/main/cagliostro-colab-ui.ipynb) `NEW`| A Customizable Stable Diffusion Web UI| [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/sd-notebook-collection/blob/main/cagliostro-colab-ui.ipynb) | "
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "tTVqCAgSmie4"
      },
      "source": [
        "# I. Install Kohya Trainer"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "_u3q60di584x"
      },
      "outputs": [],
      "source": [
        "# @title ## 1.1. Install Dependencies\n",
        "# @markdown Clone Kohya Trainer from GitHub and check for updates. Use textbox below if you want to checkout other branch or old commit. Leave it empty to stay the HEAD on main.  This will also install the required libraries.\n",
        "import os\n",
        "import zipfile\n",
        "import shutil\n",
        "import time\n",
        "from subprocess import getoutput\n",
        "from IPython.utils import capture\n",
        "from google.colab import drive\n",
        "\n",
        "%store -r\n",
        "\n",
        "# root_dir\n",
        "root_dir = \"/content\"\n",
        "deps_dir = os.path.join(root_dir, \"deps\")\n",
        "repo_dir = os.path.join(root_dir, \"kohya-trainer\")\n",
        "training_dir = os.path.join(root_dir, \"dreambooth\")\n",
        "pretrained_model = os.path.join(root_dir, \"pretrained_model\")\n",
        "vae_dir = os.path.join(root_dir, \"vae\")\n",
        "config_dir = os.path.join(training_dir, \"config\")\n",
        "\n",
        "# repo_dir\n",
        "accelerate_config = os.path.join(repo_dir, \"accelerate_config/config.yaml\")\n",
        "tools_dir = os.path.join(repo_dir, \"tools\")\n",
        "finetune_dir = os.path.join(repo_dir, \"finetune\")\n",
        "\n",
        "for store in [\n",
        "    \"root_dir\",\n",
        "    \"deps_dir\",\n",
        "    \"repo_dir\",\n",
        "    \"training_dir\",\n",
        "    \"pretrained_model\",\n",
        "    \"vae_dir\",\n",
        "    \"accelerate_config\",\n",
        "    \"tools_dir\",\n",
        "    \"finetune_dir\",\n",
        "    \"config_dir\",\n",
        "]:\n",
        "    with capture.capture_output() as cap:\n",
        "        %store {store}\n",
        "        del cap\n",
        "\n",
        "repo_url = \"https://github.com/Linaqruf/kohya-trainer\"\n",
        "bitsandytes_main_py = \"/usr/local/lib/python3.10/dist-packages/bitsandbytes/cuda_setup/main.py\"\n",
        "branch = \"\"  # @param {type: \"string\"}\n",
        "install_xformers = True  # @param {'type':'boolean'}\n",
        "mount_drive = False  # @param {type: \"boolean\"}\n",
        "verbose = False # @param {type: \"boolean\"}\n",
        "\n",
        "def read_file(filename):\n",
        "    with open(filename, \"r\") as f:\n",
        "        contents = f.read()\n",
        "    return contents\n",
        "\n",
        "\n",
        "def write_file(filename, contents):\n",
        "    with open(filename, \"w\") as f:\n",
        "        f.write(contents)\n",
        "\n",
        "\n",
        "def clone_repo(url):\n",
        "    if not os.path.exists(repo_dir):\n",
        "        os.chdir(root_dir)\n",
        "        !git clone {url} {repo_dir}\n",
        "    else:\n",
        "        os.chdir(repo_dir)\n",
        "        !git pull origin {branch} if branch else !git pull\n",
        "\n",
        "\n",
        "def ubuntu_deps(url, name, dst):\n",
        "    !wget {'-q' if not verbose else ''} --show-progress {url}\n",
        "    with zipfile.ZipFile(name, \"r\") as deps:\n",
        "        deps.extractall(dst)\n",
        "    !dpkg -i {dst}/*\n",
        "    os.remove(name)\n",
        "    shutil.rmtree(dst)\n",
        "\n",
        "\n",
        "def install_dependencies():\n",
        "    s = getoutput('nvidia-smi')\n",
        "\n",
        "    if 'T4' in s:\n",
        "        !sed -i \"s@cpu@cuda@\" library/model_util.py\n",
        "\n",
        "    !pip install {'-q' if not verbose else ''} --upgrade -r requirements.txt\n",
        "    !pip install {'-q' if not verbose else ''} torch==2.0.0+cu118 torchvision==0.15.1+cu118 torchaudio==2.0.1+cu118 torchtext==0.15.1 torchdata==0.6.0 --extra-index-url https://download.pytorch.org/whl/cu118 -U\n",
        "\n",
        "    if install_xformers:\n",
        "        !pip install {'-q' if not verbose else ''} xformers==0.0.19 triton==2.0.0 -U\n",
        "        \n",
        "    from accelerate.utils import write_basic_config\n",
        "\n",
        "    if not os.path.exists(accelerate_config):\n",
        "        write_basic_config(save_location=accelerate_config)\n",
        "\n",
        "\n",
        "def remove_bitsandbytes_message(filename):\n",
        "    welcome_message = \"\"\"\n",
        "def evaluate_cuda_setup():\n",
        "    print('')\n",
        "    print('='*35 + 'BUG REPORT' + '='*35)\n",
        "    print('Welcome to bitsandbytes. For bug reports, please submit your error trace to: https://github.com/TimDettmers/bitsandbytes/issues')\n",
        "    print('For effortless bug reporting copy-paste your error into this form: https://docs.google.com/forms/d/e/1FAIpQLScPB8emS3Thkp66nvqwmjTEgxp8Y9ufuWTzFyr9kJ5AoI47dQ/viewform?usp=sf_link')\n",
        "    print('='*80)\"\"\"\n",
        "\n",
        "    new_welcome_message = \"\"\"\n",
        "def evaluate_cuda_setup():\n",
        "    import os\n",
        "    if 'BITSANDBYTES_NOWELCOME' not in os.environ or str(os.environ['BITSANDBYTES_NOWELCOME']) == '0':\n",
        "        print('')\n",
        "        print('=' * 35 + 'BUG REPORT' + '=' * 35)\n",
        "        print('Welcome to bitsandbytes. For bug reports, please submit your error trace to: https://github.com/TimDettmers/bitsandbytes/issues')\n",
        "        print('For effortless bug reporting copy-paste your error into this form: https://docs.google.com/forms/d/e/1FAIpQLScPB8emS3Thkp66nvqwmjTEgxp8Y9ufuWTzFyr9kJ5AoI47dQ/viewform?usp=sf_link')\n",
        "        print('To hide this message, set the BITSANDBYTES_NOWELCOME variable like so: export BITSANDBYTES_NOWELCOME=1')\n",
        "        print('=' * 80)\"\"\"\n",
        "\n",
        "    contents = read_file(filename)\n",
        "    new_contents = contents.replace(welcome_message, new_welcome_message)\n",
        "    write_file(filename, new_contents)\n",
        "\n",
        "\n",
        "def main():\n",
        "    os.chdir(root_dir)\n",
        "\n",
        "    if mount_drive:\n",
        "        if not os.path.exists(\"/content/drive\"):\n",
        "            drive.mount(\"/content/drive\")\n",
        "\n",
        "    for dir in [\n",
        "        deps_dir, \n",
        "        training_dir, \n",
        "        config_dir, \n",
        "        pretrained_model, \n",
        "        vae_dir\n",
        "    ]:\n",
        "        os.makedirs(dir, exist_ok=True)\n",
        "\n",
        "    clone_repo(repo_url)\n",
        "\n",
        "    if branch:\n",
        "        os.chdir(repo_dir)\n",
        "        status = os.system(f\"git checkout {branch}\")\n",
        "        if status != 0:\n",
        "            raise Exception(\"Failed to checkout branch or commit\")\n",
        "\n",
        "    os.chdir(repo_dir)\n",
        "    \n",
        "    !apt -y update {'-qq' if not verbose else ''}\n",
        "    !apt install libunwind8-dev {'-qq' if not verbose else ''}\n",
        "\n",
        "    ubuntu_deps(\n",
        "        \"https://huggingface.co/Linaqruf/fast-repo/resolve/main/deb-libs.zip\",\n",
        "        \"deb-libs.zip\",\n",
        "        deps_dir,\n",
        "    )\n",
        "\n",
        "    install_dependencies()\n",
        "    time.sleep(3)\n",
        "    \n",
        "    remove_bitsandbytes_message(bitsandytes_main_py)\n",
        "\n",
        "    os.environ[\"LD_PRELOAD\"] = \"libtcmalloc.so\"\n",
        "    os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n",
        "    os.environ[\"BITSANDBYTES_NOWELCOME\"] = \"1\"  \n",
        "    os.environ[\"SAFETENSORS_FAST_GPU\"] = \"1\"\n",
        "\n",
        "    cuda_path = \"/usr/local/cuda-11.8/targets/x86_64-linux/lib/\"\n",
        "    ld_library_path = os.environ.get(\"LD_LIBRARY_PATH\", \"\")\n",
        "    os.environ[\"LD_LIBRARY_PATH\"] = f\"{ld_library_path}:{cuda_path}\"\n",
        "\n",
        "main()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ZmIRAxgEQESm"
      },
      "outputs": [],
      "source": [
        "# @title ## 1.2. Start `File Explorer`\n",
        "# @markdown This will work in real-time even when you run other cells\n",
        "import threading\n",
        "from google.colab import output\n",
        "from imjoy_elfinder.app import main\n",
        "\n",
        "open_in_new_tab = False  # @param {type:\"boolean\"}\n",
        "\n",
        "def start_file_explorer(root_dir=root_dir, port=8765):\n",
        "    try:\n",
        "        main([\"--root-dir=\" + root_dir, \"--port=\" + str(port)])\n",
        "    except Exception as e:\n",
        "        print(\"Error starting file explorer:\", str(e))\n",
        "\n",
        "\n",
        "def open_file_explorer(open_in_new_tab=False, root_dir=root_dir, port=8765):\n",
        "    thread = threading.Thread(target=start_file_explorer, args=[root_dir, port])\n",
        "    thread.start()\n",
        "\n",
        "    if open_in_new_tab:\n",
        "        output.serve_kernel_port_as_window(port)\n",
        "    else:\n",
        "        output.serve_kernel_port_as_iframe(port, height=\"500\")\n",
        "\n",
        "\n",
        "# Example usage\n",
        "open_file_explorer(open_in_new_tab=open_in_new_tab, root_dir=root_dir, port=8765)\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "3gob9_OwTlwh"
      },
      "source": [
        "# II. Pretrained Model Selection"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "8VT6NLv-2u6q"
      },
      "outputs": [],
      "source": [
        "# @title ## 2.1. Download Available Model\n",
        "import os\n",
        "\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "models = {\n",
        "    \"Animefull-final-pruned\": \"https://huggingface.co/Linaqruf/personal-backup/resolve/main/models/animefull-final-pruned.ckpt\",\n",
        "    \"Anything-v3-1\": \"https://huggingface.co/cag/anything-v3-1/resolve/main/anything-v3-1.safetensors\",\n",
        "    \"AnyLoRA\": \"https://huggingface.co/Linaqruf/stolen/resolve/main/pruned-models/AnyLoRA_noVae_fp16-pruned.safetensors\",\n",
        "    \"AnimePastelDream\": \"https://huggingface.co/Lykon/AnimePastelDream/resolve/main/AnimePastelDream_Soft_noVae_fp16.safetensors\",\n",
        "    \"Chillout-mix\": \"https://huggingface.co/Linaqruf/stolen/resolve/main/pruned-models/chillout_mix-pruned.safetensors\",\n",
        "    \"OpenJourney-v4\": \"https://huggingface.co/prompthero/openjourney-v4/resolve/main/openjourney-v4.ckpt\",\n",
        "    \"Stable-Diffusion-v1-5\": \"https://huggingface.co/Linaqruf/stolen/resolve/main/pruned-models/stable_diffusion_1_5-pruned.safetensors\",\n",
        "}\n",
        "\n",
        "v2_models = {\n",
        "    \"stable-diffusion-2-1-base\": \"https://huggingface.co/stabilityai/stable-diffusion-2-1-base/resolve/main/v2-1_512-ema-pruned.safetensors\",\n",
        "    \"stable-diffusion-2-1-768v\": \"https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors\",\n",
        "    \"plat-diffusion-v1-3-1\": \"https://huggingface.co/p1atdev/pd-archive/resolve/main/plat-v1-3-1.safetensors\",\n",
        "    \"replicant-v1\": \"https://huggingface.co/gsdf/Replicant-V1.0/resolve/main/Replicant-V1.0.safetensors\",\n",
        "    \"illuminati-diffusion-v1-0\": \"https://huggingface.co/IlluminatiAI/Illuminati_Diffusion_v1.0/resolve/main/illuminati_diffusion_v1.0.safetensors\",\n",
        "    \"illuminati-diffusion-v1-1\": \"https://huggingface.co/4eJIoBek/Illuminati-Diffusion-v1-1/resolve/main/illuminatiDiffusionV1_v11.safetensors\",\n",
        "    \"waifu-diffusion-1-4-anime-e2\": \"https://huggingface.co/hakurei/waifu-diffusion-v1-4/resolve/main/wd-1-4-anime_e2.ckpt\",\n",
        "    \"waifu-diffusion-1-5-e2\": \"https://huggingface.co/waifu-diffusion/wd-1-5-beta2/resolve/main/checkpoints/wd-1-5-beta2-fp32.safetensors\",\n",
        "    \"waifu-diffusion-1-5-e2-aesthetic\": \"https://huggingface.co/waifu-diffusion/wd-1-5-beta2/resolve/main/checkpoints/wd-1-5-beta2-aesthetic-fp32.safetensors\",\n",
        "}\n",
        "\n",
        "installModels = []\n",
        "installv2Models = []\n",
        "\n",
        "# @markdown ### SD1.x model\n",
        "model_name = \"AnyLoRA\"  # @param [\"\", \"Animefull-final-pruned\", \"Anything-v3-1\", \"AnyLoRA\", \"AnimePastelDream\", \"Chillout-mix\", \"OpenJourney-v4\", \"Stable-Diffusion-v1-5\"]\n",
        "# @markdown ### SD2.x model\n",
        "v2_model_name = \"\"  # @param [\"\", \"stable-diffusion-2-1-base\", \"stable-diffusion-2-1-768v\", \"plat-diffusion-v1-3-1\", \"replicant-v1\", \"illuminati-diffusion-v1-0\", \"illuminati-diffusion-v1-1\", \"waifu-diffusion-1-4-anime-e2\", \"waifu-diffusion-1-5-e2\", \"waifu-diffusion-1-5-e2-aesthetic\"]\n",
        "\n",
        "if model_name:\n",
        "    model_url = models.get(model_name)\n",
        "    if model_url:\n",
        "        installModels.append((model_name, model_url))\n",
        "\n",
        "if v2_model_name:\n",
        "    v2_model_url = v2_models.get(v2_model_name)\n",
        "    if v2_model_url:\n",
        "        installv2Models.append((v2_model_name, v2_model_url))\n",
        "\n",
        "\n",
        "def install(checkpoint_name, url):\n",
        "    ext = \"ckpt\" if url.endswith(\".ckpt\") else \"safetensors\"\n",
        "\n",
        "    hf_token = \"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\"\n",
        "    user_header = f'\"Authorization: Bearer {hf_token}\"'\n",
        "    !aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 -d {pretrained_model} -o {checkpoint_name}.{ext} \"{url}\"\n",
        "\n",
        "\n",
        "def install_checkpoint():\n",
        "    for model in installModels:\n",
        "        install(model[0], model[1])\n",
        "    for v2model in installv2Models:\n",
        "        install(v2model[0], v2model[1])\n",
        "\n",
        "\n",
        "install_checkpoint()\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "3LWn6GzNQ4j5"
      },
      "outputs": [],
      "source": [
        "# @title ## 2.2. Download Custom Model\n",
        "import os\n",
        "\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "# @markdown ### Custom model\n",
        "modelUrls = \"\"  # @param {'type': 'string'}\n",
        "\n",
        "def install(url):\n",
        "    base_name = os.path.basename(url)\n",
        "\n",
        "    if \"drive.google.com\" in url:\n",
        "        os.chdir(pretrained_model)\n",
        "        !gdown --fuzzy {url}\n",
        "    elif \"huggingface.co\" in url:\n",
        "        if \"/blob/\" in url:\n",
        "            url = url.replace(\"/blob/\", \"/resolve/\")\n",
        "        # @markdown Change this part with your own huggingface token if you need to download your private model\n",
        "        hf_token = \"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\"  # @param {type:\"string\"}\n",
        "        user_header = f'\"Authorization: Bearer {hf_token}\"'\n",
        "        !aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 -d {pretrained_model} -o {base_name} {url}\n",
        "    else:\n",
        "        !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {pretrained_model} {url}\n",
        "\n",
        "if modelUrls:\n",
        "    urls = modelUrls.split(\",\")\n",
        "    for url in urls:\n",
        "        install(url.strip())\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "qrY4KtfL6Dqp"
      },
      "outputs": [],
      "source": [
        "# @title ## 2.3. Download Available VAE (Optional)\n",
        "import os\n",
        "\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "vaes = {\n",
        "    \"none\": \"\",\n",
        "    \"anime.vae.pt\": \"https://huggingface.co/Linaqruf/personal-backup/resolve/main/vae/animevae.pt\",\n",
        "    \"waifudiffusion.vae.pt\": \"https://huggingface.co/hakurei/waifu-diffusion-v1-4/resolve/main/vae/kl-f8-anime.ckpt\",\n",
        "    \"stablediffusion.vae.pt\": \"https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.ckpt\",\n",
        "}\n",
        "install_vaes = []\n",
        "\n",
        "# @markdown Select one of the VAEs to download, select `none` for not download VAE:\n",
        "vae_name = \"anime.vae.pt\"  # @param [\"none\", \"anime.vae.pt\", \"waifudiffusion.vae.pt\", \"stablediffusion.vae.pt\"]\n",
        "\n",
        "if vae_name in vaes:\n",
        "    vae_url = vaes[vae_name]\n",
        "    if vae_url:\n",
        "        install_vaes.append((vae_name, vae_url))\n",
        "\n",
        "\n",
        "def install(vae_name, url):\n",
        "    hf_token = \"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\"\n",
        "    user_header = f'\"Authorization: Bearer {hf_token}\"'\n",
        "    !aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 -d {vae_dir} -o {vae_name} \"{url}\"\n",
        "\n",
        "\n",
        "def install_vae():\n",
        "    for vae in install_vaes:\n",
        "        install(vae[0], vae[1])\n",
        "\n",
        "\n",
        "install_vae()\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "En9UUwGNMRMM"
      },
      "source": [
        "# III. Data Acquisition\n",
        "\n",
        "You have three options for acquiring your dataset:\n",
        "\n",
        "1. Uploading it to Colab's local files.\n",
        "2. Bulk downloading images from Danbooru using the `Simple Booru Scraper`.\n",
        "3. Locating your dataset from Google Drive.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "kh7CeDqK4l3Y"
      },
      "outputs": [],
      "source": [
        "# @title ## 3.1. Locating Train Data Directory\n",
        "# @markdown Define the location of your training data. This cell will also create a folder based on your input. Regularization Images is optional and can be skipped.\n",
        "import os\n",
        "from IPython.utils import capture\n",
        "\n",
        "%store -r\n",
        "\n",
        "train_data_dir = \"/content/dreambooth/train_data\"  # @param {type:'string'}\n",
        "reg_data_dir = \"/content/dreambooth/reg_data\"  # @param {type:'string'}\n",
        "\n",
        "for dir in [train_data_dir, reg_data_dir]:\n",
        "    if dir:\n",
        "        with capture.capture_output() as cap:\n",
        "            os.makedirs(dir, exist_ok=True)\n",
        "            %store dir\n",
        "            del cap\n",
        "\n",
        "print(f\"Your train data directory : {train_data_dir}\")\n",
        "if reg_data_dir:\n",
        "    print(f\"Your reg data directory : {reg_data_dir}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "t17ZfiMB8GWZ"
      },
      "outputs": [],
      "source": [
        "# @title ## 3.2. Unzip Dataset\n",
        "\n",
        "import os\n",
        "import shutil\n",
        "from pathlib import Path\n",
        "\n",
        "#@title ## Unzip Dataset\n",
        "# @markdown Use this section if your dataset is in a `zip` file and has been uploaded somewhere. This code cell will download your dataset and automatically extract it to the `train_data_dir` if the `unzip_to` variable is empty.\n",
        "zipfile_url = \"\" #@param {type:\"string\"}\n",
        "zipfile_name = \"zipfile.zip\"\n",
        "unzip_to = \"\" #@param {type:\"string\"}\n",
        "\n",
        "hf_token = \"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\"\n",
        "user_header = f'\"Authorization: Bearer {hf_token}\"'\n",
        "\n",
        "if unzip_to:\n",
        "    os.makedirs(unzip_to, exist_ok=True)\n",
        "else:\n",
        "    unzip_to = train_data_dir\n",
        "\n",
        "\n",
        "def download_dataset(url):\n",
        "    if url.startswith(\"/content\"):\n",
        "        return url\n",
        "    elif \"drive.google.com\" in url:\n",
        "        os.chdir(root_dir)\n",
        "        !gdown --fuzzy {url}\n",
        "        return f\"{root_dir}/{zipfile_name}\"\n",
        "    elif \"huggingface.co\" in url:\n",
        "        if \"/blob/\" in url:\n",
        "            url = url.replace(\"/blob/\", \"/resolve/\")\n",
        "        !aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 -d {root_dir} -o {zipfile_name} {url}\n",
        "        return f\"{root_dir}/{zipfile_name}\"\n",
        "    else:\n",
        "        !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {root_dir} -o {zipfile_name} {url}\n",
        "        return f\"{root_dir}/{zipfile_name}\"\n",
        "\n",
        "\n",
        "def extract_dataset(zip_file, output_path):\n",
        "    if zip_file.startswith(\"/content\"):\n",
        "        !unzip -j -o {zip_file} -d \"{output_path}\"\n",
        "    else:\n",
        "        !unzip -j -o \"{zip_file}\" -d \"{output_path}\"\n",
        "\n",
        "\n",
        "def remove_files(train_dir, files_to_move):\n",
        "    for filename in os.listdir(train_dir):\n",
        "        file_path = os.path.join(train_dir, filename)\n",
        "        if filename in files_to_move:\n",
        "            if not os.path.exists(file_path):\n",
        "                shutil.move(file_path, training_dir)\n",
        "            else:\n",
        "                os.remove(file_path)\n",
        "\n",
        "\n",
        "zip_file = download_dataset(zipfile_url)\n",
        "extract_dataset(zip_file, unzip_to)\n",
        "os.remove(zip_file)\n",
        "\n",
        "files_to_move = (\n",
        "    \"meta_cap.json\",\n",
        "    \"meta_cap_dd.json\",\n",
        "    \"meta_lat.json\",\n",
        "    \"meta_clean.json\",\n",
        ")\n",
        "\n",
        "remove_files(train_data_dir, files_to_move)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "A0t1dfnU5Xkq"
      },
      "outputs": [],
      "source": [
        "#@title ## 3.3. Image Scraper (Optional)\n",
        "import os\n",
        "import html\n",
        "from IPython.utils import capture\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "#@markdown Use `gallery-dl` to scrape images from an imageboard site. Specify the `prompt(s)` by separating them with commas, e.g., `hito_komoru, touhou`.\n",
        "booru = \"Danbooru\" #@param [\"Danbooru\", \"Gelbooru\", \"Safebooru\"]\n",
        "prompt = \"\" #@param {type: \"string\"}\n",
        "#@markdown  You can also specify a `custom_url` instead of using a predefined site.\n",
        "custom_url = \"\" #@param {type: \"string\"}\n",
        "#@markdown `sub_folder` option can be used to organize the downloaded images into separate folders based on their concept or category.\n",
        "sub_folder = \"\" #@param {type: \"string\"}\n",
        "user_agent = \"gdl/1.24.5\"\n",
        "#@markdown You can limit the number of images to download by using the `--range` option followed by the desired range. For example `1-200`.\n",
        "range = \"1-200\" #@param {type: \"string\"}\n",
        "write_tags = True #@param {type: \"boolean\"}\n",
        "additional_arguments = \"--filename /O --no-part\" #@param {type: \"string\"}\n",
        "#@markdown Set `with_aria_2c` to `True` to scrape images using aria2c.\n",
        "with_aria_2c = False #@param {type: \"boolean\"}\n",
        "\n",
        "tags = prompt.split(',')\n",
        "tags = '+'.join(tags)\n",
        "\n",
        "replacement_dict = {\" \": \"\", \"(\": \"%28\", \")\": \"%29\", \":\": \"%3a\"}\n",
        "tags = ''.join(replacement_dict.get(c, c) for c in tags)\n",
        "\n",
        "if sub_folder == \"\":\n",
        "    image_dir = train_data_dir\n",
        "elif sub_folder.startswith(\"/content\"):\n",
        "    image_dir = sub_folder\n",
        "else:\n",
        "    image_dir = os.path.join(train_data_dir, sub_folder)\n",
        "    os.makedirs(image_dir, exist_ok=True)\n",
        "\n",
        "if booru == \"Danbooru\":\n",
        "    url = \"https://danbooru.donmai.us/posts?tags={}\".format(tags)\n",
        "elif booru == \"Gelbooru\":\n",
        "    url = \"https://gelbooru.com/index.php?page=post&s=list&tags={}\".format(tags)\n",
        "else:\n",
        "    url = \"https://safebooru.org/index.php?page=post&s=list&tags={}\".format(tags)\n",
        "\n",
        "valid_url = custom_url if custom_url else url\n",
        "\n",
        "def scrape(config):\n",
        "    args = \"\"\n",
        "    for k, v in config.items():\n",
        "        if k.startswith(\"_\"):\n",
        "            args += f'\"{v}\" '\n",
        "        elif isinstance(v, str):\n",
        "            args += f'--{k}=\"{v}\" '\n",
        "        elif isinstance(v, bool) and v:\n",
        "            args += f\"--{k} \"\n",
        "        elif isinstance(v, float) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "        elif isinstance(v, int) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "\n",
        "    return args\n",
        "\n",
        "def pre_process_tags(directory):\n",
        "    for item in os.listdir(directory):\n",
        "        item_path = os.path.join(directory, item)\n",
        "        if os.path.isfile(item_path) and item.endswith(\".txt\"):\n",
        "            old_path = item_path\n",
        "            new_file_name = os.path.splitext(os.path.splitext(item)[0])[0] + \".txt\"\n",
        "            new_path = os.path.join(directory, new_file_name)\n",
        "\n",
        "            os.rename(old_path, new_path)\n",
        "\n",
        "            with open(new_path, \"r\") as f:\n",
        "                contents = f.read()\n",
        "\n",
        "            contents = html.unescape(contents)\n",
        "            contents = contents.replace(\"_\", \" \")\n",
        "            contents = \", \".join(contents.split(\"\\n\"))\n",
        "\n",
        "            with open(new_path, \"w\") as f:\n",
        "                f.write(contents)\n",
        "\n",
        "        elif os.path.isdir(item_path):\n",
        "            pre_process_tags(item_path)\n",
        "\n",
        "get_url_config = {\n",
        "    \"get-urls\" : True,\n",
        "    \"range\" : range if range else None,\n",
        "    \"user-agent\" : user_agent\n",
        "}\n",
        "\n",
        "scrape_config = {\n",
        "    \"directory\" : image_dir,\n",
        "    \"write-tags\" : write_tags,\n",
        "    \"range\" : range if range else None,\n",
        "    \"user-agent\" : user_agent\n",
        "}\n",
        "\n",
        "if with_aria_2c:\n",
        "    scraper_text = os.path.join(root_dir, \"scrape_this.txt\")\n",
        "    with capture.capture_output() as cap:\n",
        "        args = scrape(get_url_config)\n",
        "        !gallery-dl \"{valid_url}\" {args} {additional_arguments}\n",
        "    with open(scraper_text, \"w\") as f:\n",
        "        f.write(cap.stdout)\n",
        "\n",
        "    os.chdir(image_dir)\n",
        "    !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -i {scraper_text}\n",
        "\n",
        "else:\n",
        "    args = scrape(scrape_config)\n",
        "    !gallery-dl \"{valid_url}\" {args} {additional_arguments}\n",
        "\n",
        "if write_tags:\n",
        "    pre_process_tags(train_data_dir)\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "T-0qKyEgTchp"
      },
      "source": [
        "# IV. Data Preprocessing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Jz2emq6vWnPu"
      },
      "outputs": [],
      "source": [
        "# @title ## 4.1. Data Cleaning\n",
        "import os\n",
        "import random\n",
        "import concurrent.futures\n",
        "from tqdm import tqdm\n",
        "from PIL import Image\n",
        "\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "test = os.listdir(train_data_dir)\n",
        "# @markdown This section will delete unnecessary files and unsupported media such as `.mp4`, `.webm`, and `.gif`. \n",
        "# @markdown Set the `convert` parameter to convert your transparent dataset with an alpha channel (RGBA) to RGB and give it a white background. \n",
        "convert = False  # @param {type:\"boolean\"}\n",
        "# @markdown You can choose to give it a `random_color` background instead of white by checking the corresponding option.\n",
        "random_color = False  # @param {type:\"boolean\"}\n",
        "# @markdown Use the `recursive` option to preprocess subfolders as well.\n",
        "recursive = False #  @param {type:\"boolean\"}\n",
        " \n",
        "\n",
        "batch_size = 32\n",
        "supported_types = [\n",
        "    \".png\",\n",
        "    \".jpg\",\n",
        "    \".jpeg\",\n",
        "    \".webp\",\n",
        "    \".bmp\",\n",
        "    \".caption\",\n",
        "    \".npz\",\n",
        "    \".txt\",\n",
        "    \".json\",\n",
        "]\n",
        "\n",
        "background_colors = [\n",
        "    (255, 255, 255),\n",
        "    (0, 0, 0),\n",
        "    (255, 0, 0),\n",
        "    (0, 255, 0),\n",
        "    (0, 0, 255),\n",
        "    (255, 255, 0),\n",
        "    (255, 0, 255),\n",
        "    (0, 255, 255),\n",
        "]\n",
        "\n",
        "def clean_directory(directory):\n",
        "    for item in os.listdir(directory):\n",
        "        file_path = os.path.join(directory, item)\n",
        "        if os.path.isfile(file_path):\n",
        "            file_ext = os.path.splitext(item)[1]\n",
        "            if file_ext not in supported_types:\n",
        "                print(f\"Deleting file {item} from {directory}\")\n",
        "                os.remove(file_path)\n",
        "        elif os.path.isdir(file_path) and recursive:\n",
        "            clean_directory(file_path)\n",
        "\n",
        "def process_image(image_path):\n",
        "    img = Image.open(image_path)\n",
        "    img_dir, image_name = os.path.split(image_path)\n",
        "\n",
        "    if img.mode in (\"RGBA\", \"LA\"):\n",
        "        if random_color:\n",
        "            background_color = random.choice(background_colors)\n",
        "        else:\n",
        "            background_color = (255, 255, 255)\n",
        "        bg = Image.new(\"RGB\", img.size, background_color)\n",
        "        bg.paste(img, mask=img.split()[-1])\n",
        "\n",
        "        if image_name.endswith(\".webp\"):\n",
        "            bg = bg.convert(\"RGB\")\n",
        "            new_image_path = os.path.join(img_dir, image_name.replace(\".webp\", \".jpg\"))\n",
        "            bg.save(new_image_path, \"JPEG\")\n",
        "            os.remove(image_path)\n",
        "            print(f\" Converted image: {image_name} to {os.path.basename(new_image_path)}\")\n",
        "        else:\n",
        "            bg.save(image_path, \"PNG\")\n",
        "            print(f\" Converted image: {image_name}\")\n",
        "    else:\n",
        "        if image_name.endswith(\".webp\"):\n",
        "            new_image_path = os.path.join(img_dir, image_name.replace(\".webp\", \".jpg\"))\n",
        "            img.save(new_image_path, \"JPEG\")\n",
        "            os.remove(image_path)\n",
        "            print(f\" Converted image: {image_name} to {os.path.basename(new_image_path)}\")\n",
        "        else:\n",
        "            img.save(image_path, \"PNG\")\n",
        "\n",
        "def find_images(directory):\n",
        "    images = []\n",
        "    for root, _, files in os.walk(directory):\n",
        "        for file in files:\n",
        "            if file.endswith(\".png\") or file.endswith(\".webp\"):\n",
        "                images.append(os.path.join(root, file))\n",
        "    return images\n",
        "\n",
        "clean_directory(train_data_dir)\n",
        "images = find_images(train_data_dir)\n",
        "num_batches = len(images) // batch_size + 1\n",
        "\n",
        "if convert:\n",
        "    with concurrent.futures.ThreadPoolExecutor() as executor:\n",
        "        for i in tqdm(range(num_batches)):\n",
        "            start = i * batch_size\n",
        "            end = start + batch_size\n",
        "            batch = images[start:end]\n",
        "            executor.map(process_image, batch)\n",
        "\n",
        "    print(\"All images have been converted\")\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "qdISafLeyklg"
      },
      "source": [
        "## 4.2. Data Annotation\n",
        "You can choose to train a model using captions. We're using [BLIP](https://huggingface.co/spaces/Salesforce/BLIP) for image captioning and [Waifu Diffusion 1.4 Tagger](https://huggingface.co/spaces/SmilingWolf/wd-v1-4-tags) for image tagging similar to Danbooru.\n",
        "- Use BLIP Captioning for: `General Images`\n",
        "- Use Waifu Diffusion 1.4 Tagger V2 for: `Anime and Manga-style Images`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "xvGx2Ikhc8iy"
      },
      "outputs": [],
      "source": [
        "#@title ### 4.2.1. BLIP Captioning\n",
        "#@markdown BLIP is a pre-training framework for unified vision-language understanding and generation, which achieves state-of-the-art results on a wide range of vision-language tasks. It can be used as a tool for image captioning, for example, `astronaut riding a horse in space`. \n",
        "import os\n",
        "\n",
        "os.chdir(finetune_dir)\n",
        "\n",
        "batch_size = 8 #@param {type:'number'}\n",
        "max_data_loader_n_workers = 2 #@param {type:'number'}\n",
        "beam_search = True #@param {type:'boolean'}\n",
        "min_length = 5 #@param {type:\"slider\", min:0, max:100, step:5.0}\n",
        "max_length = 75 #@param {type:\"slider\", min:0, max:100, step:5.0}\n",
        "#@markdown Use the `recursive` option to process subfolders as well, useful for multi-concept training.\n",
        "recursive = False #@param {type:\"boolean\"} \n",
        "#@markdown Debug while captioning, it will print your image file with generated captions.\n",
        "verbose_logging = True #@param {type:\"boolean\"}\n",
        "\n",
        "config = {\n",
        "    \"_train_data_dir\" : train_data_dir,\n",
        "    \"batch_size\" : batch_size,\n",
        "    \"beam_search\" : beam_search,\n",
        "    \"min_length\" : min_length,\n",
        "    \"max_length\" : max_length,\n",
        "    \"debug\" : verbose_logging,\n",
        "    \"caption_extension\" : \".caption\",\n",
        "    \"max_data_loader_n_workers\" : max_data_loader_n_workers,\n",
        "    \"recursive\" : recursive\n",
        "}\n",
        "\n",
        "args = \"\"\n",
        "for k, v in config.items():\n",
        "    if k.startswith(\"_\"):\n",
        "        args += f'\"{v}\" '\n",
        "    elif isinstance(v, str):\n",
        "        args += f'--{k}=\"{v}\" '\n",
        "    elif isinstance(v, bool) and v:\n",
        "        args += f\"--{k} \"\n",
        "    elif isinstance(v, float) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "    elif isinstance(v, int) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "\n",
        "final_args = f\"python make_captions.py {args}\"\n",
        "\n",
        "os.chdir(finetune_dir)\n",
        "!{final_args}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-BdXV7rAy2ag"
      },
      "outputs": [],
      "source": [
        "#@title ### 4.2.2. Waifu Diffusion 1.4 Tagger V2\n",
        "import os\n",
        "%store -r\n",
        "\n",
        "os.chdir(finetune_dir)\n",
        "\n",
        "#@markdown [Waifu Diffusion 1.4 Tagger V2](https://huggingface.co/spaces/SmilingWolf/wd-v1-4-tags) is a Danbooru-styled image classification model developed by SmilingWolf. It can also be useful for general image tagging, for example, `1girl, solo, looking_at_viewer, short_hair, bangs, simple_background`.\n",
        "batch_size = 8 #@param {type:'number'}\n",
        "max_data_loader_n_workers = 2 #@param {type:'number'}\n",
        "model = \"SmilingWolf/wd-v1-4-convnextv2-tagger-v2\" #@param [\"SmilingWolf/wd-v1-4-convnextv2-tagger-v2\", \"SmilingWolf/wd-v1-4-swinv2-tagger-v2\", \"SmilingWolf/wd-v1-4-convnext-tagger-v2\", \"SmilingWolf/wd-v1-4-vit-tagger-v2\"]\n",
        "#@markdown Use the `recursive` option to process subfolders as well, useful for multi-concept training.\n",
        "recursive = False #@param {type:\"boolean\"} \n",
        "#@markdown Debug while tagging, it will print your image file with general tags and character tags.\n",
        "verbose_logging = True #@param {type:\"boolean\"}\n",
        "#@markdown Separate `undesired_tags` with comma `(,)` if you want to remove multiple tags, e.g. `1girl,solo,smile`.\n",
        "undesired_tags = \"\" #@param {type:'string'}\n",
        "#@markdown  Adjust `general_threshold` for pruning tags (less tags, less flexible). `character_threshold` is useful if you want to train with character tags, e.g. `hakurei reimu`.\n",
        "general_threshold = 0.35 #@param {type:\"slider\", min:0, max:1, step:0.05}\n",
        "character_threshold = 0.35 #@param {type:\"slider\", min:0, max:1, step:0.05}\n",
        "\n",
        "config = {\n",
        "    \"_train_data_dir\": train_data_dir,\n",
        "    \"batch_size\": batch_size,\n",
        "    \"repo_id\": model,\n",
        "    \"recursive\": recursive,\n",
        "    \"remove_underscore\": True,\n",
        "    \"general_threshold\": general_threshold,\n",
        "    \"character_threshold\": character_threshold,\n",
        "    \"caption_extension\": \".txt\",\n",
        "    \"max_data_loader_n_workers\": max_data_loader_n_workers,\n",
        "    \"debug\": verbose_logging,\n",
        "    \"undesired_tags\": undesired_tags\n",
        "}\n",
        "\n",
        "args = \"\"\n",
        "for k, v in config.items():\n",
        "    if k.startswith(\"_\"):\n",
        "        args += f'\"{v}\" '\n",
        "    elif isinstance(v, str):\n",
        "        args += f'--{k}=\"{v}\" '\n",
        "    elif isinstance(v, bool) and v:\n",
        "        args += f\"--{k} \"\n",
        "    elif isinstance(v, float) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "    elif isinstance(v, int) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "\n",
        "final_args = f\"python tag_images_by_wd14_tagger.py {args}\"\n",
        "\n",
        "os.chdir(finetune_dir)\n",
        "!{final_args}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "_mLVURhM9PFE"
      },
      "outputs": [],
      "source": [
        "# @title ### 4.2.3. Custom Caption/Tag\n",
        "import os\n",
        "\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "# @markdown Add or remove custom tags here. You can refer to this [cheatsheet](https://rentry.org/kohyaminiguide#c-custom-tagscaption) for more information.\n",
        "extension = \".txt\"  # @param [\".txt\", \".caption\"]\n",
        "custom_tag = \"\"  # @param {type:\"string\"}\n",
        "# @markdown Use `sub_folder` option to specify a subfolder for multi-concept training. \n",
        "# @markdown > Specify `--all` to process all subfolders/`recursive`\n",
        "sub_folder = \"\" #@param {type: \"string\"}\n",
        "# @markdown Enable this to append custom tags at the end of lines.\n",
        "append = False  # @param {type:\"boolean\"}\n",
        "# @markdown Enable this if you want to remove captions/tags instead.\n",
        "remove_tag = False  # @param {type:\"boolean\"}\n",
        "recursive = False\n",
        "\n",
        "if sub_folder == \"\":\n",
        "    image_dir = train_data_dir\n",
        "elif sub_folder == \"--all\":\n",
        "    image_dir = train_data_dir\n",
        "    recursive = True\n",
        "elif sub_folder.startswith(\"/content\"):\n",
        "    image_dir = sub_folder\n",
        "else:\n",
        "    image_dir = os.path.join(train_data_dir, sub_folder)\n",
        "    os.makedirs(image_dir, exist_ok=True)\n",
        "\n",
        "def read_file(filename):\n",
        "    with open(filename, \"r\") as f:\n",
        "        contents = f.read()\n",
        "    return contents\n",
        "\n",
        "def write_file(filename, contents):\n",
        "    with open(filename, \"w\") as f:\n",
        "        f.write(contents)\n",
        "\n",
        "def process_tags(filename, custom_tag, append, remove_tag):\n",
        "    contents = read_file(filename)\n",
        "    tags = [tag.strip() for tag in contents.split(',')]\n",
        "    custom_tags = [tag.strip() for tag in custom_tag.split(',')]\n",
        "\n",
        "    for custom_tag in custom_tags:\n",
        "        custom_tag = custom_tag.replace(\"_\", \" \")\n",
        "        if remove_tag:\n",
        "            while custom_tag in tags:\n",
        "                tags.remove(custom_tag)\n",
        "        else:\n",
        "            if custom_tag not in tags:\n",
        "                if append:\n",
        "                    tags.append(custom_tag)\n",
        "                else:\n",
        "                    tags.insert(0, custom_tag)\n",
        "\n",
        "    contents = ', '.join(tags)\n",
        "    write_file(filename, contents)\n",
        "\n",
        "def process_directory(image_dir, tag, append, remove_tag, recursive):\n",
        "    for filename in os.listdir(image_dir):\n",
        "        file_path = os.path.join(image_dir, filename)\n",
        "        \n",
        "        if os.path.isdir(file_path) and recursive:\n",
        "            process_directory(file_path, tag, append, remove_tag, recursive)\n",
        "        elif filename.endswith(extension):\n",
        "            process_tags(file_path, tag, append, remove_tag)\n",
        "\n",
        "tag = custom_tag\n",
        "\n",
        "if not any(\n",
        "    [filename.endswith(extension) for filename in os.listdir(image_dir)]\n",
        "):\n",
        "    for filename in os.listdir(image_dir):\n",
        "        if filename.endswith((\".png\", \".jpg\", \".jpeg\", \".webp\", \".bmp\")):\n",
        "            open(\n",
        "                os.path.join(image_dir, filename.split(\".\")[0] + extension),\n",
        "                \"w\",\n",
        "            ).close()\n",
        "\n",
        "if custom_tag:\n",
        "    process_directory(image_dir, tag, append, remove_tag, recursive)\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "yHNbl3O_NSS0"
      },
      "source": [
        "# VII. Training Model\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "H_Q23fUEJhnC"
      },
      "outputs": [],
      "source": [
        "# @title ## 5.1. Model Config\n",
        "from google.colab import drive\n",
        "\n",
        "v2 = False  # @param {type:\"boolean\"}\n",
        "v_parameterization = False  # @param {type:\"boolean\"}\n",
        "project_name = \"\"  # @param {type:\"string\"}\n",
        "if not project_name:\n",
        "    project_name = \"last\"\n",
        "%store project_name\n",
        "pretrained_model_name_or_path = \"/content/pretrained_model/AnyLoRA.safetensors\"  # @param {type:\"string\"}\n",
        "vae = \"\"  # @param {type:\"string\"}\n",
        "output_dir = \"/content/dreambooth/output\"  # @param {'type':'string'}\n",
        "resume_path = \"\"  # @param {'type':'string'}\n",
        "\n",
        "# @markdown `output_to_drive` sets default `output_dir` to `/content/drive/MyDrive/dreambooth/output`. This will override the `output_dir` variable defined above.\n",
        "output_to_drive = False  # @param {'type':'boolean'}\n",
        "\n",
        "if output_to_drive:\n",
        "    output_dir = \"/content/drive/MyDrive/dreambooth/output\"\n",
        "\n",
        "    if not os.path.exists(\"/content/drive\"):\n",
        "        drive.mount(\"/content/drive\")\n",
        "\n",
        "sample_dir = os.path.join(output_dir, \"sample\")\n",
        "for dir in [output_dir, sample_dir]:\n",
        "    os.makedirs(dir, exist_ok=True)\n",
        "\n",
        "print(\"Project Name: \", project_name)\n",
        "print(\"Model Version: Stable Diffusion V1.x\") if not v2 else \"\"\n",
        "print(\"Model Version: Stable Diffusion V2.x\") if v2 and not v_parameterization else \"\"\n",
        "print(\"Model Version: Stable Diffusion V2.x 768v\") if v2 and v_parameterization else \"\"\n",
        "print(\n",
        "    \"Pretrained Model Path: \", pretrained_model_name_or_path\n",
        ") if pretrained_model_name_or_path else print(\"No Pretrained Model path specified.\")\n",
        "print(\"VAE Path: \", vae) if vae else print(\"No VAE path specified.\")\n",
        "print(\"Output Path: \", output_dir)\n",
        "print(\"Resume Path: \", resume_path) if resume_path else print(\n",
        "    \"No resume path specified.\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "8pWaWNl_7KL6"
      },
      "outputs": [],
      "source": [
        "# @title ## 5.2. Dataset Config\n",
        "import toml\n",
        "import glob\n",
        "\n",
        "# @markdown This configuration is designed for `one concept` training. Refer to this [guide](https://rentry.org/kohyaminiguide#b-multi-concept-training) for multi-concept training.\n",
        "dataset_repeats = 10  # @param {type:\"number\"}\n",
        "# @markdown `activation_word` is not used in training if you train with captions/tags, but it is still printed to metadata.\n",
        "activation_word = \"mksks style\"  # @param {type:\"string\"}\n",
        "caption_extension = \".txt\"  # @param [\"none\", \".txt\", \".caption\"]\n",
        "# @markdown Please refer to `4.2.3. Custom Caption/Tag (Optional)` if you want to append `activation_word` to captions/tags  **OR** specify `token_to_captions` below.\n",
        "token_to_captions = False  # @param {'type':'boolean'}\n",
        "resolution = 512  # @param {type:\"slider\", min:512, max:1024, step:128}\n",
        "flip_aug = False  # @param {type:\"boolean\"}\n",
        "keep_tokens = 0  # @param {type:\"number\"}\n",
        "\n",
        "if ',' in activation_word or ' ' in activation_word:\n",
        "    words = activation_word.replace(',', ' ').split()\n",
        "    class_token = words[-1]\n",
        "\n",
        "\n",
        "def read_file(filename):\n",
        "    with open(filename, \"r\") as f:\n",
        "        contents = f.read()\n",
        "    return contents\n",
        "\n",
        "\n",
        "def write_file(filename, contents):\n",
        "    with open(filename, \"w\") as f:\n",
        "        f.write(contents)\n",
        "\n",
        "\n",
        "def get_supported_images(folder):\n",
        "    supported_extensions = (\".png\", \".jpg\", \".jpeg\", \".webp\", \".bmp\")\n",
        "    return [file for ext in supported_extensions for file in glob.glob(f\"{folder}/*{ext}\")]\n",
        "\n",
        "\n",
        "def get_subfolders_with_supported_images(folder):\n",
        "    subfolders = [os.path.join(folder, subfolder) for subfolder in os.listdir(folder) if os.path.isdir(os.path.join(folder, subfolder))]\n",
        "    return [subfolder for subfolder in subfolders if len(get_supported_images(subfolder)) > 0]\n",
        "\n",
        "\n",
        "def process_tags(filename, custom_tag, remove_tag):\n",
        "    contents = read_file(filename)\n",
        "    tags = [tag.strip() for tag in contents.split(',')]\n",
        "    custom_tags = [tag.strip() for tag in custom_tag.split(',')]\n",
        "\n",
        "    for custom_tag in custom_tags:\n",
        "        custom_tag = custom_tag.replace(\"_\", \" \")\n",
        "        if remove_tag:\n",
        "            while custom_tag in tags:\n",
        "                tags.remove(custom_tag)\n",
        "        else:\n",
        "            if custom_tag not in tags:\n",
        "                tags.insert(0, custom_tag)\n",
        "\n",
        "    contents = ', '.join(tags)\n",
        "    write_file(filename, contents)\n",
        "\n",
        "\n",
        "def process_folder_recursively(folder):\n",
        "    for root, _, files in os.walk(folder):\n",
        "        for file in files:\n",
        "            if file.endswith(caption_extension):\n",
        "                file_path = os.path.join(root, file)\n",
        "                extracted_class_token = get_class_token_from_folder_name(root, folder)\n",
        "                train_supported_images = get_supported_images(train_data_dir)\n",
        "                tag = extracted_class_token if extracted_class_token else activation_word if train_supported_images else \"\"\n",
        "                if not tag == \"\":\n",
        "                    process_tags(file_path, tag, remove_tag=(not token_to_captions))\n",
        "\n",
        "\n",
        "def get_num_repeats(folder):\n",
        "    folder_name = os.path.basename(folder)\n",
        "    try:\n",
        "        repeats, _ = folder_name.split('_', 1)\n",
        "        num_repeats = int(repeats)\n",
        "    except ValueError:\n",
        "        num_repeats = 1\n",
        "\n",
        "    return num_repeats\n",
        "\n",
        "\n",
        "def get_class_token_from_folder_name(folder, parent_folder):\n",
        "    if folder == parent_folder:\n",
        "        return class_token\n",
        "\n",
        "    folder_name = os.path.basename(folder)\n",
        "    try:\n",
        "        _, concept = folder_name.split('_', 1)\n",
        "        return concept\n",
        "    except ValueError:\n",
        "        return \"\"\n",
        "        \n",
        "train_supported_images = get_supported_images(train_data_dir)\n",
        "train_subfolders = get_subfolders_with_supported_images(train_data_dir)\n",
        "reg_supported_images = get_supported_images(reg_data_dir)\n",
        "reg_subfolders = get_subfolders_with_supported_images(reg_data_dir)\n",
        "\n",
        "subsets = []\n",
        "\n",
        "config = {\n",
        "    \"general\": {\n",
        "        \"enable_bucket\": True,\n",
        "        \"caption_extension\": caption_extension,\n",
        "        \"shuffle_caption\": True,\n",
        "        \"keep_tokens\": keep_tokens,\n",
        "        \"bucket_reso_steps\": 64,\n",
        "        \"bucket_no_upscale\": False,\n",
        "    },\n",
        "    \"datasets\": [\n",
        "        {\n",
        "            \"resolution\": resolution,\n",
        "            \"min_bucket_reso\": 320 if resolution > 640 else 256,\n",
        "            \"max_bucket_reso\": 1280 if resolution > 640 else 1024,\n",
        "            \"caption_dropout_rate\": 0,\n",
        "            \"caption_tag_dropout_rate\": 0,\n",
        "            \"caption_dropout_every_n_epochs\": 0,\n",
        "            \"flip_aug\": flip_aug,\n",
        "            \"color_aug\": False,\n",
        "            \"face_crop_aug_range\": None,\n",
        "            \"subsets\": subsets,\n",
        "        }\n",
        "    ],\n",
        "}\n",
        "\n",
        "if token_to_captions and keep_tokens < 2:\n",
        "    keep_tokens = 1\n",
        "\n",
        "if caption_extension != \"none\":\n",
        "    process_folder_recursively(train_data_dir)\n",
        "\n",
        "if train_supported_images:\n",
        "    subsets.append({\n",
        "        \"image_dir\": train_data_dir,\n",
        "        \"class_tokens\": activation_word,\n",
        "        \"num_repeats\": dataset_repeats,\n",
        "    })\n",
        "\n",
        "for subfolder in train_subfolders:\n",
        "    num_repeats = get_num_repeats(subfolder)\n",
        "    extracted_class_token = get_class_token_from_folder_name(subfolder, train_data_dir)\n",
        "    subsets.append({\n",
        "        \"image_dir\": subfolder,\n",
        "        \"class_tokens\": extracted_class_token if extracted_class_token else None,\n",
        "        \"num_repeats\": num_repeats,\n",
        "    })\n",
        "\n",
        "if reg_supported_images:\n",
        "    subsets.append({\n",
        "        \"is_reg\": True,\n",
        "        \"image_dir\": reg_data_dir,\n",
        "        \"class_tokens\": class_token if 'class_token' in globals() else None,\n",
        "        \"num_repeats\": 1,\n",
        "    })\n",
        "\n",
        "for subfolder in reg_subfolders:\n",
        "    extracted_class_token = get_class_token_from_folder_name(subfolder, reg_data_dir)\n",
        "    subsets.append({\n",
        "        \"is_reg\": True,\n",
        "        \"image_dir\": subfolder,\n",
        "        \"class_tokens\": extracted_class_token if extracted_class_token else None,\n",
        "        \"num_repeats\": num_repeats,\n",
        "    })\n",
        "\n",
        "for subset in subsets:\n",
        "    if not glob.glob(f\"{subset['image_dir']}/*.txt\"):\n",
        "        subset[\"class_tokens\"] = activation_word\n",
        "\n",
        "dataset_config = os.path.join(config_dir, \"dataset_config.toml\")\n",
        "\n",
        "for key in config:\n",
        "    if isinstance(config[key], dict):\n",
        "        for sub_key in config[key]:\n",
        "            if config[key][sub_key] == \"\":\n",
        "                config[key][sub_key] = None\n",
        "    elif config[key] == \"\":\n",
        "        config[key] = None\n",
        "\n",
        "config_str = toml.dumps(config)\n",
        "\n",
        "with open(dataset_config, \"w\") as f:\n",
        "    f.write(config_str)\n",
        "\n",
        "print(config_str)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "UkrWPjlUfMRm"
      },
      "outputs": [],
      "source": [
        "# @title ## 5.3. Optimizer Config\n",
        "from IPython.utils import capture\n",
        "\n",
        "# @markdown `NEW` Gamma for reducing the weight of high-loss timesteps. Lower numbers have a stronger effect. The paper recommends 5. Read the paper [here](https://arxiv.org/abs/2303.09556).\n",
        "min_snr_gamma = -1 #@param {type:\"number\"}\n",
        "# @markdown `AdamW8bit` was the old `--use_8bit_adam`.\n",
        "optimizer_type = \"AdamW8bit\"  # @param [\"AdamW\", \"AdamW8bit\", \"Lion\", \"SGDNesterov\", \"SGDNesterov8bit\", \"DAdaptation\", \"AdaFactor\"]\n",
        "# @markdown Additional arguments for optimizer, e.g: `[\"decouple=true\",\"weight_decay=0.6\"]`\n",
        "optimizer_args = \"\"  # @param {'type':'string'}\n",
        "# @markdown Set `learning_rate` to `1.0` if you use `DAdaptation` optimizer, as it's a [free learning rate](https://github.com/facebookresearch/dadaptation) algorithm.\n",
        "# @markdown You probably need to specify `optimizer_args` for custom optimizer, like using `[\"decouple=true\",\"weight_decay=0.6\"]` for `DAdaptation`.\n",
        "learning_rate = 2e-6  # @param {'type':'number'}\n",
        "stop_train_text_encoder = -1 #@param {'type':'number'}\n",
        "lr_scheduler = \"constant\"  # @param [\"linear\", \"cosine\", \"cosine_with_restarts\", \"polynomial\", \"constant\", \"constant_with_warmup\", \"adafactor\"] {allow-input: false}\n",
        "lr_warmup_steps = 0  # @param {'type':'number'}\n",
        "# @markdown You can define `num_cycles` value for `cosine_with_restarts` or `power` value for `polynomial` in the field below.\n",
        "lr_scheduler_num_cycles = 0  # @param {'type':'number'}\n",
        "lr_scheduler_power = 0  # @param {'type':'number'}\n",
        "\n",
        "print(f\"  - Min-SNR Weighting: {min_snr_gamma}\") if not min_snr_gamma == -1 else \"\"\n",
        "print(f\"Using {optimizer_type} as Optimizer\")\n",
        "if optimizer_args:\n",
        "  print(f\"Optimizer Args :\", optimizer_args)\n",
        "print(\"Learning rate: \", learning_rate)\n",
        "if stop_train_text_encoder > 0:\n",
        "  print(f\"Text Encoder training stopped at {stop_train_text_encoder} steps\")\n",
        "print(\"Learning rate warmup steps: \", lr_warmup_steps)\n",
        "print(\"Learning rate Scheduler:\", lr_scheduler)\n",
        "if lr_scheduler == \"cosine_with_restarts\":\n",
        "  print(\"- lr_scheduler_num_cycles: \", lr_scheduler_num_cycles)\n",
        "elif lr_scheduler == \"polynomial\":\n",
        "  print(\"- lr_scheduler_power: \", lr_scheduler_power)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "0glNU5FxgWN1"
      },
      "outputs": [],
      "source": [
        "# @title ## 5.4. Training Config\n",
        "import toml\n",
        "import os\n",
        "\n",
        "%store -r\n",
        "enable_sample_prompt = True  # @param {type:\"boolean\"}\n",
        "sampler = \"ddim\"  # @param [\"ddim\", \"pndm\", \"lms\", \"euler\", \"euler_a\", \"heun\", \"dpm_2\", \"dpm_2_a\", \"dpmsolver\",\"dpmsolver++\", \"dpmsingle\", \"k_lms\", \"k_euler\", \"k_euler_a\", \"k_dpm_2\", \"k_dpm_2_a\"]\n",
        "noise_offset = 0.0  # @param {type:\"number\"}\n",
        "max_train_steps = 2500  # @param {type:\"number\"}\n",
        "vae_batch_size = 1  # @param {type:\"number\"}\n",
        "train_batch_size = 4  # @param {type:\"number\"}\n",
        "mixed_precision = \"fp16\"  # @param [\"no\",\"fp16\",\"bf16\"] {allow-input: false}\n",
        "save_state = False  # @param {type:\"boolean\"}\n",
        "save_precision = \"fp16\"  # @param [\"float\", \"fp16\", \"bf16\"] {allow-input: false}\n",
        "save_n_epoch_ratio = 1  # @param {type:\"number\"}\n",
        "save_model_as = \"ckpt\"  # @param [\"ckpt\", \"safetensors\", \"diffusers\", \"diffusers_safetensors\"] {allow-input: false}\n",
        "max_token_length = 225  # @param {type:\"number\"}\n",
        "clip_skip = 2  # @param {type:\"number\"}\n",
        "gradient_checkpointing = False  # @param {type:\"boolean\"}\n",
        "gradient_accumulation_steps = 1  # @param {type:\"number\"}\n",
        "seed = -1  # @param {type:\"number\"}\n",
        "logging_dir = \"/content/dreambooth/logs\"\n",
        "prior_loss_weight = 1.0\n",
        "\n",
        "os.chdir(repo_dir)\n",
        "\n",
        "sample_str = f\"\"\"\n",
        "  masterpiece, best quality, 1girl, aqua eyes, baseball cap, blonde hair, closed mouth, earrings, green background, hat, hoop earrings, jewelry, looking at viewer, shirt, short hair, simple background, solo, upper body, yellow shirt \\\n",
        "  --n lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry \\\n",
        "  --w 512 \\\n",
        "  --h 768 \\\n",
        "  --l 7 \\\n",
        "  --s 28    \n",
        "\"\"\"\n",
        "\n",
        "config = {\n",
        "    \"model_arguments\": {\n",
        "        \"v2\": v2,\n",
        "        \"v_parameterization\": v_parameterization if v2 and v_parameterization else False,\n",
        "        \"pretrained_model_name_or_path\": pretrained_model_name_or_path,\n",
        "        \"vae\": vae,\n",
        "    },\n",
        "    \"optimizer_arguments\": {\n",
        "        \"min_snr_gamma\": min_snr_gamma if not min_snr_gamma == -1 else None,\n",
        "        \"optimizer_type\": optimizer_type,\n",
        "        \"learning_rate\": learning_rate,\n",
        "        \"max_grad_norm\": 1.0,\n",
        "        \"stop_train_text_encoder\": stop_train_text_encoder if stop_train_text_encoder > 0 else None,\n",
        "        \"optimizer_args\": eval(optimizer_args) if optimizer_args else None,\n",
        "        \"lr_scheduler\": lr_scheduler,\n",
        "        \"lr_warmup_steps\": lr_warmup_steps,\n",
        "        \"lr_scheduler_num_cycles\": lr_scheduler_num_cycles if lr_scheduler == \"cosine_with_restarts\" else None,\n",
        "        \"lr_scheduler_power\": lr_scheduler_power if lr_scheduler == \"polynomial\" else None,\n",
        "    },\n",
        "    \"dataset_arguments\": {\n",
        "        \"cache_latents\": True,\n",
        "        \"debug_dataset\": False,\n",
        "        \"vae_batch_size\": vae_batch_size,\n",
        "    },\n",
        "    \"training_arguments\": {\n",
        "        \"output_dir\": output_dir,\n",
        "        \"output_name\": project_name,\n",
        "        \"save_precision\": save_precision,\n",
        "        \"save_every_n_epochs\": None,\n",
        "        \"save_n_epoch_ratio\": save_n_epoch_ratio,\n",
        "        \"save_last_n_epochs\": None,\n",
        "        \"save_state\": save_state,\n",
        "        \"save_last_n_epochs_state\": None,\n",
        "        \"resume\": resume_path,\n",
        "        \"train_batch_size\": train_batch_size,\n",
        "        \"max_token_length\": 225,\n",
        "        \"mem_eff_attn\": False,\n",
        "        \"xformers\": True,\n",
        "        \"max_train_steps\": max_train_steps,\n",
        "        \"max_data_loader_n_workers\": 8,\n",
        "        \"persistent_data_loader_workers\": True,\n",
        "        \"seed\": seed if seed > 0 else None,\n",
        "        \"gradient_checkpointing\": gradient_checkpointing,\n",
        "        \"gradient_accumulation_steps\": gradient_accumulation_steps,\n",
        "        \"mixed_precision\": mixed_precision,\n",
        "        \"clip_skip\": clip_skip if not v2 else None,\n",
        "        \"logging_dir\": logging_dir,\n",
        "        \"log_prefix\": project_name,\n",
        "        \"noise_offset\": noise_offset if noise_offset > 0 else None,\n",
        "    },\n",
        "    \"sample_prompt_arguments\": {\n",
        "        \"sample_every_n_steps\": 100 if enable_sample_prompt else 999999,\n",
        "        \"sample_every_n_epochs\": None,\n",
        "        \"sample_sampler\": sampler,\n",
        "    },\n",
        "    \"dreambooth_arguments\": {\n",
        "        \"prior_loss_weight\": 1.0,\n",
        "    },\n",
        "    \"saving_arguments\": {\n",
        "        \"save_model_as\": save_model_as\n",
        "    },\n",
        "}\n",
        "\n",
        "config_path = os.path.join(config_dir, \"config_file.toml\")\n",
        "prompt_path = os.path.join(config_dir, \"sample_prompt.txt\")\n",
        "\n",
        "for key in config:\n",
        "    if isinstance(config[key], dict):\n",
        "        for sub_key in config[key]:\n",
        "            if config[key][sub_key] == \"\":\n",
        "                config[key][sub_key] = None\n",
        "    elif config[key] == \"\":\n",
        "        config[key] = None\n",
        "\n",
        "config_str = toml.dumps(config)\n",
        "\n",
        "def write_file(filename, contents):\n",
        "    with open(filename, \"w\") as f:\n",
        "        f.write(contents)\n",
        "\n",
        "write_file(config_path, config_str)\n",
        "write_file(prompt_path, sample_str)\n",
        "    \n",
        "print(config_str)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "f2ylsv-J8pfX"
      },
      "outputs": [],
      "source": [
        "#@title ## 5.5. Start Training\n",
        "\n",
        "#@markdown Check your config here if you want to edit something: \n",
        "#@markdown - `sample_prompt` : /content/dreambooth/config/sample_prompt.txt\n",
        "#@markdown - `config_file` : /content/dreambooth/config/config_file.toml\n",
        "#@markdown - `dataset_config` : /content/dreambooth/config/dataset_config.toml\n",
        "\n",
        "#@markdown Generated sample can be seen here: /content/dreambooth/output/sample\n",
        "\n",
        "#@markdown You can import config from another session if you want.\n",
        "sample_prompt = \"/content/dreambooth/config/sample_prompt.txt\" #@param {type:'string'}\n",
        "config_file = \"/content/dreambooth/config/config_file.toml\" #@param {type:'string'}\n",
        "dataset_config = \"/content/dreambooth/config/dataset_config.toml\" #@param {type:'string'}\n",
        "\n",
        "accelerate_conf = {\n",
        "    \"config_file\" : accelerate_config,\n",
        "    \"num_cpu_threads_per_process\" : 1,\n",
        "}\n",
        "\n",
        "train_conf = {\n",
        "    \"sample_prompts\" : sample_prompt,\n",
        "    \"dataset_config\" : dataset_config,\n",
        "    \"config_file\" : config_file\n",
        "}\n",
        "\n",
        "def train(config):\n",
        "    args = \"\"\n",
        "    for k, v in config.items():\n",
        "        if k.startswith(\"_\"):\n",
        "            args += f'\"{v}\" '\n",
        "        elif isinstance(v, str):\n",
        "            args += f'--{k}=\"{v}\" '\n",
        "        elif isinstance(v, bool) and v:\n",
        "            args += f\"--{k} \"\n",
        "        elif isinstance(v, float) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "        elif isinstance(v, int) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "\n",
        "    return args\n",
        "\n",
        "accelerate_args = train(accelerate_conf)\n",
        "train_args = train(train_conf)\n",
        "final_args = f\"accelerate launch {accelerate_args} train_db.py {train_args}\"\n",
        "\n",
        "os.chdir(repo_dir)\n",
        "!{final_args}"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "vqfgyL-thgdw"
      },
      "source": [
        "# VI. Testing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "jgKi9y4w7tmV"
      },
      "outputs": [],
      "source": [
        "# @title ## 6.1. Visualize loss graph (Optional)\n",
        "import os\n",
        "\n",
        "training_logs_path = \"/content/dreambooth/logs\"  # @param {type : \"string\"}\n",
        "\n",
        "os.chdir(repo_dir)\n",
        "%load_ext tensorboard\n",
        "%tensorboard --logdir {training_logs_path}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "j1jJ4z3AXRO9"
      },
      "outputs": [],
      "source": [
        "# @title ## 6.2. Inference\n",
        "%store -r\n",
        "\n",
        "v2 = False  # @param {type:\"boolean\"}\n",
        "v_parameterization = False  # @param {type:\"boolean\"}\n",
        "prompt = \"masterpiece, best quality, 1girl, aqua eyes, baseball cap, blonde hair, closed mouth, earrings, green background, hat, hoop earrings, jewelry, looking at viewer, shirt, short hair, simple background, solo, upper body, yellow shirt\"  # @param {type: \"string\"}\n",
        "negative = \"lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry\"  # @param {type: \"string\"}\n",
        "model = \"\"  # @param {type: \"string\"}\n",
        "vae = \"\"  # @param {type: \"string\"}\n",
        "outdir = \"/content/tmp\"  # @param {type: \"string\"}\n",
        "scale = 7  # @param {type: \"slider\", min: 1, max: 40}\n",
        "sampler = \"ddim\"  # @param [\"ddim\", \"pndm\", \"lms\", \"euler\", \"euler_a\", \"heun\", \"dpm_2\", \"dpm_2_a\", \"dpmsolver\",\"dpmsolver++\", \"dpmsingle\", \"k_lms\", \"k_euler\", \"k_euler_a\", \"k_dpm_2\", \"k_dpm_2_a\"]\n",
        "steps = 28  # @param {type: \"slider\", min: 1, max: 100}\n",
        "precision = \"fp16\"  # @param [\"fp16\", \"bf16\"] {allow-input: false}\n",
        "width = 512  # @param {type: \"integer\"}\n",
        "height = 768  # @param {type: \"integer\"}\n",
        "images_per_prompt = 4  # @param {type: \"integer\"}\n",
        "batch_size = 4  # @param {type: \"integer\"}\n",
        "clip_skip = 2  # @param {type: \"slider\", min: 1, max: 40}\n",
        "seed = -1  # @param {type: \"integer\"}\n",
        "\n",
        "final_prompt = f\"{prompt} --n {negative}\"\n",
        "\n",
        "config = {\n",
        "    \"v2\": v2,\n",
        "    \"v_parameterization\": v_parameterization,\n",
        "    \"ckpt\": model,\n",
        "    \"outdir\": outdir,\n",
        "    \"xformers\": True,\n",
        "    \"vae\": vae if vae else None,\n",
        "    \"fp16\": True,\n",
        "    \"W\": width,\n",
        "    \"H\": height,\n",
        "    \"seed\": seed if seed > 0 else None,\n",
        "    \"scale\": scale,\n",
        "    \"sampler\": sampler,\n",
        "    \"steps\": steps,\n",
        "    \"max_embeddings_multiples\": 3,\n",
        "    \"batch_size\": batch_size,\n",
        "    \"images_per_prompt\": images_per_prompt,\n",
        "    \"clip_skip\": clip_skip if not v2 else None,\n",
        "    \"prompt\": final_prompt,\n",
        "}\n",
        "\n",
        "args = \"\"\n",
        "for k, v in config.items():\n",
        "    if isinstance(v, str):\n",
        "        args += f'--{k}=\"{v}\" '\n",
        "    if isinstance(v, bool) and v:\n",
        "        args += f\"--{k} \"\n",
        "    if isinstance(v, float) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "    if isinstance(v, int) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "\n",
        "final_args = f\"python gen_img_diffusers.py {args}\"\n",
        "\n",
        "os.chdir(repo_dir)\n",
        "!{final_args}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "gzBFeYOA81m5"
      },
      "outputs": [],
      "source": [
        "#@title ## 6.4. Launch Portable Web UI\n",
        "import os\n",
        "import random\n",
        "import shutil\n",
        "import zipfile\n",
        "import time\n",
        "import json\n",
        "from google.colab import drive\n",
        "from datetime import timedelta\n",
        "from subprocess import getoutput\n",
        "from IPython.display import clear_output, display, HTML\n",
        "from IPython.utils import capture\n",
        "from tqdm import tqdm\n",
        "\n",
        "webui_dir = os.path.join(root_dir, \"stable-diffusion-webui\")\n",
        "tmp_dir = os.path.join(root_dir, \"tmp\")\n",
        "patches_dir = os.path.join(root_dir, \"patches\")\n",
        "deps_dir = os.path.join(root_dir, \"deps\")\n",
        "extensions_dir = os.path.join(webui_dir, \"extensions\")\n",
        "control_dir = os.path.join(webui_dir, \"models/ControlNet\")\n",
        "\n",
        "webui_models_dir = os.path.join(webui_dir, \"models/Stable-diffusion\")\n",
        "webui_lora_dir = os.path.join(webui_dir, \"models/Lora\")\n",
        "webui_vaes_dir = os.path.join(webui_dir, \"models/VAE\")\n",
        "\n",
        "control_net_max_models_num = 2\n",
        "theme = \"ogxBGreen\"\n",
        "\n",
        "default_prompt = \"masterpiece, best quality,\"\n",
        "default_neg_prompt = \"(worst quality, low quality:1.4)\"\n",
        "default_sampler = \"DPM++ 2M Karras\"\n",
        "default_steps = 20\n",
        "default_width = 512\n",
        "default_height = 768\n",
        "default_denoising_strength = 0.55\n",
        "default_cfg_scale = 7\n",
        "\n",
        "config_file = os.path.join(webui_dir, \"config.json\")\n",
        "ui_config_file = os.path.join(webui_dir, \"ui-config.json\")\n",
        "webui_style_path = os.path.join(webui_dir, \"style.css\")\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "for dir in [patches_dir, deps_dir]:\n",
        "    os.makedirs(dir, exist_ok=True)\n",
        "\n",
        "package_url = [\n",
        "    f\"https://huggingface.co/Linaqruf/fast-repo/resolve/main/anapnoe-webui.tar.lz4\",\n",
        "    f\"https://huggingface.co/Linaqruf/fast-repo/resolve/main/anapnoe-webui-deps.tar.lz4\",\n",
        "    f\"https://huggingface.co/Linaqruf/fast-repo/resolve/main/anapnoe-webui-cache.tar.lz4\",\n",
        "]\n",
        "\n",
        "def pre_download(desc):\n",
        "    for package in tqdm(package_url, desc=desc):\n",
        "        with capture.capture_output() as cap:\n",
        "            package_name = os.path.basename(package)\n",
        "            !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {root_dir} -o {package_name} {package}\n",
        "            if package_name == f\"anapnoe-webui-deps.tar.lz4\":\n",
        "                !tar -xI lz4 -f {package_name} --overwrite-dir --directory=/usr/local/lib/python3.10/dist-packages/\n",
        "            else:\n",
        "                !tar -xI lz4 -f {package_name} --directory=/\n",
        "            os.remove(package_name)\n",
        "            del cap\n",
        "\n",
        "    if os.path.exists(\"/usr/local/lib/python3.10/dist-packages/ffmpy-0.3.0.dist-info\"):\n",
        "        shutil.rmtree(\"/usr/local/lib/python3.10/dist-packages/ffmpy-0.3.0.dist-info\")\n",
        "\n",
        "    s = getoutput(\"nvidia-smi\")\n",
        "    with capture.capture_output() as cap:\n",
        "        if not \"T4\" in s:\n",
        "            !pip uninstall -y xformers\n",
        "            !pip install -q xformers==0.0.18 triton\n",
        "        del cap\n",
        "\n",
        "\n",
        "def read_config(filename):\n",
        "    if filename.endswith(\".json\"):\n",
        "        with open(filename, \"r\") as f:\n",
        "          config = json.load(f)\n",
        "    else:\n",
        "        with open(filename, 'r') as f:\n",
        "          config = f.read()\n",
        "    return config\n",
        "\n",
        "\n",
        "def write_config(filename, config):\n",
        "    if filename.endswith(\".json\"):\n",
        "        with open(filename, \"w\") as f:\n",
        "            json.dump(config, f, indent=4)\n",
        "    else:\n",
        "        with open(filename, 'w', encoding=\"utf-8\") as f:\n",
        "            f.write(config)\n",
        "\n",
        "\n",
        "def open_theme(filename):\n",
        "    themes_folder = os.path.join(webui_dir, \"extensions-builtin/sd_theme_editor/themes\")\n",
        "    themes_file = os.path.join(themes_folder, f\"{filename}.css\")\n",
        "    webui_style_path = os.path.join(webui_dir, \"style.css\")\n",
        "\n",
        "    style_config = read_config(webui_style_path)\n",
        "    style_css_contents = style_config.split(\"/*BREAKPOINT_CSS_CONTENT*/\")[1]\n",
        "\n",
        "    theme_config = read_config(themes_file)\n",
        "    style_data = \":host{\" + theme_config + \"}\" + \"/*BREAKPOINT_CSS_CONTENT*/\" + style_css_contents\n",
        "    write_config(webui_style_path, style_data)\n",
        "\n",
        "\n",
        "def change_config(filename):\n",
        "    config = read_config(filename)\n",
        "    if not \"stable-diffusion-webui\" in config[\"disabled_extensions\"]:\n",
        "        config[\"disabled_extensions\"].append(\"stable-diffusion-webui\")\n",
        "    config[\"outdir_txt2img_samples\"] = os.path.join(tmp_dir, \"outputs/txt2img-images\")\n",
        "    config[\"outdir_img2img_samples\"] = os.path.join(tmp_dir, \"outputs/img2img-images\")\n",
        "    config[\"outdir_extras_samples\"] = os.path.join(tmp_dir, \"outputs/extras-images\")\n",
        "    config[\"outdir_txt2img_grids\"] = os.path.join(tmp_dir, \"outputs/txt2img-grids\")\n",
        "    config[\"outdir_img2img_grids\"] = os.path.join(tmp_dir, \"outputs/img2img-grids\")\n",
        "    config[\"outdir_save\"] = os.path.join(tmp_dir, \"log/images\")\n",
        "    config[\"control_net_max_models_num\"] = control_net_max_models_num\n",
        "    config[\"control_net_models_path\"] = control_dir\n",
        "    config[\"control_net_allow_script_control\"] = True\n",
        "    config[\"additional_networks_extra_lora_path\"] = webui_lora_dir\n",
        "    config[\"CLIP_stop_at_last_layers\"] = 2\n",
        "    config[\"eta_noise_seed_delta\"] = 0\n",
        "    config[\"show_progress_every_n_steps\"] = 10\n",
        "    config[\"show_progressbar\"] = True\n",
        "    config[\"quicksettings\"] = \"sd_model_checkpoint, sd_vae, CLIP_stop_at_last_layers, use_old_karras_scheduler_sigmas, always_discard_next_to_last_sigma\"\n",
        "    write_config(filename, config)\n",
        "\n",
        "\n",
        "def change_ui_config(filename):\n",
        "    config = read_config(filename)\n",
        "    config[\"txt2img/Prompt/value\"] = default_prompt\n",
        "    config[\"txt2img/Negative prompt/value\"] = default_neg_prompt\n",
        "    config[\"txt2img/Sampling method/value\"] = default_sampler\n",
        "    config[\"txt2img/Sampling steps/value\"] = default_steps\n",
        "    config[\"txt2img/Width/value\"] = default_width\n",
        "    config[\"txt2img/Height/value\"] = default_height\n",
        "    config[\"txt2img/Upscaler/value\"] = \"Latent (nearest-exact)\"\n",
        "    config[\"txt2img/Denoising strength/value\"] = default_denoising_strength\n",
        "    config[\"txt2img/CFG Scale/value\"] = default_cfg_scale\n",
        "    config[\"img2img/Prompt/value\"] = default_prompt\n",
        "    config[\"img2img/Negative prompt/value\"] = default_neg_prompt\n",
        "    config[\"img2img/Sampling method/value\"] = default_sampler\n",
        "    config[\"img2img/Sampling steps/value\"] = default_steps\n",
        "    config[\"img2img/Width/value\"] = default_width\n",
        "    config[\"img2img/Height/value\"] = default_height\n",
        "    config[\"img2img/Denoising strength/value\"] = default_denoising_strength\n",
        "    config[\"img2img/CFG Scale/value\"] = default_cfg_scale\n",
        "    write_config(filename, config)\n",
        "\n",
        "\n",
        "def update_extensions():\n",
        "    start_time = time.time()\n",
        "    extensions_updated = []\n",
        "    with tqdm(\n",
        "        total=len(os.listdir(extensions_dir)),\n",
        "        desc=\"\u001b[1;32mUpdating extensions\",\n",
        "        mininterval=0,\n",
        "    ) as pbar:\n",
        "        for dir in os.listdir(extensions_dir):\n",
        "            if os.path.isdir(os.path.join(extensions_dir, dir)):\n",
        "                os.chdir(os.path.join(extensions_dir, dir))\n",
        "                try:\n",
        "                    with capture.capture_output() as cap:\n",
        "                        !git fetch origin\n",
        "                        !git pull\n",
        "                except Exception as e:\n",
        "                    print(f\"\u001b[1;32mAn error occurred while updating {dir}: {e}\")\n",
        "\n",
        "                output = cap.stdout.strip()\n",
        "                if \"Already up to date.\" not in output:\n",
        "                    extensions_updated.append(dir)\n",
        "                pbar.update(1)\n",
        "\n",
        "    print(\"\\n\")\n",
        "    for ext in extensions_updated:\n",
        "        print(f\"\u001b[1;32m- {ext} updated to new version\")\n",
        "\n",
        "    end_time = time.time()\n",
        "    elapsed_time = int(end_time - start_time)\n",
        "\n",
        "    if elapsed_time < 60:\n",
        "        print(f\"\\n\u001b[1;32mAll extensions are up to date. Took {elapsed_time} sec\")\n",
        "    else:\n",
        "        mins, secs = divmod(elapsed_time, 60)\n",
        "        print(f\"\\n\u001b[1;32mAll extensions are up to date. Took {mins} mins {secs} sec\")\n",
        "\n",
        "\n",
        "def main():\n",
        "    start_time = time.time()\n",
        "\n",
        "    print(\"\u001b[1;32mInstalling...\\n\")\n",
        "\n",
        "    if not os.path.exists(webui_dir):\n",
        "        desc = \"\u001b[1;32mUnpacking Webui\"\n",
        "        pre_download(desc)\n",
        "    else:\n",
        "        print(\"\u001b[1;32mAlready installed, skipping...\")\n",
        "\n",
        "    with capture.capture_output() as cap:\n",
        "        os.chdir(os.path.join(webui_dir, \"repositories/stable-diffusion-stability-ai\"))\n",
        "        !git apply {patches_dir}/stablediffusion-lowram.patch\n",
        "\n",
        "        !sed -i \"s@os.path.splitext(checkpoint_.*@os.path.splitext(checkpoint_file); map_location='cuda'@\" {webui_dir}/modules/sd_models.py\n",
        "        !sed -i 's@ui.create_ui().*@ui.create_ui();shared.demo.queue(concurrency_count=999999,status_update_rate=0.1)@' {webui_dir}/webui.py\n",
        "\n",
        "        !sed -i \"s@'cpu'@'cuda'@\" {webui_dir}/modules/extras.py\n",
        "        del cap\n",
        "      \n",
        "    end_time = time.time()\n",
        "    elapsed_time = int(end_time - start_time)\n",
        "\n",
        "    change_config(config_file)\n",
        "    change_ui_config(ui_config_file)\n",
        "    open_theme(theme)\n",
        "\n",
        "    if elapsed_time < 60:\n",
        "        print(f\"\u001b[1;32mFinished unpacking. Took {elapsed_time} sec\")\n",
        "    else:\n",
        "        mins, secs = divmod(elapsed_time, 60)\n",
        "        print(f\"\u001b[1;32mFinished unpacking. Took {mins} mins {secs} sec\")\n",
        "\n",
        "    update_extensions()\n",
        "\n",
        "    #@markdown > Get <b>your</b> `ngrok_token` [here](https://dashboard.ngrok.com/get-started/your-authtoken) \n",
        "    ngrok_token = \"\" #@param {type: 'string'}\n",
        "    ngrok_region = \"ap\" #@param [\"us\", \"eu\", \"au\", \"ap\", \"sa\", \"jp\", \"in\"]\n",
        "\n",
        "    with capture.capture_output() as cap:\n",
        "      for file in os.listdir(output_dir):\n",
        "        file_path = os.path.join(output_dir, file)\n",
        "        if file_path.endswith((\".safetensors\", \".ckpt\")):\n",
        "          !ln \"{file_path}\" {webui_models_dir}\n",
        "\n",
        "      for file in os.listdir(pretrained_model):\n",
        "        file_path = os.path.join(pretrained_model, file)\n",
        "        if file_path.endswith((\".safetensors\", \".ckpt\")):\n",
        "          !ln \"{file_path}\" {webui_models_dir}\n",
        "\n",
        "      for file in os.listdir(vae_dir):\n",
        "        file_path = os.path.join(vae_dir, file)\n",
        "        if file_path.endswith(\".vae.pt\"):\n",
        "          !ln \"{file_path}\" {webui_vaes_dir}\n",
        "\n",
        "      del cap\n",
        "    model_path = os.path.join(webui_models_dir, project_name + \".\" + save_model_as)\n",
        "\n",
        "    os.chdir(webui_dir)\n",
        "\n",
        "    print(\"\u001b[1;32m\")\n",
        "\n",
        "    config = {\n",
        "        \"enable-insecure-extension-access\": True,\n",
        "        \"disable-safe-unpickle\": True,\n",
        "        \"multiple\": True if not ngrok_token else False,\n",
        "        \"ckpt\": model_path if os.path.exists(model_path) else None,\n",
        "        \"ckpt-dir\": webui_models_dir,\n",
        "        \"vae-dir\": webui_vaes_dir,\n",
        "        \"share\": True if not ngrok_token else False,\n",
        "        \"no-half-vae\": True,\n",
        "        \"lowram\": True,\n",
        "        \"gradio-queue\": True,\n",
        "        \"no-hashing\": True,\n",
        "        \"disable-console-progressbars\": True,\n",
        "        \"ngrok\": ngrok_token if ngrok_token else None,\n",
        "        \"ngrok-region\": ngrok_region if ngrok_token else None,\n",
        "        \"xformers\": True,\n",
        "        \"opt-sub-quad-attention\": True,\n",
        "        \"opt-channelslast\": True,\n",
        "        \"theme\": \"dark\"\n",
        "    }\n",
        "\n",
        "    args = \"\"\n",
        "    for k, v in config.items():\n",
        "        if k.startswith(\"_\"):\n",
        "            args += f'\"{v}\" '\n",
        "        elif isinstance(v, str):\n",
        "            args += f'--{k}=\"{v}\" '\n",
        "        elif isinstance(v, bool) and v:\n",
        "            args += f\"--{k} \"\n",
        "        elif isinstance(v, float) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "        elif isinstance(v, int) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "\n",
        "    final_args = f\"python launch.py {args}\"\n",
        "\n",
        "    os.chdir(webui_dir)\n",
        "    !{final_args}\n",
        "\n",
        "main()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "N6ckzE2GWudi"
      },
      "source": [
        "# VII. Extras"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "EHOvjWCHa-JT"
      },
      "outputs": [],
      "source": [
        "#@title ## 7.1. Convert Diffusers to Checkpoint\n",
        "import os\n",
        "%store -r\n",
        "\n",
        "os.chdir(tools_dir)\n",
        "\n",
        "#@markdown ### Conversion Config\n",
        "model_to_load = \"\" #@param {'type': 'string'}\n",
        "model_to_save = os.path.splitext(model_to_load)[0]\n",
        "convert = \"checkpoint_to_diffusers\" #@param [\"diffusers_to_checkpoint\", \"checkpoint_to_diffusers\"] {'allow-input': false}\n",
        "v2 = False #@param {type:'boolean'}\n",
        "global_step = 0 #@param {'type': 'number'}\n",
        "epoch = 0 #@param {'type': 'number'}\n",
        "use_safetensors = True #@param {'type': 'boolean'}\n",
        "save_precision_as = \"--float\" #@param [\"--fp16\",\"--bf16\",\"--float\"] {'allow-input': false}\n",
        "\n",
        "#@markdown Additional option for diffusers\n",
        "feature_extractor = True #@param {'type': 'boolean'}\n",
        "safety_checker = True #@param {'type': 'boolean'}\n",
        "\n",
        "reference_model = \"stabilityai/stable-diffusion-2-1\" if v2 else \"runwayml/stable-diffusion-v1-5\" \n",
        "model_output = f\"{model_to_save}.safetensors\" if use_safetensors else f\"{model_to_save}.ckpt\"\n",
        "\n",
        "urls = [\n",
        "    (\"preprocessor_config.json\", \"https://huggingface.co/CompVis/stable-diffusion-safety-checker/resolve/main/preprocessor_config.json\"),\n",
        "    (\"config.json\", \"https://huggingface.co/CompVis/stable-diffusion-safety-checker/resolve/main/config.json\"),\n",
        "    (\"pytorch_model.bin\", \"https://huggingface.co/CompVis/stable-diffusion-safety-checker/resolve/main/pytorch_model.bin\"),\n",
        "]\n",
        "\n",
        "diffusers_to_sd_dict = {\n",
        "    \"_model_to_load\": model_to_load,\n",
        "    \"_model_to_save\": model_output,\n",
        "    \"global_step\": global_step,\n",
        "    \"epoch\": epoch,\n",
        "    \"save_precision_as\": save_precision_as,\n",
        "}\n",
        "\n",
        "sd_to_diffusers_dict = {\n",
        "    \"_model_to_load\": model_to_load,\n",
        "    \"_model_to_save\": model_to_save,\n",
        "    \"v2\": True if v2 else False,\n",
        "    \"v1\": True if not v2 else False,\n",
        "    \"global_step\": global_step,\n",
        "    \"epoch\": epoch,\n",
        "    \"fp16\": True if save_precision_as == \"fp16\" else False,\n",
        "    \"use_safetensors\": use_safetensors,\n",
        "    \"reference_model\": reference_model\n",
        "}\n",
        "\n",
        "def convert_dict(config):\n",
        "    args = \"\"\n",
        "    for k, v in config.items():\n",
        "        if k.startswith(\"_\"):\n",
        "            args += f'\"{v}\" '\n",
        "        elif isinstance(v, str):\n",
        "            args += f'--{k}=\"{v}\" '\n",
        "        elif isinstance(v, bool) and v:\n",
        "            args += f\"--{k} \"\n",
        "        elif isinstance(v, float) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "        elif isinstance(v, int) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "\n",
        "    return args\n",
        "\n",
        "def run_script(script_name, script_args):\n",
        "    !python {script_name} {script_args}\n",
        "\n",
        "def download(output, url, save_dir):\n",
        "    !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d '{save_dir}' -o '{output}' {url}\n",
        "\n",
        "diffusers_to_sd_args = convert_dict(diffusers_to_sd_dict)\n",
        "sd_to_diffusers_args = convert_dict(sd_to_diffusers_dict)\n",
        "\n",
        "if convert == \"diffusers_to_checkpoint\":\n",
        "    if model_to_load.endswith((\"ckpt\",\"safetensors\")):\n",
        "        print(f\"{os.path.basename(model_to_load)} is not in diffusers format\")\n",
        "    else:\n",
        "        run_script(\"convert_diffusers20_original_sd.py\", diffusers_to_sd_args)\n",
        "else:\n",
        "    if not model_to_load.endswith((\"ckpt\",\"safetensors\")):\n",
        "        print(f\"{os.path.basename(model_to_load)} is not in ckpt/safetensors format\")\n",
        "    else:     \n",
        "        run_script(\"convert_diffusers20_original_sd.py\", sd_to_diffusers_args)\n",
        "\n",
        "        if feature_extractor:\n",
        "            save_dir = os.path.join(model_to_save, \"feature_extractor\")\n",
        "            os.makedirs(save_dir, exist_ok=True)\n",
        "            output, url = urls[0]\n",
        "            download(output, url, save_dir)\n",
        "            \n",
        "        if safety_checker:\n",
        "            save_dir = os.path.join(model_to_save, \"safety_checker\")\n",
        "            os.makedirs(save_dir, exist_ok=True)\n",
        "            for output, url in urls[1:]:\n",
        "                download(output, url, save_dir)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "g5Iz_ikf29LV"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "%store -r\n",
        "#@title ## 7.2. Model Pruner\n",
        "\n",
        "os.chdir(tools_dir)\n",
        "\n",
        "if not os.path.exists('prune.py'):\n",
        "    !wget https://raw.githubusercontent.com/lopho/stable-diffusion-prune/main/prune.py\n",
        "\n",
        "#@markdown Convert to Float16\n",
        "fp16 = False #@param {'type':'boolean'}\n",
        "#@markdown Use EMA for weights\n",
        "ema = False #@param {'type':'boolean'}\n",
        "#@markdown Strip CLIP weights\n",
        "no_clip = False #@param {'type':'boolean'}\n",
        "#@markdown Strip VAE weights\n",
        "no_vae = False #@param {'type':'boolean'}\n",
        "#@markdown Strip depth model weights\n",
        "no_depth = False #@param {'type':'boolean'}\n",
        "#@markdown Strip UNet weights\n",
        "no_unet = False #@param {'type':'boolean'}\n",
        "\n",
        "model_path = \"\" #@param {'type' : 'string'}\n",
        "\n",
        "config = {\n",
        "    \"fp16\": fp16,\n",
        "    \"ema\": ema,\n",
        "    \"no_clip\": no_clip,\n",
        "    \"no_vae\": no_vae,\n",
        "    \"no_depth\": no_depth,\n",
        "    \"no_unet\": no_unet,\n",
        "}\n",
        "\n",
        "suffixes = {\n",
        "    \"fp16\": \"-fp16\",\n",
        "    \"ema\": \"-ema\",\n",
        "    \"no_clip\": \"-no-clip\",\n",
        "    \"no_vae\": \"-no-vae\",\n",
        "    \"no_depth\": \"-no-depth\",\n",
        "    \"no_unet\": \"-no-unet\",\n",
        "}\n",
        "\n",
        "print(f\"Loading model from {model_path}\")\n",
        "\n",
        "dir_name = os.path.dirname(model_path)\n",
        "base_name = os.path.basename(model_path)\n",
        "output_name = base_name.split('.')[0]\n",
        "\n",
        "for option, suffix in suffixes.items():\n",
        "    if config[option]:\n",
        "        print(f\"Applying option {option}\")\n",
        "        output_name += suffix\n",
        "        \n",
        "output_name += '-pruned'\n",
        "output_path = os.path.join(dir_name, output_name + ('.ckpt' if model_path.endswith(\".ckpt\") else \".safetensors\"))\n",
        "\n",
        "args = \"\"\n",
        "for k, v in config.items():\n",
        "    if k.startswith(\"_\"):\n",
        "        args += f'\"{v}\" '\n",
        "    elif isinstance(v, str):\n",
        "        args += f'--{k}=\"{v}\" '\n",
        "    elif isinstance(v, bool) and v:\n",
        "        args += f\"--{k} \"\n",
        "    elif isinstance(v, float) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "    elif isinstance(v, int) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "\n",
        "final_args = f\"python3 prune.py {model_path} {output_path} {args}\"\n",
        "!{final_args}\n",
        "\n",
        "print(f\"Saving pruned model to {output_path}\")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "nyIl9BhNXKUq"
      },
      "source": [
        "# VIII. Deployment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "QTXsM170GUpk"
      },
      "outputs": [],
      "source": [
        "# @title ## 7.1. Upload Config\n",
        "from huggingface_hub import login\n",
        "from huggingface_hub import HfApi\n",
        "from huggingface_hub.utils import validate_repo_id, HfHubHTTPError\n",
        "\n",
        "# @markdown Login to Huggingface Hub\n",
        "# @markdown > Get **your** huggingface `WRITE` token [here](https://huggingface.co/settings/tokens)\n",
        "write_token = \"\"  # @param {type:\"string\"}\n",
        "# @markdown Fill this if you want to upload to your organization, or just leave it empty.\n",
        "orgs_name = \"\"  # @param{type:\"string\"}\n",
        "# @markdown If your model/dataset repo does not exist, it will automatically create it.\n",
        "model_name = \"your-model-name\"  # @param{type:\"string\"}\n",
        "dataset_name = \"your-dataset-name\"  # @param{type:\"string\"}\n",
        "make_private = False  # @param{type:\"boolean\"}\n",
        "\n",
        "def authenticate(write_token):\n",
        "    login(write_token, add_to_git_credential=True)\n",
        "    api = HfApi()\n",
        "    return api.whoami(write_token), api\n",
        "\n",
        "\n",
        "def create_repo(api, user, orgs_name, repo_name, repo_type, make_private=False):\n",
        "    global model_repo\n",
        "    global datasets_repo\n",
        "    \n",
        "    if orgs_name == \"\":\n",
        "        repo_id = user[\"name\"] + \"/\" + repo_name.strip()\n",
        "    else:\n",
        "        repo_id = orgs_name + \"/\" + repo_name.strip()\n",
        "\n",
        "    try:\n",
        "        validate_repo_id(repo_id)\n",
        "        api.create_repo(repo_id=repo_id, repo_type=repo_type, private=make_private)\n",
        "        print(f\"{repo_type.capitalize()} repo '{repo_id}' didn't exist, creating repo\")\n",
        "    except HfHubHTTPError as e:\n",
        "        print(f\"{repo_type.capitalize()} repo '{repo_id}' exists, skipping create repo\")\n",
        "    \n",
        "    if repo_type == \"model\":\n",
        "        model_repo = repo_id\n",
        "        print(f\"{repo_type.capitalize()} repo '{repo_id}' link: https://huggingface.co/{repo_id}\\n\")\n",
        "    else:\n",
        "        datasets_repo = repo_id\n",
        "        print(f\"{repo_type.capitalize()} repo '{repo_id}' link: https://huggingface.co/datasets/{repo_id}\\n\")\n",
        "\n",
        "user, api = authenticate(write_token)\n",
        "\n",
        "if model_name:\n",
        "    create_repo(api, user, orgs_name, model_name, \"model\", make_private)\n",
        "if dataset_name:\n",
        "    create_repo(api, user, orgs_name, dataset_name, \"dataset\", make_private)\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "Fuxghk8MnG6j"
      },
      "source": [
        "## 8.2. Upload with Huggingface Hub"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "CIeoJA-eO-8t"
      },
      "outputs": [],
      "source": [
        "# @title ### 8.2.1. Upload Model\n",
        "from huggingface_hub import HfApi\n",
        "from pathlib import Path\n",
        "\n",
        "api = HfApi()\n",
        "\n",
        "# @markdown This will be uploaded to model repo\n",
        "model_path = \"/content/dreambooth/output\"  # @param {type :\"string\"}\n",
        "path_in_repo = \"\"  # @param {type :\"string\"}\n",
        "# @markdown Now you can save your config file for future use\n",
        "config_path = \"/content/dreambooth/config\"  # @param {type :\"string\"}\n",
        "# @markdown Other Information\n",
        "commit_message = \"\"  # @param {type :\"string\"}\n",
        "\n",
        "if not commit_message:\n",
        "    commit_message = \"feat: upload \" + project_name + \" checkpoint\"\n",
        "\n",
        "if os.path.exists(model_path):\n",
        "    vae_exists = os.path.exists(os.path.join(model_path, \"vae\"))\n",
        "    unet_exists = os.path.exists(os.path.join(model_path, \"unet\"))\n",
        "    text_encoder_exists = os.path.exists(os.path.join(model_path, \"text_encoder\"))\n",
        "\n",
        "\n",
        "def upload_model(model_paths, is_folder: bool, is_config: bool):\n",
        "    path_obj = Path(model_paths)\n",
        "    trained_model = path_obj.parts[-1]\n",
        "\n",
        "    if path_in_repo:\n",
        "        trained_model = path_in_repo\n",
        "\n",
        "    if is_config:\n",
        "        if path_in_repo:\n",
        "            trained_model = f\"{path_in_repo}_config\"\n",
        "        else:\n",
        "            trained_model = f\"{project_name}_config\"\n",
        "\n",
        "    if is_folder == True:\n",
        "        print(f\"Uploading {trained_model} to https://huggingface.co/\" + model_repo)\n",
        "        print(f\"Please wait...\")\n",
        "\n",
        "        if vae_exists and unet_exists and text_encoder_exists:\n",
        "            api.upload_folder(\n",
        "                folder_path=model_paths,\n",
        "                repo_id=model_repo,\n",
        "                commit_message=commit_message,\n",
        "                ignore_patterns=\".ipynb_checkpoints\",\n",
        "            )\n",
        "        else:\n",
        "            api.upload_folder(\n",
        "                folder_path=model_paths,\n",
        "                path_in_repo=trained_model,\n",
        "                repo_id=model_repo,\n",
        "                commit_message=commit_message,\n",
        "                ignore_patterns=\".ipynb_checkpoints\",\n",
        "            )\n",
        "        print(\n",
        "            f\"Upload success, located at https://huggingface.co/\"\n",
        "            + model_repo\n",
        "            + \"/tree/main\\n\"\n",
        "        )\n",
        "    else:\n",
        "        print(f\"Uploading {trained_model} to https://huggingface.co/\" + model_repo)\n",
        "        print(f\"Please wait...\")\n",
        "\n",
        "        api.upload_file(\n",
        "            path_or_fileobj=model_paths,\n",
        "            path_in_repo=trained_model,\n",
        "            repo_id=model_repo,\n",
        "            commit_message=commit_message,\n",
        "        )\n",
        "\n",
        "        print(\n",
        "            f\"Upload success, located at https://huggingface.co/\"\n",
        "            + model_repo\n",
        "            + \"/blob/main/\"\n",
        "            + trained_model\n",
        "            + \"\\n\"\n",
        "        )\n",
        "\n",
        "\n",
        "def upload():\n",
        "    if model_path.endswith((\".ckpt\", \".safetensors\", \".pt\")):\n",
        "        upload_model(model_path, False, False)\n",
        "    else:\n",
        "        upload_model(model_path, True, False)\n",
        "\n",
        "    if config_path:\n",
        "        upload_model(config_path, True, True)\n",
        "\n",
        "\n",
        "upload()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "IW-hS9jnmf-E"
      },
      "outputs": [],
      "source": [
        "# @title ### 8.2.2. Upload Dataset\n",
        "from huggingface_hub import HfApi\n",
        "from pathlib import Path\n",
        "import shutil\n",
        "import zipfile\n",
        "import os\n",
        "\n",
        "api = HfApi()\n",
        "\n",
        "# @markdown This will be compressed to zip and  uploaded to datasets repo, leave it empty if not necessary\n",
        "train_data_path = \"/content/dreambooth/train_data\"  # @param {type :\"string\"}\n",
        "last_state_path = \"/content/dreambooth/output/last-state\"  # @param {type :\"string\"}\n",
        "# @markdown `Nerd stuff, only if you want to save training logs`\n",
        "logs_path = \"/content/dreambooth/logs\"  # @param {type :\"string\"}\n",
        "\n",
        "if project_name:\n",
        "    tmp_dataset = \"/content/dreambooth/\" + project_name + \"_dataset\"\n",
        "    tmp_last_state = \"/content/dreambooth/\" + project_name + \"_last_state\"\n",
        "\n",
        "else:\n",
        "    tmp_dataset = \"/content/dreambooth/tmp_dataset\"\n",
        "    tmp_last_state = \"/content/dreambooth/tmp_last_state\"\n",
        "\n",
        "tmp_train_data = tmp_dataset + \"/train_data\"\n",
        "dataset_zip = tmp_dataset + \".zip\"\n",
        "last_state_zip = tmp_last_state + \".zip\"\n",
        "\n",
        "# @markdown  Other Information\n",
        "commit_message = \"\"  # @param {type :\"string\"}\n",
        "\n",
        "if not commit_message:\n",
        "    commit_message = \"feat: upload \" + project_name + \" dataset and logs\"\n",
        "\n",
        "tmp_folder = [\"tmp_dataset\", \"tmp_last_state\", \"tmp_train_data\"]\n",
        "\n",
        "\n",
        "def makedirs(tmp_folders):\n",
        "    os.makedirs(tmp_folders, exist_ok=True)\n",
        "\n",
        "\n",
        "for folder in tmp_folder:\n",
        "    makedirs(folder)\n",
        "\n",
        "\n",
        "def upload_dataset(dataset_paths, is_zip: bool):\n",
        "    path_obj = Path(dataset_paths)\n",
        "    dataset_name = path_obj.parts[-1]\n",
        "\n",
        "    if is_zip:\n",
        "        print(\n",
        "            f\"Uploading {dataset_name} to https://huggingface.co/datasets/\"\n",
        "            + datasets_repo\n",
        "        )\n",
        "        print(f\"Please wait...\")\n",
        "\n",
        "        api.upload_file(\n",
        "            path_or_fileobj=dataset_paths,\n",
        "            path_in_repo=dataset_name,\n",
        "            repo_id=datasets_repo,\n",
        "            repo_type=\"dataset\",\n",
        "            commit_message=commit_message,\n",
        "        )\n",
        "        print(\n",
        "            f\"Upload success, located at https://huggingface.co/datasets/\"\n",
        "            + datasets_repo\n",
        "            + \"/blob/main/\"\n",
        "            + dataset_name\n",
        "            + \"\\n\"\n",
        "        )\n",
        "    else:\n",
        "        print(\n",
        "            f\"Uploading {dataset_name} to https://huggingface.co/datasets/\"\n",
        "            + datasets_repo\n",
        "        )\n",
        "        print(f\"Please wait...\")\n",
        "\n",
        "        api.upload_folder(\n",
        "            folder_path=dataset_paths,\n",
        "            path_in_repo=dataset_name,\n",
        "            repo_id=datasets_repo,\n",
        "            repo_type=\"dataset\",\n",
        "            commit_message=commit_message,\n",
        "            ignore_patterns=\".ipynb_checkpoints\",\n",
        "        )\n",
        "        print(\n",
        "            f\"Upload success, located at https://huggingface.co/datasets/\"\n",
        "            + datasets_repo\n",
        "            + \"/tree/main/\"\n",
        "            + dataset_name\n",
        "            + \"\\n\"\n",
        "        )\n",
        "\n",
        "\n",
        "def zip_file(tmp_folders):\n",
        "    zipfiles = tmp_folders + \".zip\"\n",
        "    with zipfile.ZipFile(zipfiles, \"w\") as zip:\n",
        "        for tmp_folders, dirs, files in os.walk(tmp_folders):\n",
        "            for file in files:\n",
        "                zip.write(os.path.join(tmp_folders, file))\n",
        "\n",
        "\n",
        "def move(src_path, dst_path, is_metadata: bool):\n",
        "    files_to_move = [\n",
        "        \"meta_cap.json\",\n",
        "        \"meta_cap_dd.json\",\n",
        "        \"meta_lat.json\",\n",
        "        \"meta_clean.json\",\n",
        "        \"meta_final.json\",\n",
        "    ]\n",
        "\n",
        "    if os.path.exists(src_path):\n",
        "        shutil.move(src_path, dst_path)\n",
        "\n",
        "    if is_metadata:\n",
        "        parent_meta_path = os.path.dirname(src_path)\n",
        "\n",
        "        for filename in os.listdir(parent_meta_path):\n",
        "            file_path = os.path.join(parent_meta_path, filename)\n",
        "            if filename in files_to_move:\n",
        "                shutil.move(file_path, dst_path)\n",
        "\n",
        "\n",
        "def upload():\n",
        "    if train_data_path:\n",
        "        move(train_data_path, tmp_train_data, False)\n",
        "        zip_file(tmp_dataset)\n",
        "        upload_dataset(dataset_zip, True)\n",
        "        os.remove(dataset_zip)\n",
        "\n",
        "    if last_state_path:\n",
        "        move(last_state_path, tmp_last_state, False)\n",
        "        zip_file(tmp_last_state)\n",
        "        upload_dataset(last_state_zip, True)\n",
        "        os.remove(last_state_zip)\n",
        "\n",
        "    if logs_path:\n",
        "        upload_dataset(logs_path, False)\n",
        "\n",
        "\n",
        "upload()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "CKZpg4keWS5c"
      },
      "source": [
        "## 8.3. Upload with GIT (Alternative)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "6nBlrOrytO9F"
      },
      "outputs": [],
      "source": [
        "# @title ### 8.3.1. Clone Repository\n",
        "\n",
        "clone_model = True  # @param {'type': 'boolean'}\n",
        "clone_dataset = True  # @param {'type': 'boolean'}\n",
        "\n",
        "!git lfs install --skip-smudge\n",
        "!export GIT_LFS_SKIP_SMUDGE=1\n",
        "\n",
        "if clone_model:\n",
        "    !git clone https://huggingface.co/{model_repo} /content/{model_name}\n",
        "\n",
        "if clone_dataset:\n",
        "    !git clone https://huggingface.co/datasets/{datasets_repo} /content/{dataset_name}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "7bJev4PzOFFB"
      },
      "outputs": [],
      "source": [
        "# @title ### 8.3.2. Commit using Git\n",
        "import os\n",
        "\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "# @markdown Choose which repo you want to commit\n",
        "commit_model = True  # @param {'type': 'boolean'}\n",
        "commit_dataset = True  # @param {'type': 'boolean'}\n",
        "# @markdown #### Other Information\n",
        "commit_message = \"\"  # @param {type :\"string\"}\n",
        "\n",
        "if not commit_message:\n",
        "    commit_message = \"feat: upload \" + project_name + \" model and dataset\"\n",
        "\n",
        "!git config --global user.email \"example@mail.com\"\n",
        "!git config --global user.name \"example\"\n",
        "\n",
        "\n",
        "def commit(repo_folder, commit_message):\n",
        "    os.chdir(os.path.join(root_dir, repo_folder))\n",
        "    !git lfs install\n",
        "    !huggingface-cli lfs-enable-largefiles .\n",
        "    !git add .\n",
        "    !git commit -m \"{commit_message}\"\n",
        "    !git push\n",
        "\n",
        "\n",
        "commit(model_name, commit_message)\n",
        "commit(dataset_name, commit_message)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}