{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "slgjeYgd6pWp"
},
"source": [
"[![visitor][visitor-badge]][visitor-stats]\n",
"[![ko-fi][ko-fi-badge]][ko-fi-link]\n",
"\n",
"# **Kohya LoRA Trainer XL**\n",
"A Colab Notebook For SDXL LoRA Training (Fine-tuning Method)\n",
"\n",
"[visitor-badge]: https://api.visitorbadge.io/api/visitors?path=Kohya%20LoRA%20Trainer%20XL&label=Visitors&labelColor=%2334495E&countColor=%231ABC9C&style=flat&labelStyle=none\n",
"[visitor-stats]: https://visitorbadge.io/status?path=Kohya%20LoRA%20Trainer%20XL\n",
"[ko-fi-badge]: https://img.shields.io/badge/Support%20me%20on%20Ko--fi-F16061?logo=ko-fi&logoColor=white&style=flat\n",
"[ko-fi-link]: https://ko-fi.com/linaqruf"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "_MxF9feWshAp"
},
"source": [
"| Notebook Name | Description | Link |\n",
"| --- | --- | --- |\n",
"| [Kohya LoRA Trainer XL](https://github.com/Linaqruf/kohya-trainer/blob/main/kohya-LoRA-trainer-XL.ipynb) | LoRA Training | [](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/kohya-LoRA-trainer-XL.ipynb) |\n",
"| [Kohya Trainer XL](https://github.com/Linaqruf/kohya-trainer/blob/main/kohya-trainer-XL.ipynb) | Native Training | [](https://colab.research.google.com/github/Linaqruf/kohya-trainer/blob/main/kohya-trainer-XL.ipynb) |\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "iIbwGFkJ0nTx"
},
"source": [
"
\n",
"NEWS: Colab's free-tier users can now train SDXL LoRA using the diffusers format instead of checkpoint as a pretrained model.
\n",
"
"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "tTVqCAgSmie4"
},
"source": [
"# **I. Prepare Environment**"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "_u3q60di584x"
},
"outputs": [],
"source": [
"# @title ## **1.1. Install Kohya Trainer**\n",
"import os\n",
"import zipfile\n",
"import shutil\n",
"import time\n",
"import requests\n",
"import torch\n",
"from subprocess import getoutput\n",
"from IPython.utils import capture\n",
"from google.colab import drive\n",
"\n",
"%store -r\n",
"\n",
"# root_dir\n",
"root_dir = \"/content\"\n",
"drive_dir = os.path.join(root_dir, \"drive/MyDrive\")\n",
"deps_dir = os.path.join(root_dir, \"deps\")\n",
"repo_dir = os.path.join(root_dir, \"kohya-trainer\")\n",
"training_dir = os.path.join(root_dir, \"LoRA\")\n",
"pretrained_model = os.path.join(root_dir, \"pretrained_model\")\n",
"vae_dir = os.path.join(root_dir, \"vae\")\n",
"lora_dir = os.path.join(root_dir, \"network_weight\")\n",
"repositories_dir = os.path.join(root_dir, \"repositories\")\n",
"config_dir = os.path.join(training_dir, \"config\")\n",
"tools_dir = os.path.join(repo_dir, \"tools\")\n",
"finetune_dir = os.path.join(repo_dir, \"finetune\")\n",
"accelerate_config = os.path.join(repo_dir, \"accelerate_config/config.yaml\")\n",
"\n",
"for store in [\"root_dir\", \"repo_dir\", \"training_dir\", \"pretrained_model\", \"vae_dir\", \"repositories_dir\", \"accelerate_config\", \"tools_dir\", \"finetune_dir\", \"config_dir\"]:\n",
" with capture.capture_output() as cap:\n",
" %store {store}\n",
" del cap\n",
"\n",
"repo_dict = {\n",
" \"qaneel/kohya-trainer (forked repo, stable, optimized for colab use)\" : \"https://github.com/qaneel/kohya-trainer\",\n",
" \"kohya-ss/sd-scripts (original repo, latest update)\" : \"https://github.com/kohya-ss/sd-scripts\",\n",
"}\n",
"\n",
"repository = \"qaneel/kohya-trainer (forked repo, stable, optimized for colab use)\" #@param [\"qaneel/kohya-trainer (forked repo, stable, optimized for colab use)\", \"kohya-ss/sd-scripts (original repo, latest update)\"] {allow-input: true}\n",
"repo_url = repo_dict[repository]\n",
"branch = \"main\" # @param {type: \"string\"}\n",
"output_to_drive = True # @param {type: \"boolean\"}\n",
"\n",
"def clone_repo(url, dir, branch):\n",
" if not os.path.exists(dir):\n",
" !git clone -b {branch} {url} {dir}\n",
"\n",
"def mount_drive(dir):\n",
" output_dir = os.path.join(training_dir, \"output\")\n",
"\n",
" if output_to_drive:\n",
" if not os.path.exists(drive_dir):\n",
" drive.mount(os.path.dirname(drive_dir))\n",
" output_dir = os.path.join(drive_dir, \"kohya-trainer/output\")\n",
"\n",
" return output_dir\n",
"\n",
"def setup_directories():\n",
" global output_dir\n",
"\n",
" output_dir = mount_drive(drive_dir)\n",
"\n",
" for dir in [training_dir, config_dir, pretrained_model, vae_dir, repositories_dir, output_dir]:\n",
" os.makedirs(dir, exist_ok=True)\n",
"\n",
"def pastebin_reader(id):\n",
" if \"pastebin.com\" in id:\n",
" url = id\n",
" if 'raw' not in url:\n",
" url = url.replace('pastebin.com', 'pastebin.com/raw')\n",
" else:\n",
" url = \"https://pastebin.com/raw/\" + id\n",
" response = requests.get(url)\n",
" response.raise_for_status()\n",
" lines = response.text.split('\\n')\n",
" return lines\n",
"\n",
"def install_repository():\n",
" global infinite_image_browser_dir, voldy, discordia_archivum_dir\n",
"\n",
" _, voldy = pastebin_reader(\"kq6ZmHFU\")[:2]\n",
"\n",
" infinite_image_browser_url = f\"https://github.com/zanllp/{voldy}-infinite-image-browsing.git\"\n",
" infinite_image_browser_dir = os.path.join(repositories_dir, f\"infinite-image-browsing\")\n",
" infinite_image_browser_deps = os.path.join(infinite_image_browser_dir, \"requirements.txt\")\n",
"\n",
" discordia_archivum_url = \"https://github.com/Linaqruf/discordia-archivum\"\n",
" discordia_archivum_dir = os.path.join(repositories_dir, \"discordia-archivum\")\n",
" discordia_archivum_deps = os.path.join(discordia_archivum_dir, \"requirements.txt\")\n",
"\n",
" clone_repo(infinite_image_browser_url, infinite_image_browser_dir, \"main\")\n",
" clone_repo(discordia_archivum_url, discordia_archivum_dir, \"main\")\n",
"\n",
" !pip install -q --upgrade -r {infinite_image_browser_deps}\n",
" !pip install python-dotenv\n",
" !pip install -q --upgrade -r {discordia_archivum_deps}\n",
"\n",
"def install_dependencies():\n",
" requirements_file = os.path.join(repo_dir, \"requirements.txt\")\n",
" model_util = os.path.join(repo_dir, \"library/model_util.py\")\n",
" gpu_info = getoutput('nvidia-smi')\n",
" t4_xformers_wheel = \"https://github.com/Linaqruf/colab-xformers/releases/download/0.0.20/xformers-0.0.20+1d635e1.d20230519-cp310-cp310-linux_x86_64.whl\"\n",
"\n",
" !apt install aria2 lz4\n",
" !wget https://github.com/camenduru/gperftools/releases/download/v1.0/libtcmalloc_minimal.so.4 -O /content/libtcmalloc_minimal.so.4\n",
" !pip install -q --upgrade -r {requirements_file}\n",
"\n",
" !pip install -q xformers==0.0.22.post7\n",
"\n",
" from accelerate.utils import write_basic_config\n",
"\n",
" if not os.path.exists(accelerate_config):\n",
" write_basic_config(save_location=accelerate_config)\n",
"\n",
"def prepare_environment():\n",
" os.environ[\"LD_PRELOAD\"] = \"/content/libtcmalloc_minimal.so.4\"\n",
" os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n",
" os.environ[\"SAFETENSORS_FAST_GPU\"] = \"1\"\n",
" os.environ[\"PYTHONWARNINGS\"] = \"ignore\"\n",
"\n",
"def main():\n",
" os.chdir(root_dir)\n",
" clone_repo(repo_url, repo_dir, branch)\n",
" os.chdir(repo_dir)\n",
" setup_directories()\n",
" install_repository()\n",
" install_dependencies()\n",
" prepare_environment()\n",
"\n",
"main()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "wrYGu-WxFbsq"
},
"outputs": [],
"source": [
"# @title ## **1.2. Download SDXL**\n",
"import os\n",
"import re\n",
"import json\n",
"import glob\n",
"import gdown\n",
"import requests\n",
"import subprocess\n",
"from IPython.utils import capture\n",
"from urllib.parse import urlparse, unquote\n",
"from pathlib import Path\n",
"from huggingface_hub import HfFileSystem\n",
"from huggingface_hub.utils import validate_repo_id, HfHubHTTPError\n",
"\n",
"%store -r\n",
"\n",
"os.chdir(root_dir)\n",
"\n",
"# @markdown Place your Huggingface token [here](https://huggingface.co/settings/tokens) to download gated models.\n",
"\n",
"HUGGINGFACE_TOKEN = \"\" #@param {type: \"string\"}\n",
"LOAD_DIFFUSERS_MODEL = True #@param {type: \"boolean\"}\n",
"SDXL_MODEL_URL = \"stablediffusionapi/pony-diffusion-v6-xl\" # @param [\"gsdf/CounterfeitXL\", \"Linaqruf/animagine-xl\", \"stabilityai/stable-diffusion-xl-base-1.0\", \"PASTE MODEL URL OR GDRIVE PATH HERE\"] {allow-input: true}\n",
"SDXL_VAE_URL = \"FP16 VAE\" # @param [\"None\", \"Original VAE\", \"FP16 VAE\", \"PASTE VAE URL OR GDRIVE PATH HERE\"] {allow-input: true}\n",
"\n",
"MODEL_URLS = {\n",
" \"gsdf/CounterfeitXL\" : \"https://huggingface.co/gsdf/CounterfeitXL/resolve/main/CounterfeitXL_%CE%B2.safetensors\",\n",
" \"Linaqruf/animagine-xl\" : \"https://huggingface.co/Linaqruf/animagine-xl/resolve/main/animagine-xl.safetensors\",\n",
" \"stabilityai/stable-diffusion-xl-base-1.0\" : \"https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors\",\n",
"}\n",
"VAE_URLS = {\n",
" \"None\" : \"\",\n",
" \"Original VAE\" : \"https://huggingface.co/stabilityai/sdxl-vae/resolve/main/sdxl_vae.safetensors\",\n",
" \"FP16 VAE\" : \"https://huggingface.co/madebyollin/sdxl-vae-fp16-fix/resolve/main/sdxl_vae.safetensors\"\n",
"}\n",
"\n",
"SDXL_MODEL_URL = MODEL_URLS.get(SDXL_MODEL_URL, SDXL_MODEL_URL)\n",
"SDXL_VAE_URL = VAE_URLS.get(SDXL_VAE_URL, SDXL_VAE_URL)\n",
"\n",
"def get_filename(url):\n",
" if any(url.endswith(ext) for ext in [\".ckpt\", \".safetensors\", \".pt\", \".pth\"]):\n",
" return os.path.basename(url)\n",
"\n",
" response = requests.get(url, stream=True)\n",
" response.raise_for_status()\n",
"\n",
" if 'content-disposition' in response.headers:\n",
" filename = re.findall('filename=\"?([^\"]+)\"?', response.headers['content-disposition'])[0]\n",
" else:\n",
" filename = unquote(os.path.basename(urlparse(url).path))\n",
"\n",
" return filename\n",
"\n",
"def aria2_download(dir, filename, url):\n",
" user_header = f\"Authorization: Bearer {HUGGINGFACE_TOKEN}\"\n",
" aria2_args = [\n",
" \"aria2c\",\n",
" \"--console-log-level=error\",\n",
" \"--summary-interval=10\",\n",
" f\"--header={user_header}\" if \"huggingface.co\" in url else \"\",\n",
" \"--continue=true\",\n",
" \"--max-connection-per-server=16\",\n",
" \"--min-split-size=1M\",\n",
" \"--split=16\",\n",
" f\"--dir={dir}\",\n",
" f\"--out={filename}\",\n",
" url\n",
" ]\n",
" subprocess.run(aria2_args)\n",
"\n",
"def download(url, dst):\n",
" print(f\"Starting downloading from {url}\")\n",
" filename = get_filename(url)\n",
" filepath = os.path.join(dst, filename)\n",
"\n",
" if \"drive.google.com\" in url:\n",
" gdown.download(url, filepath, quiet=False)\n",
" else:\n",
" if \"huggingface.co\" in url and \"/blob/\" in url:\n",
" url = url.replace(\"/blob/\", \"/resolve/\")\n",
" aria2_download(dst, filename, url)\n",
"\n",
" print(f\"Download finished: {filepath}\")\n",
" return filepath\n",
"\n",
"def all_folders_present(base_model_url, sub_folders):\n",
" fs = HfFileSystem()\n",
" existing_folders = set(fs.ls(base_model_url, detail=False))\n",
"\n",
" for folder in sub_folders:\n",
" full_folder_path = f\"{base_model_url}/{folder}\"\n",
" if full_folder_path not in existing_folders:\n",
" return False\n",
" return True\n",
"\n",
"def get_total_ram_gb():\n",
" with open('/proc/meminfo', 'r') as f:\n",
" for line in f.readlines():\n",
" if \"MemTotal\" in line:\n",
" return int(line.split()[1]) / (1024**2) # Convert to GB\n",
"\n",
"def get_gpu_name():\n",
" try:\n",
" return subprocess.check_output(\"nvidia-smi --query-gpu=name --format=csv,noheader,nounits\", shell=True).decode('ascii').strip()\n",
" except:\n",
" return None\n",
"\n",
"def main():\n",
" global model_path, vae_path, LOAD_DIFFUSERS_MODEL\n",
"\n",
" model_path, vae_path = None, None\n",
"\n",
" required_sub_folders = [\n",
" 'scheduler',\n",
" 'text_encoder',\n",
" 'text_encoder_2',\n",
" 'tokenizer',\n",
" 'tokenizer_2',\n",
" 'unet',\n",
" 'vae',\n",
" ]\n",
"\n",
" download_targets = {\n",
" \"model\": (SDXL_MODEL_URL, pretrained_model),\n",
" \"vae\": (SDXL_VAE_URL, vae_dir),\n",
" }\n",
"\n",
" total_ram = get_total_ram_gb()\n",
" gpu_name = get_gpu_name()\n",
"\n",
" # Check hardware constraints\n",
" if total_ram < 13 and gpu_name in [\"Tesla T4\", \"Tesla V100\"]:\n",
" print(\"Attempt to load diffusers model instead due to hardware constraints.\")\n",
" if not LOAD_DIFFUSERS_MODEL:\n",
" LOAD_DIFFUSERS_MODEL = True\n",
"\n",
" for target, (url, dst) in download_targets.items():\n",
" if url and not url.startswith(f\"PASTE {target.upper()} URL OR GDRIVE PATH HERE\"):\n",
" if target == \"model\" and LOAD_DIFFUSERS_MODEL:\n",
" # Code for checking and handling diffusers model\n",
" if 'huggingface.co' in url:\n",
" match = re.search(r'huggingface\\.co/([^/]+)/([^/]+)', SDXL_MODEL_URL)\n",
" if match:\n",
" username = match.group(1)\n",
" model_name = match.group(2)\n",
" url = f\"{username}/{model_name}\"\n",
" if all_folders_present(url, required_sub_folders):\n",
" print(f\"Diffusers model is loaded : {url}\")\n",
" model_path = url\n",
" else:\n",
" print(\"Repository doesn't exist or no diffusers model detected.\")\n",
" filepath = download(url, dst) # Continue with the regular download\n",
" model_path = filepath\n",
" else:\n",
" filepath = download(url, dst)\n",
"\n",
" if target == \"model\":\n",
" model_path = filepath\n",
" elif target == \"vae\":\n",
" vae_path = filepath\n",
"\n",
" print()\n",
"\n",
" if model_path:\n",
" print(f\"Selected model: {model_path}\")\n",
"\n",
" if vae_path:\n",
" print(f\"Selected VAE: {vae_path}\")\n",
"\n",
"main()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "kh7CeDqK4l3Y"
},
"outputs": [],
"source": [
"# @title ## **1.3. Directory Config**\n",
"# @markdown Specify the location of your training data in the following cell. A folder with the same name as your input will be created.\n",
"import os\n",
"\n",
"%store -r\n",
"\n",
"train_data_dir = \"/content/LoRA/bludwing\" # @param {'type' : 'string'}\n",
"%store train_data_dir\n",
"\n",
"os.makedirs(train_data_dir, exist_ok=True)\n",
"print(f\"Your train data directory : {train_data_dir}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "t17ZfiMB8GWZ"
},
"outputs": [],
"source": [
"# @title ## **2.1. Unzip Dataset**\n",
"\n",
"import os\n",
"import subprocess\n",
"import re\n",
"from urllib.parse import unquote\n",
"import requests\n",
"from zipfile import ZipFile\n",
"\n",
"# @title ## Unzip Dataset\n",
"# @markdown If your dataset is in a `zip` file and has been uploaded to a location, use this section to extract it.\n",
"# @markdown The dataset will be downloaded and automatically extracted to `train_data_dir` if `unzip_to` is empty.\n",
"\n",
"zipfile_url = \"/content/drive/MyDrive/bludwingdata.zip\" # @param {type:\"string\"}\n",
"unzip_to = \"\" # @param {type:\"string\"}\n",
"hf_token = \"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\"\n",
"\n",
"use_aria2c = True # @param {type:\"boolean\"}\n",
"preserve_folders = False # @param {type:\"boolean\"}\n",
"remove_after_unzipping = False # @param {type:\"boolean\"}\n",
"\n",
"if \"huggingface.co\" in zipfile_url and \"blob\" in zipfile_url:\n",
" zipfile_url = zipfile_url.replace(\"blob\", \"resolve\")\n",
"\n",
"if not unzip_to:\n",
" unzip_to = train_data_dir\n",
"\n",
"def get_filename_from_url(url):\n",
" if \"huggingface.co\" or \"/content/\" in url:\n",
" return os.path.basename(url)\n",
"\n",
" response = requests.head(url, allow_redirects=True)\n",
" cd = response.headers.get('content-disposition')\n",
" if cd:\n",
" fname = re.findall('filename=(.+)', cd)\n",
" if len(fname) == 0:\n",
" return \"zipfile.zip\"\n",
" return unquote(fname[0])\n",
"\n",
" return \"zipfile.zip\"\n",
"\n",
"def download_with_requests(url, output_path):\n",
" print(f\"Downloading {url} with requests...\")\n",
" response = requests.get(url, stream=True)\n",
" with open(output_path, 'wb') as file:\n",
" for chunk in response.iter_content(chunk_size=8192):\n",
" file.write(chunk)\n",
" print(f\"Downloaded to {output_path}\")\n",
" return output_path\n",
"\n",
"def download_with_aria2c(url, output_path):\n",
" print(f\"Downloading {url} with aria2c...\")\n",
" aria_args = {\n",
" 'console-log-level': 'error',\n",
" 'summary-interval': '10',\n",
" 'continue': 'true',\n",
" 'max-connection-per-server': '16',\n",
" 'min-split-size': '1M',\n",
" 'split': '16',\n",
" 'dir': os.path.dirname(output_path),\n",
" 'out': os.path.basename(output_path),\n",
" }\n",
"\n",
" if \"huggingface.co\" in url:\n",
" aria_args['header'] = f\"Authorization: Bearer {hf_token}\"\n",
"\n",
" cmd = ['aria2c'] + [f'--{k}={v}' for k, v in aria_args.items()] + [url]\n",
" subprocess.run(cmd)\n",
" print(f\"Downloaded to {output_path}\")\n",
" return output_path\n",
"\n",
"def move_files(train_dir):\n",
" for filename in os.listdir(train_dir):\n",
" file_path = os.path.join(train_dir, filename)\n",
" if filename.startswith(\"meta_\") and filename.endswith(\".json\"):\n",
" if not os.path.exists(file_path):\n",
" shutil.move(file_path, training_dir)\n",
" else:\n",
" os.remove(file_path)\n",
"\n",
"def remove_empty_dirs(path):\n",
" for dirpath, dirnames, files in os.walk(path, topdown=False): # start from leaf folders\n",
" for dirname in dirnames:\n",
" full_dir_path = os.path.join(dirpath, dirname)\n",
" if not os.listdir(full_dir_path): # Check if directory is empty\n",
" os.rmdir(full_dir_path)\n",
" print(f\"Removed empty directory: {full_dir_path}\")\n",
"\n",
"def extract_dataset(zip_file, output_path):\n",
" with ZipFile(zip_file, 'r') as zip_ref:\n",
" print(f\"Extracting {zip_file} to {output_path}...\")\n",
"\n",
" if not preserve_folders: # If we do not want to preserve folder structure\n",
" for member in zip_ref.namelist():\n",
" # Extract only the file name, discard directory structure\n",
" filename = os.path.basename(member)\n",
" if filename: # Check if file name is not empty (this skips directories)\n",
" zip_ref.extract(member, output_path)\n",
" source_path = os.path.join(output_path, member)\n",
" target_path = os.path.join(output_path, filename)\n",
" os.rename(source_path, target_path)\n",
"\n",
" remove_empty_dirs(output_path)\n",
"\n",
" else:\n",
" zip_ref.extractall(output_path)\n",
"\n",
" print(\"Extraction completed!\")\n",
"\n",
"def download_dataset(url, output_path):\n",
" if url.startswith(\"/content\"):\n",
" print(f\"Using file at {url}\")\n",
" return url\n",
"\n",
" elif \"drive.google.com\" in url:\n",
" print(\"Downloading from Google Drive...\")\n",
" cmd = ['gdown', '--id', url.split('/')[-2], '-O', output_path]\n",
" subprocess.run(cmd)\n",
" return output_path\n",
"\n",
" elif use_aria2c:\n",
" return download_with_aria2c(url, output_path)\n",
"\n",
" else:\n",
" return download_with_requests(url, output_path)\n",
"\n",
"def main():\n",
" zipfile_name = get_filename_from_url(zipfile_url)\n",
" output_path = os.path.join(root_dir, zipfile_name)\n",
"\n",
" zip_file = download_dataset(zipfile_url, output_path)\n",
"\n",
" extract_dataset(zip_file, unzip_to)\n",
"\n",
" move_files(unzip_to)\n",
"\n",
" if remove_after_unzipping and \"/content/drive\" not in zip_file:\n",
" os.remove(zip_file)\n",
" print(f\"Removed {zip_file}\")\n",
"\n",
"main()\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "5m6iIzfKniYr"
},
"outputs": [],
"source": [
"from google.colab import drive\n",
"drive.flush_and_unmount()\n",
"drive.mount('/content/drive')"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "Jz2emq6vWnPu"
},
"outputs": [],
"source": [
"# @title ## **3.1. Data Cleaning**\n",
"import os\n",
"import random\n",
"import concurrent.futures\n",
"from tqdm import tqdm\n",
"from PIL import Image\n",
"\n",
"%store -r\n",
"\n",
"os.chdir(root_dir)\n",
"\n",
"test = os.listdir(train_data_dir)\n",
"#@markdown This section removes unsupported media types such as `.mp4`, `.webm`, and `.gif`, as well as any unnecessary files.\n",
"#@markdown To convert a transparent dataset with an alpha channel (RGBA) to RGB and give it a white background, set the `convert` parameter to `True`.\n",
"convert = True # @param {type:\"boolean\"}\n",
"#@markdown Alternatively, you can give the background a `random_color` instead of white by checking the corresponding option.\n",
"random_color = False # @param {type:\"boolean\"}\n",
"recursive = False\n",
"\n",
"batch_size = 32\n",
"supported_types = [\n",
" \".png\",\n",
" \".jpg\",\n",
" \".jpeg\",\n",
" \".webp\",\n",
" \".bmp\",\n",
" \".caption\",\n",
" \".npz\",\n",
" \".txt\",\n",
" \".json\",\n",
"]\n",
"\n",
"background_colors = [\n",
" (255, 255, 255),\n",
" (0, 0, 0),\n",
" (255, 0, 0),\n",
" (0, 255, 0),\n",
" (0, 0, 255),\n",
" (255, 255, 0),\n",
" (255, 0, 255),\n",
" (0, 255, 255),\n",
"]\n",
"\n",
"def clean_directory(directory):\n",
" for item in os.listdir(directory):\n",
" file_path = os.path.join(directory, item)\n",
" if os.path.isfile(file_path):\n",
" file_ext = os.path.splitext(item)[1]\n",
" if file_ext not in supported_types:\n",
" print(f\"Deleting file {item} from {directory}\")\n",
" os.remove(file_path)\n",
" elif os.path.isdir(file_path) and recursive:\n",
" clean_directory(file_path)\n",
"\n",
"def process_image(image_path):\n",
" img = Image.open(image_path)\n",
" img_dir, image_name = os.path.split(image_path)\n",
"\n",
" if img.mode in (\"RGBA\", \"LA\"):\n",
" if random_color:\n",
" background_color = random.choice(background_colors)\n",
" else:\n",
" background_color = (255, 255, 255)\n",
" bg = Image.new(\"RGB\", img.size, background_color)\n",
" bg.paste(img, mask=img.split()[-1])\n",
"\n",
" if image_name.endswith(\".webp\"):\n",
" bg = bg.convert(\"RGB\")\n",
" new_image_path = os.path.join(img_dir, image_name.replace(\".webp\", \".jpg\"))\n",
" bg.save(new_image_path, \"JPEG\")\n",
" os.remove(image_path)\n",
" print(f\" Converted image: {image_name} to {os.path.basename(new_image_path)}\")\n",
" else:\n",
" bg.save(image_path, \"PNG\")\n",
" print(f\" Converted image: {image_name}\")\n",
" else:\n",
" if image_name.endswith(\".webp\"):\n",
" new_image_path = os.path.join(img_dir, image_name.replace(\".webp\", \".jpg\"))\n",
" img.save(new_image_path, \"JPEG\")\n",
" os.remove(image_path)\n",
" print(f\" Converted image: {image_name} to {os.path.basename(new_image_path)}\")\n",
" else:\n",
" img.save(image_path, \"PNG\")\n",
"\n",
"def find_images(directory):\n",
" images = []\n",
" for root, _, files in os.walk(directory):\n",
" for file in files:\n",
" if file.endswith(\".png\") or file.endswith(\".webp\"):\n",
" images.append(os.path.join(root, file))\n",
" return images\n",
"\n",
"clean_directory(train_data_dir)\n",
"images = find_images(train_data_dir)\n",
"num_batches = len(images) // batch_size + 1\n",
"\n",
"if convert:\n",
" with concurrent.futures.ThreadPoolExecutor() as executor:\n",
" for i in tqdm(range(num_batches)):\n",
" start = i * batch_size\n",
" end = start + batch_size\n",
" batch = images[start:end]\n",
" executor.map(process_image, batch)\n",
"\n",
" print(\"All images have been converted\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "_mLVURhM9PFE"
},
"outputs": [],
"source": [
"# @title ### **3.2.3. Custom Caption/Tag**\n",
"import os\n",
"\n",
"%store -r\n",
"\n",
"os.chdir(root_dir)\n",
"\n",
"# @markdown Add or remove custom tags here.\n",
"extension = \".txt\" # @param [\".txt\", \".caption\"]\n",
"custom_tag = \"\" # @param {type:\"string\"}\n",
"# @markdown Use `sub_folder` option to specify a subfolder for multi-concept training.\n",
"# @markdown > Specify `--all` to process all subfolders/`recursive`\n",
"sub_folder = \"\" #@param {type: \"string\"}\n",
"# @markdown Enable this to append custom tags at the end of lines.\n",
"append = False # @param {type:\"boolean\"}\n",
"# @markdown Enable this if you want to remove captions/tags instead.\n",
"remove_tag = False # @param {type:\"boolean\"}\n",
"recursive = False\n",
"\n",
"if sub_folder == \"\":\n",
" image_dir = train_data_dir\n",
"elif sub_folder == \"--all\":\n",
" image_dir = train_data_dir\n",
" recursive = True\n",
"elif sub_folder.startswith(\"/content\"):\n",
" image_dir = sub_folder\n",
"else:\n",
" image_dir = os.path.join(train_data_dir, sub_folder)\n",
" os.makedirs(image_dir, exist_ok=True)\n",
"\n",
"def read_file(filename):\n",
" with open(filename, \"r\") as f:\n",
" contents = f.read()\n",
" return contents\n",
"\n",
"def write_file(filename, contents):\n",
" with open(filename, \"w\") as f:\n",
" f.write(contents)\n",
"\n",
"def process_tags(filename, custom_tag, append, remove_tag):\n",
" contents = read_file(filename)\n",
" tags = [tag.strip() for tag in contents.split(',')]\n",
" custom_tags = [tag.strip() for tag in custom_tag.split(',')]\n",
"\n",
" for custom_tag in custom_tags:\n",
" custom_tag = custom_tag.replace(\"_\", \" \")\n",
" if remove_tag:\n",
" while custom_tag in tags:\n",
" tags.remove(custom_tag)\n",
" else:\n",
" if custom_tag not in tags:\n",
" if append:\n",
" tags.append(custom_tag)\n",
" else:\n",
" tags.insert(0, custom_tag)\n",
"\n",
" contents = ', '.join(tags)\n",
" write_file(filename, contents)\n",
"\n",
"def process_directory(image_dir, tag, append, remove_tag, recursive):\n",
" for filename in os.listdir(image_dir):\n",
" file_path = os.path.join(image_dir, filename)\n",
"\n",
" if os.path.isdir(file_path) and recursive:\n",
" process_directory(file_path, tag, append, remove_tag, recursive)\n",
" elif filename.endswith(extension):\n",
" process_tags(file_path, tag, append, remove_tag)\n",
"\n",
"tag = custom_tag\n",
"\n",
"if not any(\n",
" [filename.endswith(extension) for filename in os.listdir(image_dir)]\n",
"):\n",
" for filename in os.listdir(image_dir):\n",
" if filename.endswith((\".png\", \".jpg\", \".jpeg\", \".webp\", \".bmp\")):\n",
" open(\n",
" os.path.join(image_dir, filename.split(\".\")[0] + extension),\n",
" \"w\",\n",
" ).close()\n",
"\n",
"if custom_tag:\n",
" process_directory(image_dir, tag, append, remove_tag, recursive)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "hhgatqF3leHJ"
},
"outputs": [],
"source": [
"# @title ## **3.4. Bucketing and Latents Caching**\n",
"%store -r\n",
"\n",
"# @markdown This code will create buckets based on the `bucket_resolution` provided for multi-aspect ratio training, and then convert all images within the `train_data_dir` to latents.\n",
"bucketing_json = os.path.join(training_dir, \"meta_lat.json\")\n",
"metadata_json = os.path.join(training_dir, \"meta_clean.json\")\n",
"bucket_resolution = 1024 # @param {type:\"slider\", min:512, max:2048, step:128}\n",
"mixed_precision = \"fp16\" # @param [\"no\", \"fp16\", \"bf16\"] {allow-input: false}\n",
"skip_existing = False # @param{type:\"boolean\"}\n",
"flip_aug = True # @param{type:\"boolean\"}\n",
"# @markdown Use `clean_caption` option to clean such as duplicate tags, `women` to `girl`, etc\n",
"clean_caption = False #@param {type:\"boolean\"}\n",
"#@markdown Use the `recursive` option to process subfolders as well\n",
"recursive = True #@param {type:\"boolean\"}\n",
"\n",
"metadata_config = {\n",
" \"_train_data_dir\": train_data_dir,\n",
" \"_out_json\": metadata_json,\n",
" \"recursive\": recursive,\n",
" \"full_path\": recursive,\n",
" \"clean_caption\": clean_caption\n",
"}\n",
"\n",
"bucketing_config = {\n",
" \"_train_data_dir\": train_data_dir,\n",
" \"_in_json\": metadata_json,\n",
" \"_out_json\": bucketing_json,\n",
" \"_model_name_or_path\": vae_path if vae_path else model_path,\n",
" \"recursive\": recursive,\n",
" \"full_path\": recursive,\n",
" \"flip_aug\": flip_aug,\n",
" \"skip_existing\": skip_existing,\n",
" \"batch_size\": 4,\n",
" \"max_data_loader_n_workers\": 2,\n",
" \"max_resolution\": f\"{bucket_resolution}, {bucket_resolution}\",\n",
" \"mixed_precision\": mixed_precision,\n",
"}\n",
"\n",
"def generate_args(config):\n",
" args = \"\"\n",
" for k, v in config.items():\n",
" if k.startswith(\"_\"):\n",
" args += f'\"{v}\" '\n",
" elif isinstance(v, str):\n",
" args += f'--{k}=\"{v}\" '\n",
" elif isinstance(v, bool) and v:\n",
" args += f\"--{k} \"\n",
" elif isinstance(v, float) and not isinstance(v, bool):\n",
" args += f\"--{k}={v} \"\n",
" elif isinstance(v, int) and not isinstance(v, bool):\n",
" args += f\"--{k}={v} \"\n",
" return args.strip()\n",
"\n",
"merge_metadata_args = generate_args(metadata_config)\n",
"prepare_buckets_args = generate_args(bucketing_config)\n",
"\n",
"merge_metadata_command = f\"python merge_all_to_metadata.py {merge_metadata_args}\"\n",
"prepare_buckets_command = f\"python prepare_buckets_latents.py {prepare_buckets_args}\"\n",
"\n",
"os.chdir(finetune_dir)\n",
"!{merge_metadata_command}\n",
"time.sleep(1)\n",
"!{prepare_buckets_command}\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "GAZVkLuaRJ9e"
},
"source": [
"# **IV. Training**\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "cJgLfRtlHSjw"
},
"outputs": [],
"source": [
"import toml\n",
"\n",
"# @title ## **4.1. LoRa: Low-Rank Adaptation Config**\n",
"# @markdown Kohya's `LoRA` renamed to `LoRA-LierLa` and Kohya's `LoCon` renamed to `LoRA-C3Lier`, read [official announcement](https://github.com/kohya-ss/sd-scripts/blob/849bc24d205a35fbe1b2a4063edd7172533c1c01/README.md#naming-of-lora).\n",
"network_category = \"LoRA_LierLa\" # @param [\"LoRA_LierLa\", \"LoRA_C3Lier\", \"DyLoRA_LierLa\", \"DyLoRA_C3Lier\", \"LoCon\", \"LoHa\", \"IA3\", \"LoKR\", \"DyLoRA_Lycoris\"]\n",
"\n",
"# @markdown | network_category | network_dim | network_alpha | conv_dim | conv_alpha | unit |\n",
"# @markdown | :---: | :---: | :---: | :---: | :---: | :---: |\n",
"# @markdown | LoRA-LierLa | 32 | 1 | - | - | - |\n",
"# @markdown | LoCon/LoRA-C3Lier | 16 | 8 | 8 | 1 | - |\n",
"# @markdown | LoHa | 8 | 4 | 4 | 1 | - |\n",
"# @markdown | Other Category | ? | ? | ? | ? | - |\n",
"\n",
"# @markdown Specify `network_args` to add `optional` training args, like for specifying each 25 block weight, read [this](https://github.com/kohya-ss/sd-scripts/blob/main/train_network_README-ja.md#%E9%9A%8E%E5%B1%A4%E5%88%A5%E5%AD%A6%E7%BF%92%E7%8E%87)\n",
"network_args = \"\" # @param {'type':'string'}\n",
"\n",
"# @markdown ### **Linear Layer Config**\n",
"# @markdown Used by all `network_category`. When in doubt, set `network_dim = network_alpha`\n",
"network_dim = 16 # @param {'type':'number'}\n",
"network_alpha = 8 # @param {'type':'number'}\n",
"\n",
"# @markdown ### **Convolutional Layer Config**\n",
"# @markdown Only required if `network_category` is not `LoRA_LierLa`, as it involves training convolutional layers in addition to linear layers.\n",
"conv_dim = 32 # @param {'type':'number'}\n",
"conv_alpha = 16 # @param {'type':'number'}\n",
"\n",
"# @markdown ### **DyLoRA Config**\n",
"# @markdown Only required if `network_category` is `DyLoRA_LierLa` and `DyLoRA_C3Lier`\n",
"unit = 4 # @param {'type':'number'}\n",
"\n",
"if isinstance(network_args, str):\n",
" network_args = network_args.strip()\n",
" if network_args.startswith('[') and network_args.endswith(']'):\n",
" try:\n",
" network_args = ast.literal_eval(network_args)\n",
" except (SyntaxError, ValueError) as e:\n",
" print(f\"Error parsing network_args: {e}\\n\")\n",
" network_args = []\n",
" elif len(network_args) > 0:\n",
" print(f\"WARNING! '{network_args}' is not a valid list! Put args like this: [\\\"args=1\\\", \\\"args=2\\\"]\\n\")\n",
" network_args = []\n",
" else:\n",
" network_args = []\n",
"else:\n",
" network_args = []\n",
"\n",
"network_config = {\n",
" \"LoRA_LierLa\": {\n",
" \"module\": \"networks.lora\",\n",
" \"args\" : []\n",
" },\n",
" \"LoRA_C3Lier\": {\n",
" \"module\": \"networks.lora\",\n",
" \"args\" : [\n",
" f\"conv_dim={conv_dim}\",\n",
" f\"conv_alpha={conv_alpha}\"\n",
" ]\n",
" },\n",
" \"DyLoRA_LierLa\": {\n",
" \"module\": \"networks.dylora\",\n",
" \"args\" : [\n",
" f\"unit={unit}\"\n",
" ]\n",
" },\n",
" \"DyLoRA_C3Lier\": {\n",
" \"module\": \"networks.dylora\",\n",
" \"args\" : [\n",
" f\"conv_dim={conv_dim}\",\n",
" f\"conv_alpha={conv_alpha}\",\n",
" f\"unit={unit}\"\n",
" ]\n",
" },\n",
" \"LoCon\": {\n",
" \"module\": \"lycoris.kohya\",\n",
" \"args\" : [\n",
" f\"algo=locon\",\n",
" f\"conv_dim={conv_dim}\",\n",
" f\"conv_alpha={conv_alpha}\"\n",
" ]\n",
" },\n",
" \"LoHa\": {\n",
" \"module\": \"lycoris.kohya\",\n",
" \"args\" : [\n",
" f\"algo=loha\",\n",
" f\"conv_dim={conv_dim}\",\n",
" f\"conv_alpha={conv_alpha}\"\n",
" ]\n",
" },\n",
" \"IA3\": {\n",
" \"module\": \"lycoris.kohya\",\n",
" \"args\" : [\n",
" f\"algo=ia3\",\n",
" f\"conv_dim={conv_dim}\",\n",
" f\"conv_alpha={conv_alpha}\"\n",
" ]\n",
" },\n",
" \"LoKR\": {\n",
" \"module\": \"lycoris.kohya\",\n",
" \"args\" : [\n",
" f\"algo=lokr\",\n",
" f\"conv_dim={conv_dim}\",\n",
" f\"conv_alpha={conv_alpha}\"\n",
" ]\n",
" },\n",
" \"DyLoRA_Lycoris\": {\n",
" \"module\": \"lycoris.kohya\",\n",
" \"args\" : [\n",
" f\"algo=dylora\",\n",
" f\"conv_dim={conv_dim}\",\n",
" f\"conv_alpha={conv_alpha}\"\n",
" ]\n",
" }\n",
"}\n",
"\n",
"network_module = network_config[network_category][\"module\"]\n",
"network_args.extend(network_config[network_category][\"args\"])\n",
"\n",
"lora_config = {\n",
" \"additional_network_arguments\": {\n",
" \"no_metadata\" : False,\n",
" \"network_module\" : network_module,\n",
" \"network_dim\" : network_dim,\n",
" \"network_alpha\" : network_alpha,\n",
" \"network_args\" : network_args,\n",
" \"network_train_unet_only\" : True,\n",
" \"training_comment\" : None,\n",
" },\n",
"}\n",
"\n",
"print(toml.dumps(lora_config))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "UJMfS3uIm3DW"
},
"outputs": [],
"source": [
"pip install prodigyopt"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "JNlw3u8arwir"
},
"outputs": [],
"source": [
"import toml\n",
"import ast\n",
"\n",
"# @title ## **4.2. Optimizer Config**\n",
"# @markdown Use `Adafactor` optimizer. `RMSprop 8bit` or `Adagrad 8bit` may work. `AdamW 8bit` doesn't seem to work.\n",
"optimizer_type = \"Prodigy\" # @param [\"AdamW\",\"Prodigy\", \"AdamW8bit\", \"Lion8bit\", \"Lion\", \"SGDNesterov\", \"SGDNesterov8bit\", \"DAdaptation(DAdaptAdamPreprint)\", \"DAdaptAdaGrad\", \"DAdaptAdam\", \"DAdaptAdan\", \"DAdaptAdanIP\", \"DAdaptLion\", \"DAdaptSGD\", \"AdaFactor\"]\n",
"# @markdown Specify `optimizer_args` to add `additional` args for optimizer, e.g: `[\"weight_decay=0.6\"]`\n",
"optimizer_args = \"[ \\\"decouple=True\\\",\\\"weight_decay=0.05\\\",\\\"d_coef=0.8\\\",\\\"use_bias_correction=True\\\",\\\"safeguard_warmup=False\\\",\\\"betas=0.9,0.99\\\" ]\" # @param {'type':'string'}\n",
"# @markdown ### **Learning Rate Config**\n",
"# @markdown Different `optimizer_type` and `network_category` for some condition requires different learning rate. It's recommended to set `text_encoder_lr = 1/2 * unet_lr`\n",
"learning_rate = 1 # @param {'type':'number'}\n",
"# @markdown ### **LR Scheduler Config**\n",
"# @markdown `lr_scheduler` provides several methods to adjust the learning rate based on the number of epochs.\n",
"lr_scheduler = \"cosine\" # @param [\"linear\", \"cosine\", \"cosine_with_restarts\", \"polynomial\", \"constant\", \"constant_with_warmup\", \"adafactor\"] {allow-input: false}\n",
"lr_warmup_steps = 0 # @param {'type':'number'}\n",
"# @markdown Specify `lr_scheduler_num` with `num_cycles` value for `cosine_with_restarts` or `power` value for `polynomial`\n",
"lr_scheduler_num = 0 # @param {'type':'number'}\n",
"\n",
"if isinstance(optimizer_args, str):\n",
" optimizer_args = optimizer_args.strip()\n",
" if optimizer_args.startswith('[') and optimizer_args.endswith(']'):\n",
" try:\n",
" optimizer_args = ast.literal_eval(optimizer_args)\n",
" except (SyntaxError, ValueError) as e:\n",
" print(f\"Error parsing optimizer_args: {e}\\n\")\n",
" optimizer_args = []\n",
" elif len(optimizer_args) > 0:\n",
" print(f\"WARNING! '{optimizer_args}' is not a valid list! Put args like this: [\\\"args=1\\\", \\\"args=2\\\"]\\n\")\n",
" optimizer_args = []\n",
" else:\n",
" optimizer_args = []\n",
"else:\n",
" optimizer_args = []\n",
"\n",
"optimizer_config = {\n",
" \"optimizer_arguments\": {\n",
" \"optimizer_type\" : optimizer_type,\n",
" \"learning_rate\" : learning_rate,\n",
" \"max_grad_norm\" : 0,\n",
" \"optimizer_args\" : optimizer_args,\n",
" \"lr_scheduler\" : lr_scheduler,\n",
" \"lr_warmup_steps\" : lr_warmup_steps,\n",
" \"lr_scheduler_num_cycles\" : lr_scheduler_num if lr_scheduler == \"cosine_with_restarts\" else None,\n",
" \"lr_scheduler_power\" : lr_scheduler_num if lr_scheduler == \"polynomial\" else None,\n",
" \"lr_scheduler_type\" : None,\n",
" \"lr_scheduler_args\" : None,\n",
" \"Scale weight norms\" : 1,\n",
" },\n",
"}\n",
"\n",
"print(toml.dumps(optimizer_config))\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"cellView": "form",
"id": "YOxNM0x7dvfO"
},
"outputs": [],
"source": [
"# @title ## **4.3. Advanced Training Config** (Optional)\n",
"import toml\n",
"\n",
"\n",
"# @markdown ### **Optimizer State Config**\n",
"save_optimizer_state = False #@param {type:\"boolean\"}\n",
"load_optimizer_state = \"\" #@param {type:\"string\"}\n",
"# @markdown ### **Noise Control**\n",
"noise_control_type = \"none\" #@param [\"none\", \"noise_offset\", \"multires_noise\"]\n",
"# @markdown #### **a. Noise Offset**\n",
"# @markdown Control and easily generating darker or light images by offset the noise when fine-tuning the model. Recommended value: `0.1`. Read [Diffusion With Offset Noise](https://www.crosslabs.org//blog/diffusion-with-offset-noise)\n",
"noise_offset_num = 0 # @param {type:\"number\"}\n",
"# @markdown **[Experimental]**\n",
"# @markdown Automatically adjusts the noise offset based on the absolute mean values of each channel in the latents when used with `--noise_offset`. Specify a value around 1/10 to the same magnitude as the `--noise_offset` for best results. Set `0` to disable.\n",
"adaptive_noise_scale = 0 # @param {type:\"number\"}\n",
"# @markdown #### **b. Multires Noise**\n",
"# @markdown enable multires noise with this number of iterations (if enabled, around 6-10 is recommended)\n",
"multires_noise_iterations = 1 #@param {type:\"slider\", min:1, max:10, step:1}\n",
"multires_noise_discount = 0.1 #@param {type:\"slider\", min:0.1, max:1, step:0.1}\n",
"# @markdown ### **Caption Dropout**\n",
"caption_dropout_rate = 0 # @param {type:\"number\"}\n",
"caption_tag_dropout_rate = 0 # @param {type:\"number\"}\n",
"caption_dropout_every_n_epochs = 0 # @param {type:\"number\"}\n",
"# @markdown ### **Custom Train Function**\n",
"# @markdown Gamma for reducing the weight of high-loss timesteps. Lower numbers have a stronger effect. The paper recommends `5`. Read the paper [here](https://arxiv.org/abs/2303.09556).\n",
"min_snr_gamma = 6 #@param {type:\"number\"}\n",
"\n",
"advanced_training_config = {\n",
" \"advanced_training_config\": {\n",
" \"resume\" : load_optimizer_state,\n",
" \"save_state\" : save_optimizer_state,\n",
" \"save_last_n_epochs_state\" : save_optimizer_state,\n",
" \"noise_offset\" : noise_offset_num if noise_control_type == \"noise_offset\" else None,\n",
" \"adaptive_noise_scale\" : adaptive_noise_scale if adaptive_noise_scale and noise_control_type == \"noise_offset\" else None,\n",
" \"multires_noise_iterations\" : multires_noise_iterations if noise_control_type ==\"multires_noise\" else None,\n",
" \"multires_noise_discount\" : multires_noise_discount if noise_control_type ==\"multires_noise\" else None,\n",
" \"caption_dropout_rate\" : caption_dropout_rate,\n",
" \"caption_tag_dropout_rate\" : caption_tag_dropout_rate,\n",
" \"caption_dropout_every_n_epochs\": caption_dropout_every_n_epochs,\n",
" \"min_snr_gamma\" : min_snr_gamma if not min_snr_gamma == -1 else None,\n",
" }\n",
"}\n",
"\n",
"print(toml.dumps(advanced_training_config))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "-Z4w3lfFKLjr"
},
"outputs": [],
"source": [
"# @title ## **4.4. Training Config**\n",
"import toml\n",
"import os\n",
"from subprocess import getoutput\n",
"\n",
"%store -r\n",
"\n",
"# @markdown ### **Project Config**\n",
"project_name = \"bludwing\" # @param {type:\"string\"}\n",
"# @markdown Get your `wandb_api_key` [here](https://wandb.ai/settings) to logs with wandb.\n",
"wandb_api_key = \"\" # @param {type:\"string\"}\n",
"in_json = \"/content/LoRA/meta_lat.json\" # @param {type:\"string\"}\n",
"# @markdown ### **SDXL Config**\n",
"gradient_checkpointing = True # @param {type:\"boolean\"}\n",
"no_half_vae = True # @param {type:\"boolean\"}\n",
"#@markdown Recommended parameter for SDXL training but if you enable it, `shuffle_caption` won't work\n",
"cache_text_encoder_outputs = False # @param {type:\"boolean\"}\n",
"#@markdown These options can be used to train U-Net with different timesteps. The default values are 0 and 1000.\n",
"min_timestep = 0 # @param {type:\"number\"}\n",
"max_timestep = 1000 # @param {type:\"number\"}\n",
"# @markdown ### **Dataset Config**\n",
"num_repeats = 1 # @param {type:\"number\"}\n",
"resolution = 2048 # @param {type:\"slider\", min:512, max:2048, step:128}\n",
"keep_tokens = 0 # @param {type:\"number\"}\n",
"# @markdown ### **General Config**\n",
"num_epochs = 62 # @param {type:\"number\"}\n",
"train_batch_size = 4 # @param {type:\"number\"}\n",
"mixed_precision = \"fp16\" # @param [\"no\",\"fp16\",\"bf16\"] {allow-input: false}\n",
"seed = 1 # @param {type:\"number\"}\n",
"optimization = \"scaled dot-product attention\" # @param [\"xformers\", \"scaled dot-product attention\"]\n",
"# @markdown ### **Save Output Config**\n",
"save_precision = \"fp16\" # @param [\"float\", \"fp16\", \"bf16\"] {allow-input: false}\n",
"save_every_n_epochs = 10 # @param {type:\"number\"}\n",
"# @markdown ### **Sample Prompt Config**\n",
"enable_sample = False # @param {type:\"boolean\"}\n",
"sampler = \"euler_a\" # @param [\"ddim\", \"pndm\", \"lms\", \"euler\", \"euler_a\", \"heun\", \"dpm_2\", \"dpm_2_a\", \"dpmsolver\",\"dpmsolver++\", \"dpmsingle\", \"k_lms\", \"k_euler\", \"k_euler_a\", \"k_dpm_2\", \"k_dpm_2_a\"]\n",
"positive_prompt = \"\"\n",
"negative_prompt = \"\"\n",
"quality_prompt = \"NovelAI\" # @param [\"None\", \"Waifu Diffusion 1.5\", \"NovelAI\", \"AbyssOrangeMix\", \"Stable Diffusion XL\"] {allow-input: false}\n",
"if quality_prompt == \"NovelAI\":\n",
" positive_prompt = \"masterpiece, best quality, \"\n",
" negative_prompt = \"lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry, \"\n",
"if quality_prompt == \"AbyssOrangeMix\":\n",
" positive_prompt = \"masterpiece, best quality, \"\n",
" negative_prompt = \"(worst quality, low quality:1.4), \"\n",
"if quality_prompt == \"Stable Diffusion XL\":\n",
" negative_prompt = \"3d render, smooth, plastic, blurry, grainy, low-resolution, deep-fried, oversaturated\"\n",
"custom_prompt = \"face focus, cute, 1girl, green hair, sweater, looking at viewer, upper body, beanie, outdoors, night, turtleneck\" # @param {type:\"string\"}\n",
"# @markdown Specify `prompt_from_caption` if you want to use caption as prompt instead. Will be chosen randomly.\n",
"prompt_from_caption = \"none\" # @param [\"none\", \".txt\", \".caption\"]\n",
"if prompt_from_caption != \"none\":\n",
" custom_prompt = \"\"\n",
"num_prompt = 2 # @param {type:\"number\"}\n",
"logging_dir = os.path.join(training_dir, \"logs\")\n",
"lowram = int(next(line.split()[1] for line in open('/proc/meminfo') if \"MemTotal\" in line)) / (1024**2) < 15\n",
"\n",
"os.chdir(repo_dir)\n",
"\n",
"prompt_config = {\n",
" \"prompt\": {\n",
" \"negative_prompt\" : negative_prompt,\n",
" \"width\" : resolution,\n",
" \"height\" : resolution,\n",
" \"scale\" : 12,\n",
" \"sample_steps\" : 28,\n",
" \"subset\" : [],\n",
" }\n",
"}\n",
"\n",
"train_config = {\n",
" \"sdxl_arguments\": {\n",
" \"cache_text_encoder_outputs\" : cache_text_encoder_outputs,\n",
" \"no_half_vae\" : True,\n",
" \"min_timestep\" : min_timestep,\n",
" \"max_timestep\" : max_timestep,\n",
" \"shuffle_caption\" : True if not cache_text_encoder_outputs else False,\n",
" \"lowram\" : lowram\n",
" },\n",
" \"model_arguments\": {\n",
" \"pretrained_model_name_or_path\" : model_path,\n",
" \"vae\" : vae_path,\n",
" },\n",
" \"dataset_arguments\": {\n",
" \"debug_dataset\" : False,\n",
" \"in_json\" : in_json,\n",
" \"train_data_dir\" : train_data_dir,\n",
" \"dataset_repeats\" : num_repeats,\n",
" \"keep_tokens\" : keep_tokens,\n",
" \"resolution\" : str(resolution) + ',' + str(resolution),\n",
" \"color_aug\" : False,\n",
" \"face_crop_aug_range\" : None,\n",
" \"token_warmup_min\" : 1,\n",
" \"token_warmup_step\" : 0,\n",
" },\n",
" \"training_arguments\": {\n",
" \"output_dir\" : os.path.join(output_dir, project_name),\n",
" \"output_name\" : project_name if project_name else \"last\",\n",
" \"save_precision\" : save_precision,\n",
" \"save_every_n_epochs\" : save_every_n_epochs,\n",
" \"save_n_epoch_ratio\" : None,\n",
" \"save_last_n_epochs\" : None,\n",
" \"resume\" : None,\n",
" \"train_batch_size\" : train_batch_size,\n",
" \"max_token_length\" : 225,\n",
" \"mem_eff_attn\" : False,\n",
" \"sdpa\" : True if optimization == \"scaled dot-product attention\" else False,\n",
" \"xformers\" : True if optimization == \"xformers\" else False,\n",
" \"max_train_epochs\" : num_epochs,\n",
" \"max_data_loader_n_workers\" : 8,\n",
" \"persistent_data_loader_workers\": True,\n",
" \"seed\" : seed if seed > 0 else None,\n",
" \"gradient_checkpointing\" : gradient_checkpointing,\n",
" \"gradient_accumulation_steps\" : 1,\n",
" \"mixed_precision\" : mixed_precision,\n",
" },\n",
" \"logging_arguments\": {\n",
" \"log_with\" : \"wandb\" if wandb_api_key else \"tensorboard\",\n",
" \"log_tracker_name\" : project_name if wandb_api_key and not project_name == \"last\" else None,\n",
" \"logging_dir\" : logging_dir,\n",
" \"log_prefix\" : project_name if not wandb_api_key else None,\n",
" },\n",
" \"sample_prompt_arguments\": {\n",
" \"sample_every_n_steps\" : None,\n",
" \"sample_every_n_epochs\" : save_every_n_epochs if enable_sample else None,\n",
" \"sample_sampler\" : sampler,\n",
" },\n",
" \"saving_arguments\": {\n",
" \"save_model_as\": \"safetensors\"\n",
" },\n",
"}\n",
"\n",
"def write_file(filename, contents):\n",
" with open(filename, \"w\") as f:\n",
" f.write(contents)\n",
"\n",
"def prompt_convert(enable_sample, num_prompt, train_data_dir, prompt_config, custom_prompt):\n",
" if enable_sample:\n",
" search_pattern = os.path.join(train_data_dir, '**/*' + prompt_from_caption)\n",
" caption_files = glob.glob(search_pattern, recursive=True)\n",
"\n",
" if not caption_files:\n",
" if not custom_prompt:\n",
" custom_prompt = \"masterpiece, best quality, 1girl, aqua eyes, baseball cap, blonde hair, closed mouth, earrings, green background, hat, hoop earrings, jewelry, looking at viewer, shirt, short hair, simple background, solo, upper body, yellow shirt\"\n",
" new_prompt_config = prompt_config.copy()\n",
" new_prompt_config['prompt']['subset'] = [\n",
" {\"prompt\": positive_prompt + custom_prompt if positive_prompt else custom_prompt}\n",
" ]\n",
" else:\n",
" selected_files = random.sample(caption_files, min(num_prompt, len(caption_files)))\n",
"\n",
" prompts = []\n",
" for file in selected_files:\n",
" with open(file, 'r') as f:\n",
" prompts.append(f.read().strip())\n",
"\n",
" new_prompt_config = prompt_config.copy()\n",
" new_prompt_config['prompt']['subset'] = []\n",
"\n",
" for prompt in prompts:\n",
" new_prompt = {\n",
" \"prompt\": positive_prompt + prompt if positive_prompt else prompt,\n",
" }\n",
" new_prompt_config['prompt']['subset'].append(new_prompt)\n",
"\n",
" return new_prompt_config\n",
" else:\n",
" return prompt_config\n",
"\n",
"def eliminate_none_variable(config):\n",
" for key in config:\n",
" if isinstance(config[key], dict):\n",
" for sub_key in config[key]:\n",
" if config[key][sub_key] == \"\":\n",
" config[key][sub_key] = None\n",
" elif config[key] == \"\":\n",
" config[key] = None\n",
"\n",
" return config\n",
"\n",
"try:\n",
" train_config.update(optimizer_config)\n",
"except NameError:\n",
" raise NameError(\"'optimizer_config' dictionary is missing. Please run '4.1. Optimizer Config' cell.\")\n",
"\n",
"try:\n",
" train_config.update(lora_config)\n",
"except NameError:\n",
" raise NameError(\"'lora_config' dictionary is missing. Please run '4.1. LoRa: Low-Rank Adaptation Config' cell.\")\n",
"\n",
"advanced_training_warning = False\n",
"try:\n",
" train_config.update(advanced_training_config)\n",
"except NameError:\n",
" advanced_training_warning = True\n",
" pass\n",
"\n",
"prompt_config = prompt_convert(enable_sample, num_prompt, train_data_dir, prompt_config, custom_prompt)\n",
"\n",
"config_path = os.path.join(config_dir, \"config_file.toml\")\n",
"prompt_path = os.path.join(config_dir, \"sample_prompt.toml\")\n",
"\n",
"config_str = toml.dumps(eliminate_none_variable(train_config))\n",
"prompt_str = toml.dumps(eliminate_none_variable(prompt_config))\n",
"\n",
"write_file(config_path, config_str)\n",
"write_file(prompt_path, prompt_str)\n",
"\n",
"print(config_str)\n",
"\n",
"if advanced_training_warning:\n",
" import textwrap\n",
" error_message = \"WARNING: This is not an error message, but the [advanced_training_config] dictionary is missing. Please run the '4.2. Advanced Training Config' cell if you intend to use it, or continue to the next step.\"\n",
" wrapped_message = textwrap.fill(error_message, width=80)\n",
" print('\\033[38;2;204;102;102m' + wrapped_message + '\\033[0m\\n')\n",
" pass\n",
"\n",
"print(prompt_str)\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"background_save": true
},
"id": "p_SHtbFwHVl1"
},
"outputs": [],
"source": [
"#@title ## **4.5. Start Training**\n",
"import os\n",
"import toml\n",
"\n",
"#@markdown Check your config here if you want to edit something:\n",
"#@markdown - `sample_prompt` : /content/LoRA/config/sample_prompt.toml\n",
"#@markdown - `config_file` : /content/LoRA/config/config_file.toml\n",
"\n",
"\n",
"#@markdown You can import config from another session if you want.\n",
"\n",
"sample_prompt = \"/content/LoRA/config/sample_prompt.toml\" #@param {type:'string'}\n",
"config_file = \"/content/LoRA/config/config_file.toml\" #@param {type:'string'}\n",
"\n",
"def read_file(filename):\n",
" with open(filename, \"r\") as f:\n",
" contents = f.read()\n",
" return contents\n",
"\n",
"def train(config):\n",
" args = \"\"\n",
" for k, v in config.items():\n",
" if k.startswith(\"_\"):\n",
" args += f'\"{v}\" '\n",
" elif isinstance(v, str):\n",
" args += f'--{k}=\"{v}\" '\n",
" elif isinstance(v, bool) and v:\n",
" args += f\"--{k} \"\n",
" elif isinstance(v, float) and not isinstance(v, bool):\n",
" args += f\"--{k}={v} \"\n",
" elif isinstance(v, int) and not isinstance(v, bool):\n",
" args += f\"--{k}={v} \"\n",
"\n",
" return args\n",
"\n",
"accelerate_conf = {\n",
" \"config_file\" : \"/content/kohya-trainer/accelerate_config/config.yaml\",\n",
" \"num_cpu_threads_per_process\" : 1,\n",
"}\n",
"\n",
"train_conf = {\n",
" \"sample_prompts\" : sample_prompt if os.path.exists(sample_prompt) else None,\n",
" \"config_file\" : config_file,\n",
" \"wandb_api_key\" : wandb_api_key if wandb_api_key else None\n",
"}\n",
"\n",
"accelerate_args = train(accelerate_conf)\n",
"train_args = train(train_conf)\n",
"\n",
"final_args = f\"accelerate launch {accelerate_args} sdxl_train_network.py {train_args}\"\n",
"\n",
"os.chdir(repo_dir)\n",
"!{final_args}"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "0vhtDOMtG7qd"
},
"source": [
"## V. Testing"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "nyIl9BhNXKUq"
},
"source": [
"# **VI. Deployment**"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "uyDk2aH-fhZx"
},
"outputs": [],
"source": [
"462757"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "awv9qUK1fhWw"
},
"outputs": [],
"source": [
"778"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "qf0J9jBpfhOv"
},
"outputs": [],
"source": [
"97"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "y8CTrlPRs6x4"
},
"outputs": [],
"source": [
"2475754547"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "hI9-Ynuis6nF"
},
"outputs": [],
"source": [
"979"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "J9qaWn7rkd0M"
},
"outputs": [],
"source": [
"76896798"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "r1Yzj2wMkdoA"
},
"outputs": [],
"source": [
"67896789679"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "O2gjmhg2kdda"
},
"outputs": [],
"source": [
"78978976"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "5MgVIlkwkdSz"
},
"outputs": [],
"source": [
"695695"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "zn49N7AxkdF9"
},
"outputs": [],
"source": [
"79547959"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Bydd7AEWkc7T"
},
"outputs": [],
"source": [
"547945795"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "F06jAOXrkcwf"
},
"outputs": [],
"source": [
"947957"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "mwN0Z0PYkcmr"
},
"outputs": [],
"source": [
"94594579"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "zHavPmKukcbA"
},
"outputs": [],
"source": [
"54646546"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "RVCcCWXhkcQq"
},
"outputs": [],
"source": [
"54954945"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "j6QrlFxQkcCe"
},
"outputs": [],
"source": [
"34683468368358"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "Bs2_vuh8kbuU"
},
"outputs": [],
"source": [
"245752475247245"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"gpuType": "T4",
"provenance": []
},
"kernelspec": {
"display_name": "Python 3",
"name": "python3"
},
"language_info": {
"name": "python"
}
},
"nbformat": 4,
"nbformat_minor": 0
}