dreambooth_corridor_vastai.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "aa2c1ada",
   "metadata": {
    "id": "aa2c1ada"
   },
   "source": [
    "# Dreambooth (Corridor Digital)\n",
    "### Notebook implementation by Joe Penna, David Bielejeski\n",
    "### Adapted for those following the [Corridor Digital Dreambooth Tutorial](https://www.corridordigital.com/video/2551) from Feb 2023 but have found that the repo has changed since then.\n",
    "\n",
    "#### ***Prerequisites*** : Your training and regularization datasets are stored in zip files (one per set) in your google drive. Example : `training_img_nikop.zip`, `training_img_vmphntd.zip`, `reg_img_man.zip`, `reg_img_aesthetic.zip`. Their names won't matter but you will need their file IDs. \n",
    "\n",
    "\n",
    "More information on:\n",
    "https://github.com/JoePenna/Dreambooth-Stable-Diffusion"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7b971cc0",
   "metadata": {
    "id": "7b971cc0"
   },
   "source": [
    "## 1. Build Environment\n",
    "This might take a few minutes..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72dbc14b",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# BUILD ENV\n",
    "\n",
    "!pip install numpy==1.23.1\n",
    "!pip install pytorch-lightning==1.7.6\n",
    "!pip install csv-logger\n",
    "!pip install torchmetrics==0.11.1\n",
    "!pip install torch-fidelity==0.3.0\n",
    "!pip install albumentations==1.1.0\n",
    "!pip install opencv-python==4.7.0.72\n",
    "!pip install pudb==2019.2\n",
    "!pip install omegaconf==2.1.1\n",
    "!pip install pillow==9.4.0\n",
    "!pip install einops==0.4.1\n",
    "!pip install transformers==4.25.1\n",
    "!pip install kornia==0.6.7\n",
    "!pip install diffusers[training]==0.3.0\n",
    "!pip install captionizer==1.0.1\n",
    "!pip install -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers\n",
    "!pip install -e git+https://github.com/openai/CLIP.git@main#egg=clip\n",
    "!pip install -e .\n",
    "!pip install huggingface_hub\n",
    "!pip install gdown\n",
    "!pip install gitpython\n",
    "\n",
    "from IPython.display import clear_output\n",
    "clear_output()\n",
    "print(\"✅ Environment Built.\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "01045f19",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "## 2. Download the 1.5 model from Hugging Face\n",
    "This might take a few minutes...<br/>\n",
    "You can also provide your own v1.* model for training by uploading it and renaming it to \"model.ckpt\".  It should be in the same directory as dreambooth_runpod_joepenna.ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "411c9815",
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Download the 1.5 sd model\n",
    "from huggingface_hub import hf_hub_download\n",
    "\n",
    "downloaded_model_path = hf_hub_download(\n",
    " repo_id=\"panopstor/EveryDream\",\n",
    " filename=\"sd_v1-5_vae.ckpt\"\n",
    ")\n",
    "\n",
    "# Move the sd_v1-5_vae.ckpt to the root of this directory as \"model.ckpt\"\n",
    "actual_locations_of_model_blob = !readlink -f {downloaded_model_path}\n",
    "!mv {actual_locations_of_model_blob[-1]} model.ckpt\n",
    "clear_output()\n",
    "print(\"✅ model.ckpt successfully downloaded\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b0295dc4",
   "metadata": {},
   "source": [
    "## 3. Create Training Image Folders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d36c361",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The example folder names here reflect what was used in the Corridor tutorial, \n",
    "# rename them according to your probject's preferences.\n",
    "training_images_root = \"trainingImages\" # <-- change this (optional)\n",
    "\n",
    "subject_token_word = \"nikopueringer\" # <-- change this!\n",
    "subject_class_word = \"man\" # <-- change this!\n",
    "\n",
    "style_token_word = \"vmphntd\" # <-- change this!\n",
    "style_class_word = \"aesthetic\" # <-- change this!\n",
    "\n",
    "subject_path = f'{training_images_root}/{subject_token_word}/{subject_class_word}'\n",
    "style_path = f'{training_images_root}/{style_token_word}/{style_class_word}'\n",
    "\n",
    "# ==============================================\n",
    "# create folders \n",
    "import os\n",
    "if os.path.exists(subject_path) == False:\n",
    "  os.makedirs(subject_path)\n",
    "  print(f'{subject_path} Created.')\n",
    "else:\n",
    "  print(f'{subject_path} already exists.')\n",
    "\n",
    "if os.path.exists(style_path) == False:\n",
    "  os.makedirs(style_path)\n",
    "  print(f'{style_path} Created.')\n",
    "else:\n",
    "  print(f'{style_path} already exists.')\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "8e6fc290",
   "metadata": {},
   "source": [
    "## 4. Download Training Images\n",
    "##### With smaller datasets, you can probably drag & drop the images into the folders manually via the file manager to the left. However if they are large, you can download the zip files and extract them to the relevent folders.  Make sure `share access` to the file is set to `\"anyone with the link\"`. A GDrive share link typically looks like : https://drive.google.com/file/d/1FeHTdwDXcxoW3Nv7486FsbIm679ek5zI/view?usp=sharing\n",
    "##### You just need the File ID part, in this case _**1FeHTdwDXcxoW3Nv7486FsbIm679ek5zI**_\n",
    "##### Also to make life easier, make sure the images in the zip files are in the root and not buried under a nest of other folders. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "762f2825",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "# set file_name and file_id for your training image zip files\n",
    "# Note!: file_names do not need to match the actual filename themselves, \n",
    "# they are only used for convenience (so that a name is known when extracting).\n",
    "\n",
    "file_name_subject = 'training_images_subject.zip'   # <-- change this (optional)\n",
    "file_id_subject = '1YGEgEPD-9oKf9830VHZgcbmpolI5AoqP' # <-- change this!\n",
    "\n",
    "file_name_style = 'training_images_style.zip' # <-- change this(optional)\n",
    "file_id_style = '1b_DJ8Y-F_et7It1T57QO2Hf6pQ6zFE-6' # <-- change this!\n",
    "\n",
    "# ====================================================\n",
    "# download them\n",
    "!gdown $file_id_subject -O $file_name_subject\n",
    "!gdown $file_id_style -O $file_name_style\n",
    "\n",
    "# ====================================================\n",
    "# now extract them into the relevant locations\n",
    "import zipfile as z\n",
    "\n",
    "# extract subject training images\n",
    "zf = z.ZipFile(f'{file_name_subject}','r')\n",
    "zf.extractall(f'{subject_path}')\n",
    "zf.close()\n",
    "\n",
    "# extract style training images\n",
    "zf = z.ZipFile(f'{file_name_style}','r')\n",
    "zf.extractall(f'{style_path}')\n",
    "zf.close()\n",
    "\n",
    "# optional : delete zip files after\n",
    "# os.remove(f'{file_name_subject}')\n",
    "# os.remove(f'{file_name_style}')\n",
    "\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'{subject_path}'\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'\\033[91m folder {file_path} was found in style training images folder.  Check and remove it.\\033[0m')\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",
    "# now do the same for the style training images folder\n",
    "folder_path = f'{style_path}'\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'\\033[91m folder {file_path} was found in style training images folder.  Check and remove it.\\033[0m')\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')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f379c246",
   "metadata": {},
   "source": [
    "\n",
    "## 5. Create Regularization Image Folders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fabc915",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "# The example folder names here reflect what was used in the Corridor tutorial, \n",
    "# rename them according to your probject's preferences.\n",
    "reg_images_root = \"regImages\" # <-- change this (optional)\n",
    "\n",
    "# subject_class_word and style_class_word used below were already defined when creating the training image folders\n",
    "# will be reused here. \n",
    "reg_subject_path = f'{reg_images_root}/{subject_class_word}'\n",
    "reg_style_path = f'{reg_images_root}/{style_class_word}'\n",
    "\n",
    "# create folders \n",
    "if os.path.exists(reg_subject_path) == False:\n",
    "  os.makedirs(reg_subject_path)\n",
    "  print(f'{reg_subject_path} Created.')\n",
    "else:\n",
    "  print(f'{reg_subject_path} already exists.')\n",
    "\n",
    "if os.path.exists(reg_style_path) == False:\n",
    "  os.makedirs(reg_style_path)\n",
    "  print(f'{reg_style_path} Created.')\n",
    "else:\n",
    "  print(f'{reg_style_path} already exists.')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "17d1d11a",
   "metadata": {
    "id": "17d1d11a"
   },
   "source": [
    "## 6. Download Regularization Images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43350a1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# set file_name and file_id for your regularization image zip files\n",
    "# Note!: file_names do not need to match the actual filename themselves, \n",
    "# they are only used for convenience (so that the name is known when extracting).\n",
    "\n",
    "file_name_subject_class = 'reg_man.zip'  # <-- change this (optional)\n",
    "file_id_subject_class = '1b_JT1yCrsw3DrLvnHqrRaywEU53jA-Tj'  # <-- change this!\n",
    "\n",
    "file_name_style_class = 'reg_aesthetic.zip'  # <-- change this (optional)\n",
    "file_id_style_class = '1h1EogKPXU8NIue00VZDRscqzqWI1D43M'  # <-- change this!\n",
    "\n",
    "# ====================================================\n",
    "# download them\n",
    "!gdown $file_id_subject_class -O $file_name_subject_class\n",
    "!gdown $file_id_style_class -O $file_name_style_class\n",
    "\n",
    "# ====================================================\n",
    "# now extract them into the correct locations\n",
    "import zipfile as z\n",
    "\n",
    "# extract subject reg images\n",
    "zf = z.ZipFile(f'{file_name_subject_class}','r')\n",
    "zf.extractall(f'{reg_subject_path}')\n",
    "zf.close()\n",
    "\n",
    "# extract style reg images\n",
    "zf = z.ZipFile(f'{file_name_style_class}','r')\n",
    "zf.extractall(f'{reg_style_path}')\n",
    "zf.close()\n",
    "\n",
    "# optional : delete zip files after\n",
    "# os.remove(f'{file_name_subject_class}')\n",
    "# os.remove(f'{file_name_style_class}')\n",
    "\n",
    "# ====================================================\n",
    "# delete any non-image files & warn if any additional folders\n",
    "import os\n",
    "import shutil\n",
    "from glob import glob\n",
    "folder_path = f'{reg_subject_path}'\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'\\033[91m folder {file_path} was found in style reg images folder.  Check and remove it.\\033[0m')\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",
    "# now do the same for the style training images folder\n",
    "folder_path = f'{reg_style_path}'\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'\\033[91m folder {file_path} was found in style reg images folder.  Check and remove it.\\033[0m')\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"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ee303303",
   "metadata": {},
   "source": [
    "## 7. Setting Training Image Repeats\n",
    "\n",
    "#### In the Corridor video, the training **repeats** value was changed in the `v1-finetune_unfrozen.yaml` file. This is now moved to the file : `dreambooth_helpers/dreambooth_trainer_configurations.py`\n",
    "#### On Line 192 in the file you will find the relevant repeats value. If you prefer this to be done for you, run the next cell, otherwise skip it and move to Cell 8. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ed01274",
   "metadata": {},
   "outputs": [],
   "source": [
    "repeats = 100 # <-- change this (optional)\n",
    "\n",
    "filename = \"dreambooth_helpers/dreambooth_trainer_configurations.py\"\n",
    "file = open(filename)\n",
    "lines = file.readlines()\n",
    "\n",
    "repeats_line = lines[192-1]\n",
    "\n",
    "#print(repeats_line)\n",
    "\n",
    "pattern = '\"repeats\":'.rjust(30) # note: The argument to rjust needs to be the length of the final string + padding, not amount of padding\n",
    "#print(pattern)\n",
    "if repeats_line.startswith(pattern):    \n",
    "    # replace the line\n",
    "    new_line = ''\n",
    "    if repeats >0 and repeats <=9:\n",
    "        new_line = f'\"repeats\": {repeats},\\n'.rjust(34)\n",
    "    elif repeats >=10 and repeats <=99:\n",
    "        new_line = f'\"repeats\": {repeats},\\n'.rjust(35)\n",
    "    elif repeats >=100 and repeats <=999:\n",
    "        new_line = f'\"repeats\": {repeats},\\n'.rjust(36)\n",
    "\n",
    "    lines[192-1] = new_line\n",
    "\n",
    "    file = open(filename, \"w\")\n",
    "\n",
    "    for line in lines:\n",
    "        file.write(line)\n",
    "\n",
    "    file.close()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ad4e50df",
   "metadata": {
    "id": "ad4e50df"
   },
   "source": [
    "## 8. Training\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fa5dd66-2ca0-4819-907e-802e25583ae6",
   "metadata": {
    "id": "6fa5dd66-2ca0-4819-907e-802e25583ae6",
    "tags": []
   },
   "outputs": [],
   "source": [
    "import time\n",
    "from datetime import timedelta\n",
    "\n",
    "# Token and class will be ignored because the folder structure used for the training & regularization images will be used for token and class\n",
    "# However the token paramater is still required when launching main.py. \n",
    "token_word=\"xxx\"  \n",
    "class_word=\"yyy\"\n",
    "max_steps=3000  # <-- change this!\n",
    "save_every_x_steps=0  # <-- change this (optional)\n",
    "model_path=\"model.ckpt\"\n",
    "train_img_path=f'{training_images_root}'\n",
    "reg_img_path=f'{reg_images_root}'\n",
    "proj_name=\"myProject\"  # <-- change this (optional)\n",
    "\n",
    "\n",
    "start = time.time()\n",
    "\n",
    "# ==================================================\n",
    "# Start Training\n",
    "# ==================================================\n",
    "!python \"main.py\" \\\n",
    "--project_name \"{proj_name}\" \\\n",
    "--token \"{token_word}\" \\\n",
    "--max_training_steps {max_steps} \\\n",
    "--save_every_x_steps {save_every_x_steps} \\\n",
    "--regularization_images \"{reg_img_path}\" \\\n",
    "--training_images \"{train_img_path}\" \\\n",
    "--training_model \"{model_path}\" \\\n",
    "--flip_p 0 \\\n",
    "--learning_rate 1.0e-06\n",
    "\n",
    "\n",
    "# ==================================================\n",
    "# Show Training Time\n",
    "# ==================================================\n",
    "end = time.time()\n",
    "print(f'Elapsed Time: {timedelta(seconds=end-start)}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a90ac5c",
   "metadata": {},
   "source": [
    "# Big Important Note!\n",
    "\n",
    "The way to use your token is `<token> <class>` ie `joepenna person` and not just `joepenna`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d28d0139",
   "metadata": {},
   "source": [
    "## Generate Images With Your Trained Model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80ddb03b",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python scripts/stable_txt2img.py \\\n",
    " --ddim_eta 0.0 \\\n",
    " --n_samples 1 \\\n",
    " --n_iter 4 \\\n",
    " --scale 7.0 \\\n",
    " --ddim_steps 50 \\\n",
    " --ckpt \"/workspace/Dreambooth-Stable-Diffusion/trained_models/{file_name}\" \\\n",
    " --prompt \"joepenna person as a masterpiece portrait painting by John Singer Sargent in the style of Rembrandt\""
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "provenance": []
  },
  "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.11.3"
  },
  "vscode": {
   "interpreter": {
    "hash": "5c7b89af1651d0b8571dde13640ecdccf7d5a6204171d6ab33e7c296e100e08a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}