{
 "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": {
    "jp-MarkdownHeadingCollapsed": true,
    "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/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F0EINk5M0s-w",
    "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",
    "I recommend downloading both the standard and waifu diffusion. You can choose which one you want later.\n",
    "\n",
    "You don't need to repeat this step if you've already downloaded the models."
   ]
  },
  {
   "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": {
    "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": {
    "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",
    "\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}