Lora_Trainer_XL.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rmCPmqFL6hCQ"
   },
   "source": [
    "# 🌟 XL Lora Trainer by Hollowstrawberry\n",
    "\n",
    "❗ **Colab Premium is recommended.** Ideally you'd be changing the runtime to an A100 and use the maximum batch size.  \n",
    "However you can still train for free if you load a diffusers model, it will just take much longer.  \n",
    "\n",
    "\n",
    "This colab is based on the work of [Kohya-ss](https://github.com/kohya-ss/sd-scripts) and [Linaqruf](https://github.com/Linaqruf/kohya-trainer). Thank you!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "vJ8clWTZEu-g"
   },
   "source": [
    "### ⭕ Disclaimer\n",
    "The purpose of this document is to research bleeding-edge technologies in the field of machine learning inference.  \n",
    "Please read and follow the [Google Colab guidelines](https://research.google.com/colaboratory/faq.html) and its [Terms of Service](https://research.google.com/colaboratory/tos_v3.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dPQlB4djNm3C"
   },
   "source": [
    "| |GitHub|🇬🇧 English|🇪🇸 Spanish|\n",
    "|:--|:-:|:-:|:-:|\n",
    "| 🏠 **Homepage** | [![GitHub](https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/assets/github.svg)](https://github.com/hollowstrawberry/kohya-colab) | | |\n",
    "| 📊 **Dataset Maker** | [![GitHub](https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/assets/github.svg)](https://github.com/hollowstrawberry/kohya-colab/blob/main/Dataset_Maker.ipynb) | [![Open in Colab](https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/assets/colab-badge.svg)](https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Dataset_Maker.ipynb) | [![Abrir en Colab](https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/assets/colab-badge-spanish.svg)](https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Spanish_Dataset_Maker.ipynb) |\n",
    "| ⭐ **Lora Trainer** | [![GitHub](https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/assets/github.svg)](https://github.com/hollowstrawberry/kohya-colab/blob/main/Lora_Trainer.ipynb) | [![Open in Colab](https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/assets/colab-badge.svg)](https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Lora_Trainer.ipynb) | [![Abrir en Colab](https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/assets/colab-badge-spanish.svg)](https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Spanish_Lora_Trainer.ipynb) |\n",
    "| 🌟 **XL Lora Trainer** | [![GitHub](https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/assets/github.svg)](https://github.com/hollowstrawberry/kohya-colab/blob/main/Lora_Trainer_XL.ipynb) | [![Open in Colab](https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/assets/colab-badge.svg)](https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Lora_Trainer_XL.ipynb) |  |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "33ec890f486f41d7979ec21dc403dd6e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HTML(value='<h1>🚩 Start Here</h1>')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dc5b72656d5e4fb6bb9b4f63aa2ca70b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HTML(value='<h2>▶️ Setup</h2>')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f146abca61eb47d3951daab361f04e24",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value=\"Your project name will be the same as the folder containing your images. Spaces aren't allowed.\")"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2245c39012bf470daec8fbca8b682b76",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Text(value='', description='Project name:')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b2ff4cbd2a6140a484c2b8f7004b84ea",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value=\"The folder structure doesn't matter and is purely for comfort. Make sure to always pick the same …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e0f7abcc4bbe440db268b1feefc94f56",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Folder structure:', index=1, options=('Organize by category (MyDrive/lora_training/datas…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "29b7e1f792364bb2838e43fd01467461",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value=' Decide the model that will be downloaded and used for training. You can also choose your own by …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c350857529894338a34afba20b7f10a5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Training model:', options=('Pony Diffusion V6 XL', 'Animagine XL V3', 'Stable Diffusion …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "29603b3a292249c48e71cbdd1a8db6f4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Text(value='', description='Optional: custom training model url')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6769a1d9a2634490a5f3babde13b36db",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value='Using a diffusers model utilizes less resources. The 3 options above will work with or without di…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "798f8988e19c45fd9d2ad9b9799308cb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Checkbox(value=True, description='Load diffusers', indent=False)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d38007569bbc4bf9a05e76068c6c7d91",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value='Use wandb if you want to visualize the progress of your training over time.')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1a0049fe2ae44250abf5cdd1998bf6c8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Text(value='', description='Wandb key:')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3b7483d97c2c4c48b3975450ba878de3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HTML(value='<h2> ▶️ Processing </h2>')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "67e7d3dbb7144866a55d88af7fd069d6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "BoundedIntText(value=1024, description='Resolution:', max=1536, min=768, step=128)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5cd814fe605e402590583f93103e16b4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Text(value='.txt', description='Caption extension')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "02a3cbed381246ebb451994ce3df43d9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value='Shuffling anime tags in place improves learning and prompting. An activation tag goes at the star…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5b76b8b241b24a74ac5906e7eea04718",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Checkbox(value=True, description='Shuffle tags', indent=False)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2d1744ec65644fb399099ddfb6f3fc64",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Activation tags:', index=1, options=('0', '1', '2', '3'), value='1')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4ee26a14f47f40589c6957882753b934",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HTML(value='<h2>▶️ Steps</h2>')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3ff6de30fe3f426fa5a8b69f3fd5e4fa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value='Your images will repeat this number of times during training. I recommend that your images multip…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2716da332d4b4d149533676b7db3c53a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "IntText(value=2, description='# repeats:')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "77af4c4cd12d4f18bdf9771e133aac23",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value='Choose how long you want to train for. A good starting point is around 10 epochs or around 2000 s…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "00939dcd6c75425daf9899f305e57877",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value='One epoch is a number of steps equal to: your number of images multiplied by their repeats, divid…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a589e75c3f5a418fbe7cb79d6a8e587d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "RadioButtons(description='Preferred unit', options=('Epochs', 'Steps'), value='Epochs')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "144367212ffa4e649d9a6ea41ee268b5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "IntText(value=10, description='How many:')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f05c24a49bc94306bb98b8df839bcce2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value=\"Saving more epochs will let you compare your Lora's progress better.\")"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "372d599e316e4aa8a5fd0102e39d561a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "IntText(value=1, description='Saving every N epochs:')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9ed72d16d5d54ad2828f1aa4aa49b58c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "IntText(value=10, description='Keep only last N epochs:')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bd9c91c2294446cd9597bba5aa196848",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HTML(value='<h2>▶️ Learning </h2>')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6e4357f0bad84348a00405762b70089c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value='The learning rate is the most important for your results. If you want to train slower with lots o…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "14480d6d8c8d48b4a0c0c3505285d99e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value=\"The text encoder helps your Lora learn concepts slightly better. It is recommended to make it hal…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8308f9617d11498a8d3f259f1eecada8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatText(value=0.0003, description='UNet LR:')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b9bc2dbc769a4f8aa98025d812ab2301",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatText(value=6e-05, description='Text encoder LR:')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "66fc0da88cb148b0a2e7d596a957d17c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value=\"The scheduler is the algorithm that guides the learning rate. If you're not sure, pick `constant`…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dac060dfef2d4ff8ac954b5e0de91481",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='LR Scheduler:', index=2, options=('constant', 'cosine', 'cosine_with_restarts', 'constan…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2aac3f984e94434c8a36b5d711db7640",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "IntText(value=3, description='LR Scheduler #:')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ab7d630e6b20456091d8909b56d35db1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value='Steps spent \"warming up\" the learning rate during training for efficiency. I recommend leaving it…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a2d0206bddab483eb3f6a0d80c27cbb8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatSlider(value=0.05, continuous_update=False, description='LR Warmiup Ratio:', max=0.2, readout_format='.1f…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "39f6eec21989456c81b5a91e0ef889c7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value='Adjusts loss over time, makes learning much more efficient. The paper recommends 5.0, I recommend…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b132937809984e89a628f6710b145cf0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "FloatSlider(value=8.0, continuous_update=False, description='Min SNR Gamma:', max=16.0, min=4.0, readout_forma…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bc5b4b5097ab4b9a9b510b83097be459",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value='Multinoise may help with color balance (darker darks, lighter lights).')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "82a08a4fc9fd4255bf54ff7ec7d74c60",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Checkbox(value=True, description='Multinoise', indent=False)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5ad8a3596e7d493fb9113f5cbec54b2c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HTML(value='<h2>▶️ Structure</h2>')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "177e7d84011140a5b4f7d5c8c4a30367",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "RadioButtons(description='LoRA type:', index=1, options=('LoRA', 'LoCon'), value='LoCon')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5a4a51be742343549cecf434a7930c3c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value='Below are some recommended XL values for the following settings:')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1531bd0f7ab843868112578f72bdbc27",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HTML(value='\\n        <table>\\n            <thead>\\n                <tr>\\n                    <th>type</th>\\n …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fe237e44893943709ce79b9ed0ed13a4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value=\"More dim means larger Lora, it can hold more information but more isn't always better.\")"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fc1e9d81421d43f686e56711cfcc56b8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "IntSlider(value=16, continuous_update=False, description='network_dim:', max=32, min=1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cdaca319d7044e299676e2b3e4886567",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "IntSlider(value=8, continuous_update=False, description='network_alpha:', max=32, min=1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3920454b18de463aaf59f9f9f09a9eba",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value='The following two values only apply to the additional layers of LoCon.')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aea343e14c604d639ab53452de56b17d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "IntSlider(value=16, continuous_update=False, description='conv_dim:', max=32, min=1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "82c38eccde5247fd8c100d23bc710533",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "IntSlider(value=8, continuous_update=False, description='conv_alpha:', max=32, min=1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0f09f2c2931644d3948b150446c935f0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HTML(value='<h2>▶️ Training</h2>')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a09a78d319b14525a2e9dcf840241032",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value='Adjust these parameters depending on your colab configuration.')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "aa9dc4c90d254b689f21be812dea95cc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value=\"If you're using the free tier you should select a diffusers model at the top of this cell.\")"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "688393cba3af487c8effda729bedfb95",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HTML(value='<hr>')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2e86adeaa8644553a5e5c573f9419e8b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value='Higher batch size is often faster but uses more memory.')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a400ea74c5fd4a8b8b375dced87377b8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "IntSlider(value=2, continuous_update=False, description='Train batch size:', max=16, min=1)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "82e133b7fa9244c786b6718a03378259",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value='I have found no substantial difference between sdpa and xformers.')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b3f57a766f454d099b28d2e040e8c2f4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "RadioButtons(description='Cross attention:', options=('sdpa', 'xformers'), value='sdpa')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "826d15ce06a94a70ab06762e118a082d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value=\"If you're on an A100 you should enable bf16.\")"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b3a11d27bc2e4106be9cf6a4c64c45d0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "RadioButtons(description='Mixed precision:', options=('bf16', 'fp16'), value='bf16')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ae1ac59ed7594d91813a610984917322",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value='Caching latents to drive will add a 250KB file next to each image but will use considerably less …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "449b555fbb33443bb32348d3a3d7220d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Checkbox(value=True, description='Cache latents.', indent=False)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "962ca9e9b8714ea692bfe2a3412bee5b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Checkbox(value=True, description='Cache latents to Drive.', indent=False)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e9c43f0c349749c7a9938258c76b7d02",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value='The following option will turn off shuffle_tags and disable text encoder training.')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d3cbed485e394c1088f1bd26fc4ec579",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Checkbox(value=False, description='Cache text encoder outputs:', indent=False)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "eedbfb48cf7143a69d246c2995666669",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HTML(value='<h2>▶️ Advanced</h2>')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "57572c6bc95345efa9d5c4521d7268dc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value='The optimizer is the algorithm used for training. AdanW8Bit is the default and works great, while…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "05fae4d64e7b4b7ca911da0502d8d40f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Dropdown(description='Optimizer:', index=1, options=('AdamW8bit', 'Prodigy', 'DAdaptation', 'DadaptAdam', 'Dad…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7d6e91c6b7a34126b2385a03334869a6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HTML(value=\"\\n        Recommended args for:\\n        <ul>\\n            <li><b>AdamW8bit</b>: <div style='font-…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4af83c7ebb4d416d8c8f19dbbaccab4d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Text(value='weight_decay=0.1 betas=[0.9,0.99]', description='Optimizer args:')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a046543205604731b76880c53b1f93c1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HTML(value=\"\\n        <p> If <b>Dadapt</b> or <b>Prodigy</b> are selected and the recommended box is checked, …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1ac127f4446042458473cd5ed67ee0c9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Checkbox(value=True, description='Recommended values', indent=False)"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import os\n",
    "import re\n",
    "import toml\n",
    "from time import time\n",
    "from huggingface_hub import HfFileSystem\n",
    "from IPython.display import Markdown, display\n",
    "import ipywidgets as widgets\n",
    "\n",
    "# These carry information from past executions\n",
    "if \"model_url\" in globals():\n",
    "  old_model_url = model_url\n",
    "else:\n",
    "  old_model_url = None\n",
    "if \"dependencies_installed\" not in globals():\n",
    "  dependencies_installed = False\n",
    "if \"model_file\" not in globals():\n",
    "  model_file = None\n",
    "\n",
    "# These may be set by other cells, some are legacy\n",
    "if \"custom_dataset\" not in globals():\n",
    "  custom_dataset = None\n",
    "if \"override_dataset_config_file\" not in globals():\n",
    "  override_dataset_config_file = None\n",
    "if \"override_config_file\" not in globals():\n",
    "  override_config_file = None  \n",
    "\n",
    "COLAB = False\n",
    "SOURCE = \"https://github.com/kohya-ss/sd-scripts\"\n",
    "TAG = \"tags/v0.8.7\"\n",
    "COMMIT = None\n",
    "BETTER_EPOCH_NAMES = True\n",
    "LOAD_TRUNCATED_IMAGES = True\n",
    "try:\n",
    "  LOWRAM = int(next(line.split()[1] for line in open('/proc/meminfo') if \"MemTotal\" in line)) / (1024**2) < 15\n",
    "except:\n",
    "  LOWRAM = False\n",
    "\n",
    "display(\n",
    "    widgets.HTML(\"<h1>🚩 Start Here</h1>\"),\n",
    ")\n",
    "\n",
    "project_name_widget = widgets.Text(description='Project name:')\n",
    "folder_structure_widget = widgets.Dropdown(\n",
    "    options=[\n",
    "        \"Organize by category (MyDrive/lora_training/datasets/project_name)\", \n",
    "        \"Organize by project (MyDrive/Loras/project_name/dataset)\"\n",
    "    ],\n",
    "    value=\"Organize by project (MyDrive/Loras/project_name/dataset)\",\n",
    "    description='Folder structure:',\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "training_model_widget = widgets.Dropdown(\n",
    "    options=[\"Pony Diffusion V6 XL\", \"Animagine XL V3\", \"Stable Diffusion XL 1.0 base\"],\n",
    "    value=\"Pony Diffusion V6 XL\",\n",
    "    description='Training model:',\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "optional_custom_training_model_url_widget = widgets.Text(\n",
    "    description = 'Optional: custom training model url'\n",
    ")\n",
    "\n",
    "load_diffusers_widget = widgets.Checkbox(\n",
    "    value=True,\n",
    "    description='Load diffusers',\n",
    "    disabled=False,\n",
    "    indent=False\n",
    ")\n",
    "\n",
    "wandb_key_widget = widgets.Text(description='Wandb key:')\n",
    "\n",
    "display(\n",
    "    widgets.HTML(\"<h2>▶️ Setup</h2>\"),\n",
    "    widgets.Label(\"Your project name will be the same as the folder containing your images. Spaces aren't allowed.\"),\n",
    "    project_name_widget,\n",
    "    widgets.Label(\"The folder structure doesn't matter and is purely for comfort. Make sure to always pick the same one. I like organizing by project.\"),\n",
    "    folder_structure_widget,\n",
    "    widgets.Label(\" Decide the model that will be downloaded and used for training. You can also choose your own by pasting its download link.\"),\n",
    "    training_model_widget,\n",
    "    optional_custom_training_model_url_widget,\n",
    "    widgets.Label(\"Using a diffusers model utilizes less resources. The 3 options above will work with or without diffusers.\"),\n",
    "    load_diffusers_widget,\n",
    "    widgets.Label(\"Use wandb if you want to visualize the progress of your training over time.\"),\n",
    "    wandb_key_widget\n",
    ")\n",
    "\n",
    "display(widgets.HTML('<h2> ▶️ Processing </h2>'))\n",
    "\n",
    "resolution_widget = widgets.BoundedIntText(\n",
    "    value=1024,\n",
    "    min=768,\n",
    "    max=1536,\n",
    "    step=128,\n",
    "    description='Resolution:',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "caption_extension_widget = widgets.Text(\n",
    "    value=\".txt\",\n",
    "    description=\"Caption extension\"\n",
    ")\n",
    "\n",
    "shuffle_tags_widget = widgets.Checkbox(\n",
    "    value=True,\n",
    "    description='Shuffle tags',\n",
    "    disabled=False,\n",
    "    indent=False\n",
    ")\n",
    "\n",
    "activation_tags_widget = widgets.Dropdown(\n",
    "    options= ['0', '1' , '2' ,'3'],\n",
    "    value='1',\n",
    "    description='Activation tags:',\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "display(\n",
    "    resolution_widget,\n",
    "    caption_extension_widget,\n",
    "    widgets.Label('Shuffling anime tags in place improves learning and prompting. An activation tag goes at the start of every text file and will not be shuffled.'),\n",
    "    shuffle_tags_widget,\n",
    "    activation_tags_widget,\n",
    ")\n",
    "\n",
    "display(widgets.HTML('<h2>▶️ Steps</h2>'))\n",
    "\n",
    "num_repeats_widget = widgets.IntText(\n",
    "    value=2,\n",
    "    description='# repeats:',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "preferred_unit_widget = widgets.RadioButtons(\n",
    "    options=[\"Epochs\", \"Steps\"],\n",
    "    value='Epochs', \n",
    "    description='Preferred unit',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "how_many_widget = widgets.IntText(\n",
    "    value=10,\n",
    "    description='How many:',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "save_every_n_epochs_widget = widgets.IntText(\n",
    "    value=1,\n",
    "    description='Saving every N epochs:',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "keep_only_last_n_epochs_widget = widgets.IntText(\n",
    "    value=10,\n",
    "    description='Keep only last N epochs:',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "display(\n",
    "    widgets.Label('Your images will repeat this number of times during training. I recommend that your images multiplied by their repeats is between 200 and 400.'),\n",
    "    num_repeats_widget,\n",
    "    widgets.Label('Choose how long you want to train for. A good starting point is around 10 epochs or around 2000 steps.'),\n",
    "    widgets.Label('One epoch is a number of steps equal to: your number of images multiplied by their repeats, divided by batch size.'),\n",
    "    preferred_unit_widget,\n",
    "    how_many_widget,\n",
    "    widgets.Label(\"Saving more epochs will let you compare your Lora's progress better.\"),\n",
    "    save_every_n_epochs_widget,\n",
    "    keep_only_last_n_epochs_widget,\n",
    ")\n",
    "\n",
    "unet_lr_widget = widgets.FloatText(\n",
    "    value=3e-4,\n",
    "    description='UNet LR:',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "text_encoder_lr_widget = widgets.FloatText(\n",
    "    value=6e-5,\n",
    "    description='Text encoder LR:',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "lr_scheduler_widget = widgets.Dropdown(\n",
    "    options=[\"constant\", \"cosine\", \"cosine_with_restarts\", \"constant_with_warmup\", \"linear\", \"polynomial\"],\n",
    "    value='cosine_with_restarts',\n",
    "    description='LR Scheduler:',\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "lr_scheduler_number_widget = widgets.IntText(\n",
    "    value=3,\n",
    "    description='LR Scheduler #:',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "lr_warmup_ratio_widget = widgets.FloatSlider(\n",
    "    value=0.05,\n",
    "    min=0,\n",
    "    max=0.2,\n",
    "    step=0.01,\n",
    "    description='LR Warmiup Ratio:',\n",
    "    disabled=False,\n",
    "    continuous_update=False,\n",
    "    orientation='horizontal',\n",
    "    readout=True,\n",
    "    readout_format='.1f',\n",
    ")\n",
    "\n",
    "min_snr_gamma_widget = widgets.FloatSlider(\n",
    "    value=8.0,\n",
    "    min=4,\n",
    "    max=16.0,\n",
    "    step=0.5,\n",
    "    description='Min SNR Gamma:',\n",
    "    disabled=False,\n",
    "    continuous_update=False,\n",
    "    orientation='horizontal',\n",
    "    readout=True,\n",
    "    readout_format='.1f',\n",
    ")\n",
    "\n",
    "multinoise_widget = widgets.Checkbox(\n",
    "    value=True,\n",
    "    description='Multinoise',\n",
    "    disabled=False,\n",
    "    indent=False\n",
    ")\n",
    "\n",
    "display(\n",
    "    widgets.HTML('<h2>▶️ Learning </h2>'),\n",
    "    widgets.Label('The learning rate is the most important for your results. If you want to train slower with lots of images, or if your dim and alpha are high, move the unet to 2e-4 or lower.'),\n",
    "    widgets.Label(\"The text encoder helps your Lora learn concepts slightly better. It is recommended to make it half or a fifth of the unet. If you're training a style you can even set it to 0.\"),\n",
    "    unet_lr_widget,\n",
    "    text_encoder_lr_widget,\n",
    "    widgets.Label(\"The scheduler is the algorithm that guides the learning rate. If you're not sure, pick `constant` and ignore the number. I personally recommend `cosine_with_restarts` with 3 restarts.\"),\n",
    "    lr_scheduler_widget,\n",
    "    lr_scheduler_number_widget,\n",
    "    widgets.Label('Steps spent \"warming up\" the learning rate during training for efficiency. I recommend leaving it at 5%.'),\n",
    "    lr_warmup_ratio_widget,\n",
    "    widgets.Label('Adjusts loss over time, makes learning much more efficient. The paper recommends 5.0, I recommend 8.0 for anime. Higher value makes it less strict. Set to 0 to disable.'),\n",
    "    min_snr_gamma_widget,\n",
    "    widgets.Label('Multinoise may help with color balance (darker darks, lighter lights).'),\n",
    "    multinoise_widget\n",
    ")\n",
    "\n",
    "lora_type_widget = widgets.RadioButtons(\n",
    "    options=[\"LoRA\", \"LoCon\"],\n",
    "    value='LoCon',\n",
    "    description='LoRA type:',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "network_dim_widget = widgets.IntSlider(\n",
    "    value=16,\n",
    "    min=1,\n",
    "    max=32,\n",
    "    step=1,\n",
    "    description='network_dim:',\n",
    "    disabled=False,\n",
    "    continuous_update=False,\n",
    "    orientation='horizontal',\n",
    "    readout=True,\n",
    "    readout_format='d'\n",
    ")\n",
    "\n",
    "network_alpha_widget = widgets.IntSlider(\n",
    "    value=8,\n",
    "    min=1,\n",
    "    max=32,\n",
    "    step=1,\n",
    "    description='network_alpha:',\n",
    "    disabled=False,\n",
    "    continuous_update=False,\n",
    "    orientation='horizontal',\n",
    "    readout=True,\n",
    "    readout_format='d'\n",
    ")\n",
    "\n",
    "conv_dim_widget = widgets.IntSlider(\n",
    "    value=16,\n",
    "    min=1,\n",
    "    max=32,\n",
    "    step=1,\n",
    "    description='conv_dim:',\n",
    "    disabled=False,\n",
    "    continuous_update=False,\n",
    "    orientation='horizontal',\n",
    "    readout=True,\n",
    "    readout_format='d'\n",
    ")\n",
    "\n",
    "conv_alpha_widget = widgets.IntSlider(\n",
    "    value=8,\n",
    "    min=1,\n",
    "    max=32,\n",
    "    step=1,\n",
    "    description='conv_alpha:',\n",
    "    disabled=False,\n",
    "    continuous_update=False,\n",
    "    orientation='horizontal',\n",
    "    readout=True,\n",
    "    readout_format='d'\n",
    ")\n",
    "\n",
    "display(\n",
    "    widgets.HTML('<h2>▶️ Structure</h2>'),\n",
    "    lora_type_widget,\n",
    "    widgets.Label('Below are some recommended XL values for the following settings:'),\n",
    "    widgets.HTML(\"\"\"\n",
    "        <table>\n",
    "            <thead>\n",
    "                <tr>\n",
    "                    <th>type</th>\n",
    "                    <th>network_dim</th>\n",
    "                    <th>network_alpha</th>\n",
    "                    <th>conv_dim</th>\n",
    "                    <th>conv_alpha</th>\n",
    "                </tr>\n",
    "            </thead>\n",
    "            <tbody>\n",
    "                <tr>\n",
    "                    <th>Regular LoRA</th>\n",
    "                    <td>8</td>\n",
    "                    <td>4</td>\n",
    "                    <td>-</td>\n",
    "                    <td>-</td>\n",
    "                </tr>\n",
    "                <tr>\n",
    "                    <th>Style LoCon</th>\n",
    "                    <td>16</td>\n",
    "                    <td>8</td>\n",
    "                    <td>16</td>\n",
    "                    <td>8<td>\n",
    "                </tr>\n",
    "            </tbody>\n",
    "        <table>\n",
    "    \"\"\"),\n",
    "    widgets.Label(\"More dim means larger Lora, it can hold more information but more isn't always better.\"),\n",
    "    network_dim_widget,\n",
    "    network_alpha_widget,\n",
    "    widgets.Label(\"The following two values only apply to the additional layers of LoCon.\"),\n",
    "    conv_dim_widget,\n",
    "    conv_alpha_widget,\n",
    ")\n",
    "\n",
    "train_batch_size_widget = widgets.IntSlider(\n",
    "    value=2,\n",
    "    min=1,\n",
    "    max=16,\n",
    "    step=1,\n",
    "    description='Train batch size:',\n",
    "    disabled=False,\n",
    "    continuous_update=False,\n",
    "    orientation='horizontal',\n",
    "    readout=True,\n",
    "    readout_format='d'\n",
    ")\n",
    "\n",
    "cross_attention_widget = widgets.RadioButtons(\n",
    "    options=[\"sdpa\", \"xformers\"],\n",
    "    value=\"sdpa\",\n",
    "    description='Cross attention:',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "mixed_precision_widget = widgets.RadioButtons(\n",
    "    options=[\"bf16\", \"fp16\"],\n",
    "    value=\"bf16\",\n",
    "    description='Mixed precision:',\n",
    "    disabled=False\n",
    ")\n",
    "\n",
    "cache_latents_widget = widgets.Checkbox(\n",
    "    value=True,\n",
    "    description='Cache latents.',\n",
    "    disabled=False,\n",
    "    indent=False\n",
    ")\n",
    "\n",
    "cache_latents_to_drive_widget = widgets.Checkbox(\n",
    "    value=True,\n",
    "    description='Cache latents to Drive.',\n",
    "    disabled=False,\n",
    "    indent=False\n",
    ")\n",
    "\n",
    "cache_text_encoder_outputs_widget = widgets.Checkbox(\n",
    "    value=False,\n",
    "    description='Cache text encoder outputs:',\n",
    "    disabled=False,\n",
    "    indent=False\n",
    ")\n",
    "\n",
    "display(\n",
    "    widgets.HTML('<h2>▶️ Training</h2>'),\n",
    "    widgets.Label('Adjust these parameters depending on your colab configuration.'),\n",
    "    widgets.Label('''If you're using the free tier you should select a diffusers model at the top of this cell.'''),\n",
    "    widgets.HTML('<hr>'),\n",
    "    widgets.Label('Higher batch size is often faster but uses more memory.'),\n",
    "    train_batch_size_widget,\n",
    "    widgets.Label('I have found no substantial difference between sdpa and xformers.'),\n",
    "    cross_attention_widget,\n",
    "    widgets.Label('''If you're on an A100 you should enable bf16.'''),\n",
    "    mixed_precision_widget,\n",
    "    widgets.Label('Caching latents to drive will add a 250KB file next to each image but will use considerably less memory.'),\n",
    "    cache_latents_widget,\n",
    "    cache_latents_to_drive_widget,\n",
    "    widgets.Label('The following option will turn off shuffle_tags and disable text encoder training.'),\n",
    "    cache_text_encoder_outputs_widget\n",
    ")\n",
    "\n",
    "optimizer_widget = widgets.Dropdown(\n",
    "    options=[\"AdamW8bit\", \"Prodigy\", \"DAdaptation\", \"DadaptAdam\", \"DadaptLion\", \"AdamW\", \"Lion\", \"SGDNesterov\", \"SGDNesterov8bit\", \"AdaFactor\"],\n",
    "    value=\"Prodigy\",\n",
    "    description='Optimizer:',\n",
    "    disabled=False,\n",
    ")\n",
    "\n",
    "optimizer_args_widget = widgets.Text(\n",
    "    description=\"Optimizer args:\",\n",
    "    value=\"weight_decay=0.1 betas=[0.9,0.99]\"\n",
    ")\n",
    "\n",
    "recommended_values_widget = widgets.Checkbox(\n",
    "    value=True,\n",
    "    description='Recommended values',\n",
    "    disabled=False,\n",
    "    indent=False\n",
    ")\n",
    "\n",
    "display(\n",
    "    widgets.HTML('<h2>▶️ Advanced</h2>'),\n",
    "    widgets.Label('The optimizer is the algorithm used for training. AdanW8Bit is the default and works great, while Prodigy manages learning rate automatically and may have several advantages such as training faster due to needing less steps as well as working better for small datasets.'),\n",
    "    optimizer_widget,\n",
    "    widgets.HTML('''\n",
    "        Recommended args for:\n",
    "        <ul>\n",
    "            <li><b>AdamW8bit</b>: <div style='font-family: monospace;'>weight_decay=0.1 betas=[0.9,0.99]</div></li>\n",
    "            <li><b>Prodigy</b>: <div style='font-family: monospace;'>decouple=True weight_decay=0.01 betas=[0.9,0.999] d_coef=2 use_bias_correction=True safeguard_warmup=True</div></li>\n",
    "            <li><b>CAME</b>: <div style='font-family: monospace;'>weight_decay=0.04</div></li>\n",
    "        </ul>\n",
    "    '''),\n",
    "    optimizer_args_widget,\n",
    "    widgets.HTML('''\n",
    "        <p> If <b>Dadapt</b> or <b>Prodigy</b> are selected and the recommended box is checked, the following recommended values will override any previous settings:\n",
    "        <div style='font-family: monospace;'>unet_lr=0.75, text_encoder_lr=0.75, network_alpha=network_dim</div></p>\n",
    "        <p> If <b>CAME + REX</b> is selected and the recommended box is checked, the following recommended values will override any previous settings:\n",
    "        <div style='font-family: monospace;'>unet_lr=1e-4, text_encoder_lr=1e-6, betas=default_came_betas</div></p>\n",
    "    '''),\n",
    "    recommended_values_widget\n",
    ")    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "OglZzI_ujZq-",
    "outputId": "e8af5902-6785-4d01-d6f6-9c1559e3a6c3"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b6330de462314a69823b8a38e40b9ccb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HTML(value='<h2>▶️ Ready</h2>')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c301162d079a4a21ad0fae56eeb8dd8f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Label(value='You can now run this cell to cook your Lora. Good luck!')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "46676625a8fe4aee83d270e11cafbacd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Button(button_style='success', description='Start', icon='play', style=ButtonStyle(), tooltip='Start')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# 👩‍💻 Cool code goes here\n",
    "\n",
    "import pathlib\n",
    "from accelerate.utils import write_basic_config\n",
    "\n",
    "def install_dependencies(root_dir, repo_dir):\n",
    "    global project_name, folder_structure, training_model, optional_custom_training_model_url\n",
    "    global training_model, training_model, model_url, vae_file, model_file, resolution, caption_extension\n",
    "    global shuffle_tags, shuffle_caption, activation_tags\n",
    "    global num_repeats, preferred_unit, how_many, max_train_epochs, max_train_steps\n",
    "    global max_train_epochs, max_train_steps, preferred_unit, how_many\n",
    "    global save_every_n_epochs, keep_only_last_n_epochs, unet_lr, text_encoder_lr\n",
    "    global lr_scheduler, lr_scheduler_number, lr_warmup_ratio, lr_warmup_steps\n",
    "    global min_snr_gamma, multinoise, lora_type, network_dim, conv_alpha\n",
    "    global network_module, network_args, train_batch_size, cross_attention \n",
    "    global mixed_precision, cache_latents, cache_latents_to_drive\n",
    "    global cache_text_encoder_outputs, optimizer, optimizer_args, recommended_values\n",
    "    global images_folder\n",
    "\n",
    "    os.chdir(root_dir)\n",
    "    global SOURCE\n",
    "    !git clone {SOURCE} {repo_dir}\n",
    "    os.chdir(repo_dir)\n",
    "    global TAG\n",
    "    if TAG:\n",
    "        !git checkout {TAG}\n",
    "        !git reset --hard\n",
    "            \n",
    "    !wget https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/train_network_xl_wrapper.py -q -O train_network_xl_wrapper.py\n",
    "    !wget https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/dracula.py -q -O dracula.py\n",
    "\n",
    "    !pip install -r requirements.txt\n",
    "            \n",
    "    #if cross_attention == \"xformers\":\n",
    "    #    !pip install -q xformers==0.0.26.dev778\n",
    "    if \"CAME\" in optimizer:\n",
    "        !pip install came-pytorch\n",
    "        !wget https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/train_util.py -q -O library/train_util.py\n",
    "\n",
    "    # patch kohya for minor stuff\n",
    "    if LOWRAM:\n",
    "        !sed -i \"s@cpu@cuda@\" library/model_util.py\n",
    "    if LOAD_TRUNCATED_IMAGES:\n",
    "        !sed -i 's/from PIL import Image/from PIL import Image, ImageFile\\nImageFile.LOAD_TRUNCATED_IMAGES=True/g' library/train_util.py # fix truncated jpegs error\n",
    "    if BETTER_EPOCH_NAMES:\n",
    "        !sed -i 's/{:06d}/{:02d}/g' library/train_util.py # make epoch names shorter\n",
    "        !sed -i 's/\".\" + args.save_model_as)/\"-{:02d}.\".format(num_train_epochs) + args.save_model_as)/g' train_network.py # name of the last epoch will match the rest\n",
    "\n",
    "    if not os.path.exists(accelerate_config_file):\n",
    "        write_basic_config(save_location=accelerate_config_file)\n",
    "\n",
    "    os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n",
    "    os.environ[\"BITSANDBYTES_NOWELCOME\"] = \"1\"\n",
    "    os.environ[\"SAFETENSORS_FAST_GPU\"] = \"1\"\n",
    "\n",
    "def validate_dataset():\n",
    "    global lr_warmup_steps, lr_warmup_ratio, caption_extension, keep_tokens\n",
    "    global project_name, folder_structure, training_model, optional_custom_training_model_url\n",
    "    global training_model, training_model, model_url, vae_file, model_file, resolution, caption_extension\n",
    "    global shuffle_tags, shuffle_caption, activation_tags\n",
    "    global num_repeats, preferred_unit, how_many, max_train_epochs, max_train_steps\n",
    "    global max_train_epochs, max_train_steps, preferred_unit, how_many\n",
    "    global save_every_n_epochs, keep_only_last_n_epochs, unet_lr, text_encoder_lr\n",
    "    global lr_scheduler, lr_scheduler_number, lr_warmup_ratio, lr_warmup_steps\n",
    "    global min_snr_gamma, multinoise, lora_type, network_dim, conv_alpha\n",
    "    global network_module, network_args, train_batch_size, cross_attention \n",
    "    global mixed_precision, cache_latents, cache_latents_to_drive\n",
    "    global cache_text_encoder_outputs, optimizer, optimizer_args, recommended_values\n",
    "    global images_folder\n",
    "    \n",
    "    supported_types = (\".png\", \".jpg\", \".jpeg\", \".webp\", \".bmp\")\n",
    "\n",
    "    print(\"\\n💿 Checking dataset...\")\n",
    "    if not project_name.strip() or any(c in project_name for c in \" .()\\\"'\\\\/\"):\n",
    "        print(\"💥 Error: Please choose a valid project name.\")\n",
    "        return\n",
    "\n",
    "    if custom_dataset:\n",
    "        try:\n",
    "            datconf = toml.loads(custom_dataset)\n",
    "            datasets = [d for d in datconf[\"datasets\"][0][\"subsets\"]]\n",
    "        except:\n",
    "            print(f\"💥 Error: Your custom dataset is invalid or contains an error! Please check the original template.\")\n",
    "            return\n",
    "        reg = [d.get(\"image_dir\") for d in datasets if d.get(\"is_reg\", False)]\n",
    "        datasets_dict = {d[\"image_dir\"]: d[\"num_repeats\"] for d in datasets}\n",
    "        folders = datasets_dict.keys()\n",
    "        files = [f for folder in folders for f in os.listdir(folder)]\n",
    "        images_repeats = {folder: (len([f for f in os.listdir(folder) if f.lower().endswith(supported_types)]), datasets_dict[folder]) for folder in folders}\n",
    "    else:\n",
    "        reg = []\n",
    "        folders = [images_folder]\n",
    "        files = os.listdir(images_folder)\n",
    "        images_repeats = {images_folder: (len([f for f in files if f.lower().endswith(supported_types)]), num_repeats)}\n",
    "\n",
    "    for folder in folders:\n",
    "        if not os.path.exists(folder):\n",
    "            print(f\"💥 Error: The folder {folder.replace('/content/drive/', '')} doesn't exist.\")\n",
    "            return\n",
    "    for folder, (img, rep) in images_repeats.items():\n",
    "        if not img:\n",
    "            print(f\"💥 Error: Your {folder.replace('/content/drive/', '')} folder is empty.\")\n",
    "            return\n",
    "    for f in files:\n",
    "        if not f.lower().endswith((\".txt\", \".npz\")) and not f.lower().endswith(supported_types):\n",
    "            print(f\"💥 Error: Invalid file in dataset: \\\"{f}\\\". Aborting.\")\n",
    "            return\n",
    "\n",
    "    if not [txt for txt in files if txt.lower().endswith(\".txt\")]:\n",
    "        caption_extension = \"\"\n",
    "\n",
    "    pre_steps_per_epoch = sum(img*rep for (img, rep) in images_repeats.values())\n",
    "    steps_per_epoch = pre_steps_per_epoch/train_batch_size\n",
    "    total_steps = max_train_steps or int(max_train_epochs*steps_per_epoch)\n",
    "    estimated_epochs = int(total_steps/steps_per_epoch)\n",
    "    lr_warmup_steps = int(total_steps*lr_warmup_ratio)\n",
    "\n",
    "    for folder, (img, rep) in images_repeats.items():\n",
    "        print(\"📁\"+folder.replace(\"/content/drive/\", \"\") + (\" (Regularization)\" if folder in reg else \"\"))\n",
    "        print(f\"📈 Found {img} images with {rep} repeats, equaling {img*rep} steps.\")\n",
    "    print(f\"📉 Divide {pre_steps_per_epoch} steps by {train_batch_size} batch size to get {steps_per_epoch} steps per epoch.\")\n",
    "    if max_train_epochs:\n",
    "        print(f\"🔮 There will be {max_train_epochs} epochs, for around {total_steps} total training steps.\")\n",
    "    else:\n",
    "        print(f\"🔮 There will be {total_steps} steps, divided into {estimated_epochs} epochs and then some.\")\n",
    "\n",
    "    if total_steps > 10000:\n",
    "        print(\"💥 Error: Your total steps are too high. You probably made a mistake. Aborting...\")\n",
    "        return\n",
    "\n",
    "    return True\n",
    "\n",
    "def create_config(images_folder, output_folder, log_folder):\n",
    "    global dataset_config_file, config_file, model_file\n",
    "    global project_name, folder_structure, training_model, optional_custom_training_model_url\n",
    "    global training_model, training_model, model_url, vae_file, model_file, resolution, caption_extension\n",
    "    global shuffle_tags, shuffle_caption, activation_tags\n",
    "    global num_repeats, preferred_unit, how_many, max_train_epochs, max_train_steps\n",
    "    global max_train_epochs, max_train_steps, preferred_unit, how_many\n",
    "    global save_every_n_epochs, keep_only_last_n_epochs, unet_lr, text_encoder_lr\n",
    "    global lr_scheduler, lr_scheduler_number, lr_warmup_ratio, lr_warmup_steps\n",
    "    global min_snr_gamma, multinoise, lora_type, network_dim, conv_alpha\n",
    "    global network_module, network_args, train_batch_size, cross_attention \n",
    "    global mixed_precision, cache_latents, cache_latents_to_drive\n",
    "    global cache_text_encoder_outputs, optimizer, optimizer_args, recommended_values\n",
    "\n",
    "    if override_config_file:\n",
    "        config_file = override_config_file\n",
    "        print(f\"\\n⭕ Using custom config file {config_file}\")\n",
    "    else:\n",
    "        config_dict = {\n",
    "            \"network_arguments\": {\n",
    "                \"unet_lr\": unet_lr,\n",
    "                \"text_encoder_lr\": text_encoder_lr if not cache_text_encoder_outputs else 0,\n",
    "                \"network_dim\": network_dim,\n",
    "                \"network_alpha\": network_alpha,\n",
    "                \"network_module\": network_module,\n",
    "                \"network_args\": network_args,\n",
    "                \"network_train_unet_only\": text_encoder_lr == 0 or cache_text_encoder_outputs,\n",
    "            },\n",
    "            \"optimizer_arguments\": {\n",
    "                \"learning_rate\": unet_lr,\n",
    "                \"lr_scheduler\": lr_scheduler,\n",
    "                \"lr_scheduler_num_cycles\": lr_scheduler_num_cycles if lr_scheduler == \"cosine_with_restarts\" else None,\n",
    "                \"lr_scheduler_power\": lr_scheduler_power if lr_scheduler == \"polynomial\" else None,\n",
    "                \"lr_warmup_steps\": lr_warmup_steps if lr_scheduler != \"constant\" else None,\n",
    "                \"optimizer_type\": optimizer,\n",
    "                \"optimizer_args\": optimizer_args if optimizer_args else None,\n",
    "            },\n",
    "            \"training_arguments\": {\n",
    "                \"pretrained_model_name_or_path\": model_file,\n",
    "                \"vae\": vae_file,\n",
    "                \"max_train_steps\": max_train_steps,\n",
    "                \"max_train_epochs\": max_train_epochs,\n",
    "                \"train_batch_size\": train_batch_size,\n",
    "                \"seed\": 42,\n",
    "                \"max_token_length\": 225,\n",
    "                \"xformers\": cross_attention == \"xformers\",\n",
    "                \"sdpa\": cross_attention == \"sdpa\",\n",
    "                \"min_snr_gamma\": min_snr_gamma if min_snr_gamma > 0 else None,\n",
    "                \"lowram\": LOWRAM,\n",
    "                \"no_half_vae\": True,\n",
    "                \"gradient_checkpointing\": True,\n",
    "                \"gradient_accumulation_steps\": 1,\n",
    "                \"max_data_loader_n_workers\": 8,\n",
    "                \"persistent_data_loader_workers\": True,\n",
    "                \"mixed_precision\": mixed_precision,\n",
    "                \"full_bf16\": mixed_precision == \"bf16\",\n",
    "                \"cache_latents\": cache_latents,\n",
    "                \"cache_latents_to_disk\": cache_latents_to_drive,\n",
    "                \"cache_text_encoder_outputs\": cache_text_encoder_outputs,\n",
    "                \"min_timestep\": 0,\n",
    "                \"max_timestep\": 1000,\n",
    "                \"prior_loss_weight\": 1.0,\n",
    "                \"multires_noise_iterations\": 6 if multinoise else None,\n",
    "                \"multires_noise_discount\": 0.3 if multinoise else None,\n",
    "            },\n",
    "            \"saving_arguments\": {\n",
    "                \"save_precision\": \"fp16\",\n",
    "                \"save_model_as\": \"safetensors\",\n",
    "                \"save_every_n_epochs\": save_every_n_epochs,\n",
    "                \"save_last_n_epochs\": keep_only_last_n_epochs,\n",
    "                \"output_name\": project_name,\n",
    "                \"output_dir\": output_folder,\n",
    "                \"log_prefix\": project_name,\n",
    "                \"logging_dir\": log_folder,\n",
    "                \"wandb_api_key\": wandb_key if wandb_key else None,\n",
    "                \"log_with\": \"wandb\" if wandb_key else None,\n",
    "            }\n",
    "        }\n",
    "\n",
    "        for key in config_dict:\n",
    "            if isinstance(config_dict[key], dict):\n",
    "                config_dict[key] = {k: v for k, v in config_dict[key].items() if v is not None}\n",
    "\n",
    "        with open(config_file, \"w\") as f:\n",
    "            f.write(toml.dumps(config_dict))\n",
    "        print(f\"\\n📄 Config saved to {config_file}\")\n",
    "\n",
    "    if override_dataset_config_file:\n",
    "        dataset_config_file = override_dataset_config_file\n",
    "        print(f\"⭕ Using custom dataset config file {dataset_config_file}\")\n",
    "    else:\n",
    "        dataset_config_dict = {\n",
    "            \"general\": {\n",
    "                \"resolution\": resolution,\n",
    "                \"shuffle_caption\": shuffle_caption and not cache_text_encoder_outputs,\n",
    "                \"keep_tokens\": keep_tokens,\n",
    "                \"flip_aug\": False,\n",
    "                \"caption_extension\": caption_extension,\n",
    "                \"enable_bucket\": True,\n",
    "                \"bucket_no_upscale\": False,\n",
    "                \"bucket_reso_steps\": 64,\n",
    "                \"min_bucket_reso\": 256,\n",
    "                \"max_bucket_reso\": 4096,\n",
    "            },\n",
    "            \"datasets\": toml.loads(custom_dataset)[\"datasets\"] if custom_dataset else [\n",
    "                {\n",
    "                    \"subsets\": [\n",
    "                        {\n",
    "                            \"num_repeats\": num_repeats,\n",
    "                            \"image_dir\": images_folder,\n",
    "                            \"class_tokens\": None if caption_extension else project_name\n",
    "                        }\n",
    "                    ]\n",
    "                }\n",
    "            ]\n",
    "        }\n",
    "\n",
    "        for key in dataset_config_dict:\n",
    "            if isinstance(dataset_config_dict[key], dict):\n",
    "                dataset_config_dict[key] = {k: v for k, v in dataset_config_dict[key].items() if v is not None}\n",
    "\n",
    "        with open(dataset_config_file, \"w\") as f:\n",
    "            f.write(toml.dumps(dataset_config_dict))\n",
    "        print(f\"📄 Dataset config saved to {dataset_config_file}\")\n",
    "\n",
    "def download_model(load_diffusers):\n",
    "    global old_model_url, model_url, model_file, vae_url, vae_file\n",
    "    real_model_url = model_url.strip()\n",
    "\n",
    "    if load_diffusers:\n",
    "        if 'huggingface.co' in real_model_url:\n",
    "            match = re.search(r'huggingface\\.co/([^/]+)/([^/]+)', real_model_url)\n",
    "            if match:\n",
    "                username = match.group(1)\n",
    "                model_name = match.group(2)\n",
    "                model_file = f\"{username}/{model_name}\"\n",
    "                fs = HfFileSystem()\n",
    "                existing_folders = set(fs.ls(model_file, detail=False))\n",
    "                necessary_folders = [ \"scheduler\", \"text_encoder\", \"text_encoder_2\", \"tokenizer\", \"tokenizer_2\", \"unet\", \"vae\" ]\n",
    "                if all(f\"{model_file}/{folder}\" in existing_folders for folder in necessary_folders):\n",
    "                    print(\"🍃 Diffusers model identified.\")\n",
    "                    return True\n",
    "        print(\"🍃 Failed to load Diffusers model.\")\n",
    "\n",
    "    if not model_file:\n",
    "        if real_model_url.lower().endswith((\".ckpt\", \".safetensors\")):\n",
    "            model_file = f\"/content{real_model_url[real_model_url.rfind('/'):]}\"\n",
    "        else:\n",
    "            model_file = \"/content/downloaded_model.safetensors\"\n",
    "            if os.path.exists(model_file):\n",
    "                !rm \"{model_file}\"\n",
    "\n",
    "    if m := re.search(r\"(?:https?://)?(?:www\\.)?huggingface\\.co/[^/]+/[^/]+/blob\", model_url):\n",
    "        real_model_url = real_model_url.replace(\"blob\", \"resolve\")\n",
    "    elif m := re.search(r\"(?:https?://)?(?:www\\\\.)?civitai\\.com/models/([0-9]+)(/[A-Za-z0-9-_]+)?\", model_url):\n",
    "        if m.group(2):\n",
    "            model_file = f\"/content{m.group(2)}.safetensors\"\n",
    "        if m := re.search(r\"modelVersionId=([0-9]+)\", model_url):\n",
    "            real_model_url = f\"https://civitai.com/api/download/models/{m.group(1)}\"\n",
    "        else:\n",
    "            raise ValueError(\"optional_custom_training_model_url contains a civitai link, but the link doesn't include a modelVersionId. You can also right click the download button to copy the direct download link.\")\n",
    "\n",
    "    !aria2c \"{real_model_url}\" --console-log-level=warn -c -s 16 -x 16 -k 10M -d / -o \"{model_file}\"\n",
    "\n",
    "    if not os.path.exists(vae_file):\n",
    "        !aria2c \"{vae_url}\" --console-log-level=warn -c -s 16 -x 16 -k 10M -d / -o \"{vae_file}\"\n",
    "\n",
    "    if model_file.lower().endswith(\".safetensors\"):\n",
    "        from safetensors.torch import load_file as load_safetensors\n",
    "        try:\n",
    "            test = load_safetensors(model_file)\n",
    "            del test\n",
    "        except:\n",
    "            #if \"HeaderTooLarge\" in str(e):\n",
    "            new_model_file = os.path.splitext(model_file)[0]+\".ckpt\"\n",
    "            !mv \"{model_file}\" \"{new_model_file}\"\n",
    "            model_file = new_model_file\n",
    "            print(f\"Renamed model to {os.path.splitext(model_file)[0]}.ckpt\")\n",
    "\n",
    "    if model_file.lower().endswith(\".ckpt\"):\n",
    "        from torch import load as load_ckpt\n",
    "        try:\n",
    "            test = load_ckpt(model_file)\n",
    "            del test\n",
    "        except:\n",
    "            return False\n",
    "\n",
    "    return True\n",
    "\n",
    "def main(b):\n",
    "    global project_name, folder_structure, training_model, optional_custom_training_model_url\n",
    "    global training_model, training_model, model_url, vae_file, model_file, resolution, caption_extension\n",
    "    global shuffle_tags, shuffle_caption, activation_tags\n",
    "    global num_repeats, preferred_unit, how_many, max_train_epochs, max_train_steps\n",
    "    global max_train_epochs, max_train_steps, preferred_unit, how_many\n",
    "    global save_every_n_epochs, keep_only_last_n_epochs, unet_lr, text_encoder_lr\n",
    "    global lr_scheduler, lr_scheduler_number, lr_warmup_ratio, lr_warmup_steps\n",
    "    global min_snr_gamma, multinoise, lora_type, network_dim, conv_alpha\n",
    "    global network_module, network_args, train_batch_size, cross_attention \n",
    "    global mixed_precision, cache_latents, cache_latents_to_drive\n",
    "    global cache_text_encoder_outputs, optimizer, optimizer_args, recommended_values\n",
    "    global images_folder, root_dir, network_alpha, wandb_key, vae_url, keep_tokens\n",
    "    global lr_scheduler_num_cycles, lr_scheduler_power, output_folder\n",
    "    \n",
    "    project_name = project_name_widget.value\n",
    "    project_name = project_name.strip()\n",
    "    \n",
    "    folder_structure = folder_structure_widget.value\n",
    "    training_model = training_model_widget.value\n",
    "    optional_custom_training_model_url = optional_custom_training_model_url_widget.value\n",
    "    load_diffusers = load_diffusers_widget.value\n",
    "    wandb_key = wandb_key_widget.value\n",
    "    \n",
    "    if optional_custom_training_model_url:\n",
    "        model_url = optional_custom_training_model_url\n",
    "    elif \"Pony\" in training_model:\n",
    "        if load_diffusers:\n",
    "            model_url = \"https://huggingface.co/hollowstrawberry/67AB2F\"\n",
    "        else:\n",
    "            model_url = \"https://civitai.com/api/download/models/290640\"\n",
    "        model_file = \"/content/ponyDiffusionV6XL.safetensors\"\n",
    "    elif \"Animagine\" in training_model:\n",
    "        if load_diffusers:\n",
    "            model_url = \"https://huggingface.co/cagliostrolab/animagine-xl-3.0\"\n",
    "        else:\n",
    "            model_url = \"https://civitai.com/api/download/models/293564\"\n",
    "        model_file = \"/content/animagineXLV3.safetensors\"\n",
    "    else:\n",
    "        if load_diffusers:\n",
    "            model_url = \"https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/\"\n",
    "        else:\n",
    "            model_url = \"https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors\"\n",
    "    \n",
    "    if load_diffusers:\n",
    "        vae_file = \"stabilityai/sdxl-vae\"\n",
    "    else:\n",
    "        vae_url = \"https://huggingface.co/stabilityai/sdxl-vae/resolve/main/sdxl_vae.safetensors\"\n",
    "        vae_file = \"/content/sdxl_vae.safetensors\"\n",
    "    \n",
    "    resolution = resolution_widget.value\n",
    "    caption_extension = caption_extension_widget.value\n",
    "\n",
    "    shuffle_tags = shuffle_tags_widget.value\n",
    "    shuffle_caption = shuffle_tags\n",
    "    activation_tags = activation_tags_widget.value\n",
    "    keep_tokens = int(activation_tags)\n",
    "    \n",
    "    num_repeats = num_repeats_widget.value\n",
    "    preferred_unit = preferred_unit_widget.value\n",
    "    how_many = how_many_widget.value\n",
    "    \n",
    "    max_train_epochs = how_many if preferred_unit == \"Epochs\" else None\n",
    "    max_train_steps = how_many if preferred_unit == \"Steps\" else None\n",
    "    \n",
    "    save_every_n_epochs = save_every_n_epochs_widget.value\n",
    "    keep_only_last_n_epochs = keep_only_last_n_epochs_widget.value\n",
    "    \n",
    "    if not save_every_n_epochs:\n",
    "        save_every_n_epochs = max_train_epochs\n",
    "    if not keep_only_last_n_epochs:\n",
    "        keep_only_last_n_epochs = max_train_epochs\n",
    "    \n",
    "    unet_lr = unet_lr_widget.value\n",
    "    text_encoder_lr = text_encoder_lr_widget.value\n",
    "    \n",
    "    lr_scheduler = lr_scheduler_widget.value\n",
    "    lr_scheduler_number = lr_scheduler_number_widget.value\n",
    "    \n",
    "    lr_warmup_ratio = lr_warmup_ratio_widget.value\n",
    "    lr_warmup_steps = 0\n",
    "    \n",
    "    min_snr_gamma = min_snr_gamma_widget.value\n",
    "    \n",
    "    multinoise = multinoise_widget.value\n",
    "    \n",
    "    lora_type = lora_type_widget.value\n",
    "    \n",
    "    network_dim = network_dim_widget.value\n",
    "    network_alpha = network_alpha_widget.value\n",
    "    \n",
    "    conv_dim = conv_dim_widget.value\n",
    "    conv_alpha = conv_alpha_widget.value\n",
    "    \n",
    "    network_module = \"networks.lora\"\n",
    "    network_args = None\n",
    "    if lora_type.lower() == \"locon\":\n",
    "        network_args = [f\"conv_dim={conv_dim}\", f\"conv_alpha={conv_alpha}\"]\n",
    "\n",
    "    train_batch_size = train_batch_size_widget.value\n",
    "    cross_attention = cross_attention_widget.value\n",
    "    \n",
    "    mixed_precision = mixed_precision_widget.value\n",
    "    \n",
    "    cache_latents = cache_latents_widget.value\n",
    "    cache_latents_to_drive = cache_latents_to_drive_widget.value\n",
    "    \n",
    "    cache_text_encoder_outputs = cache_text_encoder_outputs_widget.value\n",
    "    \n",
    "    optimizer = optimizer_widget.value\n",
    "    \n",
    "    optimizer_args = optimizer_args_widget.value\n",
    "    optimizer_args = [a.strip() for a in optimizer_args.split(' ') if a]\n",
    "    \n",
    "    recommended_values = recommended_values_widget.value\n",
    "    \n",
    "    if any(opt in optimizer.lower() for opt in [\"dadapt\", \"prodigy\"]):\n",
    "        if recommended_values:\n",
    "            unet_lr = 0.75\n",
    "            text_encoder_lr = 0.75\n",
    "            network_alpha = network_dim\n",
    "    elif \"CAME\" in optimizer:\n",
    "        optimizer = \"CAME\"\n",
    "        lr_scheduler = \"REX\"\n",
    "        if recommended_values:\n",
    "            unet_lr = 1e-4\n",
    "            text_encoder_lr = 1e-6\n",
    "            for i in range(len(optimizer_args)):\n",
    "                if \"betas\" in optimizer_args[i]:\n",
    "                    optimizer_args.pop(i)\n",
    "                    break\n",
    "    \n",
    "    lr_scheduler_num_cycles = lr_scheduler_number if lr_scheduler == \"cosine_with_restarts\" else 0\n",
    "    lr_scheduler_power = lr_scheduler_number if lr_scheduler == \"polynomial\" else 0\n",
    "\n",
    "    root_dir = \"/content\" if COLAB else pathlib.Path.home() / \"Loras\"\n",
    "    deps_dir = os.path.join(root_dir, \"deps\")\n",
    "    repo_dir = os.path.join(root_dir, \"sd-scripts\")\n",
    "\n",
    "    global images_folder\n",
    "    if \"/Loras\" in folder_structure:\n",
    "        main_dir            = os.path.join(root_dir, \"drive/MyDrive/Loras\") if COLAB else root_dir\n",
    "        log_folder        = os.path.join(main_dir, \"_logs\")\n",
    "        config_folder = os.path.join(main_dir, project_name)\n",
    "        images_folder = os.path.join(main_dir, project_name, \"dataset\")\n",
    "        output_folder = os.path.join(main_dir, project_name, \"output\")\n",
    "    else:\n",
    "        main_dir            = os.path.join(root_dir, \"drive/MyDrive/lora_training\") if COLAB else root_dir\n",
    "        images_folder = os.path.join(main_dir, \"datasets\", project_name)\n",
    "        output_folder = os.path.join(main_dir, \"output\", project_name)\n",
    "        config_folder = os.path.join(main_dir, \"config\", project_name)\n",
    "        log_folder        = os.path.join(main_dir, \"log\")\n",
    "    \n",
    "    global config_file, dataset_config_file, accelerate_config_file\n",
    "    \n",
    "    config_file = os.path.join(config_folder, \"training_config.toml\")\n",
    "    dataset_config_file = os.path.join(config_folder, \"dataset_config.toml\")\n",
    "    accelerate_config_file = os.path.join(repo_dir, \"accelerate_config/config.yaml\")\n",
    "    \n",
    "    global dependencies_installed\n",
    "\n",
    "    if COLAB and not os.path.exists('/content/drive'):\n",
    "        from google.colab import drive\n",
    "        print(\"📂 Connecting to Google Drive...\")\n",
    "        drive.mount('/content/drive')\n",
    "\n",
    "    for dir in (main_dir, deps_dir, repo_dir, log_folder, images_folder, output_folder, config_folder):\n",
    "        os.makedirs(dir, exist_ok=True)\n",
    "\n",
    "    if not validate_dataset():\n",
    "        return\n",
    "\n",
    "    if not dependencies_installed:\n",
    "        print(\"\\n🏭 Installing dependencies...\\n\")\n",
    "        t0 = time()\n",
    "        install_dependencies(root_dir, repo_dir)\n",
    "        t1 = time()\n",
    "        dependencies_installed = True\n",
    "        print(f\"\\n✅ Installation finished in {int(t1-t0)} seconds.\")\n",
    "    else:\n",
    "        print(\"\\n✅ Dependencies already installed.\")\n",
    "\n",
    "    if old_model_url != model_url or not model_file or not os.path.exists(model_file):\n",
    "        print(\"\\n🔄 Downloading model...\")\n",
    "        if not download_model(load_diffusers):\n",
    "            print(\"\\n💥 Error: The model you selected is invalid or corrupted, or couldn't be downloaded. You can use a civitai or huggingface link, or any direct download link.\")\n",
    "            return\n",
    "        print()\n",
    "    else:\n",
    "        print(\"\\n🔄 Model already downloaded.\\n\")\n",
    "\n",
    "    create_config(images_folder, output_folder, log_folder)\n",
    "\n",
    "    print(\"\\n⭐ Starting trainer...\\n\")\n",
    "    os.chdir(repo_dir)\n",
    "\n",
    "    !accelerate launch --quiet --config_file={accelerate_config_file} --num_cpu_threads_per_process=1 train_network_xl_wrapper.py --dataset_config={dataset_config_file} --config_file={config_file}\n",
    "\n",
    "    if not get_ipython().__dict__['user_ns']['_exit_code']:\n",
    "        display(Markdown(\"### ✅ Done! [Go download your Lora from Google Drive](https://drive.google.com/drive/my-drive)\\n\"\n",
    "                                         \"### There will be several files, you should try the latest version (the file with the largest number next to it)\"))\n",
    "\n",
    "run_btn = widgets.Button(\n",
    "    description='Start',\n",
    "    disabled=False,\n",
    "    button_style='success', \n",
    "    tooltip='Start',\n",
    "    icon='play' \n",
    ")\n",
    "\n",
    "display(\n",
    "    widgets.HTML('<h2>▶️ Ready</h2>'),\n",
    "    widgets.Label('You can now run this cell to cook your Lora. Good luck!'),\n",
    "    run_btn\n",
    ")\n",
    "\n",
    "run_btn.on_click(main)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mBMUJ7BuvNcn"
   },
   "source": [
    "## *️⃣ Extras\n",
    "\n",
    "You can run these before starting the training."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Wd4916Eu1tb9"
   },
   "source": [
    "### 📚 Multiple folders in dataset\n",
    "Below is a template allowing you to define multiple folders in your dataset. You must include the location of each folder and you can set different number of repeats for each one. To add more folders simply copy and paste the sections starting with `[[datasets.subsets]]`.\n",
    "\n",
    "When enabling this, the number of repeats set in the main cell will be ignored, and the main folder set by the project name will also be ignored.\n",
    "\n",
    "You can make one of them a regularization folder by adding `is_reg = true`  \n",
    "You can also set different `keep_tokens`, `flip_aug`, etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "executionInfo": {
     "elapsed": 804,
     "status": "ok",
     "timestamp": 1708481414906,
     "user": {
      "displayName": "Samuel",
      "userId": "11002898838993959229"
     },
     "user_tz": 180
    },
    "id": "Y037lagnJWmn"
   },
   "outputs": [],
   "source": [
    "custom_dataset = \"\"\"\n",
    "[[datasets]]\n",
    "\n",
    "[[datasets.subsets]]\n",
    "image_dir = \"/content/drive/MyDrive/Loras/example/dataset/good_images\"\n",
    "num_repeats = 3\n",
    "\n",
    "[[datasets.subsets]]\n",
    "image_dir = \"/content/drive/MyDrive/Loras/example/dataset/normal_images\"\n",
    "num_repeats = 1\n",
    "\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "W84Jxf-U2TIU"
   },
   "outputs": [],
   "source": [
    "custom_dataset = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "cellView": "form",
    "id": "WDjkp4scvPgE"
   },
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'google.colab'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[12], line 9\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mos\u001b[39;00m\u001b[38;5;241m,\u001b[39m \u001b[38;5;21;01mzipfile\u001b[39;00m\n\u001b[1;32m      8\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mexists(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m/content/drive\u001b[39m\u001b[38;5;124m'\u001b[39m):\n\u001b[0;32m----> 9\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mgoogle\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcolab\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m drive\n\u001b[1;32m     10\u001b[0m     \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m📂 Connecting to Google Drive...\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m     11\u001b[0m     drive\u001b[38;5;241m.\u001b[39mmount(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m/content/drive\u001b[39m\u001b[38;5;124m'\u001b[39m)\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'google.colab'"
     ]
    }
   ],
   "source": [
    "#@markdown ### 📂 Unzip dataset\n",
    "#@markdown It's much slower to upload individual files to your Drive, so you may want to upload a zip if you have your dataset in your computer.\n",
    "zip = \"/content/drive/MyDrive/my_dataset.zip\" #@param {type:\"string\"}\n",
    "extract_to = \"/content/drive/MyDrive/Loras/example/dataset\" #@param {type:\"string\"}\n",
    "\n",
    "import os, zipfile\n",
    "\n",
    "if not os.path.exists('/content/drive'):\n",
    "    from google.colab import drive\n",
    "    print(\"📂 Connecting to Google Drive...\")\n",
    "    drive.mount('/content/drive')\n",
    "\n",
    "os.makedirs(extract_to, exist_ok=True)\n",
    "\n",
    "with zipfile.ZipFile(zip, 'r') as f:\n",
    "    f.extractall(extract_to)\n",
    "\n",
    "print(\"✅ Done\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "aKWlpsG0jrX3"
   },
   "outputs": [],
   "source": [
    "#@markdown ### 🔢 Count datasets\n",
    "#@markdown Google Drive makes it impossible to count the files in a folder, so this will show you the file counts in all folders and subfolders.\n",
    "folder = \"/content/drive/MyDrive/Loras\" #@param {type:\"string\"}\n",
    "\n",
    "import os\n",
    "from google.colab import drive\n",
    "\n",
    "if not os.path.exists('/content/drive'):\n",
    "    print(\"📂 Connecting to Google Drive...\\n\")\n",
    "    drive.mount('/content/drive')\n",
    "\n",
    "tree = {}\n",
    "exclude = (\"_logs\", \"/output\")\n",
    "for i, (root, dirs, files) in enumerate(os.walk(folder, topdown=True)):\n",
    "    dirs[:] = [d for d in dirs if all(ex not in d for ex in exclude)]\n",
    "    images = len([f for f in files if f.lower().endswith((\".png\", \".jpg\", \".jpeg\"))])\n",
    "    captions = len([f for f in files if f.lower().endswith(\".txt\")])\n",
    "    others = len(files) - images - captions\n",
    "    path = root[folder.rfind(\"/\")+1:]\n",
    "    tree[path] = None if not images else f\"{images:>4} images | {captions:>4} captions |\"\n",
    "    if tree[path] and others:\n",
    "        tree[path] += f\" {others:>4} other files\"\n",
    "\n",
    "pad = max(len(k) for k in tree)\n",
    "print(\"\\n\".join(f\"📁{k.ljust(pad)} | {v}\" for k, v in tree.items() if v))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YcTjh07x90Ro"
   },
   "source": [
    "# 📈 Plot training results\n",
    "You can do this after running the trainer. You don't need this unless you know what you're doing.  \n",
    "The first cell below may fail to load all your logs. Keep trying the second cell until all data has loaded."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Z_TRI3eX90Rp"
   },
   "outputs": [],
   "source": [
    "%load_ext tensorboard\n",
    "%tensorboard --logdir={log_folder}/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6rM5SLq990Rp"
   },
   "outputs": [],
   "source": [
    "from tensorboard import notebook\n",
    "notebook.display(port=6006, height=800)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "A100",
   "machine_shape": "hm",
   "provenance": []
  },
  "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.11.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}