Furry Dreambooth Runpod

dreambooth_runpod_joepenna_dogarrow_4.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 upload the regularization images, or generate them now.\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",
    "!git checkout 605faeaa2f0656bf31ab79749fb5d96b3ca2e69f"
   ]
  },
  {
   "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\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"
   ]
  },
  {
   "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",
    "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",
    "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": [
    "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. If you *really* want to skip regularization, delete `--reg_data_root \"{reg_data_root}\"` from the training block."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2be83cbd",
   "metadata": {},
   "source": [
    "# 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": {},
   "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 distintive 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"
   ]
  },
  {
   "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.\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."
   ]
  },
  {
   "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",
    "reg_data_root = \"/workspace/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 \"/workspace/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": "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!\")"
   ]
  }
 ],
 "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.7"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}