Merge pull request #41 from JoePenna/feature/google-colab-notebook

Feature/google colab notebook

Author: JoePenna

dreambooth_colab_joepenna.ipynb

@@ -0,0 +1,245 @@
+{
+ "nbformat": 4,
+ "nbformat_minor": 0,
+ "metadata": {
+  "colab": {
+   "provenance": [],
+   "collapsed_sections": []
+  },
+  "kernelspec": {
+   "name": "python3",
+   "display_name": "Python 3"
+  },
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "qeTrc2vOeiNh"
+   },
+   "outputs": [],
+   "source": [
+    "#@title BUILD ENV\n",
+    "!git clone https://github.com/JoePenna/Dreambooth-Stable-Diffusion\n",
+    "%cd Dreambooth-Stable-Diffusion\n",
+    "\n",
+    "!git checkout feature/google-colab-notebook\n",
+    "\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",
+    "\n",
+    "import os\n",
+    "os._exit(00)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "#@title # Required - Navigate back to the directory.\n",
+    "%cd Dreambooth-Stable-Diffusion"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "#@markdown Hugging Face Login\n",
+    "from huggingface_hub import notebook_login\n",
+    "\n",
+    "notebook_login()"
+   ],
+   "metadata": {
+    "id": "6tjx0HcjesFo"
+   },
+   "execution_count": 1,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "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"
+   ],
+   "metadata": {
+    "id": "O15vMMhCevib"
+   },
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "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}"
+   ],
+   "metadata": {
+    "id": "N96aedTtfBjO"
+   },
+   "execution_count": 2,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "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"
+   ],
+   "metadata": {
+    "id": "A7hmdOdOfGzs"
+   },
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "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"
+   ],
+   "metadata": {
+    "id": "m2o_fFFvfxHi"
+   },
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "source": [
+    "#@title # Prune the Model\n",
+    "\n",
+    "directory_paths = !ls -d logs/*\n",
+    "last_checkpoint_file = directory_paths[-1] + \"/checkpoints/last.ckpt\"\n",
+    "!python \"prune_ckpt.py\" --ckpt {last_checkpoint_file}\n",
+    "last_checkpoint_file_pruned = directory_paths[-1] + \"/checkpoints/last-pruned.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",
+    "!mkdir -p trained_models\n",
+    "!mv {last_checkpoint_file_pruned} 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": {
+    "id": "Ll_ZIFNUulKJ"
+   },
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "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": {
+    "id": "mkidEm4evn1J"
+   },
+   "execution_count": null,
+   "outputs": []
+  }
+ ]
+}