kohya-LoRA-dreambooth.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "slgjeYgd6pWp"
      },
      "source": [
        "![visitors](https://visitor-badge.glitch.me/badge?page_id=linaqruf.lora-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 LoRA Dreambooth**\n",
        "A Colab Notebook For LoRA Training (Dreambooth Method)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D_6E01tRBfQU"
      },
      "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, \"LoRA\")\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/LoRA/train_data\"  # @param {type:'string'}\n",
        "reg_data_dir = \"/content/LoRA/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)\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yHNbl3O_NSS0"
      },
      "source": [
        "# V. 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/LoRA/output\"  # @param {'type':'string'}\n",
        "\n",
        "# @markdown This will ignore `output_dir` defined above, and changed to `/content/drive/MyDrive/LoRA/output` by default\n",
        "output_to_drive = False  # @param {'type':'boolean'}\n",
        "\n",
        "if output_to_drive:\n",
        "    output_dir = \"/content/drive/MyDrive/LoRA/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)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "G5u_DhFeyJ6R"
      },
      "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": "iD5Ecamp4rVW"
      },
      "outputs": [],
      "source": [
        "# @title ## 5.3. LoRA and Optimizer Config\n",
        "\n",
        "# @markdown ### LoRA Config:\n",
        "network_category = \"LoRA\"  # @param [\"LoRA\", \"LoCon\", \"LoCon_Lycoris\", \"LoHa\"]\n",
        "\n",
        "# @markdown Recommended values:\n",
        "\n",
        "# @markdown | network_category | network_dim | network_alpha | conv_dim | conv_alpha |\n",
        "# @markdown | :---: | :---: | :---: | :---: | :---: |\n",
        "# @markdown | LoRA | 32 | 1 | - | - |\n",
        "# @markdown | LoCon | 16 | 8 | 8 | 1 |\n",
        "# @markdown | LoHa | 8 | 4 | 4 | 1 |\n",
        "\n",
        "# @markdown - Currently, `dropout` and `cp_decomposition` is not available in this notebook.\n",
        "\n",
        "# @markdown `conv_dim` and `conv_alpha` are needed to train `LoCon` and `LoHa`, skip it if you train normal `LoRA`. But remember, when in doubt, set `dim = alpha`.\n",
        "conv_dim = 32  # @param {'type':'number'}\n",
        "conv_alpha = 16  # @param {'type':'number'}\n",
        "# @markdown It's recommended to not set `network_dim` and `network_alpha` higher than `64`, especially for LoHa.\n",
        "# @markdown But if you want to train with higher dim/alpha so badly, try using higher learning rate. Because the model learning faster in higher dim.\n",
        "network_dim = 32  # @param {'type':'number'}\n",
        "network_alpha = 16  # @param {'type':'number'}\n",
        "# @markdown You can specify this field for resume training.\n",
        "network_weight = \"\"  # @param {'type':'string'}\n",
        "network_module = \"lycoris.kohya\" if network_category in [\"LoHa\", \"LoCon_Lycoris\"] else \"networks.lora\"\n",
        "network_args = \"\" if network_category == \"LoRA\" else [\n",
        "    f\"conv_dim={conv_dim}\", f\"conv_alpha={conv_alpha}\",\n",
        "    ]\n",
        "\n",
        "if network_category == \"LoHa\":\n",
        "  network_args.append(\"algo=loha\")\n",
        "elif network_category == \"LoCon_Lycoris\":\n",
        "  network_args.append(\"algo=lora\")\n",
        "\n",
        "# @markdown ### <br>Optimizer Config:\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 `unet_lr` to `1.0` if you use `DAdaptation` optimizer, because it's a [free learning rate](https://github.com/facebookresearch/dadaptation) algorithm.\n",
        "# @markdown However `text_encoder_lr = 1/2 * unet_lr` still applied, so you need to set `0.5` for `text_encoder_lr`.\n",
        "# @markdown Also actually you don't need to specify `learning_rate` value if both `unet_lr` and `text_encoder_lr` are defined.\n",
        "train_unet = True  # @param {'type':'boolean'}\n",
        "unet_lr = 1e-4  # @param {'type':'number'}\n",
        "train_text_encoder = True  # @param {'type':'boolean'}\n",
        "text_encoder_lr = 5e-5  # @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(\"- LoRA Config:\")\n",
        "print(f\"  - Additional network category: {network_category}\")\n",
        "print(f\"  - Loading network module: {network_module}\")\n",
        "if not network_category == \"LoRA\":\n",
        "  print(f\"  - network args: {network_args}\")\n",
        "print(f\"  - {network_module} linear_dim set to: {network_dim}\")\n",
        "print(f\"  - {network_module} linear_alpha set to: {network_alpha}\")\n",
        "if not network_category == \"LoRA\":\n",
        "  print(f\"  - {network_module} conv_dim set to: {conv_dim}\")\n",
        "  print(f\"  - {network_module} conv_alpha set to: {conv_alpha}\")\n",
        "\n",
        "if not network_weight:\n",
        "    print(\"  - No LoRA weight loaded.\")\n",
        "else:\n",
        "    if os.path.exists(network_weight):\n",
        "        print(f\"  - Loading LoRA weight: {network_weight}\")\n",
        "    else:\n",
        "        print(f\"  - {network_weight} does not exist.\")\n",
        "        network_weight = \"\"\n",
        "\n",
        "print(\"- Optimizer Config:\")\n",
        "print(f\"  - Using {optimizer_type} as Optimizer\")\n",
        "if optimizer_args:\n",
        "    print(f\"  - Optimizer Args: {optimizer_args}\")\n",
        "if train_unet and train_text_encoder:\n",
        "    print(\"  - Train UNet and Text Encoder\")\n",
        "    print(f\"    - UNet learning rate: {unet_lr}\")\n",
        "    print(f\"    - Text encoder learning rate: {text_encoder_lr}\")\n",
        "if train_unet and not train_text_encoder:\n",
        "    print(\"  - Train UNet only\")\n",
        "    print(f\"    - UNet learning rate: {unet_lr}\")\n",
        "if train_text_encoder and not train_unet:\n",
        "    print(\"  - Train Text Encoder only\")\n",
        "    print(f\"    - Text encoder learning rate: {text_encoder_lr}\")\n",
        "print(f\"  - Learning rate warmup steps: {lr_warmup_steps}\")\n",
        "print(f\"  - Learning rate Scheduler: {lr_scheduler}\")\n",
        "if lr_scheduler == \"cosine_with_restarts\":\n",
        "    print(f\"  - lr_scheduler_num_cycles: {lr_scheduler_num_cycles}\")\n",
        "elif lr_scheduler == \"polynomial\":\n",
        "    print(f\"  - lr_scheduler_power: {lr_scheduler_power}\")\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-Z4w3lfFKLjr"
      },
      "outputs": [],
      "source": [
        "# @title ## 5.4. Training Config\n",
        "\n",
        "import toml\n",
        "import os\n",
        "\n",
        "%store -r\n",
        "lowram = True  # @param {type:\"boolean\"}\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",
        "num_epochs = 10  # @param {type:\"number\"}\n",
        "train_batch_size = 6  # @param {type:\"number\"}\n",
        "mixed_precision = \"fp16\"  # @param [\"no\",\"fp16\",\"bf16\"] {allow-input: false}\n",
        "save_precision = \"fp16\"  # @param [\"float\", \"fp16\", \"bf16\"] {allow-input: false}\n",
        "save_n_epochs_type = \"save_every_n_epochs\"  # @param [\"save_every_n_epochs\", \"save_n_epoch_ratio\"] {allow-input: false}\n",
        "save_n_epochs_type_value = 1  # @param {type:\"number\"}\n",
        "save_model_as = \"safetensors\"  # @param [\"ckpt\", \"pt\", \"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/LoRA/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\n",
        "        if v2 and v_parameterization\n",
        "        else False,\n",
        "        \"pretrained_model_name_or_path\": pretrained_model_name_or_path,\n",
        "        \"vae\": vae,\n",
        "    },\n",
        "    \"additional_network_arguments\": {\n",
        "        \"no_metadata\": False,\n",
        "        \"unet_lr\": float(unet_lr) if train_unet else None,\n",
        "        \"text_encoder_lr\": float(text_encoder_lr) if train_text_encoder else None,\n",
        "        \"network_weights\": network_weight,\n",
        "        \"network_module\": network_module,\n",
        "        \"network_dim\": network_dim,\n",
        "        \"network_alpha\": network_alpha,\n",
        "        \"network_args\": network_args,\n",
        "        \"network_train_unet_only\": True if train_unet and not train_text_encoder else False,\n",
        "        \"network_train_text_encoder_only\": True if train_text_encoder and not train_unet else False,\n",
        "        \"training_comment\": None,\n",
        "    },\n",
        "    \"optimizer_arguments\": {\n",
        "        \"optimizer_type\": optimizer_type,\n",
        "        \"learning_rate\": unet_lr,\n",
        "        \"max_grad_norm\": 1.0,\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\": save_n_epochs_type_value if save_n_epochs_type == \"save_every_n_epochs\" else None,\n",
        "        \"save_n_epoch_ratio\": save_n_epochs_type_value if save_n_epochs_type == \"save_n_epoch_ratio\" else None,\n",
        "        \"save_last_n_epochs\": None,\n",
        "        \"save_state\": None,\n",
        "        \"save_last_n_epochs_state\": None,\n",
        "        \"resume\": None,\n",
        "        \"train_batch_size\": train_batch_size,\n",
        "        \"max_token_length\": 225,\n",
        "        \"mem_eff_attn\": False,\n",
        "        \"xformers\": True,\n",
        "        \"max_train_epochs\": num_epochs,\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",
        "        \"lowram\": lowram,\n",
        "    },\n",
        "    \"sample_prompt_arguments\": {\n",
        "        \"sample_every_n_steps\": None,\n",
        "        \"sample_every_n_epochs\": 1 if enable_sample_prompt else 999999,\n",
        "        \"sample_sampler\": sampler,\n",
        "    },\n",
        "    \"dreambooth_arguments\": {\n",
        "        \"prior_loss_weight\": 1.0,\n",
        "    },\n",
        "    \"saving_arguments\": {\n",
        "        \"save_model_as\": save_model_as\n",
        "    },\n",
        "}\n",
        "\n",
        "config_path = os.path.join(config_dir, \"config_file.toml\")\n",
        "prompt_path = os.path.join(config_dir, \"sample_prompt.txt\")\n",
        "\n",
        "for key in config:\n",
        "    if isinstance(config[key], dict):\n",
        "        for sub_key in config[key]:\n",
        "            if config[key][sub_key] == \"\":\n",
        "                config[key][sub_key] = None\n",
        "    elif config[key] == \"\":\n",
        "        config[key] = None\n",
        "\n",
        "config_str = toml.dumps(config)\n",
        "\n",
        "def write_file(filename, contents):\n",
        "    with open(filename, \"w\") as f:\n",
        "        f.write(contents)\n",
        "\n",
        "write_file(config_path, config_str)\n",
        "write_file(prompt_path, sample_str)\n",
        "    \n",
        "print(config_str)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "p_SHtbFwHVl1"
      },
      "outputs": [],
      "source": [
        "#@title ## 5.5. Start Training\n",
        "\n",
        "#@markdown Check your config here if you want to edit something: \n",
        "#@markdown - `sample_prompt` : /content/LoRA/config/sample_prompt.txt\n",
        "#@markdown - `config_file` : /content/LoRA/config/config_file.toml\n",
        "#@markdown - `dataset_config` : /content/LoRA/config/dataset_config.toml\n",
        "\n",
        "#@markdown Generated sample can be seen here: /content/LoRA/output/sample\n",
        "\n",
        "#@markdown You can import config from another session if you want.\n",
        "sample_prompt = \"/content/LoRA/config/sample_prompt.txt\" #@param {type:'string'}\n",
        "config_file = \"/content/LoRA/config/config_file.toml\" #@param {type:'string'}\n",
        "dataset_config = \"/content/LoRA/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_network.py \\\n",
        "  --sample_prompts={sample_prompt} \\\n",
        "  --dataset_config={dataset_config} \\\n",
        "  --config_file={config_file}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "reMcN0bM_o53"
      },
      "source": [
        "# VI. Testing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "TLSQslfFcQde"
      },
      "outputs": [],
      "source": [
        "# @title ## 6.1. Visualize loss graph (Optional)\n",
        "training_logs_path = \"/content/LoRA/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": "Rt7CKCog_4tm"
      },
      "outputs": [],
      "source": [
        "# @title ## 6.2. Interrogating LoRA Weights\n",
        "# @markdown Now you can check if your LoRA trained properly.\n",
        "\n",
        "# @markdown  If you used `clip_skip = 2` during training, the values of `lora_te_text_model_encoder_layers_11_*` will be `0.0`, this is normal. These layers are not trained at this value of `Clip Skip`.\n",
        "network_weight = \"\"  # @param {'type':'string'}\n",
        "no_verbose = True  # @param {type:\"boolean\"}\n",
        "\n",
        "import os\n",
        "import torch\n",
        "import json\n",
        "from safetensors.torch import load_file\n",
        "from safetensors.torch import safe_open\n",
        "import library.model_util as model_util\n",
        "\n",
        "print(\"Loading LoRA weight:\", network_weight)\n",
        "\n",
        "\n",
        "def main(file, verbose: bool):\n",
        "    if not verbose:\n",
        "        sd = (\n",
        "            load_file(file)\n",
        "            if os.path.splitext(file)[1] == \".safetensors\"\n",
        "            else torch.load(file, map_location=\"cuda\")\n",
        "        )\n",
        "        values = []\n",
        "\n",
        "        keys = list(sd.keys())\n",
        "        for key in keys:\n",
        "            if \"lora_up\" in key or \"lora_down\" in key:\n",
        "                values.append((key, sd[key]))\n",
        "        print(f\"number of LoRA modules: {len(values)}\")\n",
        "\n",
        "        for key, value in values:\n",
        "            value = value.to(torch.float32)\n",
        "            print(f\"{key},{torch.mean(torch.abs(value))},{torch.min(torch.abs(value))}\")\n",
        "\n",
        "    if model_util.is_safetensors(file):\n",
        "        with safe_open(file, framework=\"pt\") as f:\n",
        "            metadata = f.metadata()\n",
        "        if metadata is not None:\n",
        "            print(f\"\\nLoad metadata for: {file}\")\n",
        "            print(json.dumps(metadata, indent=4))\n",
        "    else:\n",
        "        print(\"No metadata saved, your model is not in safetensors format\")\n",
        "\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    main(network_weight, no_verbose)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "FKBrTDPrcNjP"
      },
      "outputs": [],
      "source": [
        "# @title ## 6.3. Inference\n",
        "%store -r\n",
        "\n",
        "# @markdown ### LoRA Config\n",
        "# @markdown Currently, `LoHa` and `LoCon_Lycoris` is not supported, please run `Portable Web UI` instead\n",
        "network_weight = \"\"  # @param {'type':'string'}\n",
        "network_mul = 0.7  # @param {type:\"slider\", min:-1, max:2, step:0.05}\n",
        "network_module = \"networks.lora\"\n",
        "network_args = \"\"\n",
        "\n",
        "# @markdown ### <br> General Config\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 = (\n",
        "    f\"{instance_prompt}, {prompt} --n {negative}\"\n",
        "    if instance_prompt\n",
        "    else f\"{prompt} --n {negative}\"\n",
        ")\n",
        "\n",
        "config = {\n",
        "    \"v2\": v2,\n",
        "    \"v_parameterization\": v_parameterization,\n",
        "    \"network_module\": network_module,\n",
        "    \"network_weight\": network_weight,\n",
        "    \"network_mul\": float(network_mul),\n",
        "    \"network_args\": eval(network_args) if network_args else None,\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": "dglQfu3A0oaj"
      },
      "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\", \".pt\", \".ckpt\")):\n",
        "      !ln \"{file_path}\" {webui_lora_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",
        "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",
        "  --lora-dir {webui_lora_dir} \\\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": "nyIl9BhNXKUq"
      },
      "source": [
        "# VII. Deployment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "QTXsM170GUpk"
      },
      "outputs": [],
      "source": [
        "# @title ## 7.1. Upload Config\n",
        "from huggingface_hub import login\n",
        "from huggingface_hub import HfApi\n",
        "from huggingface_hub.utils import validate_repo_id, HfHubHTTPError\n",
        "\n",
        "# @markdown Login to Huggingface Hub\n",
        "# @markdown > Get **your** huggingface `WRITE` token [here](https://huggingface.co/settings/tokens)\n",
        "write_token = \"\"  # @param {type:\"string\"}\n",
        "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": [
        "## 7.2. Upload with Huggingface Hub"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "8HpOZOglBJTz"
      },
      "outputs": [],
      "source": [
        "# @title ### 7.2.1. Upload LoRA\n",
        "from huggingface_hub import HfApi\n",
        "from pathlib import Path\n",
        "\n",
        "%store -r\n",
        "\n",
        "api = HfApi()\n",
        "\n",
        "# @markdown This will be uploaded to model repo\n",
        "model_path = \"/content/LoRA/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/LoRA/config\"  # @param {type :\"string\"}\n",
        "# @markdown Other Information\n",
        "\n",
        "commit_message = \"\"  # @param {type :\"string\"}\n",
        "\n",
        "if not commit_message:\n",
        "    commit_message = \"feat: upload \" + project_name + \" lora model\"\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",
        "        trained_model = f\"{project_name}_config\"\n",
        "\n",
        "    if is_folder:\n",
        "        print(f\"Uploading {trained_model} to https://huggingface.co/\" + model_repo)\n",
        "        print(f\"Please wait...\")\n",
        "\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 ### 7.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/LoRA/train_data\"  # @param {type :\"string\"}\n",
        "# @markdown `Nerd stuff, only if you want to save training logs`\n",
        "logs_path = \"/content/LoRA/logs\"  # @param {type :\"string\"}\n",
        "\n",
        "if project_name:\n",
        "    tmp_dataset = \"/content/LoRA/\" + project_name + \"_dataset\"\n",
        "else:\n",
        "    tmp_dataset = \"/content/LoRA/tmp_dataset\"\n",
        "\n",
        "tmp_train_data = tmp_dataset + \"/train_data\"\n",
        "dataset_zip = tmp_dataset + \".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",
        "os.makedirs(tmp_dataset, exist_ok=True)\n",
        "os.makedirs(tmp_train_data, exist_ok=True)\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):\n",
        "    zipfiles = tmp + \".zip\"\n",
        "    with zipfile.ZipFile(zipfiles, \"w\") as zip:\n",
        "        for tmp, dirs, files in os.walk(tmp):\n",
        "            for file in files:\n",
        "                zip.write(os.path.join(tmp, 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",
        "    if logs_path:\n",
        "        upload_dataset(logs_path, False)\n",
        "\n",
        "\n",
        "upload()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CKZpg4keWS5c"
      },
      "source": [
        "## 7.3. Upload with GIT (Alternative)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "6nBlrOrytO9F"
      },
      "outputs": [],
      "source": [
        "# @title ### 7.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 ### 7.3.2. Commit using Git\n",
        "import os\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 + \" lora 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
}