dreambooth_runpod_joepenna_dogarrow_6.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "aa2c1ada",
   "metadata": {
    "id": "aa2c1ada"
   },
   "source": [
    "# *Furry* Dreambooth\n",
    "### Notebook implementation by Joe Penna (@MysteryGuitarM on Twitter) - Improvements by David Bielejeski\n",
    "### Furry specific modifications by dogarrowtype#7324\n",
    "### Join us at: https://discord.gg/furrydiffusion\n",
    "\n",
    "Dreambooth is a technique that runs on top of StableDiffusion. It teaches the model the details of your specific character, and if done right, can reproduce them with high accuracy.\n",
    "\n",
    "\n",
    "## This notebook requires a GPU with 24GB of VRAM. A 3090 or an A5000 on runpod will do the job.\n",
    "Pick a runpod with a fast upload speed, to download the model at the end faster. Choose the standard pytorch template.\n",
    "\n",
    "### Instructions\n",
    "- Run each code cell, follow the instructions, but choose whether you want to use regularization images.\n",
    "\n",
    "Modified from:\n",
    "https://github.com/JoePenna/Dreambooth-Stable-Diffusion\n",
    "\n",
    "This notebook has been designed and tested to work on https://runpod.io\n",
    "\n",
    "Be careful, the predefined model is known to produce nsfw content. DO NOT use if you are below the legal age of majority in your country."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b971cc0",
   "metadata": {
    "id": "7b971cc0"
   },
   "source": [
    "## Build Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2AsGA1xpNQnb",
   "metadata": {
    "id": "2AsGA1xpNQnb"
   },
   "outputs": [],
   "source": [
    "!git clone https://github.com/JoePenna/Dreambooth-Stable-Diffusion\n",
    "%cd /workspace/Dreambooth-Stable-Diffusion\n",
    "uses_reg = False"
   ]
  },
  {
   "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 -qq diffusers[\"training\"]==0.3.0 transformers ftfy\n",
    "!pip install -qq \"ipywidgets>=7,<8\"\n",
    "!pip install huggingface_hub==0.11.0\n",
    "!pip install ipywidgets==7.7.1\n",
    "!pip install captionizer==1.0.1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3455457b-3dbc-4727-8667-6f0fa91e3543",
   "metadata": {},
   "source": [
    "## Download the model to train on top of\n",
    "Swap this link out here, if you want to use something other than ye18."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aca852ef-c16c-47cf-9d37-c212dd68a5f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://sexy.canine.wf/file/yiffy-ckpt/yiffy-e18.ckpt -O /workspace/Dreambooth-Stable-Diffusion/model.ckpt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17d1d11a",
   "metadata": {
    "id": "17d1d11a",
    "tags": []
   },
   "source": [
    "# Regularization Images (No longer recommended)\n",
    "## Skip this part unless you know what it does.\n",
    "### Expand if you really want to."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed07a5df",
   "metadata": {
    "id": "ed07a5df",
    "tags": []
   },
   "source": [
    "This is NOT the part that contains the character you want to train on. It is a separate step that helps guide the model toward what you want, but has little to do with your desired character.\n",
    "\n",
    "For furry content, ~~regularization images help. They (seem) to prime the model to focus on the area you'd like to change.~~\n",
    "\n",
    "## Regularization images seem to cause more harm than help for furry content. I no longer advise using them.\n",
    "\n",
    "### But you can still use them if you want. Just expand this block.\n",
    "\n",
    "Change \"self_generated_files_prompt\" to be closest to what you want. \n",
    "\n",
    "**For instance: change species, gender, color.**\n",
    "\n",
    "Try to keep it close to the example and simple. Don't go heavy on the detail.\n",
    "\n",
    "The prompt system here is not as advanced as Automatic1111, so don't do anything fancy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67f9ff0c-b529-4c7c-8e26-8388d70a5d91",
   "metadata": {
    "id": "67f9ff0c-b529-4c7c-8e26-8388d70a5d91",
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# GENERATE 200 images - Optional\n",
    "uses_reg = True\n",
    "self_generated_files_prompt = \"e621 safe solo anthro male canine grey fur\"\n",
    "self_generated_files_count = 200\n",
    "\n",
    "!python scripts/stable_txt2img.py \\\n",
    " --seed 10 \\\n",
    " --ddim_eta 0.0 \\\n",
    " --n_samples 1 \\\n",
    " --n_iter {self_generated_files_count} \\\n",
    " --scale 10.0 \\\n",
    " --ddim_steps 50 \\\n",
    " --ckpt model.ckpt \\\n",
    " --prompt \"{self_generated_files_prompt}\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5c235da-e7fb-482b-be43-b99bbd2a3cd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset=\"dataset1\"\n",
    "!mkdir -p regularization_images/{dataset}\n",
    "!mv /workspace/Dreambooth-Stable-Diffusion/outputs/txt2img-samples/samples/*.* regularization_images/{dataset}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d1c7e1c",
   "metadata": {
    "id": "3d1c7e1c"
   },
   "outputs": [],
   "source": [
    "# Zip up the files for downloading and reuse.\n",
    "# Download this file locally so you can reuse during another training on this dataset\n",
    "!apt update\n",
    "!apt install -y zip\n",
    "!zip -r regularization_images.zip regularization_images/{dataset}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "022a7a1d-f14b-41d1-a816-9f0730fd2d09",
   "metadata": {
    "id": "ed07a5df",
    "tags": []
   },
   "source": [
    "## OR Upload regularization images generated separately to\n",
    "`/workspace/Dreambooth-Stable-Diffusion/regularization_images/dataset1`\n",
    "\n",
    "Don't do both."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a3a8491-06e7-446b-afa7-5ef701c68c37",
   "metadata": {},
   "outputs": [],
   "source": [
    "uses_reg = True\n",
    "dataset=\"dataset1\"\n",
    "!mkdir -p regularization_images/{dataset}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26fd5a77-5cbf-4981-bd64-45f94ec5820d",
   "metadata": {},
   "source": [
    "Generate them on your own pc, gather them, and upload to the folder."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2be83cbd",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Training data\n",
    "# Upload your training images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "311f339e-8d9d-4fc1-b720-c69ada1bb08e",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p /workspace/Dreambooth-Stable-Diffusion/training_images"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69de2f51-22d1-4be4-8eef-1f63e10dacc6",
   "metadata": {
    "tags": []
   },
   "source": [
    "Upload 100-200 images of your character to\n",
    "\n",
    "```\n",
    "/workspace/Dreambooth-Stable-Diffusion/training_images\n",
    "```\n",
    "\n",
    "Be sure to upload an *even* amount of images, otherwise the training inexplicably stops at 1500 steps.\n",
    "\n",
    "## The images should be:\n",
    "\n",
    "- Solos only!\n",
    "- Square.\n",
    "- 512x512.\n",
    "- Cropped to contain the most recognizable part of your character.\n",
    "- Don't worry about legs/arms being cropped off, the most important part is the head. The model will make up the rest.\n",
    "- A variety of angles and shots. The model can become overfitted if you only do frontal shots, for instance.\n",
    "- If your character has very distinctive features or markings you want to preserve, make sure they show up in the data often.\n",
    "- Variations of the same image (alt versions/edits) are fine, but don't do too many of the same.\n",
    "- DON'T do black and white or uncolored sketches.\n",
    "- Get as close as possible to the output you'd like to see.\n",
    "\n",
    "A good idea is to use imagemagick for automated cropping. No, it's not the fancy kind that uses algorithms to focus on faces, but it totally works.\n",
    "\n",
    "You're mostly on your own for cropping, but here's some...\n",
    "## Hints for cropping commands:\n",
    "Windows:\n",
    "\n",
    "`magick *.* -gravity North -crop 1:1 -resize 512x512 -format png ..\\raw1\\image-%d.png`\n",
    "\n",
    "Linux:\n",
    "\n",
    "`find -maxdepth 1 -type f -execdir magick '{}' -gravity North -crop 1:1 -resize 512x512 ../ready1/'{}.png' \\;`\n",
    "\n",
    "https://imagemagick.org/script/download.php\n",
    "\n",
    "\n",
    "## Use BIRME, it's by far the best option for bulk crop + resize:\n",
    "https://www.birme.net/?target_width=512&target_height=512&rename=ORIGINAL-NAME_&border_width=-1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad4e50df",
   "metadata": {
    "id": "ad4e50df",
    "tags": []
   },
   "source": [
    "# Training\n",
    "\n",
    "## This is where the magic happens.\n",
    "\n",
    "Some key things to change:\n",
    "- `\"project_name\"`: Change project name to whatever you want. Avoid spaces for everything.\n",
    "- `\"max_training_steps\"`: 2000-4000 is a good range. Choose around 2000 if you have 100 images, and around 4000 if you have 200 images.\n",
    "- `\"class_word\"`: Keep as anthro. Or feral. Or... whatever kind of thing you're doing. But it should be extremely simple, and something the model already recognizes easily. Only matters if using regularization images.\n",
    "- `\"token\"`: Pick a longer, made-up word the model has never seen before. Can really be anything, but make it something like \"blaiddwolfyfromeldenring\". \"Charname+universe\" is easier to remember, imo. This is what you will type in to activate it, so... don't make it insane.\n",
    "\n",
    "Once you've got those set, you're ready to train. It should take about an hour on a 3090 for 2500 steps.\n",
    "\n",
    "Do not worry. It will stop at exactly the number of steps you set, even if the progress bar there looks otherwise. The estimate is usually wildly wrong. There is a separate internal counter that is not shown."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fa5dd66-2ca0-4819-907e-802e25583ae6",
   "metadata": {
    "id": "6fa5dd66-2ca0-4819-907e-802e25583ae6",
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Training\n",
    "\n",
    "# This is just the name for the training run, not the token.\n",
    "project_name = \"blaidd_1\"\n",
    "\n",
    "# MAX STEPS\n",
    "# How many steps do you want to train for?\n",
    "max_training_steps = 2500\n",
    "\n",
    "# Match class_word to the category of the regularization images you chose above.\n",
    "class_word = \"anthro\" # Anthro or feral are the best choices.\n",
    "\n",
    "# This is the unique token you are teaching the stable diffusion model.\n",
    "token = \"blaiddwolfyfromeldenring\" #Can be anything, but pick something rare and easy to remember/type.\n",
    "\n",
    "\n",
    "dataset=\"dataset1\"\n",
    "img_data_root = \"/workspace/Dreambooth-Stable-Diffusion/training_images/\"\n",
    "reg_data_root = \"/workspace/Dreambooth-Stable-Diffusion/regularization_images/\" + dataset\n",
    "\n",
    "!rm -rf training_images/.ipynb_checkpoints\n",
    "\n",
    "if uses_reg:\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 \"{img_data_root}\" \\\n",
    "    --max_training_steps {max_training_steps} \\\n",
    "    --class_word \"{class_word}\" \\\n",
    "    --token \"{token}\" \\\n",
    "    --no-test\n",
    "else:\n",
    "    !python \"main.py\" \\\n",
    "    --base configs/stable-diffusion/v1-finetune_unfrozen.yaml \\\n",
    "    -t \\\n",
    "    --actual_resume \"model.ckpt\" \\\n",
    "    -n \"{project_name}\" \\\n",
    "    --gpus 0, \\\n",
    "    --data_root \"{img_data_root}\" \\\n",
    "    --max_training_steps {max_training_steps} \\\n",
    "    --token \"{token}\" \\\n",
    "    --no-test"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc49d0bd",
   "metadata": {},
   "source": [
    "## Copy and name the checkpoint file\n",
    "\n",
    "Also prints the sha256sum, very helpful for verification/sharing your dreambooth model publicly. Take note of it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d65e779f",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Copy the checkpoint into the `trained_models` folder\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",
    "!sha256sum \"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": "markdown",
   "id": "a99dfbb1-5a71-4369-8ded-cb6580b222aa",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Upload checkpoint directly to Pixeldrain\n",
    "Saves you time and money if you have a slow download speed. Slightly risky, but for some, very worth it. If your rented pod has a fast internet connection, it will upload to pixeldrain very quickly.\n",
    "\n",
    "### It won't show any output while it is uploading. But if it worked, it will print out a link at at the bottom that you can use to download/share your model.\n",
    "\n",
    "Take extra care, this method can go wrong in any number of ways."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f20f277c-fecf-4abe-baa7-20010d5d9746",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://github.com/ManuelReschke/go-pd/releases/download/v1.2.1/go-pd_1.2.1_linux_amd64.tar.gz -O ./go-pd.tar.gz\n",
    "!tar -xf ./go-pd.tar.gz go-pd\n",
    "!./go-pd upload \"/workspace/Dreambooth-Stable-Diffusion/trained_models/{file_name}\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04384133-1b69-46dc-bbab-6bf4c3c96d52",
   "metadata": {},
   "source": [
    "# Make sure to delete your runpod instance when you are done! Once you have the model downloaded, you should stop the pod and delete it so you don't get charged any more credit."
   ]
  }
 ],
 "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.8"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}