dreambooth_google_colab_joepenna.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "Ab5eS5Zu2538"
   },
   "source": [
    "## Dreambooth\n",
    "#### Colab implementation by David Bielejeski. Latest information on: https://github.com/JoePenna/Dreambooth-Stable-Diffusion\n",
    "##### Before starting, make sure you have the appropriate Accelerator and GPU Type selected from the Runtime menu `Runtime > Change runtime type`.  A minimum of 24GB of VRAM is required so you should select the A100 GPU (40GB). Both the T4 and V100 have less than 20GB. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title Check GPU and VRAM available. (Optional)\n",
    "!nvidia-smi --query-gpu=name,memory.total,memory.free --format=csv,noheader\n",
    "print(\"\\n\\033[92mIf the available VRAM is equal or more than 24GB, then you are good to go.\\033[0m\") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#@title 1. Clone & Download The Repo\n",
    "!git clone https://github.com/JoePenna/Dreambooth-Stable-Diffusion\n",
    "%cd Dreambooth-Stable-Diffusion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qeTrc2vOeiNh"
   },
   "outputs": [],
   "source": [
    "#@title 2. Build The Environment\n",
    "#@markdown You might get warnings about restarting the runtime. Do this from the Runtime menu and after restarting, resume from Cell.\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 gitpython\n",
    "\n",
    "import os\n",
    "os._exit(00)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#@title 3. Just to ensure you are in the right directory.\n",
    "%cd Dreambooth-Stable-Diffusion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "O15vMMhCevib"
   },
   "outputs": [],
   "source": [
    "#@title 4. Download the 1.5 SD model with the improved VAE\n",
    "from IPython.display import clear_output\n",
    "\n",
    "from huggingface_hub import hf_hub_download\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\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "N96aedTtfBjO"
   },
   "outputs": [],
   "source": [
    "#@title 5. Download Regularization Images\n",
    "#@markdown We’ve created the following image sets\n",
    "#@markdown - `man_euler` - provided by Niko Pueringer (Corridor Digital) - euler @ 40 steps, CFG 7.5\n",
    "#@markdown - `man_unsplash` - pictures from various photographers\n",
    "#@markdown - `person_ddim`\n",
    "#@markdown - `woman_ddim` - provided by David Bielejeski - ddim @ 50 steps, CFG 10.0 <br />\n",
    "#@markdown - `artstyle` - provided by Hackmans - ddim @ 50 steps, CFG 10.0 <br />\n",
    "\n",
    "dataset=\"person_ddim\" #@param [\"man_euler\", \"man_unsplash\", \"person_ddim\", \"woman_ddim\", \"artstyle\"]\n",
    "!git clone https://github.com/djbielejeski/Stable-Diffusion-Regularization-Images-{dataset}.git\n",
    "\n",
    "!mkdir -p regularization_images/{dataset}\n",
    "!mv -v Stable-Diffusion-Regularization-Images-{dataset}/{dataset}/*.* regularization_images/{dataset}\n",
    "\n",
    "# remove temp folder now it is empty. \n",
    "!rm -rf Stable-Diffusion-Regularization-Images-{dataset}\n",
    "\n",
    "clear_output()\n",
    "print(\"✅ \\033[92mRegularization Images downloaded.\\033[0m\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "A7hmdOdOfGzs"
   },
   "outputs": [],
   "source": [
    "#@title 6. Training Images\n",
    "#@markdown ## Upload your training images\n",
    "#@markdown WARNING: Be sure to upload an even amount of images, otherwise the training inexplicably stops at 1500 steps. <br />\n",
    "#@markdown - 2-3 full body\n",
    "#@markdown - 3-5 upper body\n",
    "#@markdown - 5-12 close-up on face  <br /> <br />\n",
    "#@markdown The images should be as close as possible to the kind of images you’re trying to make (most of the time, that means no selfies).<br /><br />\n",
    "#@markdown If you get an error during uploading, just manually drag your training images into the training_images folder.\n",
    "from google.colab import files\n",
    "from IPython.display import clear_output\n",
    "\n",
    "# Create the directory\n",
    "!rm -rf training_images\n",
    "!mkdir -p training_images\n",
    "\n",
    "# Upload the files\n",
    "uploaded = files.upload()\n",
    "for filename in uploaded.keys():\n",
    " updated_file_name = filename.replace(\" \", \"_\")\n",
    " !mv \"{filename}\" \"training_images/{updated_file_name}\"\n",
    " clear_output()\n",
    "\n",
    "# Tell the user what is going on\n",
    "training_images_file_paths = !find training_images/*\n",
    "if len(training_images_file_paths) == 0:\n",
    " print(\"❌ \\033[91mno training images found. Please upload images to training_images.\\033[0m\")\n",
    "else:\n",
    " print(\"✅ \\033[92m\" + str(len(training_images_file_paths)) + \" training images found.\\033[0m\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "m2o_fFFvfxHi"
   },
   "outputs": [],
   "source": [
    "#@title 7. Final Setup & Training\n",
    "\n",
    "#@markdown This isn't used for training, just used to name the folders etc\n",
    "project_name = \"project_name\" #@param {type:\"string\"}\n",
    "\n",
    "#@markdown This is the unique token i.e. you can use a nonsensical word like zwx or your name.\n",
    "token = \"firstNamelastName\" #@param {type:\"string\"}\n",
    "\n",
    "#@markdown Match class_word to the category of the regularization images you chose above.\n",
    "class_word = \"person\" #@param [\"man\", \"person\", \"woman\"] {allow-input: true}\n",
    "\n",
    "# MAX STEPS\n",
    "#@markdown How many steps do you want to train for?\n",
    "max_training_steps = 2000 #@param {type:\"integer\"}\n",
    "\n",
    "#@markdown If you are training a person's face, set this to True\n",
    "i_am_training_a_persons_face = True #@param {type:\"boolean\"}\n",
    "flip_p_arg = 0.0 if i_am_training_a_persons_face else 0.5\n",
    "\n",
    "#@markdown Would you like to save a model every X steps? (Example: 250 would output a trained model at 250, 500, 750 steps, etc)\n",
    "save_every_x_steps = 0 #@param {type:\"integer\"}\n",
    "\n",
    "\n",
    "reg_data_root = \"/content/Dreambooth-Stable-Diffusion/regularization_images/\" + dataset\n",
    "\n",
    "!rm -rf training_images/.ipynb_checkpoints\n",
    "!python \"main.py\" \\\n",
    " --project_name \"{project_name}\" \\\n",
    " --debug False \\\n",
    " --max_training_steps {max_training_steps} \\\n",
    " --token \"{token}\" \\\n",
    " --training_model \"model.ckpt\" \\\n",
    " --training_images \"/content/Dreambooth-Stable-Diffusion/training_images\" \\\n",
    " --regularization_images \"{reg_data_root}\" \\\n",
    " --class_word \"{class_word}\" \\\n",
    " --flip_p {flip_p_arg} \\\n",
    " --save_every_x_steps {save_every_x_steps}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mkidEm4evn1J"
   },
   "outputs": [],
   "source": [
    "#@title 8. Save model into google drive\n",
    "#@markdown This is often much faster than a manual download.  it will also save you compute units. <br />\n",
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "\n",
    "# copy all ckpt files to google drive root dir\n",
    "!cp trained_models/*.ckpt /content/drive/MyDrive"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "provenance": []
  },
  "gpuClass": "premium",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11.3"
  },
  "vscode": {
   "interpreter": {
    "hash": "5c7b89af1651d0b8571dde13640ecdccf7d5a6204171d6ab33e7c296e100e08a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}