{
 "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": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "### Clone the central repository"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sBbcB4vwj_jm"
   },
   "outputs": [],
   "source": [
    "import os.path\n",
    "\n",
    "if not os.path.exists('/notebooks/stable-diffusion-webui'):\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_",
    "jp-MarkdownHeadingCollapsed": true,
    "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')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import os\n",
    "# os.kill(os.getpid(), 9) # This will crash Colab (required, everything will still be intact so dont worry)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "F0EINk5M0s-w"
   },
   "source": [
    "### Download the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Pick only one!**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "#### Standard Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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, create the directory `/notebooks/stable-diffusion/` in this notebook (pro-tip: use `mkdir -p`), then upload it to that dir (make sure it's named `model.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 = \"<enter user_token here>\"\n",
    "\n",
    "# ===============================================================================================\n",
    "\n",
    "if not os.path.exists('/notebooks/stable-diffusion/models/ldm/stable-diffusion-v1/model.ckpt'):\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/stable-diffusion-webui/model.ckpt\n",
    "else:\n",
    "    print('Standard model already downloaded')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "#### Waifu Diffusion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**HTTPS**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists('/notebooks/stable-diffusion-webui/model.ckpt'):\n",
    "    !wget https://thisanimedoesnotexist.ai/downloads/wd-v1-2-full-ema.ckpt -O /notebooks/stable-diffusion-webui/model.ckpt\n",
    "else:\n",
    "    print('Waifu Diffusion model already downloaded')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Gradient Dataset**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mounted_dataset_path = '/datasets/<name>'\n",
    "!rsync -ah --info=progress2 $mounted_dataset_path/wd-v1-2-full-ema.ckpt /notebooks/stable-diffusion-webui/model.ckpt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xt8lbdmC04ox"
   },
   "source": [
    "# Launch the WebUI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You will get a link to nnn.gradio.app, that's your WebUI. Follow it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "R-xAdMA5wxXd"
   },
   "outputs": [],
   "source": [
    "import sys\n",
    "# sys.argv = [\"--opt-split-attention\"] # low memory optimizations, not nessesary for Paperspace.\n",
    "sys.argv = [''] # if you aren't using any args you must provide an empty array.\n",
    "\n",
    "import webui\n",
    "import modules.ui\n",
    "import modules.txt2img\n",
    "import modules.img2img\n",
    "\n",
    "demo = modules.ui.create_ui(\n",
    "    txt2img=webui.wrap_gradio_gpu_call(modules.txt2img.txt2img),\n",
    "    img2img=webui.wrap_gradio_gpu_call(modules.img2img.img2img),\n",
    "    run_extras=webui.wrap_gradio_gpu_call(webui.run_extras),\n",
    "    run_pnginfo=webui.run_pnginfo\n",
    ")\n",
    "\n",
    "demo.launch(share=True)"
   ]
  }
 ],
 "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
}