{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "aa2c1ada",
   "metadata": {},
   "source": [
    "# Dreambooth\n",
    "### Notebook implementation by Joe Penna (@MysteryGuitarM on Twitter)\n",
    "https://github.com/JoePenna/Dreambooth-Stable-Diffusion\n",
    "\n",
    "### If on runpod / vast.ai / etc, spin up an A6000 or A100 pod using a Stable Diffusion template with Jupyter pre-installed."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4d54876",
   "metadata": {},
   "source": [
    "## Check GPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3d20275-7c9b-447f-b160-dc556fcaee9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# mostly for Google Colab -- you should know what you're running on Runpod / Vast.ai / etc\n",
    "from IPython.display import HTML\n",
    "from subprocess import getoutput\n",
    "s = getoutput('nvidia-smi')\n",
    "print(s)\n",
    "# or simply\n",
    "#!nvidia-smi -L"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8347baa",
   "metadata": {},
   "source": [
    "## Installation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed803d51-d42f-4b07-9582-5cf1c9b52274",
   "metadata": {},
   "outputs": [],
   "source": [
    "# GIT THE REPO\n",
    "!git clone https://github.com/JoePenna/Dreambooth-Stable-Diffusion\n",
    "\n",
    "# OR THE ORIGINAL\n",
    "# !git clone https://github.com/XavierXiao/Dreambooth-Stable-Diffusion"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b971cc0",
   "metadata": {},
   "source": [
    "## Build Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e1bc458-091b-42f4-a125-c3f0df20f29d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#BUILD ENV\n",
    "%cd /workspace/Dreambooth-Stable-Diffusion\n",
    "!pip install omegaconf einops pytorch-lightning==1.6.5 test-tube transformers kornia -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers -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 ."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17d1d11a",
   "metadata": {},
   "source": [
    "## Generate regularization images (you should change the code below)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed07a5df",
   "metadata": {},
   "source": [
    "This repo simultaneously trains both your token **and** retrains your class.\n",
    "\n",
    "From cursory testing, it does not seem like the number of reg images affects the model too much.\n",
    "\n",
    "However, it does affect your class greatly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67f9ff0c-b529-4c7c-8e26-8388d70a5d91",
   "metadata": {},
   "outputs": [],
   "source": [
    "#GENERATE 50 images\n",
    "!python scripts/stable_txt2img.py --seed 10 --ddim_eta 0.0 --n_samples 1 --n_iter 50 --scale 10.0 --ddim_steps 50  --ckpt /workspace/stable-diffusion-webui/model.ckpt \\\n",
    "--prompt \"person\"\n",
    "\n",
    "# IMPORTANT!! Change \"person\" above, and on the cell below:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00717137",
   "metadata": {},
   "source": [
    "## Upload your samples"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c90ed3ab",
   "metadata": {},
   "source": [
    "Drag your training images into /workspace/Dreambooth-Stable-Diffusion/training_samples\n",
    "\n",
    "At the moment, it seems like 10-20 images of someone's face is enough."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad4e50df",
   "metadata": {},
   "source": [
    "## Prep Training Variables"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a26964af",
   "metadata": {},
   "source": [
    "Navigate to:\n",
    "/workspace/Dreambooth-Stable-Diffusion/ldm/data/personalized.py\n",
    "\n",
    "Change \"sks\" in line 11 to whatever you want your token to be.\n",
    "\n",
    "e.g. I changed mine to\n",
    "```python\n",
    "training_templates_smallest = [\n",
    "    'joepenna {}',\n",
    "]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37612c32",
   "metadata": {},
   "source": [
    "The last line of this file trains for `3000` steps.\n",
    "/workspace/Dreambooth-Stable-Diffusion/configs/stable-diffusion/v1-finetune_unfrozen.yaml\n",
    "\n",
    "If training a person or subject, keep an eye on your project's `/images/train/samples_scaled_gs-00xxxx` generations.\n",
    "\n",
    "If training a style, keep an eye on your project's `/images/train/samples_gs-00xxxx` generations."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12f19de7",
   "metadata": {},
   "source": [
    "## Start Training (you should also change the code below)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fa5dd66-2ca0-4819-907e-802e25583ae6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# START THE TRAINING\n",
    "!python \"main.py\" \\\n",
    " --base configs/stable-diffusion/v1-finetune_unfrozen.yaml \\\n",
    " -t \\\n",
    " --actual_resume \"/workspace/stable-diffusion-webui/model.ckpt\" \\\n",
    " --reg_data_root \"/workspace/Dreambooth-Stable-Diffusion/outputs/txt2img-samples/samples\" \\\n",
    " -n \"project_name\" \\\n",
    " --gpus 0, \\\n",
    " --data_root \"/workspace/Dreambooth-Stable-Diffusion/training_imgs\" \\\n",
    " --class_word \"person\" # << match this word to the class word from regularization images above"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "982ad0c8",
   "metadata": {},
   "source": [
    "## Experimental Pruning (12GB to 2GB)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24c7167e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# EXPERIMENTAL - this version should automatically prune around 10GB from the ckpt file\n",
    "!python \"main_prune.py\" \\\n",
    " --base configs/stable-diffusion/v1-finetune_unfrozen.yaml \\\n",
    " -t \\\n",
    " --actual_resume \"/workspace/stable-diffusion-webui/model.ckpt\" \\\n",
    " --reg_data_root \"/workspace/Dreambooth-Stable-Diffusion/outputs/txt2img-samples/samples\" \\\n",
    " -n \"project_name\" \\\n",
    " --gpus 0, \\\n",
    " --data_root \"/workspace/Dreambooth-Stable-Diffusion/training_imgs\" \\\n",
    " --class_word \"person\" # << match this word to the class word from regularization images above"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ab885a3",
   "metadata": {},
   "source": [
    "## Generate Images With Trained Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b955a68e",
   "metadata": {},
   "source": [
    "Be sure to change `PROJECT_PATH` below to match the folder in:\n",
    "/workspace/Dreambooth-Stable-Diffusion/logs\n",
    "\n",
    "Also change `sks person` in the prompt to the token *and* class_word."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6e40a7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python scripts/stable_txt2img.py \\\n",
    "--ddim_eta 0.0 --n_samples 1 --n_iter 4 \\\n",
    "--scale 7.0 --ddim_steps 60 --ckpt \"/workspace/Dreambooth-Stable-Diffusion/logs/PROJECT_PATH_GOES_HERE/last.ckpt\" \\\n",
    "--prompt \"sks person as a masterpiece portrait painting by John Singer Sargent in the style of Rembrandt\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e77031d7",
   "metadata": {},
   "source": [
    "Trained generated images at /workspace/Dreambooth-Stable-Diffusion/outputs/"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.6 64-bit",
   "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
}