kohya-LoRA-trainer-XL.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/kohya-LoRA-trainer-XL.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "slgjeYgd6pWp"
      },
      "source": [
        "[![visitor][visitor-badge]][visitor-stats]\n",
        "[![ko-fi][ko-fi-badge]][ko-fi-link]\n",
        "\n",
        "# **Kohya LoRA Trainer XL**\n",
        "A Colab Notebook For SDXL LoRA Training (Fine-tuning Method)\n",
        "\n",
        "[visitor-badge]: https://api.visitorbadge.io/api/visitors?path=Kohya%20LoRA%20Trainer%20XL&label=Visitors&labelColor=%2334495E&countColor=%231ABC9C&style=flat&labelStyle=none\n",
        "[visitor-stats]: https://visitorbadge.io/status?path=Kohya%20LoRA%20Trainer%20XL\n",
        "[ko-fi-badge]: https://img.shields.io/badge/Support%20me%20on%20Ko--fi-F16061?logo=ko-fi&logoColor=white&style=flat\n",
        "[ko-fi-link]: https://ko-fi.com/linaqruf"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_MxF9feWshAp"
      },
      "source": [
        "| Notebook Name | Description | Link |\n",
        "| --- | --- | --- |\n",
        "| [Kohya LoRA Trainer XL](https://github.com/Linaqruf/kohya-trainer/blob/main/kohya-LoRA-trainer-XL.ipynb) | LoRA 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-LoRA-trainer-XL.ipynb) |\n",
        "| [Kohya Trainer XL](https://github.com/Linaqruf/kohya-trainer/blob/main/kohya-trainer-XL.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-XL.ipynb) |\n"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "<hr>\n",
        "<h4><font color=\"#4a90e2\"><b>NEWS:</b></font> <i>Colab's free-tier users can now train SDXL LoRA using the diffusers format instead of checkpoint as a pretrained model.</i></h4>\n",
        "<hr>"
      ],
      "metadata": {
        "id": "iIbwGFkJ0nTx"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tTVqCAgSmie4"
      },
      "source": [
        "# **I. Prepare Environment**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "_u3q60di584x"
      },
      "outputs": [],
      "source": [
        "# @title ## **1.1. Install Kohya Trainer**\n",
        "import os\n",
        "import zipfile\n",
        "import shutil\n",
        "import time\n",
        "import requests\n",
        "import torch\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",
        "drive_dir         = os.path.join(root_dir, \"drive/MyDrive\")\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",
        "lora_dir          = os.path.join(root_dir, \"network_weight\")\n",
        "repositories_dir  = os.path.join(root_dir, \"repositories\")\n",
        "config_dir        = os.path.join(training_dir, \"config\")\n",
        "tools_dir         = os.path.join(repo_dir, \"tools\")\n",
        "finetune_dir      = os.path.join(repo_dir, \"finetune\")\n",
        "accelerate_config = os.path.join(repo_dir, \"accelerate_config/config.yaml\")\n",
        "\n",
        "for store in [\"root_dir\", \"repo_dir\", \"training_dir\", \"pretrained_model\", \"vae_dir\", \"repositories_dir\", \"accelerate_config\", \"tools_dir\", \"finetune_dir\", \"config_dir\"]:\n",
        "    with capture.capture_output() as cap:\n",
        "        %store {store}\n",
        "        del cap\n",
        "\n",
        "repo_dict = {\n",
        "    \"qaneel/kohya-trainer (forked repo, stable, optimized for colab use)\" : \"https://github.com/qaneel/kohya-trainer\",\n",
        "    \"kohya-ss/sd-scripts (original repo, latest update)\"                    : \"https://github.com/kohya-ss/sd-scripts\",\n",
        "}\n",
        "\n",
        "repository        = \"qaneel/kohya-trainer (forked repo, stable, optimized for colab use)\" #@param [\"qaneel/kohya-trainer (forked repo, stable, optimized for colab use)\", \"kohya-ss/sd-scripts (original repo, latest update)\"] {allow-input: true}\n",
        "repo_url          = repo_dict[repository]\n",
        "branch            = \"main\"  # @param {type: \"string\"}\n",
        "output_to_drive   = True  # @param {type: \"boolean\"}\n",
        "\n",
        "def clone_repo(url, dir, branch):\n",
        "    if not os.path.exists(dir):\n",
        "       !git clone -b {branch} {url} {dir}\n",
        "\n",
        "def mount_drive(dir):\n",
        "    output_dir      = os.path.join(training_dir, \"output\")\n",
        "\n",
        "    if output_to_drive:\n",
        "        if not os.path.exists(drive_dir):\n",
        "            drive.mount(os.path.dirname(drive_dir))\n",
        "        output_dir  = os.path.join(drive_dir, \"kohya-trainer/output\")\n",
        "\n",
        "    return output_dir\n",
        "\n",
        "def setup_directories():\n",
        "    global output_dir\n",
        "\n",
        "    output_dir      = mount_drive(drive_dir)\n",
        "\n",
        "    for dir in [training_dir, config_dir, pretrained_model, vae_dir, repositories_dir, output_dir]:\n",
        "        os.makedirs(dir, exist_ok=True)\n",
        "\n",
        "def pastebin_reader(id):\n",
        "    if \"pastebin.com\" in id:\n",
        "        url = id\n",
        "        if 'raw' not in url:\n",
        "                url = url.replace('pastebin.com', 'pastebin.com/raw')\n",
        "    else:\n",
        "        url = \"https://pastebin.com/raw/\" + id\n",
        "    response = requests.get(url)\n",
        "    response.raise_for_status()\n",
        "    lines = response.text.split('\\n')\n",
        "    return lines\n",
        "\n",
        "def install_repository():\n",
        "    global infinite_image_browser_dir, voldy, discordia_archivum_dir\n",
        "\n",
        "    _, voldy = pastebin_reader(\"kq6ZmHFU\")[:2]\n",
        "\n",
        "    infinite_image_browser_url  = f\"https://github.com/zanllp/{voldy}-infinite-image-browsing.git\"\n",
        "    infinite_image_browser_dir  = os.path.join(repositories_dir, f\"infinite-image-browsing\")\n",
        "    infinite_image_browser_deps = os.path.join(infinite_image_browser_dir, \"requirements.txt\")\n",
        "\n",
        "    discordia_archivum_url = \"https://github.com/Linaqruf/discordia-archivum\"\n",
        "    discordia_archivum_dir = os.path.join(repositories_dir, \"discordia-archivum\")\n",
        "    discordia_archivum_deps = os.path.join(discordia_archivum_dir, \"requirements.txt\")\n",
        "\n",
        "    clone_repo(infinite_image_browser_url, infinite_image_browser_dir, \"main\")\n",
        "    clone_repo(discordia_archivum_url, discordia_archivum_dir, \"main\")\n",
        "\n",
        "    !pip install -q --upgrade -r {infinite_image_browser_deps}\n",
        "    !pip install python-dotenv\n",
        "    !pip install -q --upgrade -r {discordia_archivum_deps}\n",
        "\n",
        "def install_dependencies():\n",
        "    requirements_file = os.path.join(repo_dir, \"requirements.txt\")\n",
        "    model_util        = os.path.join(repo_dir, \"library/model_util.py\")\n",
        "    gpu_info          = getoutput('nvidia-smi')\n",
        "    t4_xformers_wheel = \"https://github.com/Linaqruf/colab-xformers/releases/download/0.0.20/xformers-0.0.20+1d635e1.d20230519-cp310-cp310-linux_x86_64.whl\"\n",
        "\n",
        "    !apt install aria2 lz4\n",
        "    !wget https://github.com/camenduru/gperftools/releases/download/v1.0/libtcmalloc_minimal.so.4 -O /content/libtcmalloc_minimal.so.4\n",
        "    !pip install -q --upgrade -r {requirements_file}\n",
        "\n",
        "    if '2.0.1+cu118' in torch.__version__:\n",
        "        if 'T4' in gpu_info:\n",
        "            !pip install -q {t4_xformers_wheel}\n",
        "        else:\n",
        "            !pip install -q xformers==0.0.20\n",
        "    else:\n",
        "        !pip install -q torch==2.0.0+cu118 torchvision==0.15.1+cu118 torchaudio==2.0.1+cu118 torchtext==0.15.1 torchdata==0.6.0 --extra-index-url https://download.pytorch.org/whl/cu118 -U\n",
        "        !pip install -q xformers==0.0.19 triton==2.0.0 -U\n",
        "\n",
        "    from accelerate.utils import write_basic_config\n",
        "\n",
        "    if not os.path.exists(accelerate_config):\n",
        "        write_basic_config(save_location=accelerate_config)\n",
        "\n",
        "def prepare_environment():\n",
        "    os.environ[\"LD_PRELOAD\"] = \"/content/libtcmalloc_minimal.so.4\"\n",
        "    os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n",
        "    os.environ[\"SAFETENSORS_FAST_GPU\"] = \"1\"\n",
        "    os.environ[\"PYTHONWARNINGS\"] = \"ignore\"\n",
        "\n",
        "def main():\n",
        "    os.chdir(root_dir)\n",
        "    clone_repo(repo_url, repo_dir, branch)\n",
        "    os.chdir(repo_dir)\n",
        "    setup_directories()\n",
        "    install_repository()\n",
        "    install_dependencies()\n",
        "    prepare_environment()\n",
        "\n",
        "main()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wrYGu-WxFbsq",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "# @title ## **1.2. Download SDXL**\n",
        "import os\n",
        "import re\n",
        "import json\n",
        "import glob\n",
        "import gdown\n",
        "import requests\n",
        "import subprocess\n",
        "from IPython.utils import capture\n",
        "from urllib.parse import urlparse, unquote\n",
        "from pathlib import Path\n",
        "from huggingface_hub import HfFileSystem\n",
        "from huggingface_hub.utils import validate_repo_id, HfHubHTTPError\n",
        "\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "# @markdown Place your Huggingface token [here](https://huggingface.co/settings/tokens) to download gated models.\n",
        "\n",
        "HUGGINGFACE_TOKEN     = \"\" #@param {type: \"string\"}\n",
        "LOAD_DIFFUSERS_MODEL  = True #@param {type: \"boolean\"}\n",
        "SDXL_MODEL_URL        = \"Linaqruf/animagine-xl\" # @param [\"gsdf/CounterfeitXL\", \"Linaqruf/animagine-xl\", \"stabilityai/stable-diffusion-xl-base-1.0\", \"PASTE MODEL URL OR GDRIVE PATH HERE\"] {allow-input: true}\n",
        "SDXL_VAE_URL          = \"Original VAE\" # @param [\"None\", \"Original VAE\", \"FP16 VAE\", \"PASTE VAE URL OR GDRIVE PATH HERE\"] {allow-input: true}\n",
        "\n",
        "MODEL_URLS = {\n",
        "    \"gsdf/CounterfeitXL\"        : \"https://huggingface.co/gsdf/CounterfeitXL/resolve/main/CounterfeitXL_%CE%B2.safetensors\",\n",
        "    \"Linaqruf/animagine-xl\"   : \"https://huggingface.co/Linaqruf/animagine-xl/resolve/main/animagine-xl.safetensors\",\n",
        "    \"stabilityai/stable-diffusion-xl-base-1.0\" : \"https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors\",\n",
        "}\n",
        "VAE_URLS = {\n",
        "    \"None\"                    : \"\",\n",
        "    \"Original VAE\"           : \"https://huggingface.co/stabilityai/sdxl-vae/resolve/main/sdxl_vae.safetensors\",\n",
        "    \"FP16 VAE\"           : \"https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/resolve/main/sdxl_vae.safetensors\"\n",
        "}\n",
        "\n",
        "SDXL_MODEL_URL = MODEL_URLS.get(SDXL_MODEL_URL, SDXL_MODEL_URL)\n",
        "SDXL_VAE_URL = VAE_URLS.get(SDXL_VAE_URL, SDXL_VAE_URL)\n",
        "\n",
        "def get_filename(url):\n",
        "    if any(url.endswith(ext) for ext in [\".ckpt\", \".safetensors\", \".pt\", \".pth\"]):\n",
        "        return os.path.basename(url)\n",
        "\n",
        "    response = requests.get(url, stream=True)\n",
        "    response.raise_for_status()\n",
        "\n",
        "    if 'content-disposition' in response.headers:\n",
        "        filename = re.findall('filename=\"?([^\"]+)\"?', response.headers['content-disposition'])[0]\n",
        "    else:\n",
        "        filename = unquote(os.path.basename(urlparse(url).path))\n",
        "\n",
        "    return filename\n",
        "\n",
        "def aria2_download(dir, filename, url):\n",
        "    user_header = f\"Authorization: Bearer {HUGGINGFACE_TOKEN}\"\n",
        "    aria2_args = [\n",
        "        \"aria2c\",\n",
        "        \"--console-log-level=error\",\n",
        "        \"--summary-interval=10\",\n",
        "        f\"--header={user_header}\" if \"huggingface.co\" in url else \"\",\n",
        "        \"--continue=true\",\n",
        "        \"--max-connection-per-server=16\",\n",
        "        \"--min-split-size=1M\",\n",
        "        \"--split=16\",\n",
        "        f\"--dir={dir}\",\n",
        "        f\"--out={filename}\",\n",
        "        url\n",
        "    ]\n",
        "    subprocess.run(aria2_args)\n",
        "\n",
        "def download(url, dst):\n",
        "    print(f\"Starting downloading from {url}\")\n",
        "    filename = get_filename(url)\n",
        "    filepath = os.path.join(dst, filename)\n",
        "\n",
        "    if \"drive.google.com\" in url:\n",
        "        gdown.download(url, filepath, quiet=False)\n",
        "    else:\n",
        "        if \"huggingface.co\" in url and \"/blob/\" in url:\n",
        "            url = url.replace(\"/blob/\", \"/resolve/\")\n",
        "        aria2_download(dst, filename, url)\n",
        "\n",
        "    print(f\"Download finished: {filepath}\")\n",
        "    return filepath\n",
        "\n",
        "def all_folders_present(base_model_url, sub_folders):\n",
        "    fs = HfFileSystem()\n",
        "    existing_folders = set(fs.ls(base_model_url, detail=False))\n",
        "\n",
        "    for folder in sub_folders:\n",
        "        full_folder_path = f\"{base_model_url}/{folder}\"\n",
        "        if full_folder_path not in existing_folders:\n",
        "            return False\n",
        "    return True\n",
        "\n",
        "def get_total_ram_gb():\n",
        "    with open('/proc/meminfo', 'r') as f:\n",
        "        for line in f.readlines():\n",
        "            if \"MemTotal\" in line:\n",
        "                return int(line.split()[1]) / (1024**2)  # Convert to GB\n",
        "\n",
        "def get_gpu_name():\n",
        "    try:\n",
        "        return subprocess.check_output(\"nvidia-smi --query-gpu=name --format=csv,noheader,nounits\", shell=True).decode('ascii').strip()\n",
        "    except:\n",
        "        return None\n",
        "\n",
        "def main():\n",
        "    global model_path, vae_path, LOAD_DIFFUSERS_MODEL\n",
        "\n",
        "    model_path, vae_path = None, None\n",
        "\n",
        "    required_sub_folders = [\n",
        "        'scheduler',\n",
        "        'text_encoder',\n",
        "        'text_encoder_2',\n",
        "        'tokenizer',\n",
        "        'tokenizer_2',\n",
        "        'unet',\n",
        "        'vae',\n",
        "    ]\n",
        "\n",
        "    download_targets = {\n",
        "        \"model\": (SDXL_MODEL_URL, pretrained_model),\n",
        "        \"vae\": (SDXL_VAE_URL, vae_dir),\n",
        "    }\n",
        "\n",
        "    total_ram = get_total_ram_gb()\n",
        "    gpu_name = get_gpu_name()\n",
        "\n",
        "    # Check hardware constraints\n",
        "    if total_ram < 13 and gpu_name in [\"Tesla T4\", \"Tesla V100\"]:\n",
        "        print(\"Attempt to load diffusers model instead due to hardware constraints.\")\n",
        "        if not LOAD_DIFFUSERS_MODEL:\n",
        "            LOAD_DIFFUSERS_MODEL = True\n",
        "\n",
        "    for target, (url, dst) in download_targets.items():\n",
        "        if url and not url.startswith(f\"PASTE {target.upper()} URL OR GDRIVE PATH HERE\"):\n",
        "            if target == \"model\" and LOAD_DIFFUSERS_MODEL:\n",
        "                # Code for checking and handling diffusers model\n",
        "                if 'huggingface.co' in url:\n",
        "                    match = re.search(r'huggingface\\.co/([^/]+)/([^/]+)', SDXL_MODEL_URL)\n",
        "                    if match:\n",
        "                        username = match.group(1)\n",
        "                        model_name = match.group(2)\n",
        "                        url = f\"{username}/{model_name}\"\n",
        "                if all_folders_present(url, required_sub_folders):\n",
        "                    print(f\"Diffusers model is loaded : {url}\")\n",
        "                    model_path = url\n",
        "                else:\n",
        "                    print(\"Repository doesn't exist or no diffusers model detected.\")\n",
        "                    filepath = download(url, dst)  # Continue with the regular download\n",
        "                    model_path = filepath\n",
        "            else:\n",
        "                filepath = download(url, dst)\n",
        "\n",
        "                if target == \"model\":\n",
        "                    model_path = filepath\n",
        "                elif target == \"vae\":\n",
        "                    vae_path = filepath\n",
        "\n",
        "            print()\n",
        "\n",
        "    if model_path:\n",
        "        print(f\"Selected model: {model_path}\")\n",
        "\n",
        "    if vae_path:\n",
        "        print(f\"Selected VAE: {vae_path}\")\n",
        "\n",
        "main()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "kh7CeDqK4l3Y"
      },
      "outputs": [],
      "source": [
        "# @title ## **1.3. Directory Config**\n",
        "# @markdown Specify the location of your training data in the following cell. A folder with the same name as your input will be created.\n",
        "import os\n",
        "\n",
        "%store -r\n",
        "\n",
        "train_data_dir = \"/content/LoRA/train_data\"  # @param {'type' : 'string'}\n",
        "%store train_data_dir\n",
        "\n",
        "os.makedirs(train_data_dir, exist_ok=True)\n",
        "print(f\"Your train data directory : {train_data_dir}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "qqUYtRn0RPoK"
      },
      "outputs": [],
      "source": [
        "# @title ## **1.4. Image Browser**\n",
        "import os\n",
        "import json\n",
        "import random\n",
        "import portpicker\n",
        "from IPython.utils import capture\n",
        "from IPython.display import clear_output\n",
        "from threading import Thread\n",
        "from imjoy_elfinder.app import main\n",
        "from google.colab.output import serve_kernel_port_as_iframe, serve_kernel_port_as_window\n",
        "\n",
        "%store -r\n",
        "\n",
        "# @markdown This cell allows you to view and manage your images in real-time. You can use it to:\n",
        "# @markdown - Prepare your dataset before training\n",
        "# @markdown - Monitor the sample outputs during training.\n",
        "\n",
        "root_dir      = \"/content\"\n",
        "browser_type  = \"infinite-image-browsing\" #@param [\"imjoy-elfinder\", \"infinite-image-browsing\"]\n",
        "window_height = 550 #@param {type:\"slider\", min:0, max:1000, step:1}\n",
        "\n",
        "main_app          = os.path.join(infinite_image_browser_dir, \"app.py\")\n",
        "config_file       = os.path.join(infinite_image_browser_dir, \"config.json\")\n",
        "port              = portpicker.pick_unused_port()\n",
        "\n",
        "config = {\n",
        "    \"outdir_txt2img_samples\": train_data_dir,\n",
        "}\n",
        "\n",
        "def write_file(filename, config):\n",
        "    with open(filename, 'w',) as f:\n",
        "        json.dump(config, f, indent=4)\n",
        "\n",
        "def run_app():\n",
        "    !python {main_app} --port={port} --sd_webui_config={config_file} > /dev/null 2>&1\n",
        "\n",
        "def launch():\n",
        "    os.chdir(root_dir)\n",
        "\n",
        "    thread = Thread(target=main, args=[[f\"--root-dir={root_dir}\",\n",
        "                                        f\"--port={port}\",\n",
        "                                        f\"--thumbnail\"]])\n",
        "\n",
        "    if browser_type == \"infinite-image-browsing\":\n",
        "        os.chdir(train_data_dir)\n",
        "        write_file(config_file, config)\n",
        "\n",
        "        thread = Thread(target=run_app)\n",
        "\n",
        "    thread.start()\n",
        "\n",
        "    serve_kernel_port_as_iframe(port, width='100%', height=window_height, cache_in_notebook=False)\n",
        "\n",
        "    clear_output(wait=True)\n",
        "\n",
        "launch()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "En9UUwGNMRMM"
      },
      "source": [
        "# **II. Data Gathering**\n",
        "\n",
        "You have three options for collecting your dataset:\n",
        "\n",
        "1. Upload it to Colab's local files.\n",
        "2. Use the `Simple Booru Scraper` to download images in bulk from Danbooru.\n",
        "3. Locate your dataset in Google Drive.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "t17ZfiMB8GWZ"
      },
      "outputs": [],
      "source": [
        "# @title ## **2.1. Unzip Dataset**\n",
        "\n",
        "import os\n",
        "import subprocess\n",
        "import re\n",
        "from urllib.parse import unquote\n",
        "import requests\n",
        "from zipfile import ZipFile\n",
        "\n",
        "# @title ## Unzip Dataset\n",
        "# @markdown If your dataset is in a `zip` file and has been uploaded to a location, use this section to extract it.\n",
        "# @markdown The dataset will be downloaded and automatically extracted to `train_data_dir` if `unzip_to` is empty.\n",
        "\n",
        "zipfile_url = \"\"  # @param {type:\"string\"}\n",
        "unzip_to = \"\"  # @param {type:\"string\"}\n",
        "hf_token = \"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\"\n",
        "\n",
        "use_aria2c = True # @param {type:\"boolean\"}\n",
        "preserve_folders = True # @param {type:\"boolean\"}\n",
        "remove_after_unzipping = False # @param {type:\"boolean\"}\n",
        "\n",
        "if \"huggingface.co\" in zipfile_url and \"blob\" in zipfile_url:\n",
        "    zipfile_url = zipfile_url.replace(\"blob\", \"resolve\")\n",
        "\n",
        "if not unzip_to:\n",
        "    unzip_to = train_data_dir\n",
        "\n",
        "def get_filename_from_url(url):\n",
        "    if \"huggingface.co\" or \"/content/\" in url:\n",
        "        return os.path.basename(url)\n",
        "\n",
        "    response = requests.head(url, allow_redirects=True)\n",
        "    cd = response.headers.get('content-disposition')\n",
        "    if cd:\n",
        "        fname = re.findall('filename=(.+)', cd)\n",
        "        if len(fname) == 0:\n",
        "            return \"zipfile.zip\"\n",
        "        return unquote(fname[0])\n",
        "\n",
        "    return \"zipfile.zip\"\n",
        "\n",
        "def download_with_requests(url, output_path):\n",
        "    print(f\"Downloading {url} with requests...\")\n",
        "    response = requests.get(url, stream=True)\n",
        "    with open(output_path, 'wb') as file:\n",
        "        for chunk in response.iter_content(chunk_size=8192):\n",
        "            file.write(chunk)\n",
        "    print(f\"Downloaded to {output_path}\")\n",
        "    return output_path\n",
        "\n",
        "def download_with_aria2c(url, output_path):\n",
        "    print(f\"Downloading {url} with aria2c...\")\n",
        "    aria_args = {\n",
        "        'console-log-level': 'error',\n",
        "        'summary-interval': '10',\n",
        "        'continue': 'true',\n",
        "        'max-connection-per-server': '16',\n",
        "        'min-split-size': '1M',\n",
        "        'split': '16',\n",
        "        'dir': os.path.dirname(output_path),\n",
        "        'out': os.path.basename(output_path),\n",
        "    }\n",
        "\n",
        "    if \"huggingface.co\" in url:\n",
        "        aria_args['header'] = f\"Authorization: Bearer {hf_token}\"\n",
        "\n",
        "    cmd = ['aria2c'] + [f'--{k}={v}' for k, v in aria_args.items()] + [url]\n",
        "    subprocess.run(cmd)\n",
        "    print(f\"Downloaded to {output_path}\")\n",
        "    return output_path\n",
        "\n",
        "def move_files(train_dir):\n",
        "    for filename in os.listdir(train_dir):\n",
        "        file_path = os.path.join(train_dir, filename)\n",
        "        if filename.startswith(\"meta_\") and filename.endswith(\".json\"):\n",
        "            if not os.path.exists(file_path):\n",
        "                shutil.move(file_path, training_dir)\n",
        "            else:\n",
        "                os.remove(file_path)\n",
        "\n",
        "def remove_empty_dirs(path):\n",
        "    for dirpath, dirnames, files in os.walk(path, topdown=False):  # start from leaf folders\n",
        "        for dirname in dirnames:\n",
        "            full_dir_path = os.path.join(dirpath, dirname)\n",
        "            if not os.listdir(full_dir_path):  # Check if directory is empty\n",
        "                os.rmdir(full_dir_path)\n",
        "                print(f\"Removed empty directory: {full_dir_path}\")\n",
        "\n",
        "def extract_dataset(zip_file, output_path):\n",
        "    with ZipFile(zip_file, 'r') as zip_ref:\n",
        "        print(f\"Extracting {zip_file} to {output_path}...\")\n",
        "\n",
        "        if not preserve_folders:  # If we do not want to preserve folder structure\n",
        "            for member in zip_ref.namelist():\n",
        "                # Extract only the file name, discard directory structure\n",
        "                filename = os.path.basename(member)\n",
        "                if filename:  # Check if file name is not empty (this skips directories)\n",
        "                    zip_ref.extract(member, output_path)\n",
        "                    source_path = os.path.join(output_path, member)\n",
        "                    target_path = os.path.join(output_path, filename)\n",
        "                    os.rename(source_path, target_path)\n",
        "\n",
        "            remove_empty_dirs(output_path)\n",
        "\n",
        "        else:\n",
        "            zip_ref.extractall(output_path)\n",
        "\n",
        "        print(\"Extraction completed!\")\n",
        "\n",
        "def download_dataset(url, output_path):\n",
        "    if url.startswith(\"/content\"):\n",
        "        print(f\"Using file at {url}\")\n",
        "        return url\n",
        "\n",
        "    elif \"drive.google.com\" in url:\n",
        "        print(\"Downloading from Google Drive...\")\n",
        "        cmd = ['gdown', '--id', url.split('/')[-2], '-O', output_path]\n",
        "        subprocess.run(cmd)\n",
        "        return output_path\n",
        "\n",
        "    elif use_aria2c:\n",
        "        return download_with_aria2c(url, output_path)\n",
        "\n",
        "    else:\n",
        "        return download_with_requests(url, output_path)\n",
        "\n",
        "def main():\n",
        "    zipfile_name = get_filename_from_url(zipfile_url)\n",
        "    output_path = os.path.join(root_dir, zipfile_name)\n",
        "\n",
        "    zip_file = download_dataset(zipfile_url, output_path)\n",
        "\n",
        "    extract_dataset(zip_file, unzip_to)\n",
        "\n",
        "    move_files(unzip_to)\n",
        "\n",
        "    if remove_after_unzipping and \"/content/drive\" not in zip_file:\n",
        "        os.remove(zip_file)\n",
        "        print(f\"Removed {zip_file}\")\n",
        "\n",
        "main()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "A0t1dfnU5Xkq"
      },
      "outputs": [],
      "source": [
        "#@title ## **2.2. Imageboard Scraper**\n",
        "import os\n",
        "import html\n",
        "from IPython.utils import capture\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "#@markdown Use `gallery-dl` to scrape images from an imageboard site. To specify `prompt(s)`, separate them with commas (e.g., `hito_komoru, touhou`).\n",
        "booru = \"Danbooru\" #@param [\"Danbooru\", \"Gelbooru\", \"Safebooru\"]\n",
        "prompt = \"\" #@param {type: \"string\"}\n",
        "\n",
        "#@markdown Alternatively, you can provide a `custom_url` instead of using a predefined site.\n",
        "custom_url = \"\" #@param {type: \"string\"}\n",
        "\n",
        "#@markdown Use the `sub_folder` option to organize the downloaded images into separate folders based on their concept or category.\n",
        "sub_folder = \"\" #@param {type: \"string\"}\n",
        "\n",
        "user_agent = \"gdl/1.24.5\"\n",
        "\n",
        "#@markdown You can limit the number of images to download by using the `--range` option followed by the desired range (e.g., `1-200`).\n",
        "range = \"\" #@param {type: \"string\"}\n",
        "\n",
        "write_tags = False #@param {type: \"boolean\"}\n",
        "\n",
        "additional_arguments = \"--filename /O --no-part\"\n",
        "\n",
        "tags = prompt.split(',')\n",
        "tags = '+'.join(tags)\n",
        "\n",
        "replacement_dict = {\" \": \"\", \"(\": \"%28\", \")\": \"%29\", \":\": \"%3a\"}\n",
        "tags = ''.join(replacement_dict.get(c, c) for c in tags)\n",
        "\n",
        "if sub_folder == \"\":\n",
        "    image_dir = train_data_dir\n",
        "elif sub_folder.startswith(\"/content\"):\n",
        "    image_dir = sub_folder\n",
        "else:\n",
        "    image_dir = os.path.join(train_data_dir, sub_folder)\n",
        "    os.makedirs(image_dir, exist_ok=True)\n",
        "\n",
        "if booru == \"Danbooru\":\n",
        "    url = \"https://danbooru.donmai.us/posts?tags={}\".format(tags)\n",
        "elif booru == \"Gelbooru\":\n",
        "    url = \"https://gelbooru.com/index.php?page=post&s=list&tags={}\".format(tags)\n",
        "else:\n",
        "    url = \"https://safebooru.org/index.php?page=post&s=list&tags={}\".format(tags)\n",
        "\n",
        "valid_url = custom_url if custom_url else url\n",
        "\n",
        "def scrape(config):\n",
        "    args = \"\"\n",
        "    for k, v in config.items():\n",
        "        if k.startswith(\"_\"):\n",
        "            args += f'\"{v}\" '\n",
        "        elif isinstance(v, str):\n",
        "            args += f'--{k}=\"{v}\" '\n",
        "        elif isinstance(v, bool) and v:\n",
        "            args += f\"--{k} \"\n",
        "        elif isinstance(v, float) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "        elif isinstance(v, int) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "\n",
        "    return args\n",
        "\n",
        "def pre_process_tags(directory):\n",
        "    for item in os.listdir(directory):\n",
        "        item_path = os.path.join(directory, item)\n",
        "        if os.path.isfile(item_path) and item.endswith(\".txt\"):\n",
        "            old_path = item_path\n",
        "            new_file_name = os.path.splitext(os.path.splitext(item)[0])[0] + \".txt\"\n",
        "            new_path = os.path.join(directory, new_file_name)\n",
        "\n",
        "            os.rename(old_path, new_path)\n",
        "\n",
        "            with open(new_path, \"r\") as f:\n",
        "                contents = f.read()\n",
        "\n",
        "            contents = html.unescape(contents)\n",
        "            contents = contents.replace(\"_\", \" \")\n",
        "            contents = \", \".join(contents.split(\"\\n\"))\n",
        "\n",
        "            with open(new_path, \"w\") as f:\n",
        "                f.write(contents)\n",
        "\n",
        "        elif os.path.isdir(item_path):\n",
        "            pre_process_tags(item_path)\n",
        "\n",
        "get_url_config = {\n",
        "    \"_valid_url\" : valid_url,\n",
        "    \"get-urls\" : True,\n",
        "    \"range\" : range if range else None,\n",
        "    \"user-agent\" : user_agent\n",
        "}\n",
        "\n",
        "scrape_config = {\n",
        "    \"_valid_url\" : valid_url,\n",
        "    \"directory\" : image_dir,\n",
        "    \"write-tags\" : write_tags,\n",
        "    \"range\" : range if range else None,\n",
        "    \"user-agent\" : user_agent\n",
        "}\n",
        "\n",
        "get_url_args = scrape(get_url_config)\n",
        "scrape_args = scrape(scrape_config)\n",
        "scraper_text = os.path.join(root_dir, \"scrape_this.txt\")\n",
        "\n",
        "if write_tags:\n",
        "    !gallery-dl {scrape_args} {additional_arguments}\n",
        "    pre_process_tags(train_data_dir)\n",
        "else:\n",
        "    with capture.capture_output() as cap:\n",
        "        !gallery-dl {get_url_args} {additional_arguments}\n",
        "    with open(scraper_text, \"w\") as f:\n",
        "        f.write(cap.stdout)\n",
        "\n",
        "    os.chdir(image_dir)\n",
        "    !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -i {scraper_text}\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "EVWk3PM2KKB7"
      },
      "outputs": [],
      "source": [
        "#@title ## **2.3. Journey Scraper**\n",
        "import os\n",
        "%store -r\n",
        "\n",
        "# @markdown Enter your Discord token below.\n",
        "token = \"\" #@param {type: \"string\"}\n",
        "channel_id = \"1022054094476673085\" #@param {type: \"string\"}\n",
        "# @markdown Which bot do you want to scrape? This code is optimized to only scrape non-grid images from the Journey bot, so don't worry about cropping.\n",
        "bot = \"niji\" #@param [\"niji\", \"mid\"]\n",
        "# @markdown Set the limit of messages to scrape here. (This does not limit the number of messages to download.)\n",
        "limit = 10000 #@param {type: \"number\"}\n",
        "# @markdown To specify the `include_word` and `undesired_word`, separate them with commas (e.g., hito_komoru, touhou). By default, it scrapes the newest Niji model.\n",
        "include_word = \"girl\" #@param {type:\"string\"}\n",
        "undesired_word = \"--style, --niji 4\" #@param {type:\"string\"}\n",
        "download_attachments = \"single\"\n",
        "\n",
        "def scrape(config):\n",
        "    args = \"\"\n",
        "    for k, v in config.items():\n",
        "        if k.startswith(\"_\"):\n",
        "            args += f'\"{v}\" '\n",
        "        elif isinstance(v, str):\n",
        "            args += f'--{k}=\"{v}\" '\n",
        "        elif isinstance(v, bool) and v:\n",
        "            args += f\"--{k} \"\n",
        "        elif isinstance(v, float) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "        elif isinstance(v, int) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "\n",
        "    return args\n",
        "\n",
        "scrape_config = {\n",
        "    \"token\": token,\n",
        "    \"channel_id\": channel_id,\n",
        "    \"nijijourney\": True if bot == \"niji\" else False,\n",
        "    \"midjourney\": True if bot == \"mid\" else False,\n",
        "    \"limit\": limit if limit else None,\n",
        "    \"prompt\": include_word,\n",
        "    \"single\": True,\n",
        "    \"undesired_word\": undesired_word,\n",
        "    \"download_attachments\": True,\n",
        "    \"output_folder\": train_data_dir,\n",
        "\n",
        "}\n",
        "\n",
        "scrape_args = scrape(scrape_config)\n",
        "\n",
        "os.chdir(discordia_archivum_dir)\n",
        "!python main.py {scrape_args}\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T-0qKyEgTchp"
      },
      "source": [
        "# **III. Data Preprocessing**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Jz2emq6vWnPu"
      },
      "outputs": [],
      "source": [
        "# @title ## **3.1. Data Cleaning**\n",
        "import os\n",
        "import random\n",
        "import concurrent.futures\n",
        "from tqdm import tqdm\n",
        "from PIL import Image\n",
        "\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "test = os.listdir(train_data_dir)\n",
        "#@markdown This section removes unsupported media types such as `.mp4`, `.webm`, and `.gif`, as well as any unnecessary files.\n",
        "#@markdown To convert a transparent dataset with an alpha channel (RGBA) to RGB and give it a white background, set the `convert` parameter to `True`.\n",
        "convert = False  # @param {type:\"boolean\"}\n",
        "#@markdown Alternatively, you can give the background a `random_color` instead of white by checking the corresponding option.\n",
        "random_color = False  # @param {type:\"boolean\"}\n",
        "recursive = False\n",
        "\n",
        "batch_size = 32\n",
        "supported_types = [\n",
        "    \".png\",\n",
        "    \".jpg\",\n",
        "    \".jpeg\",\n",
        "    \".webp\",\n",
        "    \".bmp\",\n",
        "    \".caption\",\n",
        "    \".npz\",\n",
        "    \".txt\",\n",
        "    \".json\",\n",
        "]\n",
        "\n",
        "background_colors = [\n",
        "    (255, 255, 255),\n",
        "    (0, 0, 0),\n",
        "    (255, 0, 0),\n",
        "    (0, 255, 0),\n",
        "    (0, 0, 255),\n",
        "    (255, 255, 0),\n",
        "    (255, 0, 255),\n",
        "    (0, 255, 255),\n",
        "]\n",
        "\n",
        "def clean_directory(directory):\n",
        "    for item in os.listdir(directory):\n",
        "        file_path = os.path.join(directory, item)\n",
        "        if os.path.isfile(file_path):\n",
        "            file_ext = os.path.splitext(item)[1]\n",
        "            if file_ext not in supported_types:\n",
        "                print(f\"Deleting file {item} from {directory}\")\n",
        "                os.remove(file_path)\n",
        "        elif os.path.isdir(file_path) and recursive:\n",
        "            clean_directory(file_path)\n",
        "\n",
        "def process_image(image_path):\n",
        "    img = Image.open(image_path)\n",
        "    img_dir, image_name = os.path.split(image_path)\n",
        "\n",
        "    if img.mode in (\"RGBA\", \"LA\"):\n",
        "        if random_color:\n",
        "            background_color = random.choice(background_colors)\n",
        "        else:\n",
        "            background_color = (255, 255, 255)\n",
        "        bg = Image.new(\"RGB\", img.size, background_color)\n",
        "        bg.paste(img, mask=img.split()[-1])\n",
        "\n",
        "        if image_name.endswith(\".webp\"):\n",
        "            bg = bg.convert(\"RGB\")\n",
        "            new_image_path = os.path.join(img_dir, image_name.replace(\".webp\", \".jpg\"))\n",
        "            bg.save(new_image_path, \"JPEG\")\n",
        "            os.remove(image_path)\n",
        "            print(f\" Converted image: {image_name} to {os.path.basename(new_image_path)}\")\n",
        "        else:\n",
        "            bg.save(image_path, \"PNG\")\n",
        "            print(f\" Converted image: {image_name}\")\n",
        "    else:\n",
        "        if image_name.endswith(\".webp\"):\n",
        "            new_image_path = os.path.join(img_dir, image_name.replace(\".webp\", \".jpg\"))\n",
        "            img.save(new_image_path, \"JPEG\")\n",
        "            os.remove(image_path)\n",
        "            print(f\" Converted image: {image_name} to {os.path.basename(new_image_path)}\")\n",
        "        else:\n",
        "            img.save(image_path, \"PNG\")\n",
        "\n",
        "def find_images(directory):\n",
        "    images = []\n",
        "    for root, _, files in os.walk(directory):\n",
        "        for file in files:\n",
        "            if file.endswith(\".png\") or file.endswith(\".webp\"):\n",
        "                images.append(os.path.join(root, file))\n",
        "    return images\n",
        "\n",
        "clean_directory(train_data_dir)\n",
        "images = find_images(train_data_dir)\n",
        "num_batches = len(images) // batch_size + 1\n",
        "\n",
        "if convert:\n",
        "    with concurrent.futures.ThreadPoolExecutor() as executor:\n",
        "        for i in tqdm(range(num_batches)):\n",
        "            start = i * batch_size\n",
        "            end = start + batch_size\n",
        "            batch = images[start:end]\n",
        "            executor.map(process_image, batch)\n",
        "\n",
        "    print(\"All images have been converted\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qdISafLeyklg"
      },
      "source": [
        "## **3.2. Data Captioning**\n",
        "\n",
        "- For general images, use BLIP captioning.\n",
        "- For anime and manga-style images, use Waifu Diffusion 1.4 Tagger V2."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "nvPyH-G_Qdha"
      },
      "outputs": [],
      "source": [
        "#@title ### **3.2.1. BLIP Captioning**\n",
        "#@markdown BLIP is a pre-training framework for unified vision-language understanding and generation, which achieves state-of-the-art results on a wide range of vision-language tasks. It can be used as a tool for image captioning, for example, `astronaut riding a horse in space`.\n",
        "import os\n",
        "\n",
        "os.chdir(finetune_dir)\n",
        "\n",
        "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",
        "config = {\n",
        "    \"_train_data_dir\"   : train_data_dir,\n",
        "    \"batch_size\"        : 8,\n",
        "    \"beam_search\"       : beam_search,\n",
        "    \"min_length\"        : min_length,\n",
        "    \"max_length\"        : max_length,\n",
        "    \"debug\"             : True,\n",
        "    \"caption_extension\" : \".caption\",\n",
        "    \"max_data_loader_n_workers\" : 2,\n",
        "    \"recursive\"         : True\n",
        "}\n",
        "\n",
        "args = \"\"\n",
        "for k, v in config.items():\n",
        "    if k.startswith(\"_\"):\n",
        "        args += f'\"{v}\" '\n",
        "    elif isinstance(v, str):\n",
        "        args += f'--{k}=\"{v}\" '\n",
        "    elif isinstance(v, bool) and v:\n",
        "        args += f\"--{k} \"\n",
        "    elif isinstance(v, float) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "    elif isinstance(v, int) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "\n",
        "final_args = f\"python make_captions.py {args}\"\n",
        "\n",
        "os.chdir(finetune_dir)\n",
        "!{final_args}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-BdXV7rAy2ag"
      },
      "outputs": [],
      "source": [
        "#@title ### **3.2.2. Waifu Diffusion 1.4 Tagger V2**\n",
        "import os\n",
        "%store -r\n",
        "\n",
        "os.chdir(finetune_dir)\n",
        "\n",
        "#@markdown [Waifu Diffusion 1.4 Tagger V2](https://huggingface.co/spaces/SmilingWolf/wd-v1-4-tags) is a Danbooru-styled image classification model developed by SmilingWolf. It can also be useful for general image tagging, for example, `1girl, solo, looking_at_viewer, short_hair, bangs, simple_background`.\n",
        "model = \"SmilingWolf/wd-v1-4-moat-tagger-v2\" #@param [\"SmilingWolf/wd-v1-4-moat-tagger-v2\", \"SmilingWolf/wd-v1-4-convnextv2-tagger-v2\", \"SmilingWolf/wd-v1-4-swinv2-tagger-v2\", \"SmilingWolf/wd-v1-4-convnext-tagger-v2\", \"SmilingWolf/wd-v1-4-vit-tagger-v2\"]\n",
        "#@markdown Separate `undesired_tags` with comma `(,)` if you want to remove multiple tags, e.g. `1girl,solo,smile`.\n",
        "undesired_tags = \"\" #@param {type:'string'}\n",
        "#@markdown Adjust `general_threshold` for pruning tags (less tags, less flexible). `character_threshold` is useful if you want to train with character tags, e.g. `hakurei reimu`.\n",
        "general_threshold = 0.35 #@param {type:\"slider\", min:0, max:1, step:0.05}\n",
        "character_threshold = 0.85 #@param {type:\"slider\", min:0, max:1, step:0.05}\n",
        "\n",
        "config = {\n",
        "    \"_train_data_dir\"           : train_data_dir,\n",
        "    \"batch_size\"                : 8,\n",
        "    \"repo_id\"                   : model,\n",
        "    \"recursive\"                 : True,\n",
        "    \"remove_underscore\"         : True,\n",
        "    \"general_threshold\"         : general_threshold,\n",
        "    \"character_threshold\"       : character_threshold,\n",
        "    \"caption_extension\"         : \".txt\",\n",
        "    \"max_data_loader_n_workers\" : 2,\n",
        "    \"debug\"                     : True,\n",
        "    \"undesired_tags\"            : undesired_tags\n",
        "}\n",
        "\n",
        "args = \"\"\n",
        "for k, v in config.items():\n",
        "    if k.startswith(\"_\"):\n",
        "        args += f'\"{v}\" '\n",
        "    elif isinstance(v, str):\n",
        "        args += f'--{k}=\"{v}\" '\n",
        "    elif isinstance(v, bool) and v:\n",
        "        args += f\"--{k} \"\n",
        "    elif isinstance(v, float) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "    elif isinstance(v, int) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "\n",
        "final_args = f\"python tag_images_by_wd14_tagger.py {args}\"\n",
        "\n",
        "os.chdir(finetune_dir)\n",
        "!{final_args}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "_mLVURhM9PFE"
      },
      "outputs": [],
      "source": [
        "# @title ### **3.2.3. Custom Caption/Tag**\n",
        "import os\n",
        "\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "# @markdown Add or remove custom tags here.\n",
        "extension   = \".txt\"  # @param [\".txt\", \".caption\"]\n",
        "custom_tag  = \"anime\"  # @param {type:\"string\"}\n",
        "# @markdown Use `sub_folder` option to specify a subfolder for multi-concept training.\n",
        "# @markdown > Specify `--all` to process all subfolders/`recursive`\n",
        "sub_folder  = \"\" #@param {type: \"string\"}\n",
        "# @markdown Enable this to append custom tags at the end of lines.\n",
        "append      = False  # @param {type:\"boolean\"}\n",
        "# @markdown Enable this if you want to remove captions/tags instead.\n",
        "remove_tag  = False  # @param {type:\"boolean\"}\n",
        "recursive   = False\n",
        "\n",
        "if sub_folder == \"\":\n",
        "    image_dir = train_data_dir\n",
        "elif sub_folder == \"--all\":\n",
        "    image_dir = train_data_dir\n",
        "    recursive = True\n",
        "elif sub_folder.startswith(\"/content\"):\n",
        "    image_dir = sub_folder\n",
        "else:\n",
        "    image_dir = os.path.join(train_data_dir, sub_folder)\n",
        "    os.makedirs(image_dir, exist_ok=True)\n",
        "\n",
        "def read_file(filename):\n",
        "    with open(filename, \"r\") as f:\n",
        "        contents = f.read()\n",
        "    return contents\n",
        "\n",
        "def write_file(filename, contents):\n",
        "    with open(filename, \"w\") as f:\n",
        "        f.write(contents)\n",
        "\n",
        "def process_tags(filename, custom_tag, append, remove_tag):\n",
        "    contents = read_file(filename)\n",
        "    tags = [tag.strip() for tag in contents.split(',')]\n",
        "    custom_tags = [tag.strip() for tag in custom_tag.split(',')]\n",
        "\n",
        "    for custom_tag in custom_tags:\n",
        "        custom_tag = custom_tag.replace(\"_\", \" \")\n",
        "        if remove_tag:\n",
        "            while custom_tag in tags:\n",
        "                tags.remove(custom_tag)\n",
        "        else:\n",
        "            if custom_tag not in tags:\n",
        "                if append:\n",
        "                    tags.append(custom_tag)\n",
        "                else:\n",
        "                    tags.insert(0, custom_tag)\n",
        "\n",
        "    contents = ', '.join(tags)\n",
        "    write_file(filename, contents)\n",
        "\n",
        "def process_directory(image_dir, tag, append, remove_tag, recursive):\n",
        "    for filename in os.listdir(image_dir):\n",
        "        file_path = os.path.join(image_dir, filename)\n",
        "\n",
        "        if os.path.isdir(file_path) and recursive:\n",
        "            process_directory(file_path, tag, append, remove_tag, recursive)\n",
        "        elif filename.endswith(extension):\n",
        "            process_tags(file_path, tag, append, remove_tag)\n",
        "\n",
        "tag = custom_tag\n",
        "\n",
        "if not any(\n",
        "    [filename.endswith(extension) for filename in os.listdir(image_dir)]\n",
        "):\n",
        "    for filename in os.listdir(image_dir):\n",
        "        if filename.endswith((\".png\", \".jpg\", \".jpeg\", \".webp\", \".bmp\")):\n",
        "            open(\n",
        "                os.path.join(image_dir, filename.split(\".\")[0] + extension),\n",
        "                \"w\",\n",
        "            ).close()\n",
        "\n",
        "if custom_tag:\n",
        "    process_directory(image_dir, tag, append, remove_tag, recursive)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "hhgatqF3leHJ"
      },
      "outputs": [],
      "source": [
        "# @title ## **3.4. Bucketing and Latents Caching**\n",
        "%store -r\n",
        "\n",
        "# @markdown This code will create buckets based on the `bucket_resolution` provided for multi-aspect ratio training, and then convert all images within the `train_data_dir` to latents.\n",
        "bucketing_json    = os.path.join(training_dir, \"meta_lat.json\")\n",
        "metadata_json     = os.path.join(training_dir, \"meta_clean.json\")\n",
        "bucket_resolution = 1024  # @param {type:\"slider\", min:512, max:1024, step:128}\n",
        "mixed_precision   = \"no\"  # @param [\"no\", \"fp16\", \"bf16\"] {allow-input: false}\n",
        "skip_existing     = False  # @param{type:\"boolean\"}\n",
        "flip_aug          = False  # @param{type:\"boolean\"}\n",
        "# @markdown Use `clean_caption` option to clean such as duplicate tags, `women` to `girl`, etc\n",
        "clean_caption     = False #@param {type:\"boolean\"}\n",
        "#@markdown Use the `recursive` option to process subfolders as well\n",
        "recursive         = True #@param {type:\"boolean\"}\n",
        "\n",
        "metadata_config = {\n",
        "    \"_train_data_dir\": train_data_dir,\n",
        "    \"_out_json\": metadata_json,\n",
        "    \"recursive\": recursive,\n",
        "    \"full_path\": recursive,\n",
        "    \"clean_caption\": clean_caption\n",
        "}\n",
        "\n",
        "bucketing_config = {\n",
        "    \"_train_data_dir\": train_data_dir,\n",
        "    \"_in_json\": metadata_json,\n",
        "    \"_out_json\": bucketing_json,\n",
        "    \"_model_name_or_path\": vae_path if vae_path else model_path,\n",
        "    \"recursive\": recursive,\n",
        "    \"full_path\": recursive,\n",
        "    \"flip_aug\": flip_aug,\n",
        "    \"skip_existing\": skip_existing,\n",
        "    \"batch_size\": 4,\n",
        "    \"max_data_loader_n_workers\": 2,\n",
        "    \"max_resolution\": f\"{bucket_resolution}, {bucket_resolution}\",\n",
        "    \"mixed_precision\": mixed_precision,\n",
        "}\n",
        "\n",
        "def generate_args(config):\n",
        "    args = \"\"\n",
        "    for k, v in config.items():\n",
        "        if k.startswith(\"_\"):\n",
        "            args += f'\"{v}\" '\n",
        "        elif isinstance(v, str):\n",
        "            args += f'--{k}=\"{v}\" '\n",
        "        elif isinstance(v, bool) and v:\n",
        "            args += f\"--{k} \"\n",
        "        elif isinstance(v, float) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "        elif isinstance(v, int) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "    return args.strip()\n",
        "\n",
        "merge_metadata_args = generate_args(metadata_config)\n",
        "prepare_buckets_args = generate_args(bucketing_config)\n",
        "\n",
        "merge_metadata_command = f\"python merge_all_to_metadata.py {merge_metadata_args}\"\n",
        "prepare_buckets_command = f\"python prepare_buckets_latents.py {prepare_buckets_args}\"\n",
        "\n",
        "os.chdir(finetune_dir)\n",
        "!{merge_metadata_command}\n",
        "time.sleep(1)\n",
        "!{prepare_buckets_command}\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GAZVkLuaRJ9e"
      },
      "source": [
        "# **IV. Training**\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "cJgLfRtlHSjw"
      },
      "outputs": [],
      "source": [
        "import toml\n",
        "\n",
        "# @title ## **4.1. LoRa: Low-Rank Adaptation Config**\n",
        "# @markdown Kohya's `LoRA` renamed to `LoRA-LierLa` and Kohya's `LoCon` renamed to `LoRA-C3Lier`, read [official announcement](https://github.com/kohya-ss/sd-scripts/blob/849bc24d205a35fbe1b2a4063edd7172533c1c01/README.md#naming-of-lora).\n",
        "network_category = \"LoRA_LierLa\"  # @param [\"LoRA_LierLa\", \"LoRA_C3Lier\", \"DyLoRA_LierLa\", \"DyLoRA_C3Lier\", \"LoCon\", \"LoHa\", \"IA3\", \"LoKR\", \"DyLoRA_Lycoris\"]\n",
        "\n",
        "# @markdown | network_category | network_dim | network_alpha | conv_dim | conv_alpha | unit |\n",
        "# @markdown | :---: | :---: | :---: | :---: | :---: | :---: |\n",
        "# @markdown | LoRA-LierLa | 32 | 1 | - | - | - |\n",
        "# @markdown | LoCon/LoRA-C3Lier | 16 | 8 | 8 | 1 | - |\n",
        "# @markdown | LoHa | 8 | 4 | 4 | 1 | - |\n",
        "# @markdown | Other Category | ? | ? | ? | ? | - |\n",
        "\n",
        "# @markdown Specify `network_args` to add `optional` training args, like for specifying each 25 block weight, read [this](https://github.com/kohya-ss/sd-scripts/blob/main/train_network_README-ja.md#%E9%9A%8E%E5%B1%A4%E5%88%A5%E5%AD%A6%E7%BF%92%E7%8E%87)\n",
        "network_args    = \"\"  # @param {'type':'string'}\n",
        "\n",
        "# @markdown ### **Linear Layer Config**\n",
        "# @markdown Used by all `network_category`. When in doubt, set `network_dim = network_alpha`\n",
        "network_dim     = 32  # @param {'type':'number'}\n",
        "network_alpha   = 16  # @param {'type':'number'}\n",
        "\n",
        "# @markdown ### **Convolutional Layer Config**\n",
        "# @markdown Only required if `network_category` is not `LoRA_LierLa`, as it involves training convolutional layers in addition to linear layers.\n",
        "conv_dim        = 32  # @param {'type':'number'}\n",
        "conv_alpha      = 16  # @param {'type':'number'}\n",
        "\n",
        "# @markdown ### **DyLoRA Config**\n",
        "# @markdown Only required if `network_category` is `DyLoRA_LierLa` and `DyLoRA_C3Lier`\n",
        "unit = 4  # @param {'type':'number'}\n",
        "\n",
        "if isinstance(network_args, str):\n",
        "    network_args = network_args.strip()\n",
        "    if network_args.startswith('[') and network_args.endswith(']'):\n",
        "        try:\n",
        "            network_args = ast.literal_eval(network_args)\n",
        "        except (SyntaxError, ValueError) as e:\n",
        "            print(f\"Error parsing network_args: {e}\\n\")\n",
        "            network_args = []\n",
        "    elif len(network_args) > 0:\n",
        "        print(f\"WARNING! '{network_args}' is not a valid list! Put args like this: [\\\"args=1\\\", \\\"args=2\\\"]\\n\")\n",
        "        network_args = []\n",
        "    else:\n",
        "        network_args = []\n",
        "else:\n",
        "    network_args = []\n",
        "\n",
        "network_config = {\n",
        "    \"LoRA_LierLa\": {\n",
        "        \"module\": \"networks.lora\",\n",
        "        \"args\"  : []\n",
        "    },\n",
        "    \"LoRA_C3Lier\": {\n",
        "        \"module\": \"networks.lora\",\n",
        "        \"args\"  : [\n",
        "            f\"conv_dim={conv_dim}\",\n",
        "            f\"conv_alpha={conv_alpha}\"\n",
        "        ]\n",
        "    },\n",
        "    \"DyLoRA_LierLa\": {\n",
        "        \"module\": \"networks.dylora\",\n",
        "        \"args\"  : [\n",
        "            f\"unit={unit}\"\n",
        "        ]\n",
        "    },\n",
        "    \"DyLoRA_C3Lier\": {\n",
        "        \"module\": \"networks.dylora\",\n",
        "        \"args\"  : [\n",
        "            f\"conv_dim={conv_dim}\",\n",
        "            f\"conv_alpha={conv_alpha}\",\n",
        "            f\"unit={unit}\"\n",
        "        ]\n",
        "    },\n",
        "    \"LoCon\": {\n",
        "        \"module\": \"lycoris.kohya\",\n",
        "        \"args\"  : [\n",
        "            f\"algo=locon\",\n",
        "            f\"conv_dim={conv_dim}\",\n",
        "            f\"conv_alpha={conv_alpha}\"\n",
        "        ]\n",
        "    },\n",
        "    \"LoHa\": {\n",
        "        \"module\": \"lycoris.kohya\",\n",
        "        \"args\"  : [\n",
        "            f\"algo=loha\",\n",
        "            f\"conv_dim={conv_dim}\",\n",
        "            f\"conv_alpha={conv_alpha}\"\n",
        "        ]\n",
        "    },\n",
        "    \"IA3\": {\n",
        "        \"module\": \"lycoris.kohya\",\n",
        "        \"args\"  : [\n",
        "            f\"algo=ia3\",\n",
        "            f\"conv_dim={conv_dim}\",\n",
        "            f\"conv_alpha={conv_alpha}\"\n",
        "        ]\n",
        "    },\n",
        "    \"LoKR\": {\n",
        "        \"module\": \"lycoris.kohya\",\n",
        "        \"args\"  : [\n",
        "            f\"algo=lokr\",\n",
        "            f\"conv_dim={conv_dim}\",\n",
        "            f\"conv_alpha={conv_alpha}\"\n",
        "        ]\n",
        "    },\n",
        "    \"DyLoRA_Lycoris\": {\n",
        "        \"module\": \"lycoris.kohya\",\n",
        "        \"args\"  : [\n",
        "            f\"algo=dylora\",\n",
        "            f\"conv_dim={conv_dim}\",\n",
        "            f\"conv_alpha={conv_alpha}\"\n",
        "        ]\n",
        "    }\n",
        "}\n",
        "\n",
        "network_module = network_config[network_category][\"module\"]\n",
        "network_args.extend(network_config[network_category][\"args\"])\n",
        "\n",
        "lora_config = {\n",
        "    \"additional_network_arguments\": {\n",
        "        \"no_metadata\"                     : False,\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,\n",
        "        \"training_comment\"                : None,\n",
        "    },\n",
        "}\n",
        "\n",
        "print(toml.dumps(lora_config))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "JNlw3u8arwir"
      },
      "outputs": [],
      "source": [
        "import toml\n",
        "import ast\n",
        "\n",
        "# @title ## **4.2. Optimizer Config**\n",
        "# @markdown Use `Adafactor` optimizer. `RMSprop 8bit` or `Adagrad 8bit` may work. `AdamW 8bit` doesn't seem to work.\n",
        "optimizer_type = \"AdaFactor\"  # @param [\"AdamW\", \"AdamW8bit\", \"Lion8bit\", \"Lion\", \"SGDNesterov\", \"SGDNesterov8bit\", \"DAdaptation(DAdaptAdamPreprint)\", \"DAdaptAdaGrad\", \"DAdaptAdam\", \"DAdaptAdan\", \"DAdaptAdanIP\", \"DAdaptLion\", \"DAdaptSGD\", \"AdaFactor\"]\n",
        "# @markdown Specify `optimizer_args` to add `additional` args for optimizer, e.g: `[\"weight_decay=0.6\"]`\n",
        "optimizer_args = \"[ \\\"scale_parameter=False\\\", \\\"relative_step=False\\\", \\\"warmup_init=False\\\" ]\"  # @param {'type':'string'}\n",
        "# @markdown ### **Learning Rate Config**\n",
        "# @markdown Different `optimizer_type` and `network_category` for some condition requires different learning rate. It's recommended to set `text_encoder_lr = 1/2 * unet_lr`\n",
        "learning_rate = 1e-4  # @param {'type':'number'}\n",
        "# @markdown ### **LR Scheduler Config**\n",
        "# @markdown `lr_scheduler` provides several methods to adjust the learning rate based on the number of epochs.\n",
        "lr_scheduler = \"constant_with_warmup\"  # @param [\"linear\", \"cosine\", \"cosine_with_restarts\", \"polynomial\", \"constant\", \"constant_with_warmup\", \"adafactor\"] {allow-input: false}\n",
        "lr_warmup_steps = 100  # @param {'type':'number'}\n",
        "# @markdown Specify `lr_scheduler_num` with `num_cycles` value for `cosine_with_restarts` or `power` value for `polynomial`\n",
        "lr_scheduler_num = 0  # @param {'type':'number'}\n",
        "\n",
        "if isinstance(optimizer_args, str):\n",
        "    optimizer_args = optimizer_args.strip()\n",
        "    if optimizer_args.startswith('[') and optimizer_args.endswith(']'):\n",
        "        try:\n",
        "            optimizer_args = ast.literal_eval(optimizer_args)\n",
        "        except (SyntaxError, ValueError) as e:\n",
        "            print(f\"Error parsing optimizer_args: {e}\\n\")\n",
        "            optimizer_args = []\n",
        "    elif len(optimizer_args) > 0:\n",
        "        print(f\"WARNING! '{optimizer_args}' is not a valid list! Put args like this: [\\\"args=1\\\", \\\"args=2\\\"]\\n\")\n",
        "        optimizer_args = []\n",
        "    else:\n",
        "        optimizer_args = []\n",
        "else:\n",
        "    optimizer_args = []\n",
        "\n",
        "optimizer_config = {\n",
        "    \"optimizer_arguments\": {\n",
        "        \"optimizer_type\"          : optimizer_type,\n",
        "        \"learning_rate\"           : learning_rate,\n",
        "        \"max_grad_norm\"           : 0,\n",
        "        \"optimizer_args\"          : optimizer_args,\n",
        "        \"lr_scheduler\"            : lr_scheduler,\n",
        "        \"lr_warmup_steps\"         : lr_warmup_steps,\n",
        "        \"lr_scheduler_num_cycles\" : lr_scheduler_num if lr_scheduler == \"cosine_with_restarts\" else None,\n",
        "        \"lr_scheduler_power\"      : lr_scheduler_num if lr_scheduler == \"polynomial\" else None,\n",
        "        \"lr_scheduler_type\"       : None,\n",
        "        \"lr_scheduler_args\"       : None,\n",
        "    },\n",
        "}\n",
        "\n",
        "print(toml.dumps(optimizer_config))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "YOxNM0x7dvfO"
      },
      "outputs": [],
      "source": [
        "# @title ## **4.3. Advanced Training Config** (Optional)\n",
        "import toml\n",
        "\n",
        "\n",
        "# @markdown ### **Optimizer State Config**\n",
        "save_optimizer_state      = False #@param {type:\"boolean\"}\n",
        "load_optimizer_state      = \"\" #@param {type:\"string\"}\n",
        "# @markdown ### **Noise Control**\n",
        "noise_control_type        = \"none\" #@param [\"none\", \"noise_offset\", \"multires_noise\"]\n",
        "# @markdown #### **a. Noise Offset**\n",
        "# @markdown Control and easily generating darker or light images by offset the noise when fine-tuning the model. Recommended value: `0.1`. Read [Diffusion With Offset Noise](https://www.crosslabs.org//blog/diffusion-with-offset-noise)\n",
        "noise_offset_num          = 0.0357  # @param {type:\"number\"}\n",
        "# @markdown **[Experimental]**\n",
        "# @markdown Automatically adjusts the noise offset based on the absolute mean values of each channel in the latents when used with `--noise_offset`. Specify a value around 1/10 to the same magnitude as the `--noise_offset` for best results. Set `0` to disable.\n",
        "adaptive_noise_scale      = 0.00357 # @param {type:\"number\"}\n",
        "# @markdown #### **b. Multires Noise**\n",
        "# @markdown enable multires noise with this number of iterations (if enabled, around 6-10 is recommended)\n",
        "multires_noise_iterations = 6 #@param {type:\"slider\", min:1, max:10, step:1}\n",
        "multires_noise_discount = 0.3 #@param {type:\"slider\", min:0.1, max:1, step:0.1}\n",
        "# @markdown ### **Caption Dropout**\n",
        "caption_dropout_rate = 0  # @param {type:\"number\"}\n",
        "caption_tag_dropout_rate = 0.5  # @param {type:\"number\"}\n",
        "caption_dropout_every_n_epochs = 0  # @param {type:\"number\"}\n",
        "# @markdown ### **Custom Train Function**\n",
        "# @markdown Gamma for reducing the weight of high-loss timesteps. Lower numbers have a stronger effect. The paper recommends `5`. Read the paper [here](https://arxiv.org/abs/2303.09556).\n",
        "min_snr_gamma             = 5 #@param {type:\"number\"}\n",
        "\n",
        "advanced_training_config = {\n",
        "    \"advanced_training_config\": {\n",
        "        \"resume\"                        : load_optimizer_state,\n",
        "        \"save_state\"                    : save_optimizer_state,\n",
        "        \"save_last_n_epochs_state\"      : save_optimizer_state,\n",
        "        \"noise_offset\"                  : noise_offset_num if noise_control_type == \"noise_offset\" else None,\n",
        "        \"adaptive_noise_scale\"          : adaptive_noise_scale if adaptive_noise_scale and noise_control_type == \"noise_offset\" else None,\n",
        "        \"multires_noise_iterations\"     : multires_noise_iterations if noise_control_type ==\"multires_noise\" else None,\n",
        "        \"multires_noise_discount\"       : multires_noise_discount if noise_control_type ==\"multires_noise\" else None,\n",
        "        \"caption_dropout_rate\"          : caption_dropout_rate,\n",
        "        \"caption_tag_dropout_rate\"      : caption_tag_dropout_rate,\n",
        "        \"caption_dropout_every_n_epochs\": caption_dropout_every_n_epochs,\n",
        "        \"min_snr_gamma\"                 : min_snr_gamma if not min_snr_gamma == -1 else None,\n",
        "    }\n",
        "}\n",
        "\n",
        "print(toml.dumps(advanced_training_config))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-Z4w3lfFKLjr",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "# @title ## **4.4. Training Config**\n",
        "import toml\n",
        "import os\n",
        "from subprocess import getoutput\n",
        "\n",
        "%store -r\n",
        "\n",
        "# @markdown ### **Project Config**\n",
        "project_name                = \"sdxl_lora\"  # @param {type:\"string\"}\n",
        "# @markdown Get your `wandb_api_key` [here](https://wandb.ai/settings) to logs with wandb.\n",
        "wandb_api_key               = \"\" # @param {type:\"string\"}\n",
        "in_json                     = \"/content/LoRA/meta_lat.json\"  # @param {type:\"string\"}\n",
        "# @markdown ### **SDXL Config**\n",
        "gradient_checkpointing      = True  # @param {type:\"boolean\"}\n",
        "no_half_vae                 = True  # @param {type:\"boolean\"}\n",
        "#@markdown Recommended parameter for SDXL training but if you enable it, `shuffle_caption` won't work\n",
        "cache_text_encoder_outputs  = False  # @param {type:\"boolean\"}\n",
        "#@markdown These options can be used to train U-Net with different timesteps. The default values are 0 and 1000.\n",
        "min_timestep                = 0 # @param {type:\"number\"}\n",
        "max_timestep                = 1000 # @param {type:\"number\"}\n",
        "# @markdown ### **Dataset Config**\n",
        "num_repeats                 = 1  # @param {type:\"number\"}\n",
        "resolution                  = 1024  # @param {type:\"slider\", min:512, max:1024, step:128}\n",
        "keep_tokens                 = 0  # @param {type:\"number\"}\n",
        "# @markdown ### **General Config**\n",
        "num_epochs                  = 10  # @param {type:\"number\"}\n",
        "train_batch_size            = 4  # @param {type:\"number\"}\n",
        "mixed_precision             = \"fp16\"  # @param [\"no\",\"fp16\",\"bf16\"] {allow-input: false}\n",
        "seed                        = -1  # @param {type:\"number\"}\n",
        "optimization                = \"scaled dot-product attention\" # @param [\"xformers\", \"scaled dot-product attention\"]\n",
        "# @markdown ### **Save Output Config**\n",
        "save_precision              = \"fp16\"  # @param [\"float\", \"fp16\", \"bf16\"] {allow-input: false}\n",
        "save_every_n_epochs         = 1  # @param {type:\"number\"}\n",
        "# @markdown ### **Sample Prompt Config**\n",
        "enable_sample               = True  # @param {type:\"boolean\"}\n",
        "sampler                     = \"euler_a\"  # @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",
        "positive_prompt             = \"\"\n",
        "negative_prompt             = \"\"\n",
        "quality_prompt              = \"NovelAI\"  # @param [\"None\", \"Waifu Diffusion 1.5\", \"NovelAI\", \"AbyssOrangeMix\", \"Stable Diffusion XL\"] {allow-input: false}\n",
        "if quality_prompt          == \"NovelAI\":\n",
        "    positive_prompt         = \"masterpiece, best quality, \"\n",
        "    negative_prompt         = \"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",
        "if quality_prompt          == \"AbyssOrangeMix\":\n",
        "    positive_prompt         = \"masterpiece, best quality, \"\n",
        "    negative_prompt         = \"(worst quality, low quality:1.4), \"\n",
        "if quality_prompt          == \"Stable Diffusion XL\":\n",
        "    negative_prompt         = \"3d render, smooth, plastic, blurry, grainy, low-resolution, deep-fried, oversaturated\"\n",
        "custom_prompt               = \"face focus, cute, 1girl, green hair, sweater, looking at viewer, upper body, beanie, outdoors, night, turtleneck\" # @param {type:\"string\"}\n",
        "# @markdown Specify `prompt_from_caption` if you want to use caption as prompt instead. Will be chosen randomly.\n",
        "prompt_from_caption         = \"none\"  # @param [\"none\", \".txt\", \".caption\"]\n",
        "if prompt_from_caption     != \"none\":\n",
        "    custom_prompt           = \"\"\n",
        "num_prompt                  = 2  # @param {type:\"number\"}\n",
        "logging_dir                 = os.path.join(training_dir, \"logs\")\n",
        "lowram                      = int(next(line.split()[1] for line in open('/proc/meminfo') if \"MemTotal\" in line)) / (1024**2) < 15\n",
        "\n",
        "os.chdir(repo_dir)\n",
        "\n",
        "prompt_config = {\n",
        "    \"prompt\": {\n",
        "        \"negative_prompt\" : negative_prompt,\n",
        "        \"width\"           : resolution,\n",
        "        \"height\"          : resolution,\n",
        "        \"scale\"           : 12,\n",
        "        \"sample_steps\"    : 28,\n",
        "        \"subset\"          : [],\n",
        "    }\n",
        "}\n",
        "\n",
        "train_config = {\n",
        "    \"sdxl_arguments\": {\n",
        "        \"cache_text_encoder_outputs\" : cache_text_encoder_outputs,\n",
        "        \"no_half_vae\"                : True,\n",
        "        \"min_timestep\"               : min_timestep,\n",
        "        \"max_timestep\"               : max_timestep,\n",
        "        \"shuffle_caption\"            : True if not cache_text_encoder_outputs else False,\n",
        "        \"lowram\"                     : lowram\n",
        "    },\n",
        "    \"model_arguments\": {\n",
        "        \"pretrained_model_name_or_path\" : model_path,\n",
        "        \"vae\"                           : vae_path,\n",
        "    },\n",
        "    \"dataset_arguments\": {\n",
        "        \"debug_dataset\"                 : False,\n",
        "        \"in_json\"                       : in_json,\n",
        "        \"train_data_dir\"                : train_data_dir,\n",
        "        \"dataset_repeats\"               : num_repeats,\n",
        "        \"keep_tokens\"                   : keep_tokens,\n",
        "        \"resolution\"                    : str(resolution) + ',' + str(resolution),\n",
        "        \"color_aug\"                     : False,\n",
        "        \"face_crop_aug_range\"           : None,\n",
        "        \"token_warmup_min\"              : 1,\n",
        "        \"token_warmup_step\"             : 0,\n",
        "    },\n",
        "    \"training_arguments\": {\n",
        "        \"output_dir\"                    : os.path.join(output_dir, project_name),\n",
        "        \"output_name\"                   : project_name if project_name else \"last\",\n",
        "        \"save_precision\"                : save_precision,\n",
        "        \"save_every_n_epochs\"           : save_every_n_epochs,\n",
        "        \"save_n_epoch_ratio\"            : None,\n",
        "        \"save_last_n_epochs\"            : None,\n",
        "        \"resume\"                        : None,\n",
        "        \"train_batch_size\"              : train_batch_size,\n",
        "        \"max_token_length\"              : 225,\n",
        "        \"mem_eff_attn\"                  : False,\n",
        "        \"sdpa\"                          : True if optimization == \"scaled dot-product attention\" else False,\n",
        "        \"xformers\"                      : True if optimization == \"xformers\" else False,\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\"   : 1,\n",
        "        \"mixed_precision\"               : mixed_precision,\n",
        "    },\n",
        "    \"logging_arguments\": {\n",
        "        \"log_with\"          : \"wandb\" if wandb_api_key else \"tensorboard\",\n",
        "        \"log_tracker_name\"  : project_name if wandb_api_key and not project_name == \"last\" else None,\n",
        "        \"logging_dir\"       : logging_dir,\n",
        "        \"log_prefix\"        : project_name if not wandb_api_key else None,\n",
        "    },\n",
        "    \"sample_prompt_arguments\": {\n",
        "        \"sample_every_n_steps\"    : None,\n",
        "        \"sample_every_n_epochs\"   : save_every_n_epochs if enable_sample else None,\n",
        "        \"sample_sampler\"          : sampler,\n",
        "    },\n",
        "    \"saving_arguments\": {\n",
        "        \"save_model_as\": \"safetensors\"\n",
        "    },\n",
        "}\n",
        "\n",
        "def write_file(filename, contents):\n",
        "    with open(filename, \"w\") as f:\n",
        "        f.write(contents)\n",
        "\n",
        "def prompt_convert(enable_sample, num_prompt, train_data_dir, prompt_config, custom_prompt):\n",
        "    if enable_sample:\n",
        "        search_pattern = os.path.join(train_data_dir, '**/*' + prompt_from_caption)\n",
        "        caption_files = glob.glob(search_pattern, recursive=True)\n",
        "\n",
        "        if not caption_files:\n",
        "            if not custom_prompt:\n",
        "                custom_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\"\n",
        "            new_prompt_config = prompt_config.copy()\n",
        "            new_prompt_config['prompt']['subset'] = [\n",
        "                {\"prompt\": positive_prompt + custom_prompt if positive_prompt else custom_prompt}\n",
        "            ]\n",
        "        else:\n",
        "            selected_files = random.sample(caption_files, min(num_prompt, len(caption_files)))\n",
        "\n",
        "            prompts = []\n",
        "            for file in selected_files:\n",
        "                with open(file, 'r') as f:\n",
        "                    prompts.append(f.read().strip())\n",
        "\n",
        "            new_prompt_config = prompt_config.copy()\n",
        "            new_prompt_config['prompt']['subset'] = []\n",
        "\n",
        "            for prompt in prompts:\n",
        "                new_prompt = {\n",
        "                    \"prompt\": positive_prompt + prompt if positive_prompt else prompt,\n",
        "                }\n",
        "                new_prompt_config['prompt']['subset'].append(new_prompt)\n",
        "\n",
        "        return new_prompt_config\n",
        "    else:\n",
        "        return prompt_config\n",
        "\n",
        "def eliminate_none_variable(config):\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",
        "    return config\n",
        "\n",
        "try:\n",
        "    train_config.update(optimizer_config)\n",
        "except NameError:\n",
        "    raise NameError(\"'optimizer_config' dictionary is missing. Please run  '4.1. Optimizer Config' cell.\")\n",
        "\n",
        "try:\n",
        "    train_config.update(lora_config)\n",
        "except NameError:\n",
        "    raise NameError(\"'lora_config' dictionary is missing. Please run  '4.1. LoRa: Low-Rank Adaptation Config' cell.\")\n",
        "\n",
        "advanced_training_warning = False\n",
        "try:\n",
        "    train_config.update(advanced_training_config)\n",
        "except NameError:\n",
        "    advanced_training_warning = True\n",
        "    pass\n",
        "\n",
        "prompt_config = prompt_convert(enable_sample, num_prompt, train_data_dir, prompt_config, custom_prompt)\n",
        "\n",
        "config_path         = os.path.join(config_dir, \"config_file.toml\")\n",
        "prompt_path         = os.path.join(config_dir, \"sample_prompt.toml\")\n",
        "\n",
        "config_str          = toml.dumps(eliminate_none_variable(train_config))\n",
        "prompt_str          = toml.dumps(eliminate_none_variable(prompt_config))\n",
        "\n",
        "write_file(config_path, config_str)\n",
        "write_file(prompt_path, prompt_str)\n",
        "\n",
        "print(config_str)\n",
        "\n",
        "if advanced_training_warning:\n",
        "    import textwrap\n",
        "    error_message = \"WARNING: This is not an error message, but the [advanced_training_config] dictionary is missing. Please run the '4.2. Advanced Training Config' cell if you intend to use it, or continue to the next step.\"\n",
        "    wrapped_message = textwrap.fill(error_message, width=80)\n",
        "    print('\\033[38;2;204;102;102m' + wrapped_message + '\\033[0m\\n')\n",
        "    pass\n",
        "\n",
        "print(prompt_str)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p_SHtbFwHVl1",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title ## **4.5. Start Training**\n",
        "import os\n",
        "import toml\n",
        "\n",
        "#@markdown Check your config here if you want to edit something:\n",
        "#@markdown - `sample_prompt` : /content/LoRA/config/sample_prompt.toml\n",
        "#@markdown - `config_file` : /content/LoRA/config/config_file.toml\n",
        "\n",
        "\n",
        "#@markdown You can import config from another session if you want.\n",
        "\n",
        "sample_prompt   = \"/content/LoRA/config/sample_prompt.toml\" #@param {type:'string'}\n",
        "config_file     = \"/content/LoRA/config/config_file.toml\" #@param {type:'string'}\n",
        "\n",
        "def read_file(filename):\n",
        "    with open(filename, \"r\") as f:\n",
        "        contents = f.read()\n",
        "    return contents\n",
        "\n",
        "def train(config):\n",
        "    args = \"\"\n",
        "    for k, v in config.items():\n",
        "        if k.startswith(\"_\"):\n",
        "            args += f'\"{v}\" '\n",
        "        elif isinstance(v, str):\n",
        "            args += f'--{k}=\"{v}\" '\n",
        "        elif isinstance(v, bool) and v:\n",
        "            args += f\"--{k} \"\n",
        "        elif isinstance(v, float) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "        elif isinstance(v, int) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "\n",
        "    return args\n",
        "\n",
        "accelerate_conf = {\n",
        "    \"config_file\" : \"/content/kohya-trainer/accelerate_config/config.yaml\",\n",
        "    \"num_cpu_threads_per_process\" : 1,\n",
        "}\n",
        "\n",
        "train_conf = {\n",
        "    \"sample_prompts\"  : sample_prompt if os.path.exists(sample_prompt) else None,\n",
        "    \"config_file\"     : config_file,\n",
        "    \"wandb_api_key\"   : wandb_api_key if wandb_api_key else None\n",
        "}\n",
        "\n",
        "accelerate_args = train(accelerate_conf)\n",
        "train_args = train(train_conf)\n",
        "\n",
        "final_args = f\"accelerate launch {accelerate_args} sdxl_train_network.py {train_args}\"\n",
        "\n",
        "os.chdir(repo_dir)\n",
        "!{final_args}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0vhtDOMtG7qd"
      },
      "source": [
        "## V. Testing"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "t4OhkX7Dz7CH"
      },
      "outputs": [],
      "source": [
        "#@title ## **5.1. Inference**\n",
        "\n",
        "import os\n",
        "%store -r\n",
        "\n",
        "# @markdown ### Model Config\n",
        "network_weights = \"\" #@param {type:'string'}\n",
        "network_mul = 0.7 # @param {type:\"slider\", min:-1, max:1, step:0.05}\n",
        "# @markdown ### Prompt Config\n",
        "prompt = \"face focus, cute, 1girl, green hair, sweater, looking at viewer, upper body, beanie, outdoors, night, turtleneck\" #@param {type:'string'}\n",
        "negative_prompt = \"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",
        "output_path = \"/content/tmp/\" #@param {type:'string'}\n",
        "resolution = \"1024, 1024\" # @param {type: \"string\"}\n",
        "optimization = \"scaled dot-product attention\" # @param [\"xformers\", \"scaled dot-product attention\"]\n",
        "conditional_resolution = \"1024, 1024\" # @param {type: \"string\"}\n",
        "steps = 28 # @param {type: \"number\"}\n",
        "sampler = \"euler_a\"  # @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",
        "scale = 7 # @param {type: \"number\"}\n",
        "seed = -1 # @param {type: \"number\"}\n",
        "images_per_prompt = 1 # @param {type: \"number\"}\n",
        "batch_size = 1 # @param {type: \"number\"}\n",
        "clip_skip = 2 # @param {type: \"number\"}\n",
        "\n",
        "os.makedirs(output_path, exist_ok=True)\n",
        "\n",
        "separators = [\"*\", \"x\", \",\"]\n",
        "\n",
        "for separator in separators:\n",
        "    if separator in resolution:\n",
        "        width, height = [value.strip() for value in resolution.split(separator)]\n",
        "        original_width, original_height = [value.strip() for value in conditional_resolution.split(separator)]\n",
        "        break\n",
        "\n",
        "network_config = {\n",
        "    \"network_module\": network_module,\n",
        "    \"network_weights\": network_weights,\n",
        "    \"network_show_meta\": True,\n",
        "    \"network_mul\": network_mul,\n",
        "}\n",
        "\n",
        "config = {\n",
        "    \"prompt\": prompt + \" --n \" + negative_prompt,\n",
        "    \"images_per_prompt\": images_per_prompt,\n",
        "    \"outdir\": output_path,\n",
        "    \"W\": width,\n",
        "    \"H\": height,\n",
        "    \"original_width\": original_width,\n",
        "    \"original_height\": original_height,\n",
        "    \"batch_size\": batch_size,\n",
        "    \"vae_batch_size\": 1,\n",
        "    \"no_half_vae\": True,\n",
        "    \"steps\": steps,\n",
        "    \"sampler\": sampler,\n",
        "    \"scale\": scale,\n",
        "    \"ckpt\": model_path,\n",
        "    \"vae\": vae_path,\n",
        "    \"seed\": seed if seed > 0 else None,\n",
        "    \"fp16\": True,\n",
        "    \"sdpa\": True if optimization == \"scaled dot-product attention\" else False,\n",
        "    \"xformers\": True if optimization == \"xformers\" else False,\n",
        "    \"opt_channels_last\": True,\n",
        "    \"clip_skip\": clip_skip,\n",
        "    \"max_embeddings_multiples\": 3,\n",
        "}\n",
        "\n",
        "if network_weights != \"\":\n",
        "    config.update(network_config)\n",
        "\n",
        "args = \"\"\n",
        "for k, v in config.items():\n",
        "    if k.startswith(\"_\"):\n",
        "        args += f'\"{v}\" '\n",
        "    elif isinstance(v, str):\n",
        "        args += f'--{k}=\"{v}\" '\n",
        "    elif isinstance(v, bool) and v:\n",
        "        args += f\"--{k} \"\n",
        "    elif isinstance(v, float) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "    elif isinstance(v, int) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "\n",
        "final_args = f\"python sdxl_gen_img.py {args}\"\n",
        "\n",
        "os.chdir(repo_dir)\n",
        "!{final_args}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nyIl9BhNXKUq"
      },
      "source": [
        "# **VI. Deployment**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "8YlP0cNQyujy"
      },
      "outputs": [],
      "source": [
        "# @title ## **6.1. Huggingface Hub config**\n",
        "from huggingface_hub import login\n",
        "from huggingface_hub import HfApi\n",
        "from huggingface_hub.utils import validate_repo_id, HfHubHTTPError\n",
        "\n",
        "# @markdown Login to Huggingface Hub\n",
        "# @markdown > Get **your** huggingface `WRITE` token [here](https://huggingface.co/settings/tokens)\n",
        "write_token = \"\"  # @param {type:\"string\"}\n",
        "# @markdown Fill this if you want to upload to your organization, or just leave it empty.\n",
        "orgs_name = \"\"  # @param{type:\"string\"}\n",
        "# @markdown If your model/dataset repo does not exist, it will automatically create it.\n",
        "model_name = \"sdxl-model\"  # @param{type:\"string\"}\n",
        "dataset_name = \"sdxl-dataset\"  # @param{type:\"string\"}\n",
        "make_private = True  # @param{type:\"boolean\"}\n",
        "\n",
        "def authenticate(write_token):\n",
        "    login(write_token, add_to_git_credential=True)\n",
        "    api = HfApi()\n",
        "    return api.whoami(write_token), api\n",
        "\n",
        "def create_repo(api, user, orgs_name, repo_name, repo_type, make_private=False):\n",
        "    global model_repo\n",
        "    global datasets_repo\n",
        "\n",
        "    if orgs_name == \"\":\n",
        "        repo_id = user[\"name\"] + \"/\" + repo_name.strip()\n",
        "    else:\n",
        "        repo_id = orgs_name + \"/\" + repo_name.strip()\n",
        "\n",
        "    try:\n",
        "        validate_repo_id(repo_id)\n",
        "        api.create_repo(repo_id=repo_id, repo_type=repo_type, private=make_private)\n",
        "        print(f\"{repo_type.capitalize()} repo '{repo_id}' didn't exist, creating repo\")\n",
        "    except HfHubHTTPError as e:\n",
        "        print(f\"{repo_type.capitalize()} repo '{repo_id}' exists, skipping create repo\")\n",
        "\n",
        "    if repo_type == \"model\":\n",
        "        model_repo = repo_id\n",
        "        print(f\"{repo_type.capitalize()} repo '{repo_id}' link: https://huggingface.co/{repo_id}\\n\")\n",
        "    else:\n",
        "        datasets_repo = repo_id\n",
        "        print(f\"{repo_type.capitalize()} repo '{repo_id}' link: https://huggingface.co/datasets/{repo_id}\\n\")\n",
        "\n",
        "user, api = authenticate(write_token)\n",
        "\n",
        "if model_name:\n",
        "    create_repo(api, user, orgs_name, model_name, \"model\", make_private)\n",
        "if dataset_name:\n",
        "    create_repo(api, user, orgs_name, dataset_name, \"dataset\", make_private)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "8HpOZOglBJTz"
      },
      "outputs": [],
      "source": [
        "# @title ## **6.2. Upload LoRA to Huggingface**\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",
        "\n",
        "# @markdown Now you can save your config file for future use\n",
        "config_path = \"/content/LoRA/config\"  # @param {type :\"string\"}\n",
        "\n",
        "# @markdown Other Information\n",
        "commit_message = \"\"  # @param {type :\"string\"}\n",
        "\n",
        "if not commit_message:\n",
        "    commit_message = f\"feat: upload {project_name} lora model\"\n",
        "\n",
        "def upload_to_hf(model_path, is_folder, is_config):\n",
        "    path_obj = Path(model_path)\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",
        "    print(f\"Uploading {trained_model} to https://huggingface.co/{model_repo}\")\n",
        "    print(\"Please wait...\")\n",
        "\n",
        "    if is_folder:\n",
        "        api.upload_folder(\n",
        "            folder_path=model_path,\n",
        "            path_in_repo=trained_model,\n",
        "            repo_id=model_repo,\n",
        "            commit_message=commit_message,\n",
        "            ignore_patterns=\".ipynb_checkpoints\",\n",
        "        )\n",
        "        print(f\"Upload success, located at https://huggingface.co/{model_repo}/tree/main\\n\")\n",
        "    else:\n",
        "        api.upload_file(\n",
        "            path_or_fileobj=model_path,\n",
        "            path_in_repo=trained_model,\n",
        "            repo_id=model_repo,\n",
        "            commit_message=commit_message,\n",
        "        )\n",
        "        print(f\"Upload success, located at https://huggingface.co/{model_repo}/blob/main/{trained_model}\\n\")\n",
        "\n",
        "def upload():\n",
        "    is_model_file = model_path.endswith((\".ckpt\", \".safetensors\", \".pt\"))\n",
        "    upload_to_hf(model_path, not is_model_file, False)\n",
        "\n",
        "    if config_path:\n",
        "        upload_to_hf(config_path, True, True)\n",
        "\n",
        "upload()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "IW-hS9jnmf-E"
      },
      "outputs": [],
      "source": [
        "# @title ## **6.3. Upload Dataset to Huggingface**\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",
        "\n",
        "# @markdown `Nerd stuff, only if you want to save training logs`\n",
        "logs_path = \"/content/LoRA/logs\"  # @param {type :\"string\"}\n",
        "\n",
        "tmp_dataset = f\"/content/LoRA/{project_name}_dataset\" if project_name else \"/content/LoRA/tmp_dataset\"\n",
        "tmp_train_data = f\"{tmp_dataset}/train_data\"\n",
        "dataset_zip = f\"{tmp_dataset}.zip\"\n",
        "\n",
        "# @markdown Other Information\n",
        "commit_message = \"\"  # @param {type :\"string\"}\n",
        "\n",
        "if not commit_message:\n",
        "    commit_message = f\"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",
        "def upload_dataset(dataset_path, is_zip):\n",
        "    path_obj = Path(dataset_path)\n",
        "    dataset_name = path_obj.parts[-1]\n",
        "\n",
        "    print(f\"Uploading {dataset_name} to https://huggingface.co/datasets/{datasets_repo}\")\n",
        "    print(\"Please wait...\")\n",
        "\n",
        "    if is_zip:\n",
        "        api.upload_file(\n",
        "            path_or_fileobj=dataset_path,\n",
        "            path_in_repo=dataset_name,\n",
        "            repo_id=datasets_repo,\n",
        "            repo_type=\"dataset\",\n",
        "            commit_message=commit_message,\n",
        "        )\n",
        "        print(f\"Upload success, located at https://huggingface.co/datasets/{datasets_repo}/blob/main/{dataset_name}\\n\")\n",
        "    else:\n",
        "        api.upload_folder(\n",
        "            folder_path=dataset_path,\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(f\"Upload success, located at https://huggingface.co/datasets/{datasets_repo}/tree/main/{dataset_name}\\n\")\n",
        "\n",
        "def zip_file(folder_path):\n",
        "    zip_path = f\"{folder_path}.zip\"\n",
        "    with zipfile.ZipFile(zip_path, \"w\") as zip_file:\n",
        "        for root, dirs, files in os.walk(folder_path):\n",
        "            for file in files:\n",
        "                zip_file.write(os.path.join(root, file))\n",
        "\n",
        "def move(src_path, dst_path, move_metadata):\n",
        "    metadata_files = [\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 move_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 metadata_files:\n",
        "                shutil.move(file_path, dst_path)\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",
        "upload()\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "T4",
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}