kohya-dreambooth.ipynb

{
  "cells": [
    {
      "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"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OQqyYzkDFA0a"
      },
      "source": [
        "| Notebook Name | Description | Link | Old Commit |\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/bc0892647cb17492a106ad1d05716e091eda13f6/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/fb96280233d3434819ba5850b2c968150c4720f7/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/21ad4942a917d3fd1ad6c03d87d16677b427254b/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/9c7f891981bee92cc7690f2094f892c46feb99e2/kohya-dreambooth.ipynb) | \n",
        "| [Fast Kohya Trainer](https://github.com/Linaqruf/kohya-trainer/blob/main/fast-kohya-trainer.ipynb) `NEW`| Easy 1-click LoRA & 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/fast-kohya-trainer.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) | "
      ]
    },
    {
      "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.9/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 {'-q' if not verbose else ''} install --upgrade -r requirements.txt\n",
        "\n",
        "    if install_xformers:\n",
        "        !pip {'-q' if not verbose else ''} install xformers==0.0.18 triton\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. Open Special `File Explorer` for Colab (Optional)\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",
        "%store -r\n",
        "\n",
        "thread = threading.Thread(target=main, args=[[\"--root-dir=/content\", \"--port=8765\"]])\n",
        "thread.start()\n",
        "\n",
        "open_in_new_tab = False  # @param {type:\"boolean\"}\n",
        "\n",
        "if open_in_new_tab:\n",
        "    output.serve_kernel_port_as_window(8765)\n",
        "else:\n",
        "    output.serve_kernel_port_as_iframe(8765, height=\"500\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3gob9_OwTlwh"
      },
      "source": [
        "# II. Pretrained Model Selection"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "wmnsZwClN1XL"
      },
      "outputs": [],
      "source": [
        "# @title ## 2.1. Download Available Model\n",
        "import os\n",
        "\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "installModels = []\n",
        "installv2Models = []\n",
        "\n",
        "# @markdown ### SD1.x model\n",
        "modelUrl = [\n",
        "    \"\",\n",
        "    \"https://huggingface.co/Linaqruf/personal-backup/resolve/main/models/animefull-final-pruned.ckpt\",\n",
        "    \"https://huggingface.co/cag/anything-v3-1/resolve/main/anything-v3-1.safetensors\",\n",
        "    \"https://huggingface.co/Lykon/AnyLoRA/resolve/main/AnyLoRA_noVae_fp16.safetensors\",\n",
        "    \"https://huggingface.co/Lykon/AnimePastelDream/resolve/main/AnimePastelDream_Soft_noVae_fp16.safetensors\",\n",
        "    \"https://huggingface.co/Linaqruf/stolen/resolve/main/pruned-models/chillout_mix-pruned.safetensors\",\n",
        "    \"https://huggingface.co/prompthero/openjourney-v4/resolve/main/openjourney-v4.ckpt\",\n",
        "    \"https://huggingface.co/Linaqruf/stolen/resolve/main/pruned-models/stable_diffusion_1_5-pruned.safetensors\",\n",
        "]\n",
        "modelList = [\n",
        "    \"\",\n",
        "    \"Animefull-final-pruned\",\n",
        "    \"Anything-v3-1\",\n",
        "    \"AnyLoRA\",\n",
        "    \"AnimePastelDream\",    \n",
        "    \"Chillout-mix\",\n",
        "    \"OpenJourney-v4\",\n",
        "    \"Stable-Diffusion-v1-5\",\n",
        "]\n",
        "modelName = \"Anything-v3-1\"  # @param [\"\", \"Animefull-final-pruned\", \"Anything-v3-1\", \"AnyLoRA\", \"AnimePastelDream\", \"Chillout-mix\", \"OpenJourney-v4\", \"Stable-Diffusion-v1-5\"]\n",
        "\n",
        "# @markdown ### SD2.x model\n",
        "v2ModelUrl = [\n",
        "    \"\",\n",
        "    \"https://huggingface.co/stabilityai/stable-diffusion-2-1-base/resolve/main/v2-1_512-ema-pruned.safetensors\",\n",
        "    \"https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors\",\n",
        "    \"https://huggingface.co/p1atdev/pd-archive/resolve/main/plat-v1-3-1.safetensors\",\n",
        "    \"https://huggingface.co/gsdf/Replicant-V1.0/resolve/main/Replicant-V1.0.safetensors\",\n",
        "    \"https://huggingface.co/IlluminatiAI/Illuminati_Diffusion_v1.0/resolve/main/illuminati_diffusion_v1.0.safetensors\",\n",
        "    \"https://huggingface.co/4eJIoBek/Illuminati-Diffusion-v1-1/resolve/main/illuminatiDiffusionV1_v11.safetensors\",\n",
        "    \"https://huggingface.co/hakurei/waifu-diffusion-v1-4/resolve/main/wd-1-4-anime_e2.ckpt\",\n",
        "    \"https://huggingface.co/waifu-diffusion/wd-1-5-beta2/resolve/main/checkpoints/wd-1-5-beta2-fp32.safetensors\",\n",
        "    \"https://huggingface.co/waifu-diffusion/wd-1-5-beta2/resolve/main/checkpoints/wd-1-5-beta2-aesthetic-fp32.safetensors\",\n",
        "]\n",
        "v2ModelList = [\n",
        "    \"\",\n",
        "    \"stable-diffusion-2-1-base\",\n",
        "    \"stable-diffusion-2-1-768v\",\n",
        "    \"plat-diffusion-v1-3-1\",\n",
        "    \"replicant-v1\",\n",
        "    \"illuminati-diffusion-v1-0\",\n",
        "    \"illuminati-diffusion-v1-1\",\n",
        "    \"waifu-diffusion-1-4-anime-e2\",\n",
        "    \"waifu-diffusion-1-5-e2\",\n",
        "    \"waifu-diffusion-1-5-e2-aesthetic\",\n",
        "]\n",
        "v2ModelName = \"\"  # @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 modelName:\n",
        "    installModels.append((modelName, modelUrl[modelList.index(modelName)]))\n",
        "if v2ModelName:\n",
        "    installv2Models.append((v2ModelName, v2ModelUrl[v2ModelList.index(v2ModelName)]))\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()"
      ]
    },
    {
      "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",
        "modelUrl = \"\"  # @param {'type': 'string'}\n",
        "\n",
        "\n",
        "def install(url):\n",
        "    base_name = os.path.basename(url)\n",
        "\n",
        "    if url.startswith(\"https://drive.google.com\"):\n",
        "        os.chdir(pretrained_model)\n",
        "        !gdown --fuzzy {url}\n",
        "    elif url.startswith(\"https://huggingface.co/\"):\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",
        "\n",
        "install(modelUrl)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "SoucgZQ6jgPQ"
      },
      "outputs": [],
      "source": [
        "# @title ## 2.3. Download Available VAE (Optional)\n",
        "import os\n",
        "\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "installVae = []\n",
        "# @markdown Select one of the VAEs to download, select `none` for not download VAE:\n",
        "vaeUrl = [\n",
        "    \"\",\n",
        "    \"https://huggingface.co/Linaqruf/personal-backup/resolve/main/vae/animevae.pt\",\n",
        "    \"https://huggingface.co/hakurei/waifu-diffusion-v1-4/resolve/main/vae/kl-f8-anime.ckpt\",\n",
        "    \"https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.ckpt\",\n",
        "]\n",
        "vaeList = [\"none\", \"anime.vae.pt\", \"waifudiffusion.vae.pt\", \"stablediffusion.vae.pt\"]\n",
        "vaeName = \"anime.vae.pt\"  # @param [\"none\", \"anime.vae.pt\", \"waifudiffusion.vae.pt\", \"stablediffusion.vae.pt\"]\n",
        "\n",
        "installVae.append((vaeName, vaeUrl[vaeList.index(vaeName)]))\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",
        "    if vaeName != \"none\":\n",
        "        for vae in installVae:\n",
        "            install(vae[0], vae[1])\n",
        "    else:\n",
        "        pass\n",
        "\n",
        "\n",
        "install_vae()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "En9UUwGNMRMM"
      },
      "source": [
        "# III. Data Acquisition\n",
        "\n",
        "You have 3 options for acquiring your dataset: \n",
        "1. Uploading it to Colab's local files.\n",
        "2. Bulk downloading images from Danbooru using the `Simple Booru Scraper`, or\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 image_dir in [train_data_dir, reg_data_dir]:\n",
        "    if image_dir:\n",
        "        with capture.capture_output() as cap:\n",
        "            os.makedirs(image_dir, exist_ok=True)\n",
        "            %store image_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": "eFFHVTWNZGbp"
      },
      "outputs": [],
      "source": [
        "# @title ## 3.2. Unzip Dataset\n",
        "import os\n",
        "import shutil\n",
        "from pathlib import Path\n",
        "\n",
        "%store -r\n",
        "\n",
        "# @markdown Specify this section if your dataset is in a `zip` file and has been uploaded somewhere. This will download your dataset and automatically extract it to the `train_data_dir` if the `unzip_to` 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",
        "        !unzip -j -o {url} -d \"{train_data_dir}\"\n",
        "    elif url.startswith(\"https://drive.google.com\"):\n",
        "        os.chdir(root_dir)\n",
        "        !gdown --fuzzy {url}\n",
        "    elif url.startswith(\"https://huggingface.co/\"):\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",
        "    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",
        "\n",
        "\n",
        "download_dataset(zipfile_url)\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "if not zipfile_url.startswith(\"/content\"):\n",
        "    !unzip -j -o \"{root_dir}/{zipfile_name}\" -d \"{unzip_to}\"\n",
        "    os.remove(f\"{root_dir}/{zipfile_name}\")\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",
        "for filename in os.listdir(train_data_dir):\n",
        "    file_path = os.path.join(train_data_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)"
      ]
    },
    {
      "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 a booru site using the specified tags\n",
        "booru = \"Danbooru\" #@param [\"Danbooru\", \"Gelbooru\", \"Safebooru\"]\n",
        "tag1 = \"\" #@param {type: \"string\"}\n",
        "tag2 = \"\" #@param {type: \"string\"}\n",
        "#@markdown This will ignore download configuration above.\n",
        "custom_url = \"\" #@param {type: \"string\"}\n",
        "#@markdown Other Config\n",
        "user_agent = \"gdl/1.24.5\" #@param {type: \"string\"}\n",
        "limit_rate = \"\" #@param {type: \"string\"}\n",
        "write_tags = True #@param {type: \"boolean\"}\n",
        "additional_arguments = \"--filename /O --no-part\" #@param {type: \"string\"}\n",
        "\n",
        "#@markdown **Experimental**: Parallel Download with `aria2c`, fast, but ignoring `write-tags`. Not a wrapper, just a simple code to extract url from `gallery-dl` and download them using `aria2c`.\n",
        "with_aria_2c = False #@param {type: \"boolean\"}\n",
        "\n",
        "tags = tag1 if not tag2 else tag1 + \"+\" + tag2\n",
        "\n",
        "tags = tags.replace(\"(\", \"%28\")\n",
        "tags = tags.replace(\")\", \"%29\")\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",
        "if with_aria_2c:\n",
        "  scraper_text = os.path.join(root_dir, \"scrape_this.txt\")\n",
        "  with capture.capture_output() as cap:\n",
        "    !gallery-dl {custom_url if custom_url else url} \\\n",
        "      --get-urls \\\n",
        "      {\"--limit_rate\" if limit_rate else \"\"} \\\n",
        "      {\"--user-agent \" + user_agent if user_agent else \"\"} \\\n",
        "      {additional_arguments}\n",
        "  with open(scraper_text, \"w\") as f:\n",
        "    f.write(cap.stdout)\n",
        "\n",
        "  os.chdir(train_data_dir)\n",
        "  !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -i {scraper_text}\n",
        "\n",
        "else:\n",
        "  !gallery-dl {custom_url if custom_url else url} \\\n",
        "    --directory \"{train_data_dir}\" \\\n",
        "    {\"--write-tags\" if write_tags else \"\"} \\\n",
        "    {\"--limit_rate\" if limit_rate else \"\"} \\\n",
        "    {\"--user-agent \" + user_agent if user_agent else \"\"} \\\n",
        "    {additional_arguments}\n",
        "\n",
        "if write_tags: \n",
        "  files = [f for f in os.listdir(train_data_dir) if f.endswith(\".txt\")]\n",
        "\n",
        "  for file in files:\n",
        "    old_path = os.path.join(train_data_dir, file)\n",
        "    new_file_name = os.path.splitext(os.path.splitext(file)[0])[0] + \".txt\"\n",
        "    new_path = os.path.join(train_data_dir, 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)"
      ]
    },
    {
      "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",
        "# @markdown ### Delete Unnecessary Files\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",
        "\n",
        "supported_types = [\n",
        "    \".png\",\n",
        "    \".jpg\",\n",
        "    \".jpeg\",\n",
        "    \".webp\",\n",
        "    \".bmp\",\n",
        "    \".caption\",\n",
        "    \".npz\",\n",
        "    \".txt\",\n",
        "    \".json\",\n",
        "]\n",
        "\n",
        "for item in test:\n",
        "    file_ext = os.path.splitext(item)[1]\n",
        "    if file_ext not in supported_types:\n",
        "        print(f\"Deleting file {item} from {train_data_dir}\")\n",
        "        os.remove(os.path.join(train_data_dir, item))\n",
        "\n",
        "# @markdown ### <br> Convert Transparent Images\n",
        "# @markdown This code will convert your transparent dataset with alpha channel (RGBA) to RGB and give it a white background.\n",
        "\n",
        "convert = False  # @param {type:\"boolean\"}\n",
        "random_color = False  # @param {type:\"boolean\"}\n",
        "\n",
        "batch_size = 32\n",
        "\n",
        "images = [\n",
        "    image\n",
        "    for image in os.listdir(train_data_dir)\n",
        "    if image.endswith(\".png\") or image.endswith(\".webp\")\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",
        "\n",
        "def process_image(image_name):\n",
        "    img = Image.open(f\"{train_data_dir}/{image_name}\")\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",
        "            bg.save(f'{train_data_dir}/{image_name.replace(\".webp\", \".jpg\")}', \"JPEG\")\n",
        "            os.remove(f\"{train_data_dir}/{image_name}\")\n",
        "            print(\n",
        "                f\" Converted image: {image_name} to {image_name.replace('.webp', '.jpg')}\"\n",
        "            )\n",
        "        else:\n",
        "            bg.save(f\"{train_data_dir}/{image_name}\", \"PNG\")\n",
        "            print(f\" Converted image: {image_name}\")\n",
        "    else:\n",
        "        if image_name.endswith(\".webp\"):\n",
        "            img.save(f'{train_data_dir}/{image_name.replace(\".webp\", \".jpg\")}', \"JPEG\")\n",
        "            os.remove(f\"{train_data_dir}/{image_name}\")\n",
        "            print(\n",
        "                f\" Converted image: {image_name} to {image_name.replace('.webp', '.jpg')}\"\n",
        "            )\n",
        "        else:\n",
        "            img.save(f\"{train_data_dir}/{image_name}\", \"PNG\")\n",
        "\n",
        "\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\")"
      ]
    },
    {
      "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](https://huggingface.co/spaces/Salesforce/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.\n",
        "#@markdown In short, it can be used as a tool for image captioning. 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",
        "\n",
        "!python make_captions.py \\\n",
        "  \"{train_data_dir}\" \\\n",
        "  --batch_size {batch_size} \\\n",
        "  {\"--beam_search\" if beam_search else \"\"} \\\n",
        "  --min_length {min_length} \\\n",
        "  --max_length {max_length} \\\n",
        "  --caption_extension .caption \\\n",
        "  --max_data_loader_n_workers {max_data_loader_n_workers}"
      ]
    },
    {
      "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 Danbooru-styled image classification developed by [SmilingWolf](https://github.com/SmilingWolf).  Can also be useful for general image tagging.\n",
        "#@markdown 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-swinv2-tagger-v2\" #@param [\"SmilingWolf/wd-v1-4-swinv2-tagger-v2\", \"SmilingWolf/wd-v1-4-convnext-tagger-v2\", \"SmilingWolf/wd-v1-4-vit-tagger-v2\"]\n",
        "remove_underscores = True #@param {type:\"boolean\"} \n",
        "#@markdown Adjust threshold for better training results.\n",
        "#@markdown - High threshold (e.g. `0.85`) for object/character training.\n",
        "#@markdown - Low threshold (e.g. `0.35`) for general/style/environment training.\n",
        "threshold = 0.35 #@param {type:\"slider\", min:0, max:1, step:0.05}\n",
        "\n",
        "!python tag_images_by_wd14_tagger.py \\\n",
        "    \"{train_data_dir}\" \\\n",
        "    --batch_size {batch_size} \\\n",
        "    --repo_id {model} \\\n",
        "    --thresh {threshold} \\\n",
        "    --caption_extension .txt \\\n",
        "    --max_data_loader_n_workers {max_data_loader_n_workers}\n",
        "\n",
        "if remove_underscores:\n",
        "  for filename in os.listdir(train_data_dir):\n",
        "      if filename.endswith('.txt'):\n",
        "          filepath = os.path.join(train_data_dir, filename)\n",
        "          with open(filepath, 'r') as f:\n",
        "              contents = f.read()\n",
        "          contents = contents.replace('_', ' ')\n",
        "          with open(filepath, 'w') as f:\n",
        "              f.write(contents)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "_mLVURhM9PFE"
      },
      "outputs": [],
      "source": [
        "# @title ### 4.2.3. Custom Caption/Tag (Optional)\n",
        "import os\n",
        "\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "# @markdown Add or remove custom tags here.\n",
        "# @markdown [Cheatsheet](https://rentry.org/kohyaminiguide#c-custom-tagscaption)\n",
        "extension = \".txt\"  # @param [\".txt\", \".caption\"]\n",
        "custom_tag = \"\"  # @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 caption/tag instead\n",
        "remove_tag = 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",
        "def write_file(filename, contents):\n",
        "    with open(filename, \"w\") as f:\n",
        "        f.write(contents)\n",
        "\n",
        "def add_tag(filename, tag, append):\n",
        "    contents = read_file(filename)\n",
        "    tag = \", \".join(tag.split())\n",
        "    tag = tag.replace(\"_\", \" \")\n",
        "    if tag in contents:\n",
        "        return\n",
        "    contents = contents.rstrip() + \", \" + tag if append else tag + \", \" + contents\n",
        "    write_file(filename, contents)\n",
        "\n",
        "def delete_tag(filename, tag):\n",
        "    contents = read_file(filename)\n",
        "    tag = \", \".join(tag.split())\n",
        "    tag = tag.replace(\"_\", \" \")\n",
        "    if tag not in contents:\n",
        "        return\n",
        "    contents = \"\".join([s.strip(\", \") for s in contents.split(tag)])\n",
        "    write_file(filename, contents)\n",
        "\n",
        "tag = custom_tag\n",
        "\n",
        "if not any(\n",
        "    [filename.endswith(extension) for filename in os.listdir(train_data_dir)]\n",
        "):\n",
        "    for filename in os.listdir(train_data_dir):\n",
        "        if filename.endswith((\".png\", \".jpg\", \".jpeg\", \".webp\", \".bmp\")):\n",
        "            open(\n",
        "                os.path.join(train_data_dir, filename.split(\".\")[0] + extension),\n",
        "                \"w\",\n",
        "            ).close()\n",
        "\n",
        "if custom_tag:\n",
        "    for filename in os.listdir(train_data_dir):\n",
        "        if filename.endswith(extension):\n",
        "            file_path = os.path.join(train_data_dir, filename)\n",
        "            if custom_tag and not remove_tag:\n",
        "                add_tag(file_path, tag, append)\n",
        "            if custom_tag and remove_tag:\n",
        "                delete_tag(file_path, tag)"
      ]
    },
    {
      "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",
        "pretrained_model_name_or_path = \"/content/pretrained_model/Anything-v3-1.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 This will ignore `output_dir` defined above, and changed to `/content/drive/MyDrive/dreambooth/output` by default\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",
        "\n",
        "# @markdown ### Dreambooth Config\n",
        "# @markdown This will only work for `one concept` training. You can still do multi-concept training but it will require more work. [Guide](https://rentry.org/kohyaminiguide#b-multi-concept-training)\n",
        "\n",
        "train_repeats = 10  # @param {type:\"number\"}\n",
        "reg_repeats = 1  # @param {type:\"number\"}\n",
        "\n",
        "# @markdown There is a lot of misinformation about the activation word, so I have made it clear in this [Rentry](https://rentry.org/kohyaminiguide#a-activation-word).\n",
        "instance_token = \"mksks\"  # @param {type:\"string\"}\n",
        "class_token = \"style\"  # @param {type:\"string\"}\n",
        "\n",
        "# @markdown Enable this option if you want to add an instance token to your caption files. This will function in the same way as the `4.2.3. Custom Caption/Tag (Optional)` cell.\n",
        "add_token_to_caption = False  # @param {type:\"boolean\"}\n",
        "\n",
        "# @markdown ### <br>General Config\n",
        "resolution = 512  # @param {type:\"slider\", min:512, max:1024, step:128}\n",
        "flip_aug = False  # @param {type:\"boolean\"}\n",
        "caption_extension = \".txt\"  # @param [\"none\", \".txt\", \".caption\"]\n",
        "caption_dropout_rate = 0  # @param {type:\"slider\", min:0, max:1, step:0.05}\n",
        "caption_dropout_every_n_epochs = 0  # @param {type:\"number\"}\n",
        "keep_tokens = 0  # @param {type:\"number\"}\n",
        "\n",
        "\n",
        "\n",
        "if add_token_to_caption and keep_tokens < 2:\n",
        "    keep_tokens = 1\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 add_tag(filename, tag):\n",
        "    contents = read_file(filename)\n",
        "    tag = \", \".join(tag.split())\n",
        "    tag = tag.replace(\"_\", \" \")\n",
        "    if tag in contents:\n",
        "        return\n",
        "    contents = tag + \", \" + contents\n",
        "    write_file(filename, contents)\n",
        "\n",
        "def delete_tag(filename, tag):\n",
        "    contents = read_file(filename)\n",
        "    tag = \", \".join(tag.split())\n",
        "    tag = tag.replace(\"_\", \" \")\n",
        "    if tag not in contents:\n",
        "        return\n",
        "    contents = \"\".join([s.strip(\", \") for s in contents.split(tag)])\n",
        "    write_file(filename, contents)\n",
        "\n",
        "if caption_extension != \"none\":\n",
        "    tag = f\"{instance_token}_{class_token}\" if 'class_token' in globals() else instance_token\n",
        "    for filename in os.listdir(train_data_dir):\n",
        "        if filename.endswith(caption_extension):\n",
        "            file_path = os.path.join(train_data_dir, filename)\n",
        "            if add_token_to_caption:\n",
        "                add_tag(file_path, tag)\n",
        "            else:\n",
        "                delete_tag(file_path, tag)\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\": caption_dropout_rate if caption_extension == \".caption\" else 0,\n",
        "            \"caption_tag_dropout_rate\": caption_dropout_rate if caption_extension == \".txt\" else 0,\n",
        "            \"caption_dropout_every_n_epochs\": caption_dropout_every_n_epochs,\n",
        "            \"flip_aug\": flip_aug,\n",
        "            \"color_aug\": False,\n",
        "            \"face_crop_aug_range\": None,\n",
        "            \"subsets\": [\n",
        "                {\n",
        "                    \"image_dir\": train_data_dir,\n",
        "                    \"class_tokens\": f\"{instance_token} {class_token}\" if 'class_token' in globals() else instance_token,\n",
        "                    \"num_repeats\": train_repeats,\n",
        "                },\n",
        "                {\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\": reg_repeats,\n",
        "                },\n",
        "            ],\n",
        "        }\n",
        "    ],\n",
        "}\n",
        "\n",
        "config_str = toml.dumps(config)\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)"
      ]
    },
    {
      "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 `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\"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",
        "train_batch_size = 1  # @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",
        "        \"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",
        "    },\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\": {\"save_model_as\": save_model_as},\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.6. 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",
        "os.chdir(repo_dir)\n",
        "!accelerate launch \\\n",
        "  --config_file={accelerate_config} \\\n",
        "  --num_cpu_threads_per_process=1 \\\n",
        "  train_db.py \\\n",
        "  --sample_prompts={sample_prompt} \\\n",
        "  --dataset_config={dataset_config} \\\n",
        "  --config_file={config_file}"
      ]
    },
    {
      "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",
        "instance_prompt = \"\"  # @param {type: \"string\"}\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\"{instance_prompt}, {prompt} --n {negative}\" if instance_prompt else 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",
        "!{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",
        "os.chdir(root_dir)\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",
        "\n",
        "webui_models_dir = os.path.join(webui_dir, \"models/Stable-diffusion\")\n",
        "# webui_lora_dir = os.path.join(extensions_dir, \"sd-webui-additional-networks/models/lora\")\n",
        "webui_vaes_dir = os.path.join(webui_dir, \"models/VAE\")\n",
        "\n",
        "commit_hash = \"\" \n",
        "os.chdir(root_dir)\n",
        "\n",
        "package_url = [\"https://huggingface.co/Linaqruf/fast-repo/resolve/main/webui.tar.lz4\",\n",
        "               \"https://huggingface.co/Linaqruf/fast-repo/resolve/main/webui-deps.tar.lz4\",\n",
        "               \"https://huggingface.co/Linaqruf/fast-repo/resolve/main/webui-cache.tar.lz4\"]\n",
        "    \n",
        "def pre_download():\n",
        "  for package in tqdm(package_url, desc='\u001b[1;32mUnpacking WebUI'):\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 == \"webui-deps.tar.lz4\":\n",
        "        !tar -xI lz4 -f {package_name} --overwrite-dir --directory=/usr/local/lib/python3.9/dist-packages/\n",
        "      else:\n",
        "        !tar -xI lz4 -f {package_name} --directory=/\n",
        "      os.remove(os.path.basename(package_name))\n",
        "      del cap\n",
        "      \n",
        "  if os.path.exists(\"/usr/local/lib/python3.9/dist-packages/ffmpy-0.3.0.dist-info\"):\n",
        "    shutil.rmtree(\"/usr/local/lib/python3.9/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 --pre xformers\n",
        "      !pip install -q --pre triton\n",
        "    del cap\n",
        "\n",
        "def read_config(filename):\n",
        "  with open(filename, \"r\") as f:\n",
        "    config = json.load(f)\n",
        "  return config\n",
        "\n",
        "def write_config(filename, config):\n",
        "  with open(filename, \"w\") as f:\n",
        "    json.dump(config, f, indent=4)\n",
        "\n",
        "start_time = time.time()\n",
        "\n",
        "print(\"\u001b[1;32mInstalling...\\n\")\n",
        "pre_download()\n",
        "\n",
        "print(\"\u001b[1;32mUsing new UI/UX from @Anapnoe...\")\n",
        "try:\n",
        "  with capture.capture_output() as cap:\n",
        "    os.chdir(webui_dir)\n",
        "    config_file = os.path.join(webui_dir, \"config.json\")\n",
        "    config = read_config(config_file)\n",
        "    if not \"stable-diffusion-webui\" in config[\"disabled_extensions\"]:\n",
        "      config[\"disabled_extensions\"].append(\"stable-diffusion-webui\")\n",
        "      write_config(config_file, config)\n",
        "      \n",
        "    !git remote set-url origin https://github.com/anapnoe/stable-diffusion-webui/\n",
        "    !git pull\n",
        "    if commit_hash:\n",
        "      !git reset --hard {commit_hash}\n",
        "      print(\"\u001b[1;32mUI Commit hash: \", commit_hash)                  \n",
        "    del cap\n",
        "except Exception as e:\n",
        "  print(\"\u001b[1;32mAn error occurred:\", e)\n",
        "\n",
        "with capture.capture_output() as cap:\n",
        "  !wget https://raw.githubusercontent.com/ddPn08/automatic1111-colab/main/patches/stablediffusion-lowram.patch -P {patches_dir}  -c\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'@\" /content/stable-diffusion-webui/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)@' /content/stable-diffusion-webui/webui.py\n",
        "  !sed -i \"s@'cpu'@'cuda'@\" /content/stable-diffusion-webui/modules/extras.py\n",
        "  del cap\n",
        "  \n",
        "with open(os.path.join(webui_dir, \"config.json\"), \"r\") as f:\n",
        "  data = json.load(f)\n",
        "\n",
        "data[\"outdir_txt2img_samples\"] = os.path.join(tmp_dir, \"outputs/txt2img-images\")\n",
        "data[\"outdir_img2img_samples\"] = os.path.join(tmp_dir, \"outputs/img2img-images\")\n",
        "data[\"outdir_extras_samples\"] = os.path.join(tmp_dir, \"outputs/extras-images\")\n",
        "data[\"outdir_txt2img_grids\"] = os.path.join(tmp_dir, \"outputs/txt2img-grids\")\n",
        "data[\"outdir_img2img_grids\"] = os.path.join(tmp_dir, \"outputs/img2img-grids\")\n",
        "data[\"outdir_save\"] = os.path.join(tmp_dir, \"log/images\")\n",
        "\n",
        "with open(os.path.join(webui_dir, \"config.json\"), \"w\") as f:\n",
        "  json.dump(data, f, indent=4)\n",
        "  \n",
        "end_time = time.time()\n",
        "elapsed_time = int(end_time - start_time)\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",
        "start_time = time.time()\n",
        "extensions_updated = []\n",
        "with tqdm(total=len(os.listdir(extensions_dir)), desc=\"\u001b[1;32mUpdating extensions\",  mininterval=0) 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",
        "      with capture.capture_output() as cap:\n",
        "        !git fetch origin\n",
        "        !git pull\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",
        "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",
        "print('\\n\u001b[1;32mAll is done! Go to the next step.')\n",
        "\n",
        "#@markdown > Get <b>your</b> token for ngrok [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",
        "\n",
        "\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",
        "!python launch.py \\\n",
        "  --enable-insecure-extension-access \\\n",
        "  --disable-safe-unpickle \\\n",
        "  {\"--ckpt \" + model_path if os.path.exists(model_path) else \"\"} \\\n",
        "  --ckpt-dir {webui_models_dir} \\\n",
        "  --vae-dir {webui_vaes_dir} \\\n",
        "  {\"--share\" if not ngrok_token else \"\"} \\\n",
        "  --no-half-vae \\\n",
        "  --lowram \\\n",
        "  --gradio-queue \\\n",
        "  --no-hashing \\\n",
        "  --disable-console-progressbars \\\n",
        "  {\"--ngrok \" + ngrok_token if ngrok_token else \"\"} \\\n",
        "  {\"--ngrok-region \" + ngrok_region if ngrok_token else \"\"} \\\n",
        "  --xformers \\\n",
        "  --opt-sub-quad-attention \\\n",
        "  --opt-channelslast \\\n",
        "  --theme dark"
      ]
    },
    {
      "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 = \"--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",
        "\n",
        "if convert == \"diffusers_to_checkpoint\":\n",
        "  model_output = f\"{model_to_save}.safetensors\" if use_safetensors else f\"{model_to_save}.ckpt\"\n",
        "  if not model_to_load.endswith(\".ckpt\") or model_to_load.endswith(\".safetensors\"):\n",
        "    !python convert_diffusers20_original_sd.py \\\n",
        "        \"{model_to_load}\" \\\n",
        "        \"{model_output}\" \\\n",
        "        --global_step {global_step} \\\n",
        "        --epoch {epoch} \\\n",
        "        {save_precision}\n",
        "else:    \n",
        "    !python convert_diffusers20_original_sd.py \\\n",
        "        \"{model_to_load}\" \\\n",
        "        \"{model_to_save}\" \\\n",
        "        {\"--fp16\" if save_precision == \"--fp16\" else \"\"} \\\n",
        "        --global_step {global_step} \\\n",
        "        --epoch {epoch} \\\n",
        "        {\"--use_safetensors\" if use_safetensors else \"\"} \\\n",
        "        {\"--v2\" if v2 else \"--v1\"} \\\n",
        "        --reference_model {reference_model} \n",
        "\n",
        "    url1 = \"https://huggingface.co/CompVis/stable-diffusion-safety-checker/resolve/main/preprocessor_config.json\"\n",
        "    url2 = \"https://huggingface.co/CompVis/stable-diffusion-safety-checker/resolve/main/config.json\"\n",
        "    url3 = \"https://huggingface.co/CompVis/stable-diffusion-safety-checker/resolve/main/pytorch_model.bin\"\n",
        "\n",
        "    if feature_extractor == True:\n",
        "      if not os.path.exists(f'{model_to_save}/feature_extractor'):\n",
        "        os.makedirs(f'{model_to_save}/feature_extractor')\n",
        "      \n",
        "      !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d '{model_to_save}/feature_extractor' -o 'preprocessor_config.json' {url1}\n",
        "\n",
        "    if safety_checker == True:\n",
        "      if not os.path.exists(f'{model_to_save}/safety_checker'):\n",
        "        os.makedirs(f'{model_to_save}/safety_checker')\n",
        "      \n",
        "      !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d '{model_to_save}/safety_checker' -o 'config.json' {url2}\n",
        "      !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d '{model_to_save}/safety_checker' -o 'pytorch_model.bin' {url3}"
      ]
    },
    {
      "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",
        "input = \"\" #@param {'type' : 'string'}\n",
        "\n",
        "print(f\"Loading model from {input}\")\n",
        "\n",
        "input_path = os.path.dirname(input)\n",
        "base_name = os.path.basename(input)\n",
        "output_name = base_name.split('.')[0]\n",
        "\n",
        "if fp16:\n",
        "    print(\"Converting to float16\")\n",
        "    output_name += '-fp16'\n",
        "if ema:\n",
        "    print(\"Using EMA for weights\")\n",
        "    output_name += '-ema'\n",
        "if no_clip:\n",
        "    print(\"Stripping CLIP weights\")\n",
        "    output_name += '-no-clip'\n",
        "if no_vae:\n",
        "    print(\"Stripping VAE weights\")\n",
        "    output_name += '-no-vae'\n",
        "if no_depth:\n",
        "    print(\"Stripping depth model weights\")\n",
        "    output_name += '-no-depth'\n",
        "if no_unet:\n",
        "    print(\"Stripping UNet weights\")\n",
        "    output_name += '-no-unet'\n",
        "output_name += '-pruned'\n",
        "output_path = os.path.join(input_path, output_name + ('.ckpt' if input.endswith(\".ckpt\") else \".safetensors\"))\n",
        "\n",
        "!python3 prune.py \"{input}\" \\\n",
        "  \"{output_path}\" \\\n",
        "  {'--fp16' if fp16 else ''} \\\n",
        "  {'--ema' if ema else ''} \\\n",
        "  {'--no-clip' if no_clip else ''} \\\n",
        "  {'--no-vae' if no_vae else ''} \\\n",
        "  {'--no-depth' if no_depth else ''} \\\n",
        "  {'--no-unet' if no_unet else ''}\n",
        "\n",
        "print(f\"Saving pruned model to {output_path}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nyIl9BhNXKUq"
      },
      "source": [
        "# VIII. Deployment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "QTXsM170GUpk"
      },
      "outputs": [],
      "source": [
        "# @title ## 8.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",
        "login(write_token, add_to_git_credential=True)\n",
        "\n",
        "api = HfApi()\n",
        "user = api.whoami(write_token)\n",
        "\n",
        "# @markdown Fill this if you want to upload to your organization, or just leave it empty.\n",
        "\n",
        "orgs_name = \"\"  # @param{type:\"string\"}\n",
        "\n",
        "# @markdown If your model/dataset repo didn't exist, it will automatically create your repo.\n",
        "model_name = \"your-model-name\"  # @param{type:\"string\"}\n",
        "dataset_name = \"your-dataset-name\"  # @param{type:\"string\"}\n",
        "make_this_model_private = False  # @param{type:\"boolean\"}\n",
        "\n",
        "if orgs_name == \"\":\n",
        "    model_repo = user[\"name\"] + \"/\" + model_name.strip()\n",
        "    datasets_repo = user[\"name\"] + \"/\" + dataset_name.strip()\n",
        "else:\n",
        "    model_repo = orgs_name + \"/\" + model_name.strip()\n",
        "    datasets_repo = orgs_name + \"/\" + dataset_name.strip()\n",
        "\n",
        "if model_name != \"\":\n",
        "    try:\n",
        "        validate_repo_id(model_repo)\n",
        "        api.create_repo(repo_id=model_repo, private=make_this_model_private)\n",
        "        print(\"Model Repo didn't exists, creating repo\")\n",
        "        print(\"Model Repo: \", model_repo, \"created!\\n\")\n",
        "\n",
        "    except HfHubHTTPError as e:\n",
        "        print(f\"Model Repo: {model_repo} exists, skipping create repo\\n\")\n",
        "\n",
        "if dataset_name != \"\":\n",
        "    try:\n",
        "        validate_repo_id(datasets_repo)\n",
        "        api.create_repo(\n",
        "            repo_id=datasets_repo, repo_type=\"dataset\", private=make_this_model_private\n",
        "        )\n",
        "        print(\"Dataset Repo didn't exists, creating repo\")\n",
        "        print(\"Dataset Repo\", datasets_repo, \"created!\\n\")\n",
        "\n",
        "    except HfHubHTTPError as e:\n",
        "        print(f\"Dataset repo: {datasets_repo} exists, skipping create repo\\n\")"
      ]
    },
    {
      "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()"
      ]
    },
    {
      "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
}