dreambooth_colab_joepenna.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#@title Load repo (if needed)\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 BUILD ENV\n",
    "!pip install omegaconf\n",
    "!pip install einops\n",
    "!pip install pytorch-lightning==1.6.5\n",
    "!pip install test-tube\n",
    "!pip install transformers\n",
    "!pip install kornia\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 setuptools==59.5.0\n",
    "!pip install pillow==9.0.1\n",
    "!pip install torchmetrics==0.6.0\n",
    "!pip install -e .\n",
    "!pip install protobuf==3.20.1\n",
    "!pip install gdown\n",
    "!pip install pydrive\n",
    "!pip install -qq diffusers[\"training\"]==0.3.0 transformers ftfy\n",
    "!pip install -qq \"ipywidgets>=7,<8\"\n",
    "!pip install huggingface_hub\n",
    "!pip install ipywidgets==7.7.1\n",
    "!pip install captionizer==1.0.1\n",
    "\n",
    "import os\n",
    "os._exit(00)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#@title # Required - Navigate back to the directory.\n",
    "%cd Dreambooth-Stable-Diffusion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "6tjx0HcjesFo"
   },
   "outputs": [],
   "source": [
    "#@markdown Hugging Face Login\n",
    "from huggingface_hub import notebook_login\n",
    "\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "O15vMMhCevib"
   },
   "outputs": [],
   "source": [
    "#@markdown Download the 1.4 sd model\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=\"CompVis/stable-diffusion-v-1-4-original\",\n",
    " filename=\"sd-v1-4.ckpt\",\n",
    " use_auth_token=True\n",
    ")\n",
    "\n",
    "# Move the sd-v1-4.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 # 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 - `blonde_woman` - provided by David Bielejeski - ddim @ 50 steps, CFG 10.0 <br />\n",
    "\n",
    "dataset=\"person_ddim\" #@param [\"man_euler\", \"man_unsplash\", \"person_ddim\", \"woman_ddim\", \"blonde_woman\"]\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}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "A7hmdOdOfGzs"
   },
   "outputs": [],
   "source": [
    "#@title # 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).\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(\"❌ no training images found. Please upload images to training_images\")\n",
    "else:\n",
    " print(\"✅ \" + str(len(training_images_file_paths)) + \" training images found\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "m2o_fFFvfxHi"
   },
   "outputs": [],
   "source": [
    "#@title # Training\n",
    "\n",
    "#@markdown This isn't used for training, just to help you remember what your trained into the model.\n",
    "project_name = \"project_name\" #@param {type:\"string\"}\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 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",
    "#@markdown This is the unique token you are incorporating into the stable diffusion model.\n",
    "token = \"firstNameLastName\" #@param {type:\"string\"}\n",
    "reg_data_root = \"/content/Dreambooth-Stable-Diffusion/regularization_images/\" + dataset\n",
    "\n",
    "!rm -rf training_images/.ipynb_checkpoints\n",
    "!python \"main.py\" \\\n",
    " --base configs/stable-diffusion/v1-finetune_unfrozen.yaml \\\n",
    " -t \\\n",
    " --actual_resume \"model.ckpt\" \\\n",
    " --reg_data_root \"{reg_data_root}\" \\\n",
    " -n \"{project_name}\" \\\n",
    " --gpus 0, \\\n",
    " --data_root \"/content/Dreambooth-Stable-Diffusion/training_images\" \\\n",
    " --max_training_steps {max_training_steps} \\\n",
    " --class_word \"{class_word}\" \\\n",
    " --token \"{token}\" \\\n",
    " --no-test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ll_ZIFNUulKJ"
   },
   "outputs": [],
   "source": [
    "#@title # Copy and name the checkpoint file\n",
    "\n",
    "directory_paths = !ls -d logs/*\n",
    "last_checkpoint_file = directory_paths[-1] + \"/checkpoints/last.ckpt\"\n",
    "training_images = !find training_images/*\n",
    "date_string = !date +\"%Y-%m-%dT%H-%M-%S\"\n",
    "file_name = date_string[-1] + \"_\" + project_name + \"_\" + str(len(training_images)) + \"_training_images_\" +  str(max_training_steps) + \"_max_training_steps_\" + token + \"_token_\" + class_word + \"_class_word.ckpt\"\n",
    "\n",
    "file_name = file_name.replace(\" \", \"_\")\n",
    "\n",
    "!mkdir -p trained_models\n",
    "!mv \"{last_checkpoint_file}\" \"trained_models/{file_name}\"\n",
    "\n",
    "print(\"Download your trained model file from trained_models/\" + file_name + \" and use in your favorite Stable Diffusion repo!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "mkidEm4evn1J"
   },
   "outputs": [],
   "source": [
    "#@title Save model in google drive\n",
    "from google.colab import drive\n",
    "drive.mount('/content/drive')\n",
    "\n",
    "!cp trained_models/{file_name} /content/drive/MyDrive/{file_name}"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3.9.0 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.9.0"
  },
  "vscode": {
   "interpreter": {
    "hash": "0f8486a44be3db540a44064ef4647f290d210aa4269e54f181de788486e3e134"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}