{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gfKvWAVnz8OB"
   },
   "source": [
    "# Voldemort's Stable Diffusion WebUI\n",
    "\n",
    "Adapted from: https://colab.research.google.com/drive/1Iy-xW9t1-OQWhb0hNxueGij8phCyluOh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Installation and Setup\n",
    "\n",
    "You must reinstall everything each time you restart the machine. If already downloaded dependencies will be auto-updated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OL82Y4rBjZIV"
   },
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Clone the central repository"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sBbcB4vwj_jm",
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os.path\n",
    "\n",
    "if not os.path.exists('/notebooks/stable-diffusion-webui'):\n",
    "    %cd /notebooks/\n",
    "    !git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui\n",
    "else: # update repo if already exists\n",
    "    print('stable-diffusion-webui already downloaded, updating')\n",
    "    %cd /notebooks/stable-diffusion-webui\n",
    "    !git pull"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "C68TUpkq0nj_",
    "tags": []
   },
   "source": [
    "### Install requirements and download repositories."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SaAJk33ppFw1",
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "%cd /notebooks/stable-diffusion-webui\n",
    "\n",
    "!pip install -r requirements.txt\n",
    "\n",
    "if not os.path.exists('/notebooks/stable-diffusion-webui/repositories'): # download repositories if they don't exist\n",
    "    !mkdir /notebooks/stable-diffusion-webui/repositories\n",
    "    !git clone https://github.com/CompVis/stable-diffusion.git /notebooks/stable-diffusion-webui/repositories/stable-diffusion\n",
    "    !git clone https://github.com/CompVis/taming-transformers.git /notebooks/stable-diffusion-webui/repositories/taming-transformers\n",
    "    !git clone https://github.com/sczhou/CodeFormer.git /notebooks/stable-diffusion-webui/repositories/CodeFormer\n",
    "else: # update repositories if they do exist\n",
    "    print('Updating dependencies')\n",
    "    for dir in os.listdir('/notebooks/stable-diffusion-webui/repositories'):\n",
    "        %cd /notebooks/stable-diffusion-webui/repositories/$dir\n",
    "        !git pull\n",
    "    \n",
    "!pip install -r /notebooks/stable-diffusion-webui/repositories/CodeFormer/requirements.txt\n",
    "\n",
    "if not os.path.exists('/notebooks/stable-diffusion-webui/GFPGANv1.3.pth'):\n",
    "    !wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth\n",
    "else:\n",
    "    print('GFPGANv1.3.pth already downloaded')\n",
    "    \n",
    "!mkdir -p /notebooks/models/\n",
    "\n",
    "# Get some storage back\n",
    "!pip cache purge\n",
    "!sudo apt-get clean"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F0EINk5M0s-w",
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Download the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I've provided a few different ways of aquiring the models. Try the torrent option first.\n",
    "\n",
    "You don't need to repeat this step if you've already downloaded the models.\n",
    "\n",
    "**Filesize and Storage Disclaimer**\n",
    "\n",
    "Paperspace free tier has only 5GB of storage space. The standard model is 4GB and waifu is 7.2GB. If you run out of storage space, you have two options:\n",
    "1. Add a payment method to your account. Storage overages are billed at \\$0.29/GB and billing occurs monthly and runs at midnight on the first of each month. With a payment method on file, Paperspace will let you use more storage and if you time it right you shouldn't actually be charged for it.\n",
    "2. Upgrade to a Pro account. They'll give you 15GB and you'll get longer runtimes and more powerful free GPUs.\n",
    "3. Create an account using my referral code `KQLRH37` or [signup link](https://console.paperspace.com/signup?R=KQLRH37). You'll get \\$10 credit that you should be able to put towards the storage overage charges."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "#### Standard Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Torrent**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!apt update\n",
    "!apt install -y aria2\n",
    "%cd /notebooks/models\n",
    "!aria2c --seed-time=0 --max-overall-upload-limit=1K \"magnet:?xt=urn:btih:3A4A612D75ED088EA542ACAC52F9F45987488D1C&tr=udp://tracker.opentrackr.org:1337\"\n",
    "!rm /notebooks/models/sd-v1-4.ckpt.aria2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Web Download**\n",
    "\n",
    "Voldy provided an alternative download if you don't want to use HuggingFace.\n",
    "\n",
    "[https://drive.google.com/file/d/1wHFgl0ivCmIZv88hVZXkb8oy9qCuaBGA/view](https://drive.google.com/file/d/1wHFgl0ivCmIZv88hVZXkb8oy9qCuaBGA/view)\n",
    "\n",
    "Download it to your computer then upload it to `/notebooks/models` (make sure it's named `sd-v1-4.ckpt`).\n",
    "\n",
    "HuggingFace is much faster and reliable but you need to get access to the repo and provide your user token."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "CT_J9L7oqLxG"
   },
   "outputs": [],
   "source": [
    "user_token = \"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\"\n",
    "\n",
    "# ===============================================================================================\n",
    "\n",
    "user_header = f\"'Authorization: Bearer {user_token}'\"\n",
    "!wget --header={user_header} https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt -O /notebooks/models/sd-v1-4.ckpt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "#### Waifu Diffusion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Torrent**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!apt update\n",
    "!apt install -y aria2\n",
    "%cd /notebooks/models\n",
    "!aria2c --seed-time=0 --max-overall-upload-limit=1K \"magnet:?xt=urn:btih:F45CECF4E9DE86DA83A78DD2CCCD7F27D5557A52&tr=udp://nyquist.localghost.org:6969\"\n",
    "!rm /notebooks/models/wd-v1-2-full-ema.ckpt.aria2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "**Web Download**\n",
    "\n",
    "Very slow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://thisanimedoesnotexist.ai/downloads/wd-v1-2-full-ema.ckpt -O /notebooks/models/wd-v1-2-full-ema.ckpt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Half-size Model**\n",
    "\n",
    "If you're on free tier try this first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!apt update\n",
    "!apt install -y aria2\n",
    "%cd /notebooks/models\n",
    "!aria2c --seed-time=0 --max-overall-upload-limit=1K \"magnet:?xt=urn:btih:153590FD7E93EE11D8DB951451056C362E3A9150&dn=wd-v1-2-full-ema-pruned.ckpt&tr=udp://glotorrents.pw:6969/announce&tr=udp://tracker.opentrackr.org:1337/announce&tr=udp://torrent.gresille.org:80/announce&tr=udp://tracker.openbittorrent.com:80&tr=udp://tracker.coppersurfer.tk:6969&tr=udp://tracker.leechers-paradise.org:6969&tr=udp://p4p.arenabg.ch:1337&tr=udp://tracker.internetwarriors.net:1337&tr=wss://tracker.openwebtorrent.com\"\n",
    "!rm /notebooks/models/wd-v1-2-full-ema.ckpt.aria2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "#### Gradient Dataset Import"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Paperspace provides a way to mount external storage to notebooks, called *Datasets*. This is an advanced method but it will save you many GB since you can reuse the same model file across different notebooks.\n",
    "\n",
    "https://docs.paperspace.com/gradient/cli/datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mounted_dataset_path = '/datasets/<name>'\n",
    "\n",
    "# ============================================================\n",
    "\n",
    "for file in os.listdir(mounted_dataset_path):\n",
    "    if file.endswith(\"ckpt\"):\n",
    "        !ln -s $mounted_dataset_path/$file /notebooks/models/$file\n",
    "        !ls -la --block-size=GB /notebooks/models/$file"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Restart the Kernel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.kill(os.getpid(), 9)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xt8lbdmC04ox"
   },
   "source": [
    "# Launch the WebUI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, enter the name of the model you want to use into `model_name` in the cell below.\n",
    "\n",
    "Then run the model. You will get a link to nnn.gradio.app, that's your WebUI. Follow it.\n",
    "\n",
    "If you have any issues, restart the kernel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "R-xAdMA5wxXd",
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Possible models:\n",
    "# sd-v1-4.ckpt (standard diffusion v1.4)\n",
    "# wd-v1-2-full-ema.ckpt (waifu-diffusion v1.2)\n",
    "# wd-v1-2-full-ema-pruned.ckpt (waifu-diffusion v1.2 half-size)\n",
    "\n",
    "model_name = ''\n",
    "\n",
    "# ========================================================================================================\n",
    "\n",
    "model_file = f'/notebooks/models/{model_name}'\n",
    "\n",
    "import sys\n",
    "import os\n",
    "\n",
    "if not os.path.exists(model_file):\n",
    "    print('Cannot find model file:', model_file)\n",
    "    sys.exit(1)\n",
    "\n",
    "if model_name == '':\n",
    "    print('You forgot to fill out the variable model_file above.')\n",
    "    sys.exit(1)\n",
    "\n",
    "!python /notebooks/stable-diffusion-webui/webui.py --share --ckpt $model_file --show-negative-prompt"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "private_outputs": true,
   "provenance": []
  },
  "gpuClass": "standard",
  "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}