kohya-trainer-XL.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/kohya-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 Trainer XL**\n",
        "A Colab Notebook For Native Training\n",
        "\n",
        "[visitor-badge]: https://api.visitorbadge.io/api/visitors?path=Kohya%20Trainer%20XL&label=Visitors&labelColor=%2334495E&countColor=%231ABC9C&style=flat&labelStyle=none\n",
        "[visitor-stats]: https://visitorbadge.io/status?path=Kohya%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\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6GbCiSTjw42j"
      },
      "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",
      "metadata": {
        "id": "tTVqCAgSmie4"
      },
      "source": [
        "# **I. Prepare Environment**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NWLbzg4vwor-"
      },
      "outputs": [],
      "source": [
        "# Make sure the GPU is A100\n",
        "!nvidia-smi"
      ]
    },
    {
      "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, \"fine_tune\")\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 -yqq\n",
        "    !pip install -q --upgrade -r {requirements_file}\n",
        "\n",
        "    !pip install -q xformers==0.0.22.post7\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[\"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": {
        "cellView": "form",
        "id": "wrYGu-WxFbsq"
      },
      "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",
        "\n",
        "%store -r\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "# @markdown Place your Huggingface [Read Token](https://huggingface.co/settings/tokens) Here. Get your SDXL access [here](https://huggingface.co/stabilityai/stable-diffusion-xl-base-0.9).\n",
        "\n",
        "HUGGINGFACE_TOKEN = \"\"#@param {type: \"string\"}\n",
        "SDXL_MODEL_URL = \"https://huggingface.co/stabilityai/stable-diffusion-xl-base-0.9/resolve/main/sd_xl_base_0.9.safetensors\" #@param {type: \"string\"}\n",
        "SDXL_VAE_URL = \"https://huggingface.co/stabilityai/sdxl-vae/resolve/main/sdxl_vae.safetensors\" #@param {type: \"string\"}\n",
        "\n",
        "def get_supported_extensions():\n",
        "    return tuple([\".ckpt\", \".safetensors\", \".pt\", \".pth\"])\n",
        "\n",
        "def get_filename(url, quiet=True):\n",
        "    extensions = get_supported_extensions()\n",
        "\n",
        "    if url.startswith(drive_dir) or url.endswith(tuple(extensions)):\n",
        "        filename = os.path.basename(url)\n",
        "    else:\n",
        "        response = requests.get(url, stream=True)\n",
        "        response.raise_for_status()\n",
        "\n",
        "        if 'content-disposition' in response.headers:\n",
        "            content_disposition = response.headers['content-disposition']\n",
        "            filename = re.findall('filename=\"?([^\"]+)\"?', content_disposition)[0]\n",
        "        else:\n",
        "            url_path = urlparse(url).path\n",
        "            filename = unquote(os.path.basename(url_path))\n",
        "\n",
        "    if filename.endswith(tuple(get_supported_extensions())):\n",
        "        return filename\n",
        "    else:\n",
        "        return None\n",
        "\n",
        "def get_most_recent_file(directory):\n",
        "    files = glob.glob(os.path.join(directory, \"*\"))\n",
        "    if not files:\n",
        "        return None\n",
        "    most_recent_file = max(files, key=os.path.getmtime)\n",
        "    basename = os.path.basename(most_recent_file)\n",
        "\n",
        "    return most_recent_file\n",
        "\n",
        "def parse_args(config):\n",
        "    args = []\n",
        "\n",
        "    for k, v in config.items():\n",
        "        if k.startswith(\"_\"):\n",
        "            args.append(f\"{v}\")\n",
        "        elif isinstance(v, str) and v is not None:\n",
        "            args.append(f'--{k}={v}')\n",
        "        elif isinstance(v, bool) and v:\n",
        "            args.append(f\"--{k}\")\n",
        "        elif isinstance(v, float) and not isinstance(v, bool):\n",
        "            args.append(f\"--{k}={v}\")\n",
        "        elif isinstance(v, int) and not isinstance(v, bool):\n",
        "            args.append(f\"--{k}={v}\")\n",
        "\n",
        "    return args\n",
        "\n",
        "def aria2_download(dir, filename, url):\n",
        "    # hf_token    = \"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\" if not HUGGINGFACE_TOKEN else HUGGINGFACE_TOKEN\n",
        "    user_header = f\"Authorization: Bearer {HUGGINGFACE_TOKEN}\"\n",
        "\n",
        "    aria2_config = {\n",
        "        \"console-log-level\"         : \"error\",\n",
        "        \"summary-interval\"          : 10,\n",
        "        \"header\"                    : user_header if \"huggingface.co\" in url else None,\n",
        "        \"continue\"                  : True,\n",
        "        \"max-connection-per-server\" : 16,\n",
        "        \"min-split-size\"            : \"1M\",\n",
        "        \"split\"                     : 16,\n",
        "        \"dir\"                       : dir,\n",
        "        \"out\"                       : filename,\n",
        "        \"_url\"                      : url,\n",
        "    }\n",
        "    aria2_args = parse_args(aria2_config)\n",
        "    subprocess.run([\"aria2c\", *aria2_args])\n",
        "\n",
        "def gdown_download(url, dst, filepath):\n",
        "    if \"/uc?id/\" in url:\n",
        "        return gdown.download(url, filepath, quiet=False)\n",
        "    elif \"/file/d/\" in url:\n",
        "        return gdown.download(url=url, output=filepath, quiet=False, fuzzy=True)\n",
        "    elif \"/drive/folders/\" in url:\n",
        "        os.chdir(dst)\n",
        "        return gdown.download_folder(url, quiet=True, use_cookies=False)\n",
        "\n",
        "def download(url, dst):\n",
        "    filename = get_filename(url, quiet=False)\n",
        "    filepath = os.path.join(dst, filename)\n",
        "\n",
        "    if \"drive.google.com\" in url:\n",
        "        gdown = gdown_download(url, dst, filepath)\n",
        "    elif url.startswith(\"/content/drive/MyDrive/\"):\n",
        "        # Path(filepath).write_bytes(Path(url).read_bytes())\n",
        "        return url\n",
        "    else:\n",
        "        if \"huggingface.co\" in url:\n",
        "            if \"/blob/\" in url:\n",
        "                url = url.replace(\"/blob/\", \"/resolve/\")\n",
        "        aria2_download(dst, filename, url)\n",
        "\n",
        "def get_filepath(url, dst):\n",
        "    extensions = get_supported_extensions()\n",
        "    filename = get_filename(url)\n",
        "\n",
        "    if not filename.endswith(extensions):\n",
        "        most_recent_file = get_most_recent_file(dst)\n",
        "        filename = os.path.basename(most_recent_file)\n",
        "\n",
        "    filepath = os.path.join(dst, filename)\n",
        "\n",
        "    return filepath\n",
        "\n",
        "def main():\n",
        "    global model_path, vae_path\n",
        "\n",
        "    model_path = vae_path = None\n",
        "\n",
        "    download_targets = {\n",
        "        \"model\" : (SDXL_MODEL_URL, pretrained_model),\n",
        "        \"vae\"   : (SDXL_VAE_URL, vae_dir),\n",
        "    }\n",
        "    selected_files = {}\n",
        "\n",
        "    for target, (url, dst) in download_targets.items():\n",
        "        if url and f\"PASTE {target.upper()} URL OR GDRIVE PATH HERE\" not in url:\n",
        "            downloader = download(url, dst)\n",
        "            selected_files[target] = get_filepath(url, dst)\n",
        "\n",
        "            if target == \"model\":\n",
        "                model_path = selected_files[\"model\"] if not downloader else downloader\n",
        "            elif target == \"vae\":\n",
        "                vae_path = selected_files[\"vae\"] if not downloader else downloader\n",
        "\n",
        "    for category, path in {\n",
        "        \"model\": model_path,\n",
        "        \"vae\": vae_path,\n",
        "    }.items():\n",
        "        if path is not None and os.path.exists(path):\n",
        "            print(f\"Selected {category}: {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/fine_tune/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  = \"sd-webui-infinite-image-browsing\" #@param [\"imjoy-elfinder\", \"sd-webui-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 == \"sd-webui-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 shutil\n",
        "from pathlib import Path\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. The dataset will be downloaded and automatically extracted to `train_data_dir` if `unzip_to` is empty.\n",
        "zipfile_url  = \"https://huggingface.co/datasets/Linaqruf/hitokomoru-lora-dataset/resolve/main/hitokomoru_dataset.zip\" #@param {type:\"string\"}\n",
        "zipfile_name = \"zipfile.zip\"\n",
        "unzip_to     = \"\" #@param {type:\"string\"}\n",
        "\n",
        "hf_token     = \"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\"\n",
        "user_header  = f'\"Authorization: Bearer {hf_token}\"'\n",
        "\n",
        "if unzip_to:\n",
        "    os.makedirs(unzip_to, exist_ok=True)\n",
        "else:\n",
        "    unzip_to = train_data_dir\n",
        "\n",
        "def download_dataset(url):\n",
        "    if url.startswith(\"/content\"):\n",
        "        return url\n",
        "    elif \"drive.google.com\" in url:\n",
        "        os.chdir(root_dir)\n",
        "        !gdown --fuzzy {url}\n",
        "        return f\"{root_dir}/{zipfile_name}\"\n",
        "    elif \"huggingface.co\" in url:\n",
        "        if \"/blob/\" in url:\n",
        "            url = url.replace(\"/blob/\", \"/resolve/\")\n",
        "        !aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 -d {root_dir} -o {zipfile_name} {url}\n",
        "        return f\"{root_dir}/{zipfile_name}\"\n",
        "    else:\n",
        "        !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {root_dir} -o {zipfile_name} {url}\n",
        "        return f\"{root_dir}/{zipfile_name}\"\n",
        "\n",
        "def extract_dataset(zip_file, output_path):\n",
        "    !unzip -j -o {zip_file} -d \"{output_path}\"\n",
        "\n",
        "def remove_files(train_dir, files_to_move):\n",
        "    for filename in os.listdir(train_dir):\n",
        "        file_path = os.path.join(train_dir, filename)\n",
        "        if filename in files_to_move:\n",
        "            if not os.path.exists(file_path):\n",
        "                shutil.move(file_path, training_dir)\n",
        "            else:\n",
        "                os.remove(file_path)\n",
        "\n",
        "zip_file = download_dataset(zipfile_url)\n",
        "extract_dataset(zip_file, unzip_to)\n",
        "os.remove(zip_file)\n",
        "\n",
        "files_to_move = (\n",
        "    \"meta_cap.json\",\n",
        "    \"meta_cap_dd.json\",\n",
        "    \"meta_lat.json\",\n",
        "    \"meta_clean.json\",\n",
        ")\n",
        "\n",
        "remove_files(train_data_dir, files_to_move)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "A0t1dfnU5Xkq"
      },
      "outputs": [],
      "source": [
        "#@title ## **2.2. Scrape Dataset**\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 = \"hitokomoru\" #@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 = \"1-200\" #@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": "Plit55i32LEp"
      },
      "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 = \"\" #@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.55 #@param {type:\"slider\", min:0, max:1, step:0.05}\n",
        "character_threshold = 0.35 #@param {type:\"slider\", min:0, max:1, step:0.05}\n",
        "\n",
        "config = {\n",
        "    \"_train_data_dir\"           : train_data_dir,\n",
        "    \"batch_size\"                : 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",
        "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     = True #@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\": model_path,\n",
        "    \"recursive\": recursive,\n",
        "    \"full_path\": recursive,\n",
        "    \"flip_aug\": flip_aug,\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": "yHNbl3O_NSS0"
      },
      "source": [
        "# **IV. Training**\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "hJ0uGlqCH8SD"
      },
      "outputs": [],
      "source": [
        "# @title ## **4.1. Optimizer Config**\n",
        "import toml\n",
        "import ast\n",
        "\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 = 4e-7  # @param {'type':'number'}\n",
        "train_text_encoder = False  # @param {type:\"boolean\"}\n",
        "# train_text_encoder = False  # @param {'type':'boolean'}\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",
        "        \"train_text_encoder\"      : train_text_encoder,\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": "7guAPchHJFqU"
      },
      "outputs": [],
      "source": [
        "# @title ## **4.3. Advanced Training Config** (Optional)\n",
        "import toml\n",
        "\n",
        "\n",
        "# @markdown ### **Resume With Optimizer State**\n",
        "optimizer_state_path      = \"\" #@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.1  # @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.01 # @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 ### **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             = -1 #@param {type:\"number\"}\n",
        "\n",
        "advanced_training_config = {\n",
        "    \"advanced_training_config\": {\n",
        "        \"resume\"                    : optimizer_state_path,\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",
        "        \"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": {
        "cellView": "form",
        "id": "2y5wg-y4IKba"
      },
      "outputs": [],
      "source": [
        "# @title ## **4.3. Training Config**\n",
        "import toml\n",
        "import os\n",
        "import random\n",
        "from subprocess import getoutput\n",
        "\n",
        "%store -r\n",
        "\n",
        "# @markdown ### **Project Config**\n",
        "project_name            = \"sdxl_finetune\"  # @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/fine_tune/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 = True  # @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",
        "# @markdown Please refer to `3.2.3. Custom Caption/Tag (Optional)` if you want to append `activation_word` to captions/tags\n",
        "resolution              = 1024  # @param {type:\"slider\", min:512, max:1024, step:128}\n",
        "keep_tokens             = 0  # @param {type:\"number\"}\n",
        "# @markdown ### **General Config**\n",
        "max_train_steps         = 2500  # @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",
        "# @markdown ### **Save Output Config**\n",
        "save_precision          = \"fp16\"  # @param [\"float\", \"fp16\", \"bf16\"] {allow-input: false}\n",
        "save_every_n_steps      = 1000  # @param {type:\"number\"}\n",
        "save_optimizer_state    = False  # @param {type:\"boolean\"}\n",
        "save_model_as           = \"safetensors\" #@param [\"ckpt\", \"safetensors\", \"diffusers\", \"diffusers_safetensors\"]\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              = \"None\"  # @param [\"None\", \"Waifu Diffusion 1.5\", \"NovelAI\", \"AbyssOrangeMix\", \"Stable Diffusion XL\"] {allow-input: false}\n",
        "if quality_prompt          == \"Waifu Diffusion 1.5\":\n",
        "    positive_prompt         = \"(exceptional, best aesthetic, new, newest, best quality, masterpiece, extremely detailed, anime, waifu:1.2), \"\n",
        "    negative_prompt         = \"lowres, ((bad anatomy)), ((bad hands)), missing finger, extra digits, fewer digits, blurry, ((mutated hands and fingers)), (poorly drawn face), ((mutation)), ((deformed face)), (ugly), ((bad proportions)), ((extra limbs)), extra face, (double head), (extra head), ((extra feet)), monster, logo, cropped, worst quality, jpeg, humpbacked, long body, long neck, ((jpeg artifacts)), deleted, old, oldest, ((censored)), ((bad aesthetic)), (mosaic censoring, bar censor, blur censor), \"\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               = \"1girl, aqua eyes, baseball cap, blonde hair, closed mouth, earrings, green background, hat, hoop earrings, jewelry, looking at viewer, shirt, short hair, simple background, solo, upper body, yellow shirt\" # @param {type:\"string\"}\n",
        "# @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",
        "sample_interval             = 100  # @param {type:\"number\"}\n",
        "logging_dir             = \"/content/fine_tune/logs\"\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\"           : 7,\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",
        "        # \"enable_bucket\"              : True,\n",
        "        \"no_half_vae\"                : no_half_vae,\n",
        "        # \"cache_latents\"              : True,\n",
        "        # \"cache_latents_to_disk\"      : True,\n",
        "        # \"vae_batch_size\"             : 4,\n",
        "        \"min_timestep\"               : min_timestep,\n",
        "        \"max_timestep\"               : max_timestep,\n",
        "        \"shuffle_caption\"            : True if not cache_text_encoder_outputs else False,\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",
        "        \"caption_dropout_rate\"          : 0,\n",
        "        \"caption_tag_dropout_rate\"      : 0,\n",
        "        \"caption_dropout_every_n_epochs\": 0,\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\"                    : output_dir,\n",
        "        \"output_name\"                   : project_name if project_name else \"last\",\n",
        "        \"save_precision\"                : save_precision,\n",
        "        \"save_every_n_steps\"            : save_every_n_steps,\n",
        "        \"save_n_epoch_ratio\"            : None,\n",
        "        \"save_last_n_epochs\"            : None,\n",
        "        \"save_state\"                    : None,\n",
        "        \"save_last_n_epochs_state\"      : None,\n",
        "        \"resume\"                        : None,\n",
        "        \"train_batch_size\"              : train_batch_size,\n",
        "        \"max_token_length\"              : 225,\n",
        "        \"mem_eff_attn\"                  : False,\n",
        "        \"xformers\"                      : True,\n",
        "        \"max_train_steps\"               : max_train_steps,\n",
        "        \"max_data_loader_n_workers\"     : 8,\n",
        "        \"persistent_data_loader_workers\": True,\n",
        "        \"seed\"                          : seed if seed > 0 else None,\n",
        "        \"gradient_checkpointing\"        : gradient_checkpointing,\n",
        "        \"gradient_accumulation_steps\"   : 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\"    : sample_interval,\n",
        "        \"sample_every_n_epochs\"   : 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",
        "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)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "DNYGVTkaiQPf"
      },
      "outputs": [],
      "source": [
        "#@title ## **4.4. 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/fine_tune/config/sample_prompt.toml\n",
        "#@markdown - `config_file` : /content/fine_tune/config/config_file.toml\n",
        "\n",
        "#@markdown You can import config from another session if you want.\n",
        "\n",
        "sample_prompt   = \"/content/fine_tune/config/sample_prompt.toml\" #@param {type:'string'}\n",
        "config_file     = \"/content/fine_tune/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\" : accelerate_config,\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.py {train_args}\"\n",
        "\n",
        "os.chdir(repo_dir)\n",
        "!{final_args}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "7Jzy__Ty9NAk"
      },
      "outputs": [],
      "source": [
        "#@title ## **5.1. Inference**\n",
        "\n",
        "import os\n",
        "%store -r\n",
        "\n",
        "# @markdown ### Model Config\n",
        "ckpt_path = \"/content/fine_tune/output/sdxl_finetune.safetensors\" #@param {type:'string'}\n",
        "# @markdown ### Prompt Config\n",
        "prompt = \"1girl, aqua eyes, baseball cap, blonde hair, closed mouth, earrings, green background, hat, hoop earrings, jewelry, looking at viewer, shirt, short hair, simple background, solo, upper body, yellow shirt\" #@param {type:'string'}\n",
        "negative_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",
        "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\": ckpt_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",
        "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": [
        "# **V. Deployment**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "QTXsM170GUpk"
      },
      "outputs": [],
      "source": [
        "# @title ## **5.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 = \"\"  # @param{type:\"string\"}\n",
        "dataset_name = \"\"  # @param{type:\"string\"}\n",
        "make_private = False  # @param{type:\"boolean\"}\n",
        "\n",
        "def authenticate(write_token):\n",
        "    login(write_token, add_to_git_credential=True)\n",
        "    api = HfApi()\n",
        "    return api.whoami(write_token), api\n",
        "\n",
        "\n",
        "def create_repo(api, user, orgs_name, repo_name, repo_type, make_private=False):\n",
        "    global model_repo\n",
        "    global datasets_repo\n",
        "\n",
        "    if orgs_name == \"\":\n",
        "        repo_id = user[\"name\"] + \"/\" + repo_name.strip()\n",
        "    else:\n",
        "        repo_id = orgs_name + \"/\" + repo_name.strip()\n",
        "\n",
        "    try:\n",
        "        validate_repo_id(repo_id)\n",
        "        api.create_repo(repo_id=repo_id, repo_type=repo_type, private=make_private)\n",
        "        print(f\"{repo_type.capitalize()} repo '{repo_id}' didn't exist, creating repo\")\n",
        "    except HfHubHTTPError as e:\n",
        "        print(f\"{repo_type.capitalize()} repo '{repo_id}' exists, skipping create repo\")\n",
        "\n",
        "    if repo_type == \"model\":\n",
        "        model_repo = repo_id\n",
        "        print(f\"{repo_type.capitalize()} repo '{repo_id}' link: https://huggingface.co/{repo_id}\\n\")\n",
        "    else:\n",
        "        datasets_repo = repo_id\n",
        "        print(f\"{repo_type.capitalize()} repo '{repo_id}' link: https://huggingface.co/datasets/{repo_id}\\n\")\n",
        "\n",
        "user, api = authenticate(write_token)\n",
        "\n",
        "if model_name:\n",
        "    create_repo(api, user, orgs_name, model_name, \"model\", make_private)\n",
        "if dataset_name:\n",
        "    create_repo(api, user, orgs_name, dataset_name, \"dataset\", make_private)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "CIeoJA-eO-8t"
      },
      "outputs": [],
      "source": [
        "# @title ## **5.2. Upload Checkpoint to Huggingface**\n",
        "from huggingface_hub import HfApi\n",
        "from pathlib import Path\n",
        "\n",
        "api = HfApi()\n",
        "\n",
        "# @markdown This will be uploaded to model repo\n",
        "model_path = \"/content/fine_tune/output/sdxl_finetune.safetensors\"  # @param {type :\"string\"}\n",
        "path_in_repo = \"\"  # @param {type :\"string\"}\n",
        "# @markdown Now you can save your config file for future use\n",
        "config_path = \"/content/fine_tune/config\"  # @param {type :\"string\"}\n",
        "# @markdown Other Information\n",
        "commit_message = \"\"  # @param {type :\"string\"}\n",
        "\n",
        "if not commit_message:\n",
        "    commit_message = \"feat: upload \" + project_name + \" checkpoint\"\n",
        "\n",
        "if os.path.exists(model_path):\n",
        "    vae_exists = os.path.exists(os.path.join(model_path, \"vae\"))\n",
        "    unet_exists = os.path.exists(os.path.join(model_path, \"unet\"))\n",
        "    text_encoder_exists = os.path.exists(os.path.join(model_path, \"text_encoder\"))\n",
        "\n",
        "\n",
        "def upload_model(model_paths, is_folder: bool, is_config: bool):\n",
        "    path_obj = Path(model_paths)\n",
        "    trained_model = path_obj.parts[-1]\n",
        "\n",
        "    if path_in_repo:\n",
        "        trained_model = path_in_repo\n",
        "\n",
        "    if is_config:\n",
        "        if path_in_repo:\n",
        "            trained_model = f\"{path_in_repo}_config\"\n",
        "        else:\n",
        "            trained_model = f\"{project_name}_config\"\n",
        "\n",
        "    if is_folder == True:\n",
        "        print(f\"Uploading {trained_model} to https://huggingface.co/\" + model_repo)\n",
        "        print(f\"Please wait...\")\n",
        "\n",
        "        if vae_exists and unet_exists and text_encoder_exists:\n",
        "            api.upload_folder(\n",
        "                folder_path=model_paths,\n",
        "                repo_id=model_repo,\n",
        "                commit_message=commit_message,\n",
        "                ignore_patterns=\".ipynb_checkpoints\",\n",
        "            )\n",
        "        else:\n",
        "            api.upload_folder(\n",
        "                folder_path=model_paths,\n",
        "                path_in_repo=trained_model,\n",
        "                repo_id=model_repo,\n",
        "                commit_message=commit_message,\n",
        "                ignore_patterns=\".ipynb_checkpoints\",\n",
        "            )\n",
        "        print(\n",
        "            f\"Upload success, located at https://huggingface.co/\"\n",
        "            + model_repo\n",
        "            + \"/tree/main\\n\"\n",
        "        )\n",
        "    else:\n",
        "        print(f\"Uploading {trained_model} to https://huggingface.co/\" + model_repo)\n",
        "        print(f\"Please wait...\")\n",
        "\n",
        "        api.upload_file(\n",
        "            path_or_fileobj=model_paths,\n",
        "            path_in_repo=trained_model,\n",
        "            repo_id=model_repo,\n",
        "            commit_message=commit_message,\n",
        "        )\n",
        "\n",
        "        print(\n",
        "            f\"Upload success, located at https://huggingface.co/\"\n",
        "            + model_repo\n",
        "            + \"/blob/main/\"\n",
        "            + trained_model\n",
        "            + \"\\n\"\n",
        "        )\n",
        "\n",
        "\n",
        "def upload():\n",
        "    if model_path.endswith((\".ckpt\", \".safetensors\", \".pt\")):\n",
        "        upload_model(model_path, False, False)\n",
        "    else:\n",
        "        upload_model(model_path, True, False)\n",
        "\n",
        "    if config_path:\n",
        "        upload_model(config_path, True, True)\n",
        "\n",
        "\n",
        "upload()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "IW-hS9jnmf-E"
      },
      "outputs": [],
      "source": [
        "# @title ## **5.3. Upload Dataset to Huggingface** (optional)\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/fine_tune/train_data\"  # @param {type :\"string\"}\n",
        "meta_lat_path = \"/content/fine_tune/meta_lat.json\"  # @param {type :\"string\"}\n",
        "last_state_path = \"/content/fine_tune/output/last-state\"  # @param {type :\"string\"}\n",
        "# @markdown `Nerd stuff, only if you want to save training logs`\n",
        "logs_path = \"/content/fine_tune/logs\"  # @param {type :\"string\"}\n",
        "\n",
        "if project_name:\n",
        "    tmp_dataset = \"/content/fine_tune/\" + project_name + \"_dataset\"\n",
        "    tmp_last_state = \"/content/fine_tune/\" + project_name + \"_last_state\"\n",
        "\n",
        "else:\n",
        "    tmp_dataset = \"/content/fine_tune/tmp_dataset\"\n",
        "    tmp_last_state = \"/content/fine_tune/tmp_last_state\"\n",
        "\n",
        "tmp_train_data = tmp_dataset + \"/train_data\"\n",
        "dataset_zip = tmp_dataset + \".zip\"\n",
        "last_state_zip = tmp_last_state + \".zip\"\n",
        "\n",
        "# @markdown  Other Information\n",
        "commit_message = \"\"  # @param {type :\"string\"}\n",
        "\n",
        "if not commit_message:\n",
        "    commit_message = \"feat: upload \" + project_name + \" dataset and logs\"\n",
        "\n",
        "tmp_folder = [\"tmp_dataset\", \"tmp_last_state\", \"tmp_train_data\"]\n",
        "\n",
        "\n",
        "def makedirs(tmp_folders):\n",
        "    os.makedirs(tmp_folders, exist_ok=True)\n",
        "\n",
        "\n",
        "for folder in tmp_folder:\n",
        "    makedirs(folder)\n",
        "\n",
        "\n",
        "def upload_dataset(dataset_paths, is_zip: bool):\n",
        "    path_obj = Path(dataset_paths)\n",
        "    dataset_name = path_obj.parts[-1]\n",
        "\n",
        "    if is_zip:\n",
        "        print(\n",
        "            f\"Uploading {dataset_name} to https://huggingface.co/datasets/\"\n",
        "            + datasets_repo\n",
        "        )\n",
        "        print(f\"Please wait...\")\n",
        "\n",
        "        api.upload_file(\n",
        "            path_or_fileobj=dataset_paths,\n",
        "            path_in_repo=dataset_name,\n",
        "            repo_id=datasets_repo,\n",
        "            repo_type=\"dataset\",\n",
        "            commit_message=commit_message,\n",
        "        )\n",
        "        print(\n",
        "            f\"Upload success, located at https://huggingface.co/datasets/\"\n",
        "            + datasets_repo\n",
        "            + \"/blob/main/\"\n",
        "            + dataset_name\n",
        "            + \"\\n\"\n",
        "        )\n",
        "    else:\n",
        "        print(\n",
        "            f\"Uploading {dataset_name} to https://huggingface.co/datasets/\"\n",
        "            + datasets_repo\n",
        "        )\n",
        "        print(f\"Please wait...\")\n",
        "\n",
        "        api.upload_folder(\n",
        "            folder_path=dataset_paths,\n",
        "            path_in_repo=dataset_name,\n",
        "            repo_id=datasets_repo,\n",
        "            repo_type=\"dataset\",\n",
        "            commit_message=commit_message,\n",
        "            ignore_patterns=\".ipynb_checkpoints\",\n",
        "        )\n",
        "        print(\n",
        "            f\"Upload success, located at https://huggingface.co/datasets/\"\n",
        "            + datasets_repo\n",
        "            + \"/tree/main/\"\n",
        "            + dataset_name\n",
        "            + \"\\n\"\n",
        "        )\n",
        "\n",
        "\n",
        "def zip_file(tmp_folders):\n",
        "    zipfiles = tmp_folders + \".zip\"\n",
        "    with zipfile.ZipFile(zipfiles, \"w\") as zip:\n",
        "        for tmp_folders, dirs, files in os.walk(tmp_folders):\n",
        "            for file in files:\n",
        "                zip.write(os.path.join(tmp_folders, file))\n",
        "\n",
        "\n",
        "def move(src_path, dst_path, is_metadata: bool):\n",
        "    files_to_move = [\n",
        "        \"meta_cap.json\",\n",
        "        \"meta_cap_dd.json\",\n",
        "        \"meta_lat.json\",\n",
        "        \"meta_clean.json\",\n",
        "        \"meta_final.json\",\n",
        "    ]\n",
        "\n",
        "    if os.path.exists(src_path):\n",
        "        shutil.move(src_path, dst_path)\n",
        "\n",
        "    if is_metadata:\n",
        "        parent_meta_path = os.path.dirname(src_path)\n",
        "\n",
        "        for filename in os.listdir(parent_meta_path):\n",
        "            file_path = os.path.join(parent_meta_path, filename)\n",
        "            if filename in files_to_move:\n",
        "                shutil.move(file_path, dst_path)\n",
        "\n",
        "\n",
        "def upload():\n",
        "    if train_data_path and meta_lat_path:\n",
        "        move(train_data_path, tmp_train_data, False)\n",
        "        move(meta_lat_path, tmp_dataset, True)\n",
        "        zip_file(tmp_dataset)\n",
        "        upload_dataset(dataset_zip, True)\n",
        "        os.remove(dataset_zip)\n",
        "\n",
        "    if last_state_path:\n",
        "        if os.path.exists(last_state_path):\n",
        "            move(last_state_path, tmp_last_state, False)\n",
        "            zip_file(tmp_last_state)\n",
        "            upload_dataset(last_state_zip, True)\n",
        "            os.remove(last_state_zip)\n",
        "\n",
        "    if logs_path:\n",
        "        upload_dataset(logs_path, False)\n",
        "\n",
        "\n",
        "upload()"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "gpuType": "A100",
      "include_colab_link": true,
      "machine_shape": "hm",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}