!sdxl_kohya_vastai_no_config.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "40fba68a-cb1f-44bb-a8ef-a3dc2b696010",
   "metadata": {},
   "source": [
    "# SDXL Dreambooth Training With Kohya_SS SD-Scripts - No-config version\n",
    "# Run each cell in sequence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9888623-b6f4-4bd9-aa2d-e1234345991d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. INSTALL DEPENDENCIES\n",
    "\n",
    "!wget https://raw.githubusercontent.com/yushan777/kohya_ss_vastai/main/colors.py -q\n",
    "from colors import bcolors\n",
    "from IPython.display import clear_output\n",
    "\n",
    "print(f\"{bcolors.BOLD}{bcolors.GREEN}Installing prerequisites...{bcolors.ENDC}\")\n",
    "# !sudo apt-get update -y && apt-get install -y -qq libgl1\n",
    "# !sudo dpkg --configure -a\n",
    "# !sudo apt update -y && sudo apt install -y -qq python3-tk\n",
    "# clear_output(wait=True)\n",
    "\n",
    "print(f\"{bcolors.BOLD}{bcolors.GREEN}Installing main dependencies - this will take a few minutes...{bcolors.ENDC}\")\n",
    "%pip install gdown\n",
    "!chmod +x setup.sh\n",
    "!./setup.sh\n",
    "\n",
    "#revert huggingface_hub to 0.24.5 for compatibility\n",
    "%pip install huggingface_hub==0.24.5\n",
    "\n",
    "print(f\"{bcolors.BOLD}{bcolors.GREEN}Finished installing dependencies.{bcolors.ENDC}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67c2137a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2. Configure Accelerate With defaults and bf16 precision\n",
    "!accelerate config default --mixed_precision \"bf16\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "197e5f03-f37d-4a75-bd3c-b4ffd5b62833",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3. VARIABLES\n",
    "# The following variables can be changed according to your project preferences\n",
    "token_word = \"ohwx\"\n",
    "class_word = \"person\" \n",
    "training_repeats = 40\n",
    "training_root_dir = \"training_images\" \n",
    "regularization_root_dir = \"reg_images\"\n",
    "models_dir = \"training_models\"\n",
    "project_name = \"myProject\"\n",
    "output_dir = \"trained_models\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5db7f292-afe1-4cdd-b639-edb2d14d6bd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 4. CREATE DATASET FOLDERS & MODEL FOLDER\n",
    "\n",
    "# ================================================================\n",
    "# your training and reg image subfolders will be named according to what is set in previous cell: \n",
    "# so if you stuck with the defaults, they would be for example :\n",
    "# \"training_images/40_ohwx person\"\n",
    "# \"reg_images/1_person\"\n",
    "training_dir = f'{training_root_dir}/{training_repeats}_{token_word} {class_word}'\n",
    "reg_dir = f'{regularization_root_dir}/1_{class_word}'\n",
    "\n",
    "# back to parent folder\n",
    "%cd /workspace/kohya_ss\n",
    "\n",
    "import os\n",
    "if os.path.exists(training_dir) == False:\n",
    "  os.makedirs(training_dir)\n",
    "  print(f'{training_dir} Created.')\n",
    "else:\n",
    "  print(f'{training_dir} already exists.')\n",
    "\n",
    "if os.path.exists(reg_dir) == False:\n",
    "  os.makedirs(reg_dir)\n",
    "  print(f'{reg_dir} Created.')\n",
    "else:\n",
    "  print(f'{reg_dir} already exists.')\n",
    "\n",
    "if os.path.exists(models_dir) == False:\n",
    "  os.makedirs(models_dir)\n",
    "  print(f'{models_dir} Created.')\n",
    "else:\n",
    "  print(f'{models_dir} already exists.')\n",
    "\n",
    "if os.path.exists(output_dir) == False:\n",
    "  os.makedirs(output_dir)\n",
    "  print(f'{output_dir} Created.')\n",
    "else:\n",
    "  print(f'{output_dir} already exists.')\n",
    "\n",
    "# Create prompt file use for samples during training. \n",
    "# first 2 prompts are the same but one has CFG=1 and the other CFG=7\n",
    "# if subject likeness is strong and regular even at CFG=1 then it is an good indicator that it is overfitted\n",
    "# third prompt is to check how well it responds to styling. \n",
    "lines = [\n",
    "    f\"a photo of {token_word} {class_word} --w 1024 --h 1024 --l 7, --s 20 --d 1234567890\\n\",\n",
    "    f\"a photo of {token_word} {class_word} --w 1024 --h 1024 --l 1, --s 20 --d 1234567890\\n\",\n",
    "    f\"a portrait of {token_word} {class_word} in the style of Rembrandt --w 1024 --h 1024 --l 6, --s 20 --d 1234567890\\n\"\n",
    "]\n",
    "\n",
    "# Create and write to the text file\n",
    "with open(\"prompt.txt\", \"w\") as file:\n",
    "    file.writelines(lines)\n",
    "\n",
    "print(\"File 'prompt.txt' has been created.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16769aec",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 5. Download and unzip Training Images to Training Image folder\n",
    "import shutil\n",
    "import zipfile\n",
    "\n",
    "# back to parent folder\n",
    "%cd /workspace/kohya_ss\n",
    "\n",
    "# =================================================================\n",
    "# Use YOUR OWN Google Drive file IDs for your image zips\n",
    "training_images_file_ID = '1BIixbqMYW5xxxxxxxxxxxxxxxxxxx' \n",
    "# =================================================================\n",
    "\n",
    "# download training images from google drive, rename to train.zip\n",
    "!gdown '{training_images_file_ID}' -O train.zip\n",
    "\n",
    "# move train.zip to training images sub folder\n",
    "shutil.move('train.zip', f'{training_dir}')\n",
    "\n",
    "# extract zip contents to current folder and delete zip\n",
    "%cd $training_dir\n",
    "with zipfile.ZipFile('train.zip', 'r') as train_ref:\n",
    "    train_ref.extractall()\n",
    "\n",
    "# delete zip\n",
    "!rm train.zip\n",
    "\n",
    "# back to parent folder\n",
    "%cd /workspace/kohya_ss"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3628554f",
   "metadata": {},
   "source": [
    "## For Regularization Images, you have 2 choices: \n",
    "1) Download your own regularization images (google drive) : use Cell 6a. <br/>\n",
    "or\n",
    "2) Download a set (if class is appropriate) from my github repo. : use Cell 6b. <br/> \n",
    "Sets available : person_ddim, woman_ddim, man_ddim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02061da7-5923-4b79-ac3b-daa493d20337",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 6a. Download YOUR OWN Reg Images (Google Drive)\n",
    "import shutil\n",
    "import zipfile\n",
    "\n",
    "# back to parent folder\n",
    "%cd /workspace/kohya_ss\n",
    "\n",
    "# =================================================================\n",
    "# Use YOUR OWN Google Drive file IDs for your image zips\n",
    "reg_images_file_ID = '1CIqOhLBfzxxxxxxxxxxxxxxxxxx'\n",
    "# =================================================================\n",
    "\n",
    "# download reg images from google drive, rename to reg.zip\n",
    "!gdown '{reg_images_file_ID}' -O reg.zip\n",
    "\n",
    "# move train.zip to reg images sub folder\n",
    "shutil.move('reg.zip', f'{reg_dir}')\n",
    "\n",
    "# extract zip contents to current folder and delete zip\n",
    "%cd $reg_dir\n",
    "with zipfile.ZipFile('reg.zip', 'r') as reg_ref:\n",
    "    reg_ref.extractall()\n",
    "\n",
    "# delete zip\n",
    "!rm reg.zip\n",
    "\n",
    "# back to parent folder\n",
    "%cd /workspace/kohya_ss\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9f43fe0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 6b. Download one of my pre-made set of regularization images\n",
    "!pip install github-clone\n",
    "import shutil\n",
    "import os\n",
    "\n",
    "# back to parent folder\n",
    "%cd /workspace/kohya_ss\n",
    "\n",
    "print(\"This might take a few minutes.\")\n",
    "orig_dir_name = ''\n",
    "if class_word == 'person':\n",
    "    !ghclone https://github.com/yushan777/SD-Regularization-Images/tree/main/sdxl_person_ddim_1000\n",
    "    orig_dir_name = 'sdxl_person_ddim_1000'\n",
    "elif class_word == 'woman':\n",
    "    !ghclone https://github.com/yushan777/SD-Regularization-Images/tree/main/sdxl_woman_ddim_1000\n",
    "    orig_dir_name = 'sdxl_woman_ddim_1000'\n",
    "elif class_word == 'man':\n",
    "    !ghclone https://github.com/yushan777/SD-Regularization-Images/tree/main/sdxl_man_ddim_1000\n",
    "    orig_dir_name = 'sdxl_man_ddim_1000'\n",
    "    \n",
    "\n",
    "if len(orig_dir_name) > 0:\n",
    "    # first make sure reg dir is empty, by emptying it\n",
    "    for filename in os.listdir(reg_dir):\n",
    "        file_path = os.path.join(reg_dir, filename)\n",
    "        try:\n",
    "            if os.path.isfile(file_path) or os.path.islink(file_path):\n",
    "                os.unlink(file_path)\n",
    "            elif os.path.isdir(file_path):\n",
    "                shutil.rmtree(file_path)\n",
    "        except Exception as e:\n",
    "            print('Failed to delete %s. Reason: %s' % (file_path, e))\n",
    "            \n",
    "\n",
    "    # now move images to reg folder\n",
    "    print(\"Moving images to {reg_dir}\")\n",
    "    file_names = os.listdir(orig_dir_name)\n",
    "        \n",
    "    for file_name in file_names:\n",
    "        shutil.move(os.path.join(orig_dir_name, file_name), reg_dir)\n",
    "\n",
    "    # back to parent folder\n",
    "    %cd /workspace/kohya_ss\n",
    "\n",
    "    print(f\"{bcolors.GREEN}Finished downloading regularization images.{bcolors.ENDC}\")\n",
    "else:\n",
    "    print(f\"{bcolors.RED}Your class {class_word} does not match any available pre-made sets. Nothing downloaded.{bcolors.ENDC}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a446f2c4-984a-48db-be94-3d0bfba0f57c",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# 7. Check Dataset Folders\n",
    "# ====================================================\n",
    "# remove any non-image files & warn if any additional folders exist\n",
    "import os\n",
    "import shutil\n",
    "from glob import glob\n",
    "folder_path = f'{training_dir}'\n",
    "\n",
    "# Get a list of all files in the folder\n",
    "files = glob(folder_path + '/*', recursive=False)\n",
    "\n",
    "# Iterate over the files and delete the ones that are not JPG or PNG\n",
    "for file_path in files:\n",
    "    if not (file_path.endswith('.jpg') or file_path.endswith('.png')):\n",
    "        if os.path.isfile(file_path):\n",
    "            os.remove(file_path)\n",
    "        elif os.path.isdir(file_path):\n",
    "            print(f'{bcolors.BOLD}{bcolors.RED} Unexpected folder: \\'{file_path}\\' was found in training images folder path.  Check and remove it.{bcolors.ENDC}')\n",
    "\n",
    "# force remove hidden .ipynb_checkpoints folder in images folder. \n",
    "if os.path.exists(f'{folder_path}/.ipynb_checkpoints'):\n",
    "    shutil.rmtree(f'{folder_path}/.ipynb_checkpoints')\n",
    "\n",
    "# ====================================================\n",
    "# delete any non-image files & warn if any additional folders\n",
    "\n",
    "folder_path = f'{reg_dir}'\n",
    "\n",
    "# Get a list of all files in the folder\n",
    "files = glob(folder_path + '/*', recursive=False)\n",
    "\n",
    "# Iterate over the files and delete the ones that are not JPG or PNG\n",
    "for file_path in files:\n",
    "    if not (file_path.endswith('.jpg') or file_path.endswith('.png')):\n",
    "        if os.path.isfile(file_path):\n",
    "            os.remove(file_path)\n",
    "        elif os.path.isdir(file_path):\n",
    "            print(f'{bcolors.BOLD}{bcolors.RED} Unexpected folder: \\'{file_path}\\' was found in training images folder path.  Check and remove it.{bcolors.ENDC}')\n",
    "\n",
    "# force remove hidden .ipynb_checkpoints folder in images folder. \n",
    "if os.path.exists(f'{folder_path}/.ipynb_checkpoints'):\n",
    "    shutil.rmtree(f'{folder_path}/.ipynb_checkpoints')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3dd2358-b766-46cf-a75b-e46b6d14a1b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 8. Download SDXL 1.0 (0.9VAE) base model (the one with 0.9vae baked in instead of 1.0vae\n",
    "from huggingface_hub import hf_hub_download\n",
    "\n",
    "# Specify the repository ID, the filename, and the desired download directory\n",
    "repo_id = 'stabilityai/stable-diffusion-xl-base-1.0'  \n",
    "filename = 'sd_xl_base_1.0_0.9vae.safetensors'  \n",
    "download_directory = 'training_models'  \n",
    "\n",
    "# Download the file to the training_models dir, as a regular file (not symlink)\n",
    "base_model = hf_hub_download(repo_id=repo_id, filename=filename, local_dir=download_directory)\n",
    "print(base_model)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75e3acb2-df65-4515-b00e-4d002be7f49f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 9. START TRAINING\n",
    "\n",
    "max_steps = 2000\n",
    "learning_rate = 1e-5 \n",
    "learning_rate_te1 = 1e-7\n",
    "learning_rate_te2 = 1e-7\n",
    "save_every_n_steps = 1000\n",
    "\n",
    "print(base_model)\n",
    "!accelerate launch --num_cpu_threads_per_process=2 \"sd-scripts/sdxl_train.py\" \\\n",
    "  --pretrained_model_name_or_path=\"{base_model}\" \\\n",
    "  --train_data_dir=\"{training_root_dir}\" \\\n",
    "  --reg_data_dir=\"{regularization_root_dir}\" \\\n",
    "  --output_dir=\"{output_dir}\" \\\n",
    "  --output_name=\"{project_name}\" \\\n",
    "  --save_model_as=\"safetensors\" \\\n",
    "  --train_batch_size=2 \\\n",
    "  --max_train_steps={max_steps} \\\n",
    "  --save_every_n_steps={save_every_n_steps} \\\n",
    "  --optimizer_type=\"adafactor\" \\\n",
    "  --optimizer_args scale_parameter=False relative_step=False warmup_init=False \\\n",
    "  --xformers \\\n",
    "  --cache_latents \\\n",
    "  --lr_scheduler=\"constant_with_warmup\" \\\n",
    "  --lr_warmup_steps=100 \\\n",
    "  --learning_rate=\"{learning_rate}\" \\\n",
    "  --learning_rate_te1=\"{learning_rate_te1}\" \\\n",
    "  --learning_rate_te2=\"{learning_rate_te2}\" \\\n",
    "  --max_grad_norm=0.0 \\\n",
    "  --train_text_encoder \\\n",
    "  --resolution=\"1024,1024\" \\\n",
    "  --save_precision=\"bf16\" \\\n",
    "  --save_n_epoch_ratio=1 \\\n",
    "  --max_data_loader_n_workers=1 \\\n",
    "  --persistent_data_loader_workers \\\n",
    "  --mixed_precision=\"bf16\" \\\n",
    "  --full_bf16 \\\n",
    "  --logging_dir=\"logs\" \\\n",
    "  --log_prefix=\"last\" \\\n",
    "  --gradient_checkpointing \\\n",
    "  --min_snr_gamma=0 \\\n",
    "  --noise_offset=0.0357 \\\n",
    "  --sample_sampler=\"euler_a\" \\\n",
    "  --sample_prompts=\"prompt.txt\" \\\n",
    "  --sample_every_n_steps=100\n",
    "\n",
    "# Training finished...\n",
    "# List all checkpoints\n",
    "print(\"Trained models:\\n\")\n",
    "!ls trained_models -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e44072bf-090f-4adf-bece-8f9d3b5db145",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 10. (Optional) Extract LoRA from checkpoint\n",
    "\n",
    "trained_model_name='myProject.safetensors' #this must match one of your trained checkpoints\n",
    "lora_name='xxxxxxxxxxxxx-loRA-64.safetensors' #this can be anything. \n",
    "\n",
    "!python3 sd-scripts/networks/extract_lora_from_models.py \\\n",
    "--sdxl \\\n",
    "--model_org='{base_model}' \\\n",
    "--model_tuned='{output_dir}/{trained_model_name}' \\\n",
    "--save_to='{output_dir}/{lora_name}' \\\n",
    "--save_precision='bf16' \\\n",
    "--min_diff=0.001 \\\n",
    "--dim=64 \\\n",
    "--device='cuda'\n",
    "\n",
    "# more info about arguments at:\n",
    "# https://github.com/kohya-ss/sd-scripts/blob/95ae56bd22c285ccb2fe5fca96d92f39842bb99b/networks/extract_lora_from_models.py#L211\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ce071cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 11. (Optional) Upload to Hugging Face \n",
    "# you will need an existing repo that you have access rights to and your hugging face access token\n",
    "\n",
    "!huggingface-cli login --token hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx\n",
    "\n",
    "%cd output # or a subfolder of output/\n",
    "\n",
    "from huggingface_hub import HfApi\n",
    "# Variables\n",
    "filename = 'xxxxxxxxx.safetensors'\n",
    "target_subfolder = 'xxx'\n",
    "target_filename = 'xxxxxxxxx.safetensors'\n",
    "\n",
    "api = HfApi()\n",
    "api.upload_file(\n",
    "    path_or_fileobj=filename,  # Assuming the file is in the current working directory\n",
    "    path_in_repo=f\"{target_subfolder}/{filename}\",  # create the path in the repository\n",
    "    repo_id=\"username/repo\",\n",
    "    repo_type=\"model\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f9ef7f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Empty System Trash\n",
    "# If you have deleted models etc or large downloads, they will still be in the trash\n",
    "# you will need to run this to release storage space\n",
    "rm -rf ~/.local/share/Trash/*"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}