240215_sdxl_LoRA_trainer_XL.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/MushroomFleet/unsorted-projects/blob/main/240215_sdxl_LoRA_trainer_XL.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "slgjeYgd6pWp"
      },
      "source": [
        "[![visitor][visitor-badge]][visitor-stats]\n",
        "[![ko-fi][ko-fi-badge]][ko-fi-link]\n",
        "\n",
        "# **Kohya LoRA Trainer XL**\n",
        "A Colab Notebook For SDXL LoRA Training (Fine-tuning Method)\n",
        "*Johnsons Fork*\n",
        "\n",
        "[visitor-badge]: https://api.visitorbadge.io/api/visitors?path=Kohya%20LoRA%20Trainer%20XL&label=Visitors&labelColor=%2334495E&countColor=%231ABC9C&style=flat&labelStyle=none\n",
        "[visitor-stats]: https://visitorbadge.io/status?path=Kohya%20LoRA%20Trainer%20XL\n",
        "[ko-fi-badge]: https://img.shields.io/badge/Support%20me%20on%20Ko--fi-F16061?logo=ko-fi&logoColor=white&style=flat\n",
        "[ko-fi-link]: https://ko-fi.com/linaqruf"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_MxF9feWshAp"
      },
      "source": [
        "| Notebook Name | Description | Link |\n",
        "| --- | --- | --- |\n",
        "| [Kohya LoRA Trainer XL](https://github.com/Linaqruf/kohya-trainer/blob/main/kohya-LoRA-trainer-XL.ipynb) | LoRA Training | [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/kohya-LoRA-trainer-XL.ipynb) |\n",
        "| [Kohya Trainer XL](https://github.com/Linaqruf/kohya-trainer/blob/main/kohya-trainer-XL.ipynb) | Native Training | [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/kohya-trainer-XL.ipynb) |\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tTVqCAgSmie4"
      },
      "source": [
        "# **I. Prepare Environment**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "v6uDIVmL03SC",
        "outputId": "6918c873-f24c-4a58-fcfb-bdae5c03d928"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mon Feb  5 02:37:41 2024       \n",
            "+---------------------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 535.104.05             Driver Version: 535.104.05   CUDA Version: 12.2     |\n",
            "|-----------------------------------------+----------------------+----------------------+\n",
            "| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                                         |                      |               MIG M. |\n",
            "|=========================================+======================+======================|\n",
            "|   0  NVIDIA A100-SXM4-40GB          Off | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   32C    P0              44W / 400W |      2MiB / 40960MiB |      0%      Default |\n",
            "|                                         |                      |             Disabled |\n",
            "+-----------------------------------------+----------------------+----------------------+\n",
            "                                                                                         \n",
            "+---------------------------------------------------------------------------------------+\n",
            "| Processes:                                                                            |\n",
            "|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |\n",
            "|        ID   ID                                                             Usage      |\n",
            "|=======================================================================================|\n",
            "|  No running processes found                                                           |\n",
            "+---------------------------------------------------------------------------------------+\n"
          ]
        }
      ],
      "source": [
        "# You can run this notebook in Tesla T4 but make sure you're using High RAM setting. (Colab Pro gated)\n",
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_u3q60di584x",
        "outputId": "e9b1bf32-b75f-4937-db70-73409d328158"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into '/content/kohya-trainer'...\n",
            "remote: Enumerating objects: 2441, done.\u001b[K\n",
            "remote: Counting objects: 100% (2441/2441), done.\u001b[K\n",
            "remote: Compressing objects: 100% (941/941), done.\u001b[K\n",
            "remote: Total 2441 (delta 1633), reused 2168 (delta 1497), pack-reused 0\u001b[K\n",
            "Receiving objects: 100% (2441/2441), 4.12 MiB | 17.27 MiB/s, done.\n",
            "Resolving deltas: 100% (1633/1633), done.\n",
            "Mounted at /content/drive\n",
            "Cloning into '/content/repositories/infinite-image-browsing'...\n",
            "remote: Enumerating objects: 8404, done.\u001b[K\n",
            "remote: Counting objects: 100% (2800/2800), done.\u001b[K\n",
            "remote: Compressing objects: 100% (798/798), done.\u001b[K\n",
            "remote: Total 8404 (delta 2160), reused 2497 (delta 1946), pack-reused 5604\u001b[K\n",
            "Receiving objects: 100% (8404/8404), 18.84 MiB | 5.72 MiB/s, done.\n",
            "Resolving deltas: 100% (6044/6044), done.\n",
            "Cloning into '/content/repositories/discordia-archivum'...\n",
            "remote: Enumerating objects: 54, done.\u001b[K\n",
            "remote: Counting objects: 100% (54/54), done.\u001b[K\n",
            "remote: Compressing objects: 100% (45/45), done.\u001b[K\n",
            "remote: Total 54 (delta 26), reused 21 (delta 7), pack-reused 0\u001b[K\n",
            "Receiving objects: 100% (54/54), 16.04 KiB | 16.04 MiB/s, done.\n",
            "Resolving deltas: 100% (26/26), done.\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m92.1/92.1 kB\u001b[0m \u001b[31m2.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m60.7/60.7 kB\u001b[0m \u001b[31m5.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.5/4.5 MB\u001b[0m \u001b[31m99.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m49.5/49.5 kB\u001b[0m \u001b[31m7.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m71.5/71.5 kB\u001b[0m \u001b[31m10.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m58.3/58.3 kB\u001b[0m \u001b[31m9.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m116.3/116.3 kB\u001b[0m \u001b[31m17.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "lida 0.0.10 requires kaleido, which is not installed.\n",
            "lida 0.0.10 requires python-multipart, which is not installed.\n",
            "imageio 2.31.6 requires pillow<10.1.0,>=8.3.2, but you have pillow 10.2.0 which is incompatible.\n",
            "tensorflow-probability 0.22.0 requires typing-extensions<4.6.0, but you have typing-extensions 4.9.0 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[0mRequirement already satisfied: python-dotenv in /usr/local/lib/python3.10/dist-packages (1.0.1)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.1/1.1 MB\u001b[0m \u001b[31m21.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m75.4/75.4 kB\u001b[0m \u001b[31m11.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m74.5/74.5 kB\u001b[0m \u001b[31m11.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Building wheel for discord-protos (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "Reading package lists... Done\n",
            "Building dependency tree... Done\n",
            "Reading state information... Done\n",
            "The following additional packages will be installed:\n",
            "  libaria2-0 libc-ares2\n",
            "The following NEW packages will be installed:\n",
            "  aria2 libaria2-0 libc-ares2 lz4\n",
            "0 upgraded, 4 newly installed, 0 to remove and 32 not upgraded.\n",
            "Need to get 1,603 kB of archives.\n",
            "After this operation, 5,676 kB of additional disk space will be used.\n",
            "Get:1 http://archive.ubuntu.com/ubuntu jammy-updates/main amd64 libc-ares2 amd64 1.18.1-1ubuntu0.22.04.2 [45.0 kB]\n",
            "Get:2 http://archive.ubuntu.com/ubuntu jammy/universe amd64 libaria2-0 amd64 1.36.0-1 [1,086 kB]\n",
            "Get:3 http://archive.ubuntu.com/ubuntu jammy/universe amd64 aria2 amd64 1.36.0-1 [381 kB]\n",
            "Get:4 http://archive.ubuntu.com/ubuntu jammy/main amd64 lz4 amd64 1.9.3-2build2 [90.0 kB]\n",
            "Fetched 1,603 kB in 2s (866 kB/s)\n",
            "Selecting previously unselected package libc-ares2:amd64.\n",
            "(Reading database ... 121730 files and directories currently installed.)\n",
            "Preparing to unpack .../libc-ares2_1.18.1-1ubuntu0.22.04.2_amd64.deb ...\n",
            "Unpacking libc-ares2:amd64 (1.18.1-1ubuntu0.22.04.2) ...\n",
            "Selecting previously unselected package libaria2-0:amd64.\n",
            "Preparing to unpack .../libaria2-0_1.36.0-1_amd64.deb ...\n",
            "Unpacking libaria2-0:amd64 (1.36.0-1) ...\n",
            "Selecting previously unselected package aria2.\n",
            "Preparing to unpack .../aria2_1.36.0-1_amd64.deb ...\n",
            "Unpacking aria2 (1.36.0-1) ...\n",
            "Selecting previously unselected package lz4.\n",
            "Preparing to unpack .../lz4_1.9.3-2build2_amd64.deb ...\n",
            "Unpacking lz4 (1.9.3-2build2) ...\n",
            "Setting up libc-ares2:amd64 (1.18.1-1ubuntu0.22.04.2) ...\n",
            "Setting up lz4 (1.9.3-2build2) ...\n",
            "Setting up libaria2-0:amd64 (1.36.0-1) ...\n",
            "Setting up aria2 (1.36.0-1) ...\n",
            "Processing triggers for man-db (2.10.2-1) ...\n",
            "Processing triggers for libc-bin (2.35-0ubuntu3.4) ...\n",
            "/sbin/ldconfig.real: /usr/local/lib/libtbbmalloc_proxy.so.2 is not a symbolic link\n",
            "\n",
            "/sbin/ldconfig.real: /usr/local/lib/libtbbbind.so.3 is not a symbolic link\n",
            "\n",
            "/sbin/ldconfig.real: /usr/local/lib/libtbb.so.12 is not a symbolic link\n",
            "\n",
            "/sbin/ldconfig.real: /usr/local/lib/libtbbmalloc.so.2 is not a symbolic link\n",
            "\n",
            "/sbin/ldconfig.real: /usr/local/lib/libtbbbind_2_0.so.3 is not a symbolic link\n",
            "\n",
            "/sbin/ldconfig.real: /usr/local/lib/libtbbbind_2_5.so.3 is not a symbolic link\n",
            "\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m219.1/219.1 kB\u001b[0m \u001b[31m5.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m495.8/495.8 kB\u001b[0m \u001b[31m37.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.2/7.2 MB\u001b[0m \u001b[31m97.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.2/1.2 MB\u001b[0m \u001b[31m74.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m53.1/53.1 kB\u001b[0m \u001b[31m8.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m61.8/61.8 MB\u001b[0m \u001b[31m26.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m41.6/41.6 kB\u001b[0m \u001b[31m5.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m825.8/825.8 kB\u001b[0m \u001b[31m68.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m97.1/97.1 MB\u001b[0m \u001b[31m16.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m79.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m236.8/236.8 kB\u001b[0m \u001b[31m32.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.1/2.1 MB\u001b[0m \u001b[31m94.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.6/1.6 MB\u001b[0m \u001b[31m89.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.5/1.5 MB\u001b[0m \u001b[31m66.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m624.7/624.7 kB\u001b[0m \u001b[31m57.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m68.8/68.8 kB\u001b[0m \u001b[31m9.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m13.6/13.6 MB\u001b[0m \u001b[31m79.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m235.4/235.4 kB\u001b[0m \u001b[31m32.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Installing backend dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m42.2/42.2 kB\u001b[0m \u001b[31m6.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m79.5/79.5 kB\u001b[0m \u001b[31m11.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m47.8/47.8 MB\u001b[0m \u001b[31m29.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m204.0/204.0 kB\u001b[0m \u001b[31m28.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m104.3/104.3 MB\u001b[0m \u001b[31m15.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m8.4/8.4 MB\u001b[0m \u001b[31m110.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m377.0/377.0 kB\u001b[0m \u001b[31m44.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m70.4/70.4 kB\u001b[0m \u001b[31m9.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.8/7.8 MB\u001b[0m \u001b[31m105.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m840.2/840.2 kB\u001b[0m \u001b[31m68.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m196.4/196.4 kB\u001b[0m \u001b[31m25.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m257.5/257.5 kB\u001b[0m \u001b[31m33.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.8/1.8 MB\u001b[0m \u001b[31m93.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m247.3/247.3 kB\u001b[0m \u001b[31m32.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m43.1/43.1 kB\u001b[0m \u001b[31m6.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m56.7/56.7 kB\u001b[0m \u001b[31m9.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m62.7/62.7 kB\u001b[0m \u001b[31m10.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.6/1.6 MB\u001b[0m \u001b[31m86.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m117.0/117.0 kB\u001b[0m \u001b[31m18.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m98.9/98.9 kB\u001b[0m \u001b[31m14.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.8/3.8 MB\u001b[0m \u001b[31m104.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m115.0/115.0 kB\u001b[0m \u001b[31m17.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m247.1/247.1 kB\u001b[0m \u001b[31m32.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.8/4.8 MB\u001b[0m \u001b[31m113.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m82.1/82.1 kB\u001b[0m \u001b[31m12.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.6/3.6 MB\u001b[0m \u001b[31m100.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m468.9/468.9 kB\u001b[0m \u001b[31m52.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Building wheel for dadaptation (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for lycoris-lora (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for fairscale (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for contexttimer (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for elfinder-client (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for iopath (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for antlr4-python3-runtime (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for pathtools (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.3/2.3 GB\u001b[0m \u001b[31m894.0 kB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m6.1/6.1 MB\u001b[0m \u001b[31m101.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.4/4.4 MB\u001b[0m \u001b[31m53.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.0/2.0 MB\u001b[0m \u001b[31m83.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.6/4.6 MB\u001b[0m \u001b[31m101.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m63.3/63.3 MB\u001b[0m \u001b[31m25.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m153.0/153.0 kB\u001b[0m \u001b[31m4.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h  Installing build dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Getting requirements to build wheel ... \u001b[?25l\u001b[?25hdone\n",
            "  Installing backend dependencies ... \u001b[?25l\u001b[?25hdone\n",
            "  Preparing metadata (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "  Building wheel for lit (pyproject.toml) ... \u001b[?25l\u001b[?25hdone\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m108.2/108.2 MB\u001b[0m \u001b[31m14.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h"
          ]
        }
      ],
      "source": [
        "# @title ## **1.1. Install Kohya Trainer**\n",
        "import os\n",
        "import zipfile\n",
        "import shutil\n",
        "import time\n",
        "import requests\n",
        "import torch\n",
        "from subprocess import getoutput\n",
        "from IPython.utils import capture\n",
        "from google.colab import drive\n",
        "\n",
        "%store -r\n",
        "\n",
        "# root_dir\n",
        "root_dir          = \"/content\"\n",
        "drive_dir         = os.path.join(root_dir, \"drive/MyDrive\")\n",
        "deps_dir          = os.path.join(root_dir, \"deps\")\n",
        "repo_dir          = os.path.join(root_dir, \"kohya-trainer\")\n",
        "training_dir      = os.path.join(root_dir, \"LoRA\")\n",
        "pretrained_model  = os.path.join(root_dir, \"pretrained_model\")\n",
        "vae_dir           = os.path.join(root_dir, \"vae\")\n",
        "lora_dir          = os.path.join(root_dir, \"network_weight\")\n",
        "repositories_dir  = os.path.join(root_dir, \"repositories\")\n",
        "config_dir        = os.path.join(training_dir, \"config\")\n",
        "tools_dir         = os.path.join(repo_dir, \"tools\")\n",
        "finetune_dir      = os.path.join(repo_dir, \"finetune\")\n",
        "accelerate_config = os.path.join(repo_dir, \"accelerate_config/config.yaml\")\n",
        "\n",
        "for store in [\"root_dir\", \"repo_dir\", \"training_dir\", \"pretrained_model\", \"vae_dir\", \"repositories_dir\", \"accelerate_config\", \"tools_dir\", \"finetune_dir\", \"config_dir\"]:\n",
        "    with capture.capture_output() as cap:\n",
        "        %store {store}\n",
        "        del cap\n",
        "\n",
        "repo_dict = {\n",
        "    \"qaneel/kohya-trainer (forked repo, stable, optimized for colab use)\" : \"https://github.com/qaneel/kohya-trainer\",\n",
        "    \"kohya-ss/sd-scripts (original repo, latest update)\"                    : \"https://github.com/kohya-ss/sd-scripts\",\n",
        "}\n",
        "\n",
        "repository        = \"qaneel/kohya-trainer (forked repo, stable, optimized for colab use)\" #@param [\"qaneel/kohya-trainer (forked repo, stable, optimized for colab use)\", \"kohya-ss/sd-scripts (original repo, latest update)\"] {allow-input: true}\n",
        "repo_url          = repo_dict[repository]\n",
        "branch            = \"main\"  # @param {type: \"string\"}\n",
        "output_to_drive   = True  # @param {type: \"boolean\"}\n",
        "\n",
        "def clone_repo(url, dir, branch):\n",
        "    if not os.path.exists(dir):\n",
        "       !git clone -b {branch} {url} {dir}\n",
        "\n",
        "def mount_drive(dir):\n",
        "    output_dir      = os.path.join(training_dir, \"output\")\n",
        "\n",
        "    if output_to_drive:\n",
        "        if not os.path.exists(drive_dir):\n",
        "            drive.mount(os.path.dirname(drive_dir))\n",
        "        output_dir  = os.path.join(drive_dir, \"kohya-trainer/output\")\n",
        "\n",
        "    return output_dir\n",
        "\n",
        "def setup_directories():\n",
        "    global output_dir\n",
        "\n",
        "    output_dir      = mount_drive(drive_dir)\n",
        "\n",
        "    for dir in [training_dir, config_dir, pretrained_model, vae_dir, repositories_dir, output_dir]:\n",
        "        os.makedirs(dir, exist_ok=True)\n",
        "\n",
        "def pastebin_reader(id):\n",
        "    if \"pastebin.com\" in id:\n",
        "        url = id\n",
        "        if 'raw' not in url:\n",
        "                url = url.replace('pastebin.com', 'pastebin.com/raw')\n",
        "    else:\n",
        "        url = \"https://pastebin.com/raw/\" + id\n",
        "    response = requests.get(url)\n",
        "    response.raise_for_status()\n",
        "    lines = response.text.split('\\n')\n",
        "    return lines\n",
        "\n",
        "def install_repository():\n",
        "    global infinite_image_browser_dir, voldy, discordia_archivum_dir\n",
        "\n",
        "    _, voldy = pastebin_reader(\"kq6ZmHFU\")[:2]\n",
        "\n",
        "    infinite_image_browser_url  = f\"https://github.com/zanllp/{voldy}-infinite-image-browsing.git\"\n",
        "    infinite_image_browser_dir  = os.path.join(repositories_dir, f\"infinite-image-browsing\")\n",
        "    infinite_image_browser_deps = os.path.join(infinite_image_browser_dir, \"requirements.txt\")\n",
        "\n",
        "    discordia_archivum_url = \"https://github.com/Linaqruf/discordia-archivum\"\n",
        "    discordia_archivum_dir = os.path.join(repositories_dir, \"discordia-archivum\")\n",
        "    discordia_archivum_deps = os.path.join(discordia_archivum_dir, \"requirements.txt\")\n",
        "\n",
        "    clone_repo(infinite_image_browser_url, infinite_image_browser_dir, \"main\")\n",
        "    clone_repo(discordia_archivum_url, discordia_archivum_dir, \"main\")\n",
        "\n",
        "    !pip install -q --upgrade -r {infinite_image_browser_deps}\n",
        "    !pip install python-dotenv\n",
        "    !pip install -q --upgrade -r {discordia_archivum_deps}\n",
        "\n",
        "def install_dependencies():\n",
        "    requirements_file = os.path.join(repo_dir, \"requirements.txt\")\n",
        "    model_util        = os.path.join(repo_dir, \"library/model_util.py\")\n",
        "    gpu_info          = getoutput('nvidia-smi')\n",
        "    t4_xformers_wheel = \"https://github.com/Linaqruf/colab-xformers/releases/download/0.0.20/xformers-0.0.20+1d635e1.d20230519-cp310-cp310-linux_x86_64.whl\"\n",
        "\n",
        "    !apt install aria2 lz4\n",
        "    !pip install -q --upgrade -r {requirements_file}\n",
        "\n",
        "    if '2.0.1+cu118' in torch.__version__:\n",
        "        if 'T4' in gpu_info:\n",
        "            !pip install -q {t4_xformers_wheel}\n",
        "        else:\n",
        "            !pip install -q xformers==0.0.20\n",
        "    else:\n",
        "        !pip install -q torch==2.0.0+cu118 torchvision==0.15.1+cu118 torchaudio==2.0.1+cu118 torchtext==0.15.1 torchdata==0.6.0 --extra-index-url https://download.pytorch.org/whl/cu118 -U\n",
        "        !pip install -q xformers==0.0.19 triton==2.0.0 -U\n",
        "\n",
        "    from accelerate.utils import write_basic_config\n",
        "\n",
        "    if not os.path.exists(accelerate_config):\n",
        "        write_basic_config(save_location=accelerate_config)\n",
        "\n",
        "def prepare_environment():\n",
        "    os.environ[\"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",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "wrYGu-WxFbsq",
        "outputId": "d70f284b-a728-45b7-dd3f-733c5a290b6d"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Selected model: /content/pretrained_model/sd_xl_base_1.0.safetensors\n",
            "Selected vae: /content/vae/sdxl_vae.safetensors\n"
          ]
        }
      ],
      "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-1.0/resolve/main/sd_xl_base_1.0.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",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "kh7CeDqK4l3Y",
        "outputId": "7d4f849a-e490-41ca-e922-843cd12f109f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Stored 'train_data_dir' (str)\n",
            "Your train data directory : /content/drive/MyDrive/24-datasets/ZenkaiXL\n"
          ]
        }
      ],
      "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/drive/MyDrive/yourdataset\"  # @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",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hhgatqF3leHJ",
        "outputId": "cca111f6-7039-4004-db38-eb74e8089f4f"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Found 381 images.\n",
            "Creating a new metadata file\n",
            "Merging tags and captions into metadata json.\n",
            "100% 381/381 [04:58<00:00,  1.28it/s]\n",
            "No captions found for any of the 381 images\n",
            "All 381 images have tags\n",
            "Writing metadata: /content/LoRA/meta_clean.json\n",
            "Done!\n",
            "found 381 images.\n",
            "loading existing metadata: /content/LoRA/meta_clean.json\n",
            "load VAE: /content/pretrained_model/sd_xl_base_1.0.safetensors\n",
            "100% 381/381 [04:21<00:00,  1.46it/s]\n",
            "bucket 0 (1024, 1024): 381\n",
            "mean ar error: 0.0\n",
            "writing metadata: /content/LoRA/meta_lat.json\n",
            "done!\n"
          ]
        }
      ],
      "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     = False #@param {type:\"boolean\"}\n",
        "#@markdown Use the `recursive` option to process subfolders as well\n",
        "recursive         = True #@param {type:\"boolean\"}\n",
        "\n",
        "metadata_config = {\n",
        "    \"_train_data_dir\": train_data_dir,\n",
        "    \"_out_json\": metadata_json,\n",
        "    \"recursive\": recursive,\n",
        "    \"full_path\": recursive,\n",
        "    \"clean_caption\": clean_caption\n",
        "}\n",
        "\n",
        "bucketing_config = {\n",
        "    \"_train_data_dir\": train_data_dir,\n",
        "    \"_in_json\": metadata_json,\n",
        "    \"_out_json\": bucketing_json,\n",
        "    \"_model_name_or_path\": 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": "GAZVkLuaRJ9e"
      },
      "source": [
        "# **IV. Training**\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "cJgLfRtlHSjw",
        "outputId": "07a3b3b3-d674-441e-8a1f-98efae9ef94b"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[additional_network_arguments]\n",
            "no_metadata = false\n",
            "network_module = \"networks.lora\"\n",
            "network_dim = 16\n",
            "network_alpha = 8\n",
            "network_args = [ \"conv_dim=8\", \"conv_alpha=1\",]\n",
            "network_train_unet_only = true\n",
            "\n"
          ]
        }
      ],
      "source": [
        "import toml\n",
        "\n",
        "# @title ## **4.1. LoRa: Low-Rank Adaptation Config**\n",
        "# @markdown Kohya's `LoRA` renamed to `LoRA-LierLa` and Kohya's `LoCon` renamed to `LoRA-C3Lier`, read [official announcement](https://github.com/kohya-ss/sd-scripts/blob/849bc24d205a35fbe1b2a4063edd7172533c1c01/README.md#naming-of-lora).\n",
        "network_category = \"LoRA_C3Lier\"  # @param [\"LoRA_LierLa\", \"LoRA_C3Lier\", \"DyLoRA_LierLa\", \"DyLoRA_C3Lier\", \"LoCon\", \"LoHa\", \"IA3\", \"LoKR\", \"DyLoRA_Lycoris\"]\n",
        "\n",
        "# @markdown | network_category | network_dim | network_alpha | conv_dim | conv_alpha | unit |\n",
        "# @markdown | :---: | :---: | :---: | :---: | :---: | :---: |\n",
        "# @markdown | LoRA-LierLa | 32 | 1 | - | - | - |\n",
        "# @markdown | LoCon/LoRA-C3Lier | 16 | 8 | 8 | 1 | - |\n",
        "# @markdown | LoHa | 8 | 4 | 4 | 1 | - |\n",
        "# @markdown | Other Category | ? | ? | ? | ? | - |\n",
        "\n",
        "# @markdown Specify `network_args` to add `optional` training args, like for specifying each 25 block weight, read [this](https://github.com/kohya-ss/sd-scripts/blob/main/train_network_README-ja.md#%E9%9A%8E%E5%B1%A4%E5%88%A5%E5%AD%A6%E7%BF%92%E7%8E%87)\n",
        "network_args    = \"\"  # @param {'type':'string'}\n",
        "\n",
        "# @markdown ### **Linear Layer Config**\n",
        "# @markdown Used by all `network_category`. When in doubt, set `network_dim = network_alpha`\n",
        "network_dim     = 16  # @param {'type':'number'}\n",
        "network_alpha   = 8  # @param {'type':'number'}\n",
        "\n",
        "# @markdown ### **Convolutional Layer Config**\n",
        "# @markdown Only required if `network_category` is not `LoRA_LierLa`, as it involves training convolutional layers in addition to linear layers.\n",
        "conv_dim        = 8  # @param {'type':'number'}\n",
        "conv_alpha      = 1  # @param {'type':'number'}\n",
        "\n",
        "# @markdown ### **DyLoRA Config**\n",
        "# @markdown Only required if `network_category` is `DyLoRA_LierLa` and `DyLoRA_C3Lier`\n",
        "unit = 4  # @param {'type':'number'}\n",
        "\n",
        "if isinstance(network_args, str):\n",
        "    network_args = network_args.strip()\n",
        "    if network_args.startswith('[') and network_args.endswith(']'):\n",
        "        try:\n",
        "            network_args = ast.literal_eval(network_args)\n",
        "        except (SyntaxError, ValueError) as e:\n",
        "            print(f\"Error parsing network_args: {e}\\n\")\n",
        "            network_args = []\n",
        "    elif len(network_args) > 0:\n",
        "        print(f\"WARNING! '{network_args}' is not a valid list! Put args like this: [\\\"args=1\\\", \\\"args=2\\\"]\\n\")\n",
        "        network_args = []\n",
        "    else:\n",
        "        network_args = []\n",
        "else:\n",
        "    network_args = []\n",
        "\n",
        "network_config = {\n",
        "    \"LoRA_LierLa\": {\n",
        "        \"module\": \"networks.lora\",\n",
        "        \"args\"  : []\n",
        "    },\n",
        "    \"LoRA_C3Lier\": {\n",
        "        \"module\": \"networks.lora\",\n",
        "        \"args\"  : [\n",
        "            f\"conv_dim={conv_dim}\",\n",
        "            f\"conv_alpha={conv_alpha}\"\n",
        "        ]\n",
        "    },\n",
        "    \"DyLoRA_LierLa\": {\n",
        "        \"module\": \"networks.dylora\",\n",
        "        \"args\"  : [\n",
        "            f\"unit={unit}\"\n",
        "        ]\n",
        "    },\n",
        "    \"DyLoRA_C3Lier\": {\n",
        "        \"module\": \"networks.dylora\",\n",
        "        \"args\"  : [\n",
        "            f\"conv_dim={conv_dim}\",\n",
        "            f\"conv_alpha={conv_alpha}\",\n",
        "            f\"unit={unit}\"\n",
        "        ]\n",
        "    },\n",
        "    \"LoCon\": {\n",
        "        \"module\": \"lycoris.kohya\",\n",
        "        \"args\"  : [\n",
        "            f\"algo=locon\",\n",
        "            f\"conv_dim={conv_dim}\",\n",
        "            f\"conv_alpha={conv_alpha}\"\n",
        "        ]\n",
        "    },\n",
        "    \"LoHa\": {\n",
        "        \"module\": \"lycoris.kohya\",\n",
        "        \"args\"  : [\n",
        "            f\"algo=loha\",\n",
        "            f\"conv_dim={conv_dim}\",\n",
        "            f\"conv_alpha={conv_alpha}\"\n",
        "        ]\n",
        "    },\n",
        "    \"IA3\": {\n",
        "        \"module\": \"lycoris.kohya\",\n",
        "        \"args\"  : [\n",
        "            f\"algo=ia3\",\n",
        "            f\"conv_dim={conv_dim}\",\n",
        "            f\"conv_alpha={conv_alpha}\"\n",
        "        ]\n",
        "    },\n",
        "    \"LoKR\": {\n",
        "        \"module\": \"lycoris.kohya\",\n",
        "        \"args\"  : [\n",
        "            f\"algo=lokr\",\n",
        "            f\"conv_dim={conv_dim}\",\n",
        "            f\"conv_alpha={conv_alpha}\"\n",
        "        ]\n",
        "    },\n",
        "    \"DyLoRA_Lycoris\": {\n",
        "        \"module\": \"lycoris.kohya\",\n",
        "        \"args\"  : [\n",
        "            f\"algo=dylora\",\n",
        "            f\"conv_dim={conv_dim}\",\n",
        "            f\"conv_alpha={conv_alpha}\"\n",
        "        ]\n",
        "    }\n",
        "}\n",
        "\n",
        "network_module = network_config[network_category][\"module\"]\n",
        "network_args.extend(network_config[network_category][\"args\"])\n",
        "\n",
        "lora_config = {\n",
        "    \"additional_network_arguments\": {\n",
        "        \"no_metadata\"                     : False,\n",
        "        \"network_module\"                  : network_module,\n",
        "        \"network_dim\"                     : network_dim,\n",
        "        \"network_alpha\"                   : network_alpha,\n",
        "        \"network_args\"                    : network_args,\n",
        "        \"network_train_unet_only\"         : True,\n",
        "        \"training_comment\"                : None,\n",
        "    },\n",
        "}\n",
        "\n",
        "print(toml.dumps(lora_config))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JNlw3u8arwir",
        "outputId": "43a9b9d4-e500-4cd6-cf31-510b309b1251"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[optimizer_arguments]\n",
            "optimizer_type = \"AdaFactor\"\n",
            "learning_rate = 0.003\n",
            "max_grad_norm = 0\n",
            "optimizer_args = [ \"scale_parameter=False\", \"relative_step=False\", \"warmup_init=False\",]\n",
            "lr_scheduler = \"constant_with_warmup\"\n",
            "lr_warmup_steps = 100\n",
            "\n"
          ]
        }
      ],
      "source": [
        "import toml\n",
        "import ast\n",
        "\n",
        "# @title ## **4.2. Optimizer Config**\n",
        "# @markdown Use `Adafactor` optimizer. `RMSprop 8bit` or `Adagrad 8bit` may work. `AdamW 8bit` doesn't seem to work.\n",
        "optimizer_type = \"AdaFactor\"  # @param [\"AdamW\", \"AdamW8bit\", \"Lion8bit\", \"Lion\", \"SGDNesterov\", \"SGDNesterov8bit\", \"DAdaptation(DAdaptAdamPreprint)\", \"DAdaptAdaGrad\", \"DAdaptAdam\", \"DAdaptAdan\", \"DAdaptAdanIP\", \"DAdaptLion\", \"DAdaptSGD\", \"AdaFactor\"]\n",
        "# @markdown Specify `optimizer_args` to add `additional` args for optimizer, e.g: `[\"weight_decay=0.6\"]`\n",
        "optimizer_args = \"[ \\\"scale_parameter=False\\\", \\\"relative_step=False\\\", \\\"warmup_init=False\\\" ]\"  # @param {'type':'string'}\n",
        "# @markdown ### **Learning Rate Config**\n",
        "# @markdown Different `optimizer_type` and `network_category` for some condition requires different learning rate. It's recommended to set `text_encoder_lr = 1/2 * unet_lr`\n",
        "learning_rate = 3e-3  # @param {'type':'number'}\n",
        "# @markdown ### **LR Scheduler Config**\n",
        "# @markdown `lr_scheduler` provides several methods to adjust the learning rate based on the number of epochs.\n",
        "lr_scheduler = \"constant_with_warmup\"  # @param [\"linear\", \"cosine\", \"cosine_with_restarts\", \"polynomial\", \"constant\", \"constant_with_warmup\", \"adafactor\"] {allow-input: false}\n",
        "lr_warmup_steps = 100  # @param {'type':'number'}\n",
        "# @markdown Specify `lr_scheduler_num` with `num_cycles` value for `cosine_with_restarts` or `power` value for `polynomial`\n",
        "lr_scheduler_num = 0  # @param {'type':'number'}\n",
        "\n",
        "if isinstance(optimizer_args, str):\n",
        "    optimizer_args = optimizer_args.strip()\n",
        "    if optimizer_args.startswith('[') and optimizer_args.endswith(']'):\n",
        "        try:\n",
        "            optimizer_args = ast.literal_eval(optimizer_args)\n",
        "        except (SyntaxError, ValueError) as e:\n",
        "            print(f\"Error parsing optimizer_args: {e}\\n\")\n",
        "            optimizer_args = []\n",
        "    elif len(optimizer_args) > 0:\n",
        "        print(f\"WARNING! '{optimizer_args}' is not a valid list! Put args like this: [\\\"args=1\\\", \\\"args=2\\\"]\\n\")\n",
        "        optimizer_args = []\n",
        "    else:\n",
        "        optimizer_args = []\n",
        "else:\n",
        "    optimizer_args = []\n",
        "\n",
        "optimizer_config = {\n",
        "    \"optimizer_arguments\": {\n",
        "        \"optimizer_type\"          : optimizer_type,\n",
        "        \"learning_rate\"           : learning_rate,\n",
        "        \"max_grad_norm\"           : 0,\n",
        "        \"optimizer_args\"          : optimizer_args,\n",
        "        \"lr_scheduler\"            : lr_scheduler,\n",
        "        \"lr_warmup_steps\"         : lr_warmup_steps,\n",
        "        \"lr_scheduler_num_cycles\" : lr_scheduler_num if lr_scheduler == \"cosine_with_restarts\" else None,\n",
        "        \"lr_scheduler_power\"      : lr_scheduler_num if lr_scheduler == \"polynomial\" else None,\n",
        "        \"lr_scheduler_type\"       : None,\n",
        "        \"lr_scheduler_args\"       : None,\n",
        "    },\n",
        "}\n",
        "\n",
        "print(toml.dumps(optimizer_config))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "YOxNM0x7dvfO",
        "outputId": "feecfc75-a8ff-4a72-fe3c-e1634284db49"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[advanced_training_config]\n",
            "multires_noise_iterations = 6\n",
            "multires_noise_discount = 0.3\n",
            "min_snr_gamma = 3\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# @title ## **4.3. Advanced Training Config** (Optional)\n",
        "import toml\n",
        "\n",
        "# @markdown ### **Noise Control**\n",
        "noise_control_type        = \"multires_noise\" #@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             = 3 #@param {type:\"number\"}\n",
        "\n",
        "advanced_training_config = {\n",
        "    \"advanced_training_config\": {\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",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-Z4w3lfFKLjr",
        "outputId": "e5f8127e-7702-46da-8dd6-9b64ded92587"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[sdxl_arguments]\n",
            "cache_text_encoder_outputs = true\n",
            "no_half_vae = true\n",
            "min_timestep = 0\n",
            "max_timestep = 1000\n",
            "shuffle_caption = false\n",
            "\n",
            "[model_arguments]\n",
            "pretrained_model_name_or_path = \"/content/pretrained_model/sd_xl_base_1.0.safetensors\"\n",
            "vae = \"/content/vae/sdxl_vae.safetensors\"\n",
            "\n",
            "[dataset_arguments]\n",
            "debug_dataset = false\n",
            "in_json = \"/content/LoRA/meta_lat.json\"\n",
            "train_data_dir = \"/content/drive/MyDrive/24-datasets/ZenkaiXL\"\n",
            "dataset_repeats = 30\n",
            "keep_tokens = 0\n",
            "resolution = \"1024,1024\"\n",
            "color_aug = false\n",
            "token_warmup_min = 1\n",
            "token_warmup_step = 0\n",
            "\n",
            "[training_arguments]\n",
            "output_dir = \"/content/drive/MyDrive/kohya-trainer/output\"\n",
            "output_name = \"ZenkaiXL-v0.2\"\n",
            "save_precision = \"fp16\"\n",
            "save_every_n_epochs = 1\n",
            "train_batch_size = 30\n",
            "max_token_length = 225\n",
            "mem_eff_attn = false\n",
            "sdpa = true\n",
            "xformers = false\n",
            "max_train_epochs = 2\n",
            "max_data_loader_n_workers = 8\n",
            "persistent_data_loader_workers = true\n",
            "gradient_checkpointing = true\n",
            "gradient_accumulation_steps = 1\n",
            "mixed_precision = \"fp16\"\n",
            "\n",
            "[logging_arguments]\n",
            "log_with = \"tensorboard\"\n",
            "logging_dir = \"/content/LoRA/logs\"\n",
            "log_prefix = \"ZenkaiXL-v0.2\"\n",
            "\n",
            "[sample_prompt_arguments]\n",
            "sample_sampler = \"euler_a\"\n",
            "\n",
            "[saving_arguments]\n",
            "save_model_as = \"safetensors\"\n",
            "\n",
            "[optimizer_arguments]\n",
            "optimizer_type = \"AdaFactor\"\n",
            "learning_rate = 0.003\n",
            "max_grad_norm = 0\n",
            "optimizer_args = [ \"scale_parameter=False\", \"relative_step=False\", \"warmup_init=False\",]\n",
            "lr_scheduler = \"constant_with_warmup\"\n",
            "lr_warmup_steps = 100\n",
            "\n",
            "[additional_network_arguments]\n",
            "no_metadata = false\n",
            "network_module = \"networks.lora\"\n",
            "network_dim = 16\n",
            "network_alpha = 8\n",
            "network_args = [ \"conv_dim=8\", \"conv_alpha=1\",]\n",
            "network_train_unet_only = true\n",
            "\n",
            "[advanced_training_config]\n",
            "multires_noise_iterations = 6\n",
            "multires_noise_discount = 0.3\n",
            "min_snr_gamma = 3\n",
            "\n",
            "[prompt]\n",
            "width = 1024\n",
            "height = 1024\n",
            "scale = 7\n",
            "sample_steps = 28\n",
            "subset = []\n",
            "\n"
          ]
        }
      ],
      "source": [
        "# @title ## **4.4. Training Config**\n",
        "import toml\n",
        "import os\n",
        "from subprocess import getoutput\n",
        "\n",
        "%store -r\n",
        "\n",
        "# @markdown ### **Project Config**\n",
        "project_name                = \"ENTER-PROJECT-NAME-HERE\"  # @param {type:\"string\"}\n",
        "# @markdown Get your `wandb_api_key` [here](https://wandb.ai/settings) to logs with wandb.\n",
        "wandb_api_key               = \"\" # @param {type:\"string\"}\n",
        "in_json                     = \"/content/LoRA/meta_lat.json\"  # @param {type:\"string\"}\n",
        "# @markdown ### **SDXL Config**\n",
        "gradient_checkpointing      = True  # @param {type:\"boolean\"}\n",
        "no_half_vae                 = True  # @param {type:\"boolean\"}\n",
        "#@markdown Recommended parameter for SDXL training but if you enable it, `shuffle_caption` won't work\n",
        "cache_text_encoder_outputs  = 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                 = 30  # @param {type:\"number\"}\n",
        "resolution                  = 1024  # @param {type:\"slider\", min:512, max:1024, step:128}\n",
        "keep_tokens                 = 0  # @param {type:\"number\"}\n",
        "# @markdown ### **General Config**\n",
        "num_epochs                  = 2  # @param {type:\"number\"}\n",
        "train_batch_size            = 30  # @param {type:\"number\"}\n",
        "mixed_precision             = \"fp16\"  # @param [\"no\",\"fp16\",\"bf16\"] {allow-input: false}\n",
        "seed                        = -1  # @param {type:\"number\"}\n",
        "optimization                = \"scaled dot-product attention\" # @param [\"xformers\", \"scaled dot-product attention\"]\n",
        "# @markdown ### **Save Output Config**\n",
        "save_precision              = \"fp16\"  # @param [\"float\", \"fp16\", \"bf16\"] {allow-input: false}\n",
        "save_every_n_epochs         = 1  # @param {type:\"number\"}\n",
        "# @markdown ### **Sample Prompt Config**\n",
        "enable_sample               = False  # @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",
        "logging_dir                 = os.path.join(training_dir, \"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\"                : True,\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",
        "        \"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_epochs\"           : save_every_n_epochs,\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",
        "        \"sdpa\"                          : True if optimization == \"scaled dot-product attention\" else False,\n",
        "        \"xformers\"                      : True if optimization == \"xformers\" else False,\n",
        "        \"max_train_epochs\"              : num_epochs,\n",
        "        \"max_data_loader_n_workers\"     : 8,\n",
        "        \"persistent_data_loader_workers\": True,\n",
        "        \"seed\"                          : seed if seed > 0 else None,\n",
        "        \"gradient_checkpointing\"        : gradient_checkpointing,\n",
        "        \"gradient_accumulation_steps\"   : 1,\n",
        "        \"mixed_precision\"               : mixed_precision,\n",
        "    },\n",
        "    \"logging_arguments\": {\n",
        "        \"log_with\"          : \"wandb\" if wandb_api_key else \"tensorboard\",\n",
        "        \"log_tracker_name\"  : project_name if wandb_api_key and not project_name == \"last\" else None,\n",
        "        \"logging_dir\"       : logging_dir,\n",
        "        \"log_prefix\"        : project_name if not wandb_api_key else None,\n",
        "    },\n",
        "    \"sample_prompt_arguments\": {\n",
        "        \"sample_every_n_steps\"    : None,\n",
        "        \"sample_every_n_epochs\"   : save_every_n_epochs if enable_sample else None,\n",
        "        \"sample_sampler\"          : sampler,\n",
        "    },\n",
        "    \"saving_arguments\": {\n",
        "        \"save_model_as\": \"safetensors\"\n",
        "    },\n",
        "}\n",
        "\n",
        "def write_file(filename, contents):\n",
        "    with open(filename, \"w\") as f:\n",
        "        f.write(contents)\n",
        "\n",
        "def prompt_convert(enable_sample, num_prompt, train_data_dir, prompt_config, custom_prompt):\n",
        "    if enable_sample:\n",
        "        search_pattern = os.path.join(train_data_dir, '**/*' + prompt_from_caption)\n",
        "        caption_files = glob.glob(search_pattern, recursive=True)\n",
        "\n",
        "        if not caption_files:\n",
        "            if not custom_prompt:\n",
        "                custom_prompt = \"masterpiece, best quality, 1girl, aqua eyes, baseball cap, blonde hair, closed mouth, earrings, green background, hat, hoop earrings, jewelry, looking at viewer, shirt, short hair, simple background, solo, upper body, yellow shirt\"\n",
        "            new_prompt_config = prompt_config.copy()\n",
        "            new_prompt_config['prompt']['subset'] = [\n",
        "                {\"prompt\": positive_prompt + custom_prompt if positive_prompt else custom_prompt}\n",
        "            ]\n",
        "        else:\n",
        "            selected_files = random.sample(caption_files, min(num_prompt, len(caption_files)))\n",
        "\n",
        "            prompts = []\n",
        "            for file in selected_files:\n",
        "                with open(file, 'r') as f:\n",
        "                    prompts.append(f.read().strip())\n",
        "\n",
        "            new_prompt_config = prompt_config.copy()\n",
        "            new_prompt_config['prompt']['subset'] = []\n",
        "\n",
        "            for prompt in prompts:\n",
        "                new_prompt = {\n",
        "                    \"prompt\": positive_prompt + prompt if positive_prompt else prompt,\n",
        "                }\n",
        "                new_prompt_config['prompt']['subset'].append(new_prompt)\n",
        "\n",
        "        return new_prompt_config\n",
        "    else:\n",
        "        return prompt_config\n",
        "\n",
        "def eliminate_none_variable(config):\n",
        "    for key in config:\n",
        "        if isinstance(config[key], dict):\n",
        "            for sub_key in config[key]:\n",
        "                if config[key][sub_key] == \"\":\n",
        "                    config[key][sub_key] = None\n",
        "        elif config[key] == \"\":\n",
        "            config[key] = None\n",
        "\n",
        "    return config\n",
        "\n",
        "try:\n",
        "    train_config.update(optimizer_config)\n",
        "except NameError:\n",
        "    raise NameError(\"'optimizer_config' dictionary is missing. Please run  '4.1. Optimizer Config' cell.\")\n",
        "\n",
        "try:\n",
        "    train_config.update(lora_config)\n",
        "except NameError:\n",
        "    raise NameError(\"'lora_config' dictionary is missing. Please run  '4.1. LoRa: Low-Rank Adaptation Config' cell.\")\n",
        "\n",
        "advanced_training_warning = False\n",
        "try:\n",
        "    train_config.update(advanced_training_config)\n",
        "except NameError:\n",
        "    advanced_training_warning = True\n",
        "    pass\n",
        "\n",
        "prompt_config = prompt_convert(enable_sample, num_prompt, train_data_dir, prompt_config, custom_prompt)\n",
        "\n",
        "config_path         = os.path.join(config_dir, \"config_file.toml\")\n",
        "prompt_path         = os.path.join(config_dir, \"sample_prompt.toml\")\n",
        "\n",
        "config_str          = toml.dumps(eliminate_none_variable(train_config))\n",
        "prompt_str          = toml.dumps(eliminate_none_variable(prompt_config))\n",
        "\n",
        "write_file(config_path, config_str)\n",
        "write_file(prompt_path, prompt_str)\n",
        "\n",
        "print(config_str)\n",
        "\n",
        "if advanced_training_warning:\n",
        "    import textwrap\n",
        "    error_message = \"WARNING: This is not an error message, but the [advanced_training_config] dictionary is missing. Please run the '4.2. Advanced Training Config' cell if you intend to use it, or continue to the next step.\"\n",
        "    wrapped_message = textwrap.fill(error_message, width=80)\n",
        "    print('\\033[38;2;204;102;102m' + wrapped_message + '\\033[0m\\n')\n",
        "    pass\n",
        "\n",
        "print(prompt_str)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "p_SHtbFwHVl1",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "b15de6f7-2b00-4ab4-8848-bb0934c6181b"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Loading settings from /content/LoRA/config/config_file.toml...\n",
            "/content/LoRA/config/config_file\n",
            "prepare tokenizers\n",
            "Downloading vocab.json: 100% 961k/961k [00:00<00:00, 3.90MB/s]\n",
            "Downloading merges.txt: 100% 525k/525k [00:00<00:00, 1.07MB/s]\n",
            "Downloading (…)cial_tokens_map.json: 100% 389/389 [00:00<00:00, 1.97MB/s]\n",
            "Downloading tokenizer_config.json: 100% 905/905 [00:00<00:00, 4.97MB/s]\n",
            "Downloading vocab.json: 100% 862k/862k [00:00<00:00, 3.59MB/s]\n",
            "Downloading merges.txt: 100% 525k/525k [00:00<00:00, 713kB/s]\n",
            "Downloading (…)cial_tokens_map.json: 100% 389/389 [00:00<00:00, 2.33MB/s]\n",
            "Downloading tokenizer_config.json: 100% 904/904 [00:00<00:00, 5.30MB/s]\n",
            "update token length: 225\n",
            "Training with captions.\n",
            "loading existing metadata: /content/LoRA/meta_lat.json\n",
            "metadata has bucket info, enable bucketing / メタデータにbucket情報があるためbucketを有効にします\n",
            "using bucket info in metadata / メタデータ内のbucket情報を使います\n",
            "[Dataset 0]\n",
            "  batch_size: 30\n",
            "  resolution: (1024, 1024)\n",
            "  enable_bucket: True\n",
            "  min_bucket_reso: None\n",
            "  max_bucket_reso: None\n",
            "  bucket_reso_steps: None\n",
            "  bucket_no_upscale: None\n",
            "\n",
            "  [Subset 0 of Dataset 0]\n",
            "    image_dir: \"/content/drive/MyDrive/24-datasets/ZenkaiXL\"\n",
            "    image_count: 381\n",
            "    num_repeats: 30\n",
            "    shuffle_caption: False\n",
            "    keep_tokens: 0\n",
            "    caption_dropout_rate: 0.0\n",
            "    caption_dropout_every_n_epoches: 0\n",
            "    caption_tag_dropout_rate: 0.0\n",
            "    color_aug: False\n",
            "    flip_aug: False\n",
            "    face_crop_aug_range: None\n",
            "    random_crop: False\n",
            "    token_warmup_min: 1,\n",
            "    token_warmup_step: 0,\n",
            "    metadata_file: /content/LoRA/meta_lat.json\n",
            "\n",
            "\n",
            "[Dataset 0]\n",
            "loading image sizes.\n",
            "100% 381/381 [00:00<00:00, 2010100.41it/s]\n",
            "make buckets\n",
            "number of images (including repeats) / 各bucketの画像枚数（繰り返し回数を含む）\n",
            "bucket 0: resolution (1024, 1024), count: 11430\n",
            "mean ar error (without repeats): 0.0\n",
            "Warning: SDXL has been trained with noise_offset=0.0357, but noise_offset is disabled due to multires_noise_iterations / SDXLはnoise_offset=0.0357で学習されていますが、multires_noise_iterationsが有効になっているためnoise_offsetは無効になります\n",
            "preparing accelerator\n",
            "loading model for process 0/1\n",
            "load StableDiffusion checkpoint: /content/pretrained_model/sd_xl_base_1.0.safetensors\n",
            "building U-Net\n",
            "loading U-Net from checkpoint\n",
            "U-Net:  <All keys matched successfully>\n",
            "building text encoders\n",
            "loading text encoders from checkpoint\n",
            "text encoder 1: <All keys matched successfully>\n",
            "text encoder 2: <All keys matched successfully>\n",
            "building VAE\n",
            "loading VAE from checkpoint\n",
            "VAE: <All keys matched successfully>\n",
            "load VAE: /content/vae/sdxl_vae.safetensors\n",
            "additional VAE loaded\n",
            "Enable SDPA for U-Net\n",
            "import network module: networks.lora\n",
            "move vae and unet to cpu to save memory\n",
            "[Dataset 0]\n",
            "caching text encoder outputs.\n",
            "checking cache existence...\n",
            "100% 381/381 [00:00<00:00, 2403052.37it/s]\n",
            "caching text encoder outputs...\n",
            "100% 13/13 [00:01<00:00,  9.91it/s]\n",
            "move vae and unet back to original device\n",
            "create LoRA network. base dim (rank): 16, alpha: 8\n",
            "neuron dropout: p=None, rank dropout: p=None, module dropout: p=None\n",
            "apply LoRA to Conv2d with kernel size (3,3). dim (rank): 8, alpha: 1.0\n",
            "create LoRA for Text Encoder 1:\n",
            "create LoRA for Text Encoder 2:\n",
            "create LoRA for Text Encoder: 264 modules.\n",
            "create LoRA for U-Net: 788 modules.\n",
            "enable LoRA for U-Net\n",
            "prepare optimizer, data loader etc.\n",
            "use Adafactor optimizer | {'scale_parameter': False, 'relative_step': False, 'warmup_init': False}\n",
            "override steps. steps for 2 epochs is / 指定エポックまでのステップ数: 762\n",
            "running training / 学習開始\n",
            "  num train images * repeats / 学習画像の数×繰り返し回数: 11430\n",
            "  num reg images / 正則化画像の数: 0\n",
            "  num batches per epoch / 1epochのバッチ数: 381\n",
            "  num epochs / epoch数: 2\n",
            "  batch size per device / バッチサイズ: 30\n",
            "  gradient accumulation steps / 勾配を合計するステップ数 = 1\n",
            "  total optimization steps / 学習ステップ数: 762\n",
            "steps:   0% 0/762 [00:00<?, ?it/s]\n",
            "epoch 1/2\n",
            "steps:  50% 381/762 [55:30<55:30,  8.74s/it, loss=0.101]\n",
            "saving checkpoint: /content/drive/MyDrive/kohya-trainer/output/ZenkaiXL-v0.2-000001.safetensors\n",
            "\n",
            "epoch 2/2\n",
            "steps: 100% 762/762 [1:51:01<00:00,  8.74s/it, loss=0.0994]\n",
            "saving checkpoint: /content/drive/MyDrive/kohya-trainer/output/ZenkaiXL-v0.2.safetensors\n",
            "model saved.\n",
            "steps: 100% 762/762 [1:51:03<00:00,  8.74s/it, loss=0.0994]\n"
          ]
        }
      ],
      "source": [
        "#@title ## **4.5. Start Training**\n",
        "import os\n",
        "import toml\n",
        "\n",
        "#@markdown Check your config here if you want to edit something:\n",
        "#@markdown - `sample_prompt` : /content/LoRA/config/sample_prompt.toml\n",
        "#@markdown - `config_file` : /content/LoRA/config/config_file.toml\n",
        "\n",
        "\n",
        "#@markdown You can import config from another session if you want.\n",
        "\n",
        "sample_prompt   = \"/content/LoRA/config/sample_prompt.toml\" #@param {type:'string'}\n",
        "config_file     = \"/content/LoRA/config/config_file.toml\" #@param {type:'string'}\n",
        "\n",
        "def read_file(filename):\n",
        "    with open(filename, \"r\") as f:\n",
        "        contents = f.read()\n",
        "    return contents\n",
        "\n",
        "def train(config):\n",
        "    args = \"\"\n",
        "    for k, v in config.items():\n",
        "        if k.startswith(\"_\"):\n",
        "            args += f'\"{v}\" '\n",
        "        elif isinstance(v, str):\n",
        "            args += f'--{k}=\"{v}\" '\n",
        "        elif isinstance(v, bool) and v:\n",
        "            args += f\"--{k} \"\n",
        "        elif isinstance(v, float) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "        elif isinstance(v, int) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "\n",
        "    return args\n",
        "\n",
        "accelerate_conf = {\n",
        "    \"config_file\" : \"/content/kohya-trainer/accelerate_config/config.yaml\",\n",
        "    \"num_cpu_threads_per_process\" : 1,\n",
        "}\n",
        "\n",
        "train_conf = {\n",
        "    \"sample_prompts\"  : sample_prompt if os.path.exists(sample_prompt) else None,\n",
        "    \"config_file\"     : config_file,\n",
        "}\n",
        "\n",
        "accelerate_args = train(accelerate_conf)\n",
        "train_args = train(train_conf)\n",
        "\n",
        "final_args = f\"accelerate launch {accelerate_args} sdxl_train_network.py {train_args}\"\n",
        "\n",
        "os.chdir(repo_dir)\n",
        "!{final_args}"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "machine_shape": "hm",
      "provenance": [],
      "gpuType": "A100",
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}