notebook/train.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e9980ac8-8e1b-4515-907a-5ecfa7849c77",
   "metadata": {},
   "source": [
    "# KOHYA LoRA TRAINER XL "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a955ca4c-e2fa-4638-9a8a-e944ade9f851",
   "metadata": {},
   "source": [
    "## Install Kohya Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8ffd6aa3-892d-4fd6-8377-0fa58d5659cb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wed Mar 20 10:32:38 2024       \n",
      "+---------------------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 535.129.03             Driver Version: 535.129.03   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 RTX A4500               On  | 00000000:81:00.0 Off |                  Off |\n",
      "| 30%   27C    P8              15W / 200W |      2MiB / 20470MiB |      0%      Default |\n",
      "|                                         |                      |                  N/A |\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": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f909916d-4ee9-45f1-b8d1-11d43ba2a92f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "83 packages can be upgraded. Run 'apt list --upgradable' to see them.\n",
      "aria2 is already the newest version (1.36.0-1).\n",
      "0 upgraded, 0 newly installed, 0 to remove and 83 not upgraded.\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpython -m pip install --upgrade pip\u001b[0m\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpython -m pip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "import os, shutil\n",
    "from pathlib import Path\n",
    "\n",
    "root_dir = Path(\"/workspace\")\n",
    "repo_dir = root_dir / \"kohya-trainer\"\n",
    "training_dir = root_dir / \"fine_tune\"\n",
    "model_dir = root_dir / \"model\"\n",
    "vae_dir = root_dir / \"vae\"\n",
    "lora_dir = root_dir / \"network_weight\"\n",
    "config_dir = training_dir / \"config\"\n",
    "output_dir = training_dir / \"outputs\"\n",
    "tools_dir = repo_dir / \"tools\"\n",
    "finetune_dir = repo_dir / \"finetune\"\n",
    "accelerate_config = repo_dir / \"accelerate_config\" / \"config.yaml\"\n",
    "\n",
    "# Uncomment if you want to use latest commit\n",
    "# repo_url = \"https://github.com/kohya-ss/sd-scripts\" \n",
    "repo_url = \"https://github.com/cagliostrolab/sd-scripts\"\n",
    "\n",
    "derrian_repo_url = \"https://github.com/derrian-distro/LoRA_Easy_Training_Scripts\"\n",
    "derrian_repo_dir = repo_dir / \"LoRA_Easy_Training_Scripts\"\n",
    "\n",
    "HUGGINGFACE_TOKEN = \"\"\n",
    "\n",
    "def clone_repo(url, dir, branch=\"main\"):\n",
    "    dir = Path(dir)\n",
    "    if not dir.exists():\n",
    "        !git clone -b {branch} {url} {dir}\n",
    "\n",
    "def install_dependencies():\n",
    "    !apt update -yqq\n",
    "    !apt install aria2 -yqq\n",
    "    !pip install -q --upgrade xformers==0.0.21 accelerate==0.25.0 transformers==4.36.2 diffusers[torch]==0.25.0 ftfy==6.1.1 opencv-python==4.7.0.68 einops==0.6.0 pytorch-lightning==1.9.0 safetensors==0.3.1 toml==0.10.2 voluptuous==0.13.1 huggingface-hub==0.20.1 wandb==0.15.7 invisible-watermark==0.2.0 open-clip-torch==2.20.0 tensorflow==2.10.1 bitsandbytes==0.35.0 -e .\n",
    "    # !pip install -q --upgrade -r requirements.txt\n",
    "    \n",
    "    !rm $accelerate_config\n",
    "    from accelerate.utils import write_basic_config\n",
    "\n",
    "    write_basic_config(save_location=accelerate_config)\n",
    "\n",
    "    os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n",
    "    os.environ[\"PYTHONWARNINGS\"] = \"ignore\"\n",
    "    \n",
    "    os.chdir(derrian_repo_dir / \"custom_scheduler\")\n",
    "    !pip install -q --upgrade -e .\n",
    "    \n",
    "# initialize function\n",
    "\n",
    "def get_filename(url, bearer_token):\n",
    "    headers = {\"Authorization\": f\"Bearer {bearer_token}\"}\n",
    "    with requests.get(url, headers=headers, stream=True) as response:\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(Path(url_path).name)\n",
    "\n",
    "    return filename\n",
    "\n",
    "def parse_args(config, aria=False):\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",
    "            if aria:\n",
    "                args.append(f\"--{k}={v}\")\n",
    "            else:\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, token):\n",
    "    user_header = f\"Authorization: Bearer {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\"                       : str(dir),\n",
    "        \"out\"                       : filename,\n",
    "        \"_url\"                      : url,\n",
    "    }\n",
    "    aria2_args = parse_args(aria2_config, aria=True)\n",
    "    subprocess.run([\"aria2c\", *aria2_args])\n",
    "    \n",
    "def download(url, dst, token):\n",
    "    filename = get_filename(url, token)\n",
    "    dst = Path(dst)\n",
    "    filepath = dst / filename\n",
    "\n",
    "    if url.startswith(\"/workspace\"):\n",
    "        return url\n",
    "    else:\n",
    "        if \"/blob/\" in url:\n",
    "            url = url.replace(\"/blob/\", \"/resolve/\")\n",
    "                \n",
    "        aria2_download(dst, filename, url, token)\n",
    "\n",
    "    return filepath\n",
    "    \n",
    "def main():\n",
    "    os.chdir(root_dir)\n",
    "    clone_repo(repo_url, repo_dir)\n",
    "    clone_repo(derrian_repo_url, derrian_repo_dir)\n",
    "    \n",
    "    os.chdir(repo_dir)\n",
    "    for dir in [training_dir, config_dir, model_dir, vae_dir, output_dir]:\n",
    "        dir.mkdir(parents=True, exist_ok=True)\n",
    "    install_dependencies()\n",
    "    \n",
    "main()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30bebe28-e96e-476c-8e54-3c79879f558a",
   "metadata": {},
   "source": [
    "## Download SDXL"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "dd221ad9-66cf-4ad4-987a-b5911b370955",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[#3d846b 0B/6.4GiB(0%) CN:1 DL:0B] [FileAlloc:#3d846b 87MiB/6.4GiB(1%)]\n",
      "[#3d846b 0B/6.4GiB(0%) CN:1 DL:0B] [FileAlloc:#3d846b 809MiB/6.4GiB(12%)]\n",
      "[#3d846b 0B/6.4GiB(0%) CN:1 DL:0B] [FileAlloc:#3d846b 1.4GiB/6.4GiB(22%)]\n",
      "[#3d846b 0B/6.4GiB(0%) CN:1 DL:0B] [FileAlloc:#3d846b 2.1GiB/6.4GiB(33%)]\n",
      "[#3d846b 0B/6.4GiB(0%) CN:1 DL:0B] [FileAlloc:#3d846b 2.7GiB/6.4GiB(43%)]\n",
      "[#3d846b 0B/6.4GiB(0%) CN:1 DL:0B] [FileAlloc:#3d846b 3.4GiB/6.4GiB(53%)]\n",
      "[#3d846b 0B/6.4GiB(0%) CN:1 DL:0B] [FileAlloc:#3d846b 4.0GiB/6.4GiB(63%)]\n",
      "[#3d846b 0B/6.4GiB(0%) CN:1 DL:0B] [FileAlloc:#3d846b 4.7GiB/6.4GiB(73%)]\n",
      "[#3d846b 0B/6.4GiB(0%) CN:1 DL:0B] [FileAlloc:#3d846b 5.4GiB/6.4GiB(83%)]\n",
      "[#3d846b 0B/6.4GiB(0%) CN:1 DL:0B] [FileAlloc:#3d846b 6.0GiB/6.4GiB(94%)]\n",
      " *** Download Progress Summary as of Wed Mar 20 10:36:14 2024 *** \n",
      "===============================================================================\n",
      "[#3d846b 0B/6.4GiB(0%) CN:1 DL:0B]\n",
      "FILE: /workspace/model/animagine-xl-3.0-base.safetensors\n",
      "-------------------------------------------------------------------------------\n",
      "\n",
      "[#3d846b 0B/6.4GiB(0%) CN:1 DL:0B]\n",
      "[#3d846b 24MiB/6.4GiB(0%) CN:16 DL:25MiB ETA:4m20s]\n",
      "[#3d846b 148MiB/6.4GiB(2%) CN:16 DL:74MiB ETA:1m26s]\n",
      "[#3d846b 289MiB/6.4GiB(4%) CN:16 DL:96MiB ETA:1m5s]\n",
      "[#3d846b 433MiB/6.4GiB(6%) CN:16 DL:108MiB ETA:56s]\n",
      "[#3d846b 574MiB/6.4GiB(8%) CN:16 DL:115MiB ETA:52s]\n",
      "[#3d846b 712MiB/6.4GiB(10%) CN:16 DL:119MiB ETA:49s]\n",
      "[#3d846b 825MiB/6.4GiB(12%) CN:16 DL:118MiB ETA:48s]\n",
      "[#3d846b 0.9GiB/6.4GiB(14%) CN:16 DL:118MiB ETA:47s]\n",
      "[#3d846b 1.0GiB/6.4GiB(15%) CN:16 DL:117MiB ETA:47s]\n",
      "[#3d846b 1.1GiB/6.4GiB(17%) CN:16 DL:117MiB ETA:46s]\n",
      " *** Download Progress Summary as of Wed Mar 20 10:36:25 2024 *** \n",
      "===============================================================================\n",
      "[#3d846b 1.2GiB/6.4GiB(19%) CN:16 DL:125MiB ETA:42s]\n",
      "FILE: /workspace/model/animagine-xl-3.0-base.safetensors\n",
      "-------------------------------------------------------------------------------\n",
      "\n",
      "[#3d846b 1.2GiB/6.4GiB(19%) CN:16 DL:125MiB ETA:42s]\n",
      "[#3d846b 1.3GiB/6.4GiB(21%) CN:16 DL:124MiB ETA:41s]\n",
      "[#3d846b 1.4GiB/6.4GiB(22%) CN:16 DL:124MiB ETA:41s]\n",
      "[#3d846b 1.6GiB/6.4GiB(24%) CN:16 DL:122MiB ETA:40s]\n",
      "[#3d846b 1.7GiB/6.4GiB(26%) CN:16 DL:119MiB ETA:40s]\n",
      "[#3d846b 1.8GiB/6.4GiB(28%) CN:16 DL:116MiB ETA:40s]\n",
      "[#3d846b 1.9GiB/6.4GiB(29%) CN:16 DL:116MiB ETA:39s]\n",
      "[#3d846b 2.0GiB/6.4GiB(31%) CN:16 DL:114MiB ETA:39s]\n",
      "[#3d846b 2.1GiB/6.4GiB(33%) CN:16 DL:114MiB ETA:38s]\n",
      "[#3d846b 2.2GiB/6.4GiB(34%) CN:16 DL:114MiB ETA:37s]\n",
      "[#3d846b 2.3GiB/6.4GiB(36%) CN:16 DL:113MiB ETA:36s]\n",
      " *** Download Progress Summary as of Wed Mar 20 10:36:36 2024 *** \n",
      "===============================================================================\n",
      "[#3d846b 2.4GiB/6.4GiB(38%) CN:16 DL:113MiB ETA:36s]\n",
      "FILE: /workspace/model/animagine-xl-3.0-base.safetensors\n",
      "-------------------------------------------------------------------------------\n",
      "\n",
      "[#3d846b 2.4GiB/6.4GiB(38%) CN:16 DL:113MiB ETA:36s]\n",
      "[#3d846b 2.5GiB/6.4GiB(39%) CN:16 DL:112MiB ETA:35s]\n",
      "[#3d846b 2.7GiB/6.4GiB(42%) CN:16 DL:114MiB ETA:33s]\n",
      "[#3d846b 2.8GiB/6.4GiB(43%) CN:16 DL:115MiB ETA:32s]\n",
      "[#3d846b 2.9GiB/6.4GiB(45%) CN:16 DL:115MiB ETA:31s]\n",
      "[#3d846b 3.0GiB/6.4GiB(47%) CN:16 DL:117MiB ETA:29s]\n",
      "[#3d846b 3.1GiB/6.4GiB(49%) CN:16 DL:118MiB ETA:28s]\n",
      "[#3d846b 3.2GiB/6.4GiB(50%) CN:16 DL:118MiB ETA:27s]\n",
      "[#3d846b 3.3GiB/6.4GiB(52%) CN:16 DL:118MiB ETA:26s]\n",
      "[#3d846b 3.5GiB/6.4GiB(54%) CN:16 DL:118MiB ETA:25s]\n",
      "[#3d846b 3.6GiB/6.4GiB(55%) CN:16 DL:119MiB ETA:24s]\n",
      " *** Download Progress Summary as of Wed Mar 20 10:36:47 2024 *** \n",
      "===============================================================================\n",
      "[#3d846b 3.7GiB/6.4GiB(57%) CN:16 DL:116MiB ETA:24s]\n",
      "FILE: /workspace/model/animagine-xl-3.0-base.safetensors\n",
      "-------------------------------------------------------------------------------\n",
      "\n",
      "[#3d846b 3.7GiB/6.4GiB(57%) CN:16 DL:116MiB ETA:24s]\n",
      "[#3d846b 3.8GiB/6.4GiB(59%) CN:16 DL:114MiB ETA:23s]\n",
      "[#3d846b 3.9GiB/6.4GiB(60%) CN:16 DL:113MiB ETA:22s]\n",
      "[#3d846b 4.0GiB/6.4GiB(62%) CN:16 DL:114MiB ETA:21s]\n",
      "[#3d846b 4.1GiB/6.4GiB(64%) CN:16 DL:113MiB ETA:20s]\n",
      "[#3d846b 4.2GiB/6.4GiB(66%) CN:16 DL:113MiB ETA:19s]\n",
      "[#3d846b 4.3GiB/6.4GiB(67%) CN:16 DL:113MiB ETA:18s]\n",
      "[#3d846b 4.4GiB/6.4GiB(69%) CN:16 DL:112MiB ETA:18s]\n",
      "[#3d846b 4.5GiB/6.4GiB(70%) CN:16 DL:111MiB ETA:17s]\n",
      "[#3d846b 4.6GiB/6.4GiB(72%) CN:16 DL:111MiB ETA:16s]\n",
      "[#3d846b 4.8GiB/6.4GiB(74%) CN:16 DL:111MiB ETA:15s]\n",
      " *** Download Progress Summary as of Wed Mar 20 10:36:57 2024 *** \n",
      "===============================================================================\n",
      "[#3d846b 4.9GiB/6.4GiB(75%) CN:16 DL:111MiB ETA:14s]\n",
      "FILE: /workspace/model/animagine-xl-3.0-base.safetensors\n",
      "-------------------------------------------------------------------------------\n",
      "\n",
      "[#3d846b 4.9GiB/6.4GiB(75%) CN:16 DL:111MiB ETA:14s]\n",
      "[#3d846b 5.0GiB/6.4GiB(77%) CN:16 DL:110MiB ETA:13s]\n",
      "[#3d846b 5.1GiB/6.4GiB(79%) CN:16 DL:109MiB ETA:12s]\n",
      "[#3d846b 5.2GiB/6.4GiB(80%) CN:16 DL:109MiB ETA:11s]\n",
      "[#3d846b 5.3GiB/6.4GiB(82%) CN:16 DL:109MiB ETA:10s]\n",
      "[#3d846b 5.4GiB/6.4GiB(84%) CN:16 DL:109MiB ETA:9s]\n",
      "[#3d846b 5.5GiB/6.4GiB(85%) CN:16 DL:109MiB ETA:8s]\n",
      "[#3d846b 5.6GiB/6.4GiB(87%) CN:16 DL:113MiB ETA:7s]\n",
      "[#3d846b 5.8GiB/6.4GiB(89%) CN:16 DL:114MiB ETA:5s]\n",
      "[#3d846b 5.9GiB/6.4GiB(91%) CN:16 DL:115MiB ETA:4s]\n",
      "[#3d846b 6.0GiB/6.4GiB(93%) CN:16 DL:117MiB ETA:3s]\n",
      " *** Download Progress Summary as of Wed Mar 20 10:37:08 2024 *** \n",
      "===============================================================================\n",
      "[#3d846b 6.1GiB/6.4GiB(95%) CN:16 DL:118MiB ETA:2s]\n",
      "FILE: /workspace/model/animagine-xl-3.0-base.safetensors\n",
      "-------------------------------------------------------------------------------\n",
      "\n",
      "[#3d846b 6.1GiB/6.4GiB(95%) CN:16 DL:118MiB ETA:2s]\n",
      "[#3d846b 6.2GiB/6.4GiB(97%) CN:16 DL:119MiB ETA:1s]\n",
      "[#3d846b 6.3GiB/6.4GiB(98%) CN:16 DL:119MiB]\n",
      "[#3d846b 6.4GiB/6.4GiB(99%) CN:1 DL:115MiB]\n",
      "\n",
      "Download Results:\n",
      "gid   |stat|avg speed  |path/URI\n",
      "======+====+===========+=======================================================\n",
      "3d846b|OK  |   113MiB/s|/workspace/model/animagine-xl-3.0-base.safetensors\n",
      "\n",
      "Status Legend:\n",
      "(OK):download completed.\n",
      "\n",
      "Download Results:\n",
      "gid   |stat|avg speed  |path/URI\n",
      "======+====+===========+=======================================================\n",
      "4bb8ab|OK  |       0B/s|/workspace/vae/sdxl_vae.safetensors\n",
      "\n",
      "Status Legend:\n",
      "(OK):download completed.\n",
      "Selected model: /workspace/model/animagine-xl-3.0-base.safetensors\n",
      "Selected vae: /workspace/vae/sdxl_vae.safetensors\n"
     ]
    }
   ],
   "source": [
    "import os, re, requests, subprocess\n",
    "from urllib.parse import urlparse, unquote\n",
    "from pathlib import Path\n",
    "\n",
    "model_path = Path()\n",
    "vae_path = Path()\n",
    "\n",
    "SDXL_MODEL_URL    = \"https://huggingface.co/cagliostrolab/animagine-xl-3.0-base/resolve/main/animagine-xl-3.0-base.safetensors\"\n",
    "SDXL_VAE_URL      = \"https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/resolve/main/sdxl_vae.safetensors\"\n",
    "\n",
    "def main():\n",
    "    global model_path, vae_path\n",
    "    \n",
    "    os.chdir(root_dir)\n",
    "\n",
    "    download_targets = {\n",
    "        \"model\": (SDXL_MODEL_URL, model_dir),\n",
    "        \"vae\": (SDXL_VAE_URL, vae_dir),\n",
    "    }\n",
    "    selected_files = {}\n",
    "\n",
    "    for target, (url, dst) in download_targets.items():\n",
    "        if url.startswith(\"/workspace\"):\n",
    "            selected_files[target] = Path(url)\n",
    "        else:\n",
    "            selected_files[target] = download(url, dst, HUGGINGFACE_TOKEN)\n",
    "\n",
    "    model_path = selected_files.get(\"model\", model_path)\n",
    "    vae_path = selected_files.get(\"vae\", vae_path)\n",
    "    \n",
    "    for category, path in {\"model\": model_path, \"vae\": vae_path}.items():\n",
    "        if path and path.exists():\n",
    "            print(f\"Selected {category}: {path}\")\n",
    "\n",
    "main()\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1656b23-2980-4d4f-b258-19f229f8584b",
   "metadata": {},
   "source": [
    "## Directory Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2860b5c0-0709-4154-a622-5576db80cab1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Your train data directory : /workspace/train_data/animagine-xl-3.1\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "project_name = \"animagine-xl-3.1\"\n",
    "train_data_dir = root_dir / \"train_data\" / \"animagine-xl-3.1\"\n",
    "\n",
    "train_data_dir.mkdir(parents=True, exist_ok=True)\n",
    "print(f\"Your train data directory : {train_data_dir}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c21a92e-49b1-4893-8ee2-148ca2975c7b",
   "metadata": {},
   "source": [
    "## Data Gathering"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7345ea37-09a0-4243-866d-a768c439e1ec",
   "metadata": {},
   "source": [
    "## Unzip Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "78a818d2-6fdd-4d05-bc80-9b2278314db2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[#edb4c4 0B/4.4GiB(0%) CN:1 DL:0B] [FileAlloc:#edb4c4 885MiB/4.4GiB(19%)]\n",
      "[#edb4c4 0B/4.4GiB(0%) CN:1 DL:0B] [FileAlloc:#edb4c4 2.0GiB/4.4GiB(45%)]\n",
      "[#edb4c4 0B/4.4GiB(0%) CN:1 DL:0B] [FileAlloc:#edb4c4 3.1GiB/4.4GiB(70%)]\n",
      "[#edb4c4 0B/4.4GiB(0%) CN:1 DL:0B] [FileAlloc:#edb4c4 4.2GiB/4.4GiB(96%)]\n",
      "[#edb4c4 48MiB/4.4GiB(1%) CN:16 DL:114MiB ETA:39s]\n",
      "[#edb4c4 250MiB/4.4GiB(5%) CN:16 DL:175MiB ETA:24s]\n",
      "[#edb4c4 455MiB/4.4GiB(9%) CN:16 DL:188MiB ETA:21s]\n",
      "[#edb4c4 656MiB/4.4GiB(14%) CN:16 DL:192MiB ETA:20s]\n",
      "[#edb4c4 859MiB/4.4GiB(18%) CN:16 DL:195MiB ETA:18s]\n",
      "[#edb4c4 1.0GiB/4.4GiB(23%) CN:16 DL:196MiB ETA:17s]\n",
      " *** Download Progress Summary as of Sat Dec  2 19:41:16 2023 *** \n",
      "===============================================================================\n",
      "[#edb4c4 1.2GiB/4.4GiB(27%) CN:16 DL:199MiB ETA:16s]\n",
      "FILE: /workspace/aesthetic-beta-raw.zip\n",
      "-------------------------------------------------------------------------------\n",
      "\n",
      "[#edb4c4 1.2GiB/4.4GiB(27%) CN:16 DL:199MiB ETA:16s]\n",
      "[#edb4c4 1.4GiB/4.4GiB(32%) CN:16 DL:201MiB ETA:15s]\n",
      "[#edb4c4 1.6GiB/4.4GiB(36%) CN:16 DL:200MiB ETA:14s]\n",
      "[#edb4c4 1.8GiB/4.4GiB(41%) CN:16 DL:200MiB ETA:13s]\n",
      "[#edb4c4 2.0GiB/4.4GiB(45%) CN:16 DL:206MiB ETA:11s]\n",
      "[#edb4c4 2.2GiB/4.4GiB(50%) CN:16 DL:205MiB ETA:11s]\n",
      "[#edb4c4 2.4GiB/4.4GiB(54%) CN:16 DL:206MiB ETA:10s]\n",
      "[#edb4c4 2.6GiB/4.4GiB(59%) CN:16 DL:205MiB ETA:9s]\n",
      "[#edb4c4 2.8GiB/4.4GiB(63%) CN:16 DL:205MiB ETA:8s]\n",
      "[#edb4c4 3.0GiB/4.4GiB(67%) CN:16 DL:203MiB ETA:7s]\n",
      "[#edb4c4 3.2GiB/4.4GiB(72%) CN:16 DL:203MiB ETA:6s]\n",
      " *** Download Progress Summary as of Sat Dec  2 19:41:27 2023 *** \n",
      "===============================================================================\n",
      "[#edb4c4 3.4GiB/4.4GiB(76%) CN:16 DL:202MiB ETA:5s]\n",
      "FILE: /workspace/aesthetic-beta-raw.zip\n",
      "-------------------------------------------------------------------------------\n",
      "\n",
      "[#edb4c4 3.4GiB/4.4GiB(76%) CN:16 DL:202MiB ETA:5s]\n",
      "[#edb4c4 3.6GiB/4.4GiB(80%) CN:16 DL:203MiB ETA:4s]\n",
      "[#edb4c4 3.8GiB/4.4GiB(85%) CN:16 DL:202MiB ETA:3s]\n",
      "[#edb4c4 3.9GiB/4.4GiB(89%) CN:16 DL:202MiB ETA:2s]\n",
      "[#edb4c4 4.1GiB/4.4GiB(94%) CN:16 DL:202MiB ETA:1s]\n",
      "[#edb4c4 4.3GiB/4.4GiB(98%) CN:16 DL:201MiB]\n",
      "\n",
      "Download Results:\n",
      "gid   |stat|avg speed  |path/URI\n",
      "======+====+===========+=======================================================\n",
      "edb4c4|OK  |   197MiB/s|/workspace/aesthetic-beta-raw.zip\n",
      "\n",
      "Status Legend:\n",
      "(OK):download completed.\n"
     ]
    }
   ],
   "source": [
    "import os, zipfile, shutil\n",
    "from pathlib import Path\n",
    "\n",
    "src_url = \"\"\n",
    "dst_dir = \"\"\n",
    "\n",
    "if not dst_dir:\n",
    "    dst_dir = train_data_dir\n",
    "\n",
    "dst_dir = Path(dst_dir)\n",
    "dst_dir.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "def extract_zipfile(zip_file, output_path):\n",
    "    with zipfile.ZipFile(zip_file, \"r\") as zip_ref:\n",
    "        zip_ref.extractall(output_path)\n",
    "\n",
    "def main():\n",
    "    zip_file = download(src_url, root_dir, HUGGINGFACE_TOKEN)\n",
    "    extract_zipfile(zip_file, dst_dir)\n",
    "    os.remove(zip_file)\n",
    "\n",
    "main()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68409b1c-7172-4513-8b2e-b6733857da92",
   "metadata": {},
   "source": [
    "## WD Tagger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5d62e2a-c1cc-479a-a7c7-42b888d0c51f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.chdir(finetune_dir)\n",
    "\n",
    "models = [\"moat\", \"convnextv2\", \"swinv2\", \"convnext\", \"vit\"]\n",
    "model = models[1]\n",
    "\n",
    "tagger_config = {\n",
    "    \"_train_data_dir\" : train_data_dir,\n",
    "    \"batch_size\" : 24,\n",
    "    \"repo_id\" : f\"SmilingWolf/wd-v1-4-{model}-tagger-v2\",\n",
    "    \"recursive\" : True,\n",
    "    \"remove_underscore\" : True,\n",
    "    \"general_threshold\" : 0.35,\n",
    "    \"character_threshold\" : 1,\n",
    "    \"caption_extension\" : \".txt\",\n",
    "    \"max_data_loader_n_workers\" : 8, \n",
    "    \"force_download\" : True, \n",
    "    \"undesired_tags\" : \"\"\n",
    "}\n",
    "\n",
    "tagger_args = ' '.join(parse_args(tagger_config))\n",
    "final_args = f\"python tag_images_by_wd14_tagger.py {tagger_args}\"\n",
    "\n",
    "os.chdir(finetune_dir)\n",
    "! {final_args}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8867485e-39ce-40e6-90bc-f786be7d4138",
   "metadata": {},
   "source": [
    "## Aspect Ratio Bucketing and Caching latents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62857e6e-3ece-440d-b255-7c5d1ba48cb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "raw_metadata = training_dir / f\"{project_name}_clean.json\"\n",
    "processed_metadata = training_dir / f\"{project_name}_lat.json\"\n",
    "recursive = True\n",
    "resolution = 1024\n",
    "\n",
    "metadata_config = {\n",
    "    \"_train_data_dir\": train_data_dir,\n",
    "    \"_out_json\": raw_metadata,\n",
    "    \"recursive\": recursive,\n",
    "    \"full_path\": recursive,\n",
    "}\n",
    "\n",
    "bucketing_config = {\n",
    "    \"_train_data_dir\": train_data_dir,\n",
    "    \"_in_json\": raw_metadata,\n",
    "    \"_out_json\": processed_metadata,\n",
    "    \"_model_name_or_path\": vae_path if vae_path else model_path,\n",
    "    \"recursive\": recursive,\n",
    "    \"full_path\": recursive,\n",
    "    \"flip_aug\": False,\n",
    "    \"max_bucket_reso\" : int(resolution * 2),\n",
    "    \"min_bucket_reso\" : int(resolution / 2),\n",
    "    \"bucket_no_upscale\" : False, \n",
    "    \"bucket_reso_steps\" : 64, \n",
    "    \"batch_size\": 8,\n",
    "    \"skip_existing\": True,\n",
    "    \"max_data_loader_n_workers\": 1,\n",
    "    \"max_resolution\": \", \".join([str(resolution)] * 2),\n",
    "    \"mixed_precision\": \"fp16\",\n",
    "}\n",
    "\n",
    "merge_metadata_args = ' '.join(parse_args(metadata_config))\n",
    "prepare_buckets_args = ' '.join(parse_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",
    "if not Path(\"merge_all_to_metadata.py\").exists():\n",
    "    !wget https://raw.githubusercontent.com/qaneel/kohya-trainer/main/finetune/merge_all_to_metadata.py\n",
    "!{merge_metadata_command}\n",
    "!{prepare_buckets_command}\n",
    "\n",
    "# 20:23 18696\n",
    "# 20:28 21064 \n",
    "# 20:33 23168\n",
    "# 20:38 25576 \n",
    "# 20:43 28032  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35904ab7-9e26-4b94-a5f5-53ffa5480877",
   "metadata": {},
   "source": [
    "## Optimizer Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "ea5ed0fa-52a0-4640-9845-5ee38a81f3fc",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "estimated_num_steps =  90990\n",
      "effective_batch_size =  96\n",
      "[optimizer_arguments]\n",
      "optimizer_type = \"AdamW\"\n",
      "learning_rate = 1e-5\n",
      "train_text_encoder = true\n",
      "optimizer_args = [ \"weight_decay=0.1\", \"betas=0.9,0.99\",]\n",
      "lr_scheduler = \"cosine_with_restarts\"\n",
      "lr_scheduler_num_cycles = 10\n",
      "lr_scheduler_type = \"LoraEasyCustomOptimizer.CustomOptimizers.CosineAnnealingWarmupRestarts\"\n",
      "lr_scheduler_args = [ \"min_lr=1e-06\", \"gamma=0.9\", \"first_cycle_steps=9099\",]\n",
      "max_grad_norm = 1.0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import torch\n",
    "import toml\n",
    "\n",
    "def count_images_in_directory(directory, extensions):\n",
    "    count = 0\n",
    "    for root, dirs, files in os.walk(directory):\n",
    "        for file in files:\n",
    "            if any(file.endswith(ext) for ext in extensions):\n",
    "                count += 1\n",
    "    return count\n",
    "\n",
    "IMAGE_EXTENSIONS = [\".png\", \".jpg\", \".jpeg\", \".webp\", \".bmp\", \".PNG\", \".JPG\", \".JPEG\", \".WEBP\", \".BMP\"]\n",
    "\n",
    "train_data_count = count_images_in_directory(train_data_dir, IMAGE_EXTENSIONS)\n",
    "learning_rate = 1e-5\n",
    "\n",
    "gradient_accumulation_steps = 3 \n",
    "batch_size = 16\n",
    "num_epochs = 10\n",
    "percentage = 5 / 100  # 5%\n",
    "\n",
    "num_gpus = torch.cuda.device_count()\n",
    "effective_batch_size = batch_size * gradient_accumulation_steps * num_gpus\n",
    "num_steps = int(train_data_count / effective_batch_size * num_epochs)\n",
    "cycle_steps = int(num_steps // num_epochs)\n",
    "\n",
    "warmup_steps = int(num_steps * percentage)\n",
    "\n",
    "print(\"estimated_num_steps = \", num_steps)\n",
    "print(\"effective_batch_size = \", effective_batch_size)\n",
    "# min_lr = learning_rate * 0.1\n",
    "min_lr = 1e-6\n",
    "\n",
    "optimizer_config = {\n",
    "    \"optimizer_arguments\": {\n",
    "        \"optimizer_type\" : \"AdamW\",\n",
    "        \"learning_rate\" : learning_rate,\n",
    "        \"train_text_encoder\" : True,\n",
    "        # \"learning_rate_te1\" : learning_rate / 2,\n",
    "        # \"learning_rate_te2\" : learning_rate / 2,\n",
    "        \"optimizer_args\" : [\"weight_decay=0.1\", \"betas=0.9,0.99\"],\n",
    "        \"lr_scheduler\" : \"cosine_with_restarts\",\n",
    "        \"lr_warmup_steps\" : None, # warmup_steps,  # Updated to use the computed warmup steps\n",
    "        \"lr_scheduler_num_cycles\" : num_epochs,  # cosine_with_restarts\n",
    "        \"lr_scheduler_power\" : None,  # polynomial\n",
    "        \"lr_scheduler_type\" : \"LoraEasyCustomOptimizer.CustomOptimizers.CosineAnnealingWarmupRestarts\",\n",
    "        \"lr_scheduler_args\" : [f\"min_lr={min_lr}\", \"gamma=0.9\", f\"first_cycle_steps={cycle_steps}\"],\n",
    "        \"max_grad_norm\" : 1.0\n",
    "    },\n",
    "}\n",
    "\n",
    "print(toml.dumps(optimizer_config))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4812407-f888-4cf5-8730-d211aabe4046",
   "metadata": {},
   "source": [
    "## Advanced Training Config\n",
    "1. Specify `optimizer_state_path` to resume training with Optimizer State\n",
    "2. You can't use both `noise_offset` and `multires_noise` at the same time\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "eb3d2c24-fc34-44dc-85b9-09051ddb90b0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[advanced_training_config]\n",
      "resume = \"\"\n",
      "resume_from_huggingface = false\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import toml\n",
    "\n",
    "optimizer_state_path      = \"\" \n",
    "\n",
    "advanced_training_config = {\n",
    "    \"advanced_training_config\": {\n",
    "        \"resume\" : optimizer_state_path,\n",
    "        \"resume_from_huggingface\": False,\n",
    "        # \"noise_offset\" : 0.0357,\n",
    "        # \"adaptive_noise_scale\" : 0.00357,\n",
    "        # \"multires_noise_iterations\" : 6, \n",
    "        # \"multires_noise_discount\" : 0.3, \n",
    "        # \"min_snr_gamma\" : 5\n",
    "    }\n",
    "}\n",
    "\n",
    "print(toml.dumps(advanced_training_config))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54a49ac6-bf82-400a-b9c7-d59d892f0ca3",
   "metadata": {},
   "source": [
    "## Deployment Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "13a4659c-60d1-418e-9f4a-a733923195be",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[save_to_hub_config]\n",
      "huggingface_repo_id = \"\"\n",
      "huggingface_repo_type = \"model\"\n",
      "huggingface_path_in_repo = \"model/animagine-xl-3.1_20240320_104740\"\n",
      "huggingface_token = \"\"\n",
      "async_upload = true\n",
      "save_state_to_huggingface = true\n",
      "huggingface_repo_visibility = \"private\"\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import toml\n",
    "from datetime import datetime\n",
    "WRITE_TOKEN = \"\"\n",
    "\n",
    "current_datetime = datetime.now().strftime(\"%Y%m%d_%H%M%S\")\n",
    "\n",
    "deployment_config = {\n",
    "    \"save_to_hub_config\": {\n",
    "        \"huggingface_repo_id\" : \"\",\n",
    "        \"huggingface_repo_type\" : \"model\", \n",
    "        \"huggingface_path_in_repo\" : f\"model/{project_name}_{current_datetime}\",\n",
    "        # \"resume_from_huggingface\"\n",
    "        \"huggingface_token\" : WRITE_TOKEN,\n",
    "        \"async_upload\" : True,\n",
    "        \"save_state_to_huggingface\" : True,\n",
    "        \"huggingface_repo_visibility\" : \"private\",\n",
    "    }\n",
    "}\n",
    "print(toml.dumps(deployment_config))\n",
    "\n",
    "if WRITE_TOKEN == \"\":\n",
    "    del deployment_config"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00afaeff-7fbb-4462-983c-03b210fa2e41",
   "metadata": {},
   "source": [
    "# Training Config \n",
    "1. Get your `wandb_api_key` here: https://wandb.ai/settings\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "7fc406d6-32c4-4eb1-bd6f-c9005020e4c3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[sdxl_arguments]\n",
      "cache_text_encoder_outputs = false\n",
      "no_half_vae = false\n",
      "min_timestep = 0\n",
      "max_timestep = 1000\n",
      "\n",
      "[model_arguments]\n",
      "pretrained_model_name_or_path = \"/workspace/model/animagine-xl-3.0-base.safetensors\"\n",
      "vae = \"/workspace/vae/sdxl_vae.safetensors\"\n",
      "\n",
      "[dataset_arguments]\n",
      "shuffle_caption = true\n",
      "debug_dataset = false\n",
      "in_json = \"/workspace/fine_tune/animagine-xl-3.1_lat.json\"\n",
      "train_data_dir = \"/workspace/train_data/animagine-xl-3.1\"\n",
      "dataset_repeats = 1\n",
      "keep_tokens_separator = \"|||\"\n",
      "resolution = \"1024, 1024\"\n",
      "caption_dropout_rate = 0\n",
      "caption_tag_dropout_rate = 0\n",
      "caption_dropout_every_n_epochs = 0\n",
      "token_warmup_min = 1\n",
      "token_warmup_step = 0\n",
      "\n",
      "[training_arguments]\n",
      "output_dir = \"/workspace/fine_tune/outputs/animagine-xl-3.1\"\n",
      "output_name = \"animagine-xl-3.1\"\n",
      "save_precision = \"fp16\"\n",
      "save_every_n_steps = 1000\n",
      "save_last_n_steps = true\n",
      "save_state = true\n",
      "save_last_n_steps_state = true\n",
      "train_batch_size = 16\n",
      "max_token_length = 225\n",
      "mem_eff_attn = false\n",
      "xformers = true\n",
      "sdpa = false\n",
      "max_train_epochs = 10\n",
      "max_data_loader_n_workers = 8\n",
      "persistent_data_loader_workers = true\n",
      "gradient_checkpointing = true\n",
      "gradient_accumulation_steps = 3\n",
      "mixed_precision = \"fp16\"\n",
      "ddp_gradient_as_bucket_view = true\n",
      "ddp_static_graph = true\n",
      "ddp_timeout = 100000\n",
      "\n",
      "[logging_arguments]\n",
      "log_with = \"wandb\"\n",
      "log_tracker_name = \"animagine-xl-3.1\"\n",
      "logging_dir = \"/workspace/fine_tune/logs\"\n",
      "\n",
      "[sample_prompt_arguments]\n",
      "sample_every_n_steps = 100\n",
      "sample_sampler = \"euler_a\"\n",
      "\n",
      "[saving_arguments]\n",
      "save_model_as = \"safetensors\"\n",
      "\n",
      "[optimizer_arguments]\n",
      "optimizer_type = \"AdamW\"\n",
      "learning_rate = 1e-5\n",
      "train_text_encoder = true\n",
      "optimizer_args = [ \"weight_decay=0.1\", \"betas=0.9,0.99\",]\n",
      "lr_scheduler = \"cosine_with_restarts\"\n",
      "lr_scheduler_num_cycles = 10\n",
      "lr_scheduler_type = \"LoraEasyCustomOptimizer.CustomOptimizers.CosineAnnealingWarmupRestarts\"\n",
      "lr_scheduler_args = [ \"min_lr=1e-06\", \"gamma=0.9\", \"first_cycle_steps=9099\",]\n",
      "max_grad_norm = 1.0\n",
      "\n",
      "[advanced_training_config]\n",
      "resume_from_huggingface = false\n",
      "\n",
      "\u001b[38;2;204;102;102mWARNING: This is not an error message, but the [deployment_config] dictionary is\n",
      "missing. Please run the 'Deployment Training Config' cell if you intend to use\n",
      "it, or continue to the next step.\u001b[0m\n",
      "\n",
      "[prompt]\n",
      "negative_prompt = \"lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, displeasing, jpeg artifacts, signature, watermark, username, blurry\"\n",
      "scale = 7\n",
      "sample_steps = 28\n",
      "[[prompt.subset]]\n",
      "prompt = \"1girl, makima \\\\(chainsaw man\\\\), chainsaw man, v, looking at viewer, upper body, outdoors, night, masterpiece, best quality, very aesthetic, absurdres\"\n",
      "width = 896\n",
      "height = 1152\n",
      "\n",
      "[[prompt.subset]]\n",
      "prompt = \"1girl, maomao \\\\(kusuriya no hitorigoto\\\\), kusuriya no hitorigoto, waving, looking at viewer, upper body, outdoors, night, masterpiece, best quality\"\n",
      "width = 832\n",
      "height = 1216\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import toml\n",
    "\n",
    "wandb_api_key = \"\" \n",
    "\n",
    "seed = 42\n",
    "resolution = 1024\n",
    "\n",
    "prompt_config = {\n",
    "    \"prompt\": {\n",
    "        \"negative_prompt\" : \"lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, displeasing, jpeg artifacts, signature, watermark, username, blurry\",\n",
    "        \"scale\"           : 7,\n",
    "        \"sample_steps\"    : 28,\n",
    "        \"subset\"          : [\n",
    "            {\n",
    "                \"prompt\" : \"1girl, makima \\(chainsaw man\\), chainsaw man, v, looking at viewer, upper body, outdoors, night, masterpiece, best quality, very aesthetic, absurdres\",\n",
    "                \"width\"  : 896,\n",
    "                \"height\" : 1152,                \n",
    "            },\n",
    "            {\n",
    "                \"prompt\" : \"1girl, maomao \\(kusuriya no hitorigoto\\), kusuriya no hitorigoto, waving, looking at viewer, upper body, outdoors, night, masterpiece, best quality\",\n",
    "                \"width\"  : 832,\n",
    "                \"height\" : 1216,                \n",
    "            },\n",
    "        ],\n",
    "    }\n",
    "}\n",
    "\n",
    "train_config = {\n",
    "    \"sdxl_arguments\": {\n",
    "        \"cache_text_encoder_outputs\" : False,\n",
    "        \"no_half_vae\" : False,\n",
    "        \"min_timestep\" : 0,\n",
    "        \"max_timestep\" : 1000,\n",
    "    },\n",
    "    \"model_arguments\": {\n",
    "        \"pretrained_model_name_or_path\" : str(model_path),\n",
    "        \"vae\" : str(vae_path),\n",
    "    },\n",
    "    \"dataset_arguments\": {\n",
    "        \"shuffle_caption\" : True,\n",
    "        \"debug_dataset\" : False,\n",
    "        \"in_json\" : str(training_dir / f\"{project_name}_lat.json\"),\n",
    "        \"train_data_dir\" : str(train_data_dir),\n",
    "        \"dataset_repeats\" : 1,\n",
    "        \"keep_tokens\" : None,\n",
    "        \"keep_tokens_separator\" : \"|||\",\n",
    "        \"resolution\" : \", \".join([str(resolution)] * 2),\n",
    "        \"caption_dropout_rate\" : 0,\n",
    "        \"caption_tag_dropout_rate\" : 0,\n",
    "        \"caption_dropout_every_n_epochs\": 0,\n",
    "        \"token_warmup_min\" : 1,\n",
    "        \"token_warmup_step\" : 0,\n",
    "    },\n",
    "    \"training_arguments\": {\n",
    "        \"output_dir\" : str(output_dir / project_name),\n",
    "        \"output_name\" : project_name,\n",
    "        \"save_precision\" : \"fp16\",\n",
    "        # \"save_every_n_epochs\" : 1,\n",
    "        \"save_every_n_steps\" : 1000,\n",
    "        \"save_n_epoch_ratio\" : None,\n",
    "        # \"save_last_n_epochs\" : True,\n",
    "        \"save_last_n_steps\" : True,\n",
    "        \"save_state\" : True,\n",
    "        # \"save_last_n_epochs_state\" : True,\n",
    "        \"save_last_n_steps_state\" : True,\n",
    "        \"train_batch_size\" : batch_size,\n",
    "        \"max_token_length\" : 225,\n",
    "        \"mem_eff_attn\" : False,\n",
    "        \"xformers\" : True,\n",
    "        \"sdpa\" : False, \n",
    "        \"max_train_epochs\" : num_epochs,\n",
    "        # \"max_train_steps\": 132590 - 41259 - 62000,\n",
    "        \"max_data_loader_n_workers\" : 8,\n",
    "        \"persistent_data_loader_workers\": True,\n",
    "        \"seed\" : None,\n",
    "        \"gradient_checkpointing\" : True,\n",
    "        \"gradient_accumulation_steps\" : gradient_accumulation_steps,\n",
    "        \"mixed_precision\" : \"fp16\",\n",
    "        \"ddp_gradient_as_bucket_view\": True,\n",
    "        \"ddp_static_graph\": True,\n",
    "        \"ddp_timeout\": 100000,\n",
    "        \n",
    "    },\n",
    "    \"logging_arguments\": {\n",
    "        \"log_with\" : \"wandb\",\n",
    "        \"log_tracker_name\" : project_name,\n",
    "        \"logging_dir\" : str(training_dir / \"logs\"),\n",
    "    },\n",
    "    \"sample_prompt_arguments\": {\n",
    "        \"sample_every_n_steps\" : 100,\n",
    "        \"sample_every_n_epochs\" : None,\n",
    "        \"sample_sampler\" : \"euler_a\",\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 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  '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",
    "deployment_config_warning = False\n",
    "try:\n",
    "    train_config.update(deployment_config)\n",
    "except NameError:\n",
    "    deployment_config_warning = True\n",
    "    pass\n",
    "\n",
    "config_path         = config_dir / f\"{project_name}_config_file.toml\"\n",
    "prompt_path         = config_dir / f\"{project_name}_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 '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",
    "if deployment_config_warning:\n",
    "    import textwrap\n",
    "    error_message = \"WARNING: This is not an error message, but the [deployment_config] dictionary is missing. Please run the 'Deployment 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": "markdown",
   "id": "a92fa91d-3c2e-4a20-a6ce-9f75e7280192",
   "metadata": {},
   "source": [
    "## Start Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c4e38ac-4996-419e-a647-b3e182281af9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import toml\n",
    "\n",
    "sample_prompt   = f\"/workspace/fine_tune/config/{project_name}_sample_prompt.toml\"\n",
    "config_file     = f\"/workspace/fine_tune/config/{project_name}_config_file.toml\"\n",
    "\n",
    "script_names = [\"sdxl_train.py\", \"sdxl_train_network.py\"]\n",
    "script_name = script_names[0]\n",
    "\n",
    "accelerate_conf = {\n",
    "    \"config_file\" : str(accelerate_config),\n",
    "    \"num_cpu_threads_per_process\" : 1,\n",
    "    \"num_processes\" : 2, \n",
    "    \"multi_gpu\" : True,\n",
    "    \"num_machines\" : 1, \n",
    "    \"gpu_ids\" : \"0,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",
    "train_args = ' '.join(parse_args(train_conf))\n",
    "accelerate_args = ' '.join(parse_args(accelerate_conf))\n",
    "\n",
    "final_args = f\"accelerate launch {accelerate_args} {repo_dir}/{script_name} {train_args}\"\n",
    "\n",
    "os.chdir(repo_dir)\n",
    "!{final_args}\n",
    "print(final_args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3807d6b4-3dcc-4514-97d2-39266a5e13d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "def sleep_with_countdown(interval):\n",
    "    for remaining in range(interval, 0, -1):\n",
    "        mins, secs = divmod(remaining, 60)\n",
    "        timer = '{:02d}:{:02d}'.format(mins, secs)\n",
    "        print(timer, end=\"\\r\")\n",
    "        time.sleep(1)\n",
    "\n",
    "# Set the interval for 10 minutes (600 seconds)\n",
    "interval = 600\n",
    "\n",
    "# Display the countdown while the program is paused\n",
    "print(\"Program is paused. Time remaining:\")\n",
    "sleep_with_countdown(interval)\n",
    "\n",
    "print(\"Program resumes after countdown.\")\n",
    "\n",
    "RUNPOD_POD_ID = \"\"\n",
    "!runpodctl remove pod $RUNPOD_POD_ID\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}