EveryDream-trainer/dreambooth_runpod_joepenna.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Dreambooth\n",
    "### Notebook implementation by Joe Penna (@MysteryGuitarM on Twitter) - Improvements by David Bielejeski and Kane Wallmann\n",
    "\n",
    "### Instructions\n",
    "- Sign up for RunPod here: https://runpod.io/?ref=n8yfwyum\n",
    "    - Note: That's my personal referral link. Please don't use it if we are mortal enemies.\n",
    "\n",
    "- Click *Deploy* on either `SECURE CLOUD` or `COMMUNITY CLOUD`\n",
    "\n",
    "- Follow the rest of the instructions in this video: https://www.youtube.com/watch?v=7m__xadX0z0#t=5m33.1s\n",
    "\n",
    "Latest information on:\n",
    "https://github.com/JoePenna/Dreambooth-Stable-Diffusion"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Build Environment"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2AsGA1xpNQnb",
   "metadata": {
    "id": "2AsGA1xpNQnb"
   },
   "outputs": [],
   "source": [
    "# If running on Vast.AI, copy the code in this cell into a new notebook. Run it, then launch the `dreambooth_runpod_joepenna.ipynb` notebook from the jupyter interface.\n",
    "!git clone https://github.com/JoePenna/Dreambooth-Stable-Diffusion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e1bc458-091b-42f4-a125-c3f0df20f29d",
   "metadata": {
    "id": "9e1bc458-091b-42f4-a125-c3f0df20f29d",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dae11c10",
   "metadata": {
    "id": "dae11c10"
   },
   "outputs": [],
   "source": [
    "## Login to stable diffusion\n",
    "from huggingface_hub import notebook_login\n",
    "\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "## Download the 1.4 sd model\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",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "## 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"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17d1d11a",
   "metadata": {
    "id": "17d1d11a"
   },
   "source": [
    "# Regularization Images"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed07a5df",
   "metadata": {
    "id": "ed07a5df"
   },
   "source": [
    "Training teaches your new model both your token **but** re-trains your class simultaneously.\n",
    "\n",
    "From cursory testing, it does not seem like reg images affect the model too much. However, they do affect your class greatly, which will in turn affect your generations.\n",
    "\n",
    "The regularization images should be of the format `class_xyz.jpg` (e.g. man_001.jpg, woman_001.jpg, etc.). You can create your own or use the pre-generated ones below.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "mxPL2O0OLvBW",
   "metadata": {
    "id": "mxPL2O0OLvBW"
   },
   "source": [
    "## Download pre-generated regularization images\n",
    "\n",
    "I have prepared a repository which contains 1000 men, woman and person regularization images named correctly. This appears to work quite well at avoiding distortions if you are training photos of people. These were taken from djbielejeski/Stable-Diffusion-Regularization-Images-{dataset}.git repositories:\n",
    "\n",
    "* man_euler - provided by Niko Pueringer (Corridor Digital) - euler @ 40 steps, CFG 7.5\n",
    "* person_ddim\n",
    "* woman_ddim - provided by David Bielejeski - ddim @ 50 steps, CFG 10.0\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e7EydXCjOV1v",
   "metadata": {
    "id": "e7EydXCjOV1v"
   },
   "outputs": [],
   "source": [
    "# Grab the existing regularization images\n",
    "!git clone https://github.com/kanewallmann/Stable-Diffusion-Regularization-Images.git\n",
    "\n",
    "!mkdir -p outputs/txt2img-samples/regularization\n",
    "!mv -v Stable-Diffusion-Regularization-Images/*.* outputs/txt2img-samples/regularization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "zshrC_JuMXmM",
   "metadata": {
    "id": "zshrC_JuMXmM"
   },
   "source": [
    "# Upload your training images\n",
    "Upload 10-20 images of someone to\n",
    "\n",
    "```\n",
    "/workspace/Dreambooth-Stable-Diffusion/training_samples\n",
    "```\n",
    "\n",
    "WARNING: Be sure to upload an *even* amount of images, otherwise the training inexplicably stops at 1500 steps.\n",
    "\n",
    "*   2-3 full body\n",
    "*   3-5 upper body\n",
    "*   5-12 close-up on face\n",
    "\n",
    "The images should be:\n",
    "\n",
    "- as close as possible to the kind of images you're trying to make (most of the time, that means no selfies).\n",
    "- named in the following format `identifier class_xyz.jpg` (example: \"kwallmann man_001.jpg\", \"bholly dog_001.jpg\", etc.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#@markdown Add here the URLs to the images of the subject you are adding\n",
    "urls = [\n",
    " \"https://i.imgur.com/test1.png\",\n",
    " \"https://i.imgur.com/test2.png\",\n",
    " \"https://i.imgur.com/test3.png\",\n",
    " \"https://i.imgur.com/test4.png\",\n",
    " \"https://i.imgur.com/test5.png\",\n",
    " # You can add additional images here -- about 20-30 images in different \n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "#@title Download and check the images you have just added\n",
    "import os\n",
    "import requests\n",
    "from io import BytesIO\n",
    "from PIL import Image\n",
    "\n",
    "\n",
    "def image_grid(imgs, rows, cols):\n",
    " assert len(imgs) == rows*cols\n",
    "\n",
    " w, h = imgs[0].size\n",
    " grid = Image.new('RGB', size=(cols*w, rows*h))\n",
    " grid_w, grid_h = grid.size\n",
    "\n",
    " for i, img in enumerate(imgs):\n",
    "  grid.paste(img, box=(i%cols*w, i//cols*h))\n",
    " return grid\n",
    "\n",
    "def download_image(url):\n",
    " try:\n",
    "  response = requests.get(url)\n",
    " except:\n",
    "  return None\n",
    " return Image.open(BytesIO(response.content)).convert(\"RGB\")\n",
    "\n",
    "images = list(filter(None,[download_image(url) for url in urls]))\n",
    "save_path = \"./training_samples\"\n",
    "if not os.path.exists(save_path):\n",
    " os.mkdir(save_path)\n",
    "[image.save(f\"{save_path}/{i}.png\", format=\"png\") for i, image in enumerate(images)]\n",
    "image_grid(images, 1, len(images))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad4e50df",
   "metadata": {
    "id": "ad4e50df"
   },
   "source": [
    "## Training\n",
    "\n",
    "If training a person or subject, keep an eye on your project's `logs/{folder}/images/train/samples_scaled_gs-00xxxx` generations.\n",
    "\n",
    "If training a style, keep an eye on your project's `logs/{folder}/images/train/samples_gs-00xxxx` generations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fa5dd66-2ca0-4819-907e-802e25583ae6",
   "metadata": {
    "id": "6fa5dd66-2ca0-4819-907e-802e25583ae6",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# START THE TRAINING\n",
    "project_name = \"project_name\"\n",
    "\n",
    "# MAX STEPS\n",
    "max_training_steps = 1000\n",
    "\n",
    "reg_data_root = \"/workspace/Dreambooth-Stable-Diffusion/outputs/txt2img-samples/regularization\"\n",
    "\n",
    "!rm -rf training_samples/.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 \"/workspace/Dreambooth-Stable-Diffusion/training_samples\" \\\n",
    " --max_training_steps {max_training_steps} \\\n",
    " --no-test"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc49d0bd",
   "metadata": {},
   "source": [
    "## Pruning (12GB to 2GB)\n",
    "We are working on having this happen automatically (TODO: PR's welcome)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "directory_paths = !ls -d logs/*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# This version should automatically prune around 10GB from the ckpt file\n",
    "last_checkpoint_file = directory_paths[-1] + \"/checkpoints/last.ckpt\"\n",
    "!python \"prune_ckpt.py\" --ckpt {last_checkpoint_file}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "last_checkpoint_file_pruned = directory_paths[-1] + \"/checkpoints/last-pruned.ckpt\"\n",
    "training_samples = !ls training_samples\n",
    "date_string = !date +\"%Y-%m-%dT%H-%M-%S\"\n",
    "file_name = date_string[-1] + \"_\" + project_name + \"_\" + str(len(training_samples)) + \"_training_images_\" +  str(max_training_steps) + \"_max_training_steps_\" + class_word + \"_class_word.ckpt\"\n",
    "!mkdir -p trained_models\n",
    "!mv {last_checkpoint_file_pruned} trained_models/{file_name}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download your trained model file from `trained_models` and use in your favorite Stable Diffusion repo!"
   ]
  },
  {
   "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 (ipykernel)",
   "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.10.6"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}