Created
December 11, 2025 05:17
-
-
Save ssghost/76d46d3be1d0d06905aae030b572697f to your computer and use it in GitHub Desktop.
Lora_Trainer.ipynb
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| { | |
| "cells": [ | |
| { | |
| "cell_type": "markdown", | |
| "metadata": { | |
| "id": "view-in-github", | |
| "colab_type": "text" | |
| }, | |
| "source": [ | |
| "<a href=\"https://colab.research.google.com/gist/ssghost/76d46d3be1d0d06905aae030b572697f/lora_trainer.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>" | |
| ] | |
| }, | |
| { | |
| "cell_type": "markdown", | |
| "metadata": { | |
| "id": "rmCPmqFL6hCQ" | |
| }, | |
| "source": [ | |
| "# ⭐ Lora Trainer by Hollowstrawberry\n", | |
| "\n", | |
| "This is based on the work of [Kohya-ss](https://github.com/kohya-ss/sd-scripts) and [Linaqruf](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/kohya-LoRA-dreambooth.ipynb). Thank you!\n" | |
| ] | |
| }, | |
| { | |
| "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. \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** | [](https://github.com/hollowstrawberry/kohya-colab) | | |\n", | |
| "| 📊 **Dataset Maker** | [](https://github.com/hollowstrawberry/kohya-colab/blob/main/Dataset_Maker.ipynb) | [](https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Dataset_Maker.ipynb) | [](https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Spanish_Dataset_Maker.ipynb) |\n", | |
| "| ⭐ **Lora Trainer** | [](https://github.com/hollowstrawberry/kohya-colab/blob/main/Lora_Trainer.ipynb) | [](https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Lora_Trainer.ipynb) | [](https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Spanish_Lora_Trainer.ipynb) |\n", | |
| "| 🌟 **XL Lora Trainer** | [](https://github.com/hollowstrawberry/kohya-colab/blob/main/Lora_Trainer_XL.ipynb) | [](https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Lora_Trainer_XL.ipynb) | |\n", | |
| "| 🌟 **Legacy XL Trainer** | [](https://github.com/hollowstrawberry/kohya-colab/blob/main/Lora_Trainer_XL_Legacy.ipynb) | [](https://colab.research.google.com/github/hollowstrawberry/kohya-colab/blob/main/Lora_Trainer_XL_Legacy.ipynb) | |" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": null, | |
| "metadata": { | |
| "cellView": "form", | |
| "id": "OglZzI_ujZq-" | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "import os\n", | |
| "import re\n", | |
| "import toml\n", | |
| "from time import time\n", | |
| "from IPython.display import Markdown, display\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", | |
| "if \"optimizer\" not in globals():\n", | |
| " optimizer = \"AdamW8bit\"\n", | |
| "if \"optimizer_args\" not in globals():\n", | |
| " optimizer_args = None\n", | |
| "if \"continue_from_lora\" not in globals():\n", | |
| " continue_from_lora = \"\"\n", | |
| "if \"weighted_captions\" not in globals():\n", | |
| " weighted_captions = False\n", | |
| "if \"adjust_tags\" not in globals():\n", | |
| " adjust_tags = False\n", | |
| "if \"keep_tokens_weight\" not in globals():\n", | |
| " keep_tokens_weight = 1.0\n", | |
| "\n", | |
| "COLAB = True # low ram\n", | |
| "XFORMERS = True\n", | |
| "SOURCE = \"https://github.com/uYouUs/sd-scripts\"\n", | |
| "COMMIT = None\n", | |
| "BETTER_EPOCH_NAMES = True\n", | |
| "LOAD_TRUNCATED_IMAGES = True\n", | |
| "\n", | |
| "#@title ## 🚩 Start Here\n", | |
| "\n", | |
| "#@markdown ### ▶️ Setup\n", | |
| "#@markdown Your project name will be the same as the folder containing your images. Spaces aren't allowed.\n", | |
| "project_name = \"\" #@param {type:\"string\"}\n", | |
| "project_name = project_name.strip()\n", | |
| "#@markdown 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 = \"Organize by project (MyDrive/Loras/project_name/dataset)\" #@param [\"Organize by category (MyDrive/lora_training/datasets/project_name)\", \"Organize by project (MyDrive/Loras/project_name/dataset)\"]\n", | |
| "#@markdown Decide the model that will be downloaded and used for training. These options should produce clean and consistent results. You can also choose your own by pasting its download link.\n", | |
| "training_model = \"Anime (animefull-final-pruned-fp16.safetensors)\" #@param [\"Anime (animefull-final-pruned-fp16.safetensors)\", \"AnyLora (AnyLoRA_noVae_fp16-pruned.ckpt)\", \"Stable Diffusion (sd-v1-5-pruned-noema-fp16.safetensors)\"]\n", | |
| "optional_custom_training_model_url = \"\" #@param {type:\"string\"}\n", | |
| "custom_model_is_based_on_sd2 = False #@param {type:\"boolean\"}\n", | |
| "\n", | |
| "if optional_custom_training_model_url:\n", | |
| " model_url = optional_custom_training_model_url\n", | |
| "elif \"AnyLora\" in training_model:\n", | |
| " model_url = \"https://huggingface.co/Lykon/AnyLoRA/resolve/main/AnyLoRA_noVae_fp16-pruned.ckpt\"\n", | |
| "elif \"Anime\" in training_model:\n", | |
| " model_url = \"https://huggingface.co/hollowstrawberry/stable-diffusion-guide/resolve/main/models/animefull-final-pruned-fp16.safetensors\"\n", | |
| "else:\n", | |
| " model_url = \"https://huggingface.co/hollowstrawberry/stable-diffusion-guide/resolve/main/models/sd-v1-5-pruned-noema-fp16.safetensors\"\n", | |
| "\n", | |
| "#@markdown ### ▶️ Processing\n", | |
| "#@markdown Resolution of 512 is standard for Stable Diffusion 1.5. Higher resolution training is much slower but can lead to better details. <p>\n", | |
| "#@markdown Images will be automatically scaled while training to produce the best results, so you don't need to crop or resize anything yourself.\n", | |
| "resolution = 512 #@param {type:\"slider\", min:512, max:1024, step:128}\n", | |
| "#@markdown This option will train your images both normally and flipped, for no extra cost, to learn more from them. Turn it on specially if you have less than 20 images. <p>\n", | |
| "#@markdown **Turn it off if you care about asymmetrical elements in your Lora**.\n", | |
| "flip_aug = False #@param {type:\"boolean\"}\n", | |
| "#markdown Leave empty for no captions.\n", | |
| "caption_extension = \".txt\" #param {type:\"string\"}\n", | |
| "#@markdown 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 = True #@param {type:\"boolean\"}\n", | |
| "shuffle_caption = shuffle_tags\n", | |
| "activation_tags = \"1\" #@param [0,1,2,3]\n", | |
| "keep_tokens = int(activation_tags)\n", | |
| "\n", | |
| "#@markdown ### ▶️ Steps <p>\n", | |
| "#@markdown 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 = 10 #@param {type:\"number\"}\n", | |
| "#@markdown Choose how long you want to train for. A good starting point is around 10 epochs or around 2000 steps.<p>\n", | |
| "#@markdown One epoch is a number of steps equal to: your number of images multiplied by their repeats, divided by batch size. <p>\n", | |
| "preferred_unit = \"Epochs\" #@param [\"Epochs\", \"Steps\"]\n", | |
| "how_many = 10 #@param {type:\"number\"}\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", | |
| "#@markdown Saving more epochs will let you compare your Lora's progress better.\n", | |
| "save_every_n_epochs = 1 #@param {type:\"number\"}\n", | |
| "keep_only_last_n_epochs = 10 #@param {type:\"number\"}\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", | |
| "#@markdown Increasing the batch size makes training faster, but may make learning worse. Recommended 2 or 3.\n", | |
| "train_batch_size = 2 #@param {type:\"slider\", min:1, max:8, step:1}\n", | |
| "\n", | |
| "#@markdown ### ▶️ Learning\n", | |
| "#@markdown 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. <p>\n", | |
| "#@markdown 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 = 5e-4 #@param {type:\"number\"}\n", | |
| "text_encoder_lr = 1e-4 #@param {type:\"number\"}\n", | |
| "#@markdown 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 = \"cosine_with_restarts\" #@param [\"constant\", \"cosine\", \"cosine_with_restarts\", \"constant_with_warmup\", \"linear\", \"polynomial\"]\n", | |
| "lr_scheduler_number = 3 #@param {type:\"number\"}\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", | |
| "#@markdown Steps spent \"warming up\" the learning rate during training for efficiency. I recommend leaving it at 5%.\n", | |
| "lr_warmup_ratio = 0.05 #@param {type:\"slider\", min:0.0, max:0.5, step:0.01}\n", | |
| "lr_warmup_steps = 0\n", | |
| "#@markdown New feature that adjusts loss over time, makes learning much more efficient, and training can be done with about half as many epochs. Uses a value of 5.0 as recommended by [the paper](https://arxiv.org/abs/2303.09556).\n", | |
| "min_snr_gamma = True #@param {type:\"boolean\"}\n", | |
| "min_snr_gamma_value = 5.0 if min_snr_gamma else None\n", | |
| "\n", | |
| "#@markdown ### ▶️ Structure\n", | |
| "#@markdown LoRA is the classic type and good for a variety of purposes. LoCon is good with artstyles as it has more layers to learn more aspects of the dataset.\n", | |
| "lora_type = \"LoRA\" #@param [\"LoRA\", \"LoCon\"]\n", | |
| "\n", | |
| "#@markdown Below are some recommended values for the following settings:\n", | |
| "\n", | |
| "#@markdown | type | network_dim | network_alpha | conv_dim | conv_alpha |\n", | |
| "#@markdown | :---: | :---: | :---: | :---: | :---: |\n", | |
| "#@markdown | LoRA | 16 | 8 | | |\n", | |
| "#@markdown | LoCon | 16 | 8 | 8 | 4 |\n", | |
| "\n", | |
| "#@markdown More dim means larger Lora, it can hold more information but more isn't always better. A dim between 8-32 is recommended, and alpha equal to half the dim.\n", | |
| "network_dim = 16 #@param {type:\"slider\", min:1, max:128, step:1}\n", | |
| "network_alpha = 8 #@param {type:\"slider\", min:1, max:128, step:1}\n", | |
| "#@markdown The following two values only apply to the additional layers of LoCon.\n", | |
| "conv_dim = 8 #@param {type:\"slider\", min:1, max:64, step:1}\n", | |
| "conv_alpha = 4 #@param {type:\"slider\", min:1, max:64, step:1}\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", | |
| "#@markdown ### ▶️ Ready\n", | |
| "#@markdown You can now run this cell to cook your Lora. Good luck! <p>\n", | |
| "\n", | |
| "\n", | |
| "# 👩💻 Cool code goes here\n", | |
| "\n", | |
| "if optimizer.lower() == \"prodigy\" or \"dadapt\" in optimizer.lower():\n", | |
| " if override_values_for_dadapt_and_prodigy:\n", | |
| " unet_lr = 0.5\n", | |
| " text_encoder_lr = 0.5\n", | |
| " lr_scheduler = \"constant_with_warmup\"\n", | |
| " lr_warmup_ratio = 0.05\n", | |
| " network_alpha = network_dim\n", | |
| "\n", | |
| " if not optimizer_args:\n", | |
| " optimizer_args = [\"decouple=True\",\"weight_decay=0.01\",\"betas=[0.9,0.999]\"]\n", | |
| " if optimizer == \"Prodigy\":\n", | |
| " optimizer_args.extend([\"d_coef=2\",\"use_bias_correction=True\"])\n", | |
| " if lr_warmup_ratio > 0:\n", | |
| " optimizer_args.append(\"safeguard_warmup=True\")\n", | |
| " else:\n", | |
| " optimizer_args.append(\"safeguard_warmup=False\")\n", | |
| "\n", | |
| "root_dir = \"/content\" if COLAB else \"~/Loras\"\n", | |
| "deps_dir = os.path.join(root_dir, \"deps\")\n", | |
| "repo_dir = os.path.join(root_dir, \"kohya-trainer\")\n", | |
| "\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", | |
| "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", | |
| "def install_dependencies():\n", | |
| " os.chdir(root_dir)\n", | |
| " !git clone {SOURCE} {repo_dir}\n", | |
| " os.chdir(repo_dir)\n", | |
| " if COMMIT:\n", | |
| " !git reset --hard {COMMIT}\n", | |
| " !wget https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/train_network_wrapper.py -q -O train_network_wrapper.py\n", | |
| " !wget https://raw.githubusercontent.com/hollowstrawberry/kohya-colab/main/dracula.py -q -O dracula.py\n", | |
| "\n", | |
| " !apt -y update -qq\n", | |
| " !apt -y install aria2 -qq\n", | |
| " !pip install -U torch==2.4 xformers triton torchvision==0.19 --index-url https://download.pytorch.org/whl/cu121\n", | |
| " !pip install accelerate==0.25.0 transformers==4.36.2 diffusers[torch]==0.25.0 ftfy==6.1.1 \\\n", | |
| " opencv-python==4.8.1.78 einops==0.7.0 pytorch-lightning==1.9.0 bitsandbytes==0.43.0 \\\n", | |
| " prodigyopt==1.0 lion-pytorch==0.0.6 tensorboard safetensors==0.4.2 altair==4.2.2 \\\n", | |
| " easygui==0.98.3 toml==0.10.2 voluptuous==0.13.1 huggingface-hub==0.20.1 imagesize==1.4.1 rich==13.7.1 numpy==1.26.4\n", | |
| " !pip install -e .\n", | |
| "\n", | |
| " # patch kohya for minor stuff\n", | |
| " if COLAB:\n", | |
| " !sed -i \"s@cpu@cuda@\" library/model_util.py # low ram\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", | |
| " from accelerate.utils import write_basic_config\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, keep_tokens_weight, weighted_captions, adjust_tags\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", | |
| " if continue_from_lora and not (continue_from_lora.endswith(\".safetensors\") and os.path.exists(continue_from_lora)):\n", | |
| " print(f\"💥 Error: Invalid path to existing Lora. Example: /content/drive/MyDrive/Loras/example.safetensors\")\n", | |
| " return\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", | |
| " if adjust_tags:\n", | |
| " print(f\"\\n📎 Weighted tags: {'ON' if weighted_captions else 'OFF'}\")\n", | |
| " if weighted_captions:\n", | |
| " print(f\"📎 Will use {keep_tokens_weight} weight on {keep_tokens} activation tag(s)\")\n", | |
| " print(\"📎 Adjusting tags...\")\n", | |
| " adjust_weighted_tags(folders, keep_tokens, keep_tokens_weight, weighted_captions)\n", | |
| "\n", | |
| " return True\n", | |
| "\n", | |
| "def adjust_weighted_tags(folders, keep_tokens: int, keep_tokens_weight: float, weighted_captions: bool):\n", | |
| " weighted_tag = re.compile(r\"\\((.+?):[.\\d]+\\)(,|$)\")\n", | |
| " for folder in folders:\n", | |
| " for txt in [f for f in os.listdir(folder) if f.lower().endswith(\".txt\")]:\n", | |
| " with open(os.path.join(folder, txt), 'r') as f:\n", | |
| " content = f.read()\n", | |
| " # reset previous changes\n", | |
| " content = content.replace('\\\\', '')\n", | |
| " content = weighted_tag.sub(r'\\1\\2', content)\n", | |
| " if weighted_captions:\n", | |
| " # re-apply changes\n", | |
| " content = content.replace(r'(', r'\\(').replace(r')', r'\\)').replace(r':', r'\\:')\n", | |
| " if keep_tokens_weight > 1:\n", | |
| " tags = [s.strip() for s in content.split(\",\")]\n", | |
| " for i in range(min(keep_tokens, len(tags))):\n", | |
| " tags[i] = f'({tags[i]}:{keep_tokens_weight})'\n", | |
| " content = \", \".join(tags)\n", | |
| " with open(os.path.join(folder, txt), 'w') as f:\n", | |
| " f.write(content)\n", | |
| "\n", | |
| "def create_config():\n", | |
| " global dataset_config_file, config_file, model_file\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", | |
| " \"additional_network_arguments\": {\n", | |
| " \"unet_lr\": unet_lr,\n", | |
| " \"text_encoder_lr\": text_encoder_lr,\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\": True if text_encoder_lr == 0 else None,\n", | |
| " \"network_weights\": continue_from_lora if continue_from_lora else None\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", | |
| " \"max_train_steps\": max_train_steps,\n", | |
| " \"max_train_epochs\": max_train_epochs,\n", | |
| " \"save_every_n_epochs\": save_every_n_epochs,\n", | |
| " \"save_last_n_epochs\": keep_only_last_n_epochs,\n", | |
| " \"train_batch_size\": train_batch_size,\n", | |
| " \"noise_offset\": None,\n", | |
| " \"clip_skip\": 2,\n", | |
| " \"min_snr_gamma\": min_snr_gamma_value,\n", | |
| " \"weighted_captions\": weighted_captions,\n", | |
| " \"seed\": 42,\n", | |
| " \"max_token_length\": 225,\n", | |
| " \"xformers\": XFORMERS,\n", | |
| " \"lowram\": COLAB,\n", | |
| " \"max_data_loader_n_workers\": 8,\n", | |
| " \"persistent_data_loader_workers\": True,\n", | |
| " \"save_precision\": \"fp16\",\n", | |
| " \"mixed_precision\": \"fp16\",\n", | |
| " \"output_dir\": output_folder,\n", | |
| " \"logging_dir\": log_folder,\n", | |
| " \"output_name\": project_name,\n", | |
| " \"log_prefix\": project_name,\n", | |
| " },\n", | |
| " \"model_arguments\": {\n", | |
| " \"pretrained_model_name_or_path\": model_file,\n", | |
| " \"v2\": custom_model_is_based_on_sd2,\n", | |
| " \"v_parameterization\": True if custom_model_is_based_on_sd2 else None,\n", | |
| " },\n", | |
| " \"saving_arguments\": {\n", | |
| " \"save_model_as\": \"safetensors\",\n", | |
| " },\n", | |
| " \"dreambooth_arguments\": {\n", | |
| " \"prior_loss_weight\": 1.0,\n", | |
| " },\n", | |
| " \"dataset_arguments\": {\n", | |
| " \"cache_latents\": True,\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,\n", | |
| " \"keep_tokens\": keep_tokens,\n", | |
| " \"flip_aug\": flip_aug,\n", | |
| " \"caption_extension\": caption_extension,\n", | |
| " \"enable_bucket\": True,\n", | |
| " \"bucket_reso_steps\": 64,\n", | |
| " \"bucket_no_upscale\": False,\n", | |
| " \"min_bucket_reso\": 320 if resolution > 640 else 256,\n", | |
| " \"max_bucket_reso\": 1280 if resolution > 640 else 1024,\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():\n", | |
| " global old_model_url, model_url, model_file\n", | |
| " real_model_url = model_url.strip()\n", | |
| "\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 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():\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()\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():\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()\n", | |
| "\n", | |
| " print(\"\\n⭐ Starting trainer...\\n\")\n", | |
| " os.chdir(repo_dir)\n", | |
| "\n", | |
| " !accelerate launch --config_file={accelerate_config_file} --num_cpu_threads_per_process=1 train_network_wrapper.py --dataset_config={dataset_config_file} --config_file={config_file}\n", | |
| "\n", | |
| "main()\n" | |
| ] | |
| }, | |
| { | |
| "cell_type": "markdown", | |
| "metadata": { | |
| "id": "mBMUJ7BuvNcn" | |
| }, | |
| "source": [ | |
| "## *️⃣ Extras\n", | |
| "\n", | |
| "You can run these before starting the training." | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": null, | |
| "metadata": { | |
| "cellView": "form", | |
| "id": "sy9jU2yrdYar" | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "#@markdown ### 🔮 Optimizer\n", | |
| "#@markdown If you run this cell you will change the optimizer used for training. The default will be `AdamW8bit` otherwise, which is recommended.<p>\n", | |
| "#@markdown * Dadapt and Prodigy manage learning rate automatically and are very good for small datasets. You can use it without changing anything else here.<p>\n", | |
| "#@markdown For Dadapt and Prodigy, the following values will be overriden if the box is checked:<p>\n", | |
| "#@markdown `learning_rate=0.5`, `network_alpha=network_dim`, `lr_scheduler=\"constant_with_warmup\"`, `lr_warmup_ratio=0.05`<p>\n", | |
| "#@markdown For Dadapt and Prodigy, if `optimizer_args` is left empty the default will be `decouple=True, weight_decay=0.01, betas=[0.9,0.999]`<p>\n", | |
| "#@markdown And additionally for Prodigy: `d_coef=2, use_bias_correction=True, safeguard_warmup=True`<p>\n", | |
| "optimizer = \"Prodigy\" #@param [\"AdamW8bit\", \"Prodigy\", \"DAdaptation\", \"DadaptAdam\", \"DadaptLion\", \"AdamW\", \"Lion\", \"SGDNesterov\", \"SGDNesterov8bit\", \"AdaFactor\"]\n", | |
| "optimizer_args = \"\" #@param {type:\"string\"}\n", | |
| "splitter = \", \" if \", \" in optimizer_args else \",\"\n", | |
| "optimizer_args = [a.strip() for a in optimizer_args.split(splitter) if a]\n", | |
| "override_values_for_dadapt_and_prodigy = True #@param {type:\"boolean\"}" | |
| ] | |
| }, | |
| { | |
| "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": null, | |
| "metadata": { | |
| "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": null, | |
| "metadata": { | |
| "id": "W84Jxf-U2TIU" | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "custom_dataset = None" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": null, | |
| "metadata": { | |
| "cellView": "form", | |
| "id": "-Yq5mNvcCy2l" | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "#@markdown ### 🤓 Other\n", | |
| "#@markdown These are kept here for the few people that use them.\n", | |
| "\n", | |
| "#@markdown Weighted captions is a new feature that allows you to use (parentheses) to give more weight to certain tags in your dataset, same as in your webui prompts. <p>\n", | |
| "#@markdown Normal parentheses in your tags such as `(series names)` will need to be escaped like `\\(series names\\)`\n", | |
| "weighted_captions = False #@param {type:\"boolean\"}\n", | |
| "\n", | |
| "#markdown By enabling `adjust_tags`, you will let this colab modify your tags before running to automatically adjust to `weighted_captions` being on or off. <p>\n", | |
| "#markdown Then, you may increase `activation_tag_weight` to improve how effective your activation tag is.\n", | |
| "adjust_tags = False #param {type:\"boolean\"}\n", | |
| "activation_tag_weight = \"1.0\" #param [\"1.0\",\"1.1\",\"1.2\"]\n", | |
| "keep_tokens_weight = float(activation_tag_weight)\n", | |
| "\n", | |
| "#@markdown Here you can write a path in your Google Drive to load an existing Lora file to continue training on.<p>\n", | |
| "#@markdown **Warning:** It's not the same as one long training session. The epochs start from scratch, and it may have worse results.\n", | |
| "continue_from_lora = \"\" #@param {type:\"string\"}\n", | |
| "if continue_from_lora and not continue_from_lora.startswith(\"/content/drive/MyDrive\"):\n", | |
| " import os\n", | |
| " continue_from_lora = os.path.join(\"/content/drive/MyDrive\", continue_from_lora)\n" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": null, | |
| "metadata": { | |
| "cellView": "form", | |
| "id": "WDjkp4scvPgE" | |
| }, | |
| "outputs": [], | |
| "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": "cDyqB2ytNN08" | |
| }, | |
| "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": "gdogfLJ_NN08" | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "%load_ext tensorboard\n", | |
| "%tensorboard --logdir={log_folder}/" | |
| ] | |
| }, | |
| { | |
| "cell_type": "code", | |
| "execution_count": null, | |
| "metadata": { | |
| "id": "4NC0QOaXNN08" | |
| }, | |
| "outputs": [], | |
| "source": [ | |
| "from tensorboard import notebook\n", | |
| "notebook.display(port=6006, height=800)" | |
| ] | |
| } | |
| ], | |
| "metadata": { | |
| "accelerator": "GPU", | |
| "colab": { | |
| "provenance": [], | |
| "include_colab_link": true | |
| }, | |
| "gpuClass": "standard", | |
| "kernelspec": { | |
| "display_name": "Python 3", | |
| "name": "python3" | |
| }, | |
| "language_info": { | |
| "name": "python" | |
| } | |
| }, | |
| "nbformat": 4, | |
| "nbformat_minor": 0 | |
| } |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment