koi_colab_backend.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Colab-Backend for [KOI](https://github.com/nousr/koi)\n",
        "---\n",
        "\n",
        "This notebook serves as one of the many ways you connect `koi` to a GPU backend!\n",
        "\n",
        "I also hope that it will serve as a good \"getting started\" guide and walk you through all the steps necessary to get everything up and running!\n",
        "\n",
        "\n",
        "### Notebook & Plug-In by [nousr](https://twitter.com/nousr_)\n",
        "\n",
        "---\n",
        "\n",
        "*StableDiffusion is a model created by CompVis in conjunction with [StabilityAI](stability.ai). By using this notebook you are also agreeing to any binding agreements that are associated with the StableDiffusion-V1 model.*"
      ],
      "metadata": {
        "id": "356TbPIlNtLQ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Dependencies\n",
        "\n",
        "First we need to install a few things...\n",
        "\n",
        "1. We install the koi package to ensure we have the proper packages\n",
        "2. We install `diffusers` from source to get the latest *img2img* pipeline\n",
        "3. Finally we install `ngrok` and `flask-ngrok` *(**note:** this are not necessary if you are running the server locally)*\n",
        "\n",
        "\n",
        "`Ngrok` & `Flask` is what makes it possible to use google colab as our gpu backend. In short, flask will handle our our server and ngrok will provide us a public IP that we can use to talk to from our local machine. \n"
      ],
      "metadata": {
        "id": "epp6uF59tGM7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!git clone https://github.com/nousr/koi.git && pip install -e koi\n",
        "!pip install git+https://github.com/huggingface/diffusers.git\n",
        "!pip install ngrok\n",
        "!pip install flask-ngrok\n",
        "!sudo apt install net-tools "
      ],
      "metadata": {
        "id": "cIR9LTKNIT83"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# ⏰ NOTE ⏰\n",
        "\n",
        "Before you continue you need to refresh the notebook, otherwise you will get some mysterious errors!\n",
        "\n",
        "You can do this by going to the top menu bar in colab and navigating to `Runtime` > `Restart Runtime`. After it refreshes you can continue!\n",
        "\n",
        "> *note:* you do not need to re-run the first cell after restarting the runtime\n",
        "\n",
        "(thanks to @thefacesblur on twitter for helping me debug this)"
      ],
      "metadata": {
        "id": "7tP2VGKNIPr8"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Setup\n",
        "\n",
        "Before we continue you must login to huggingface to use stable diffusion through the `diffusers` module.\n",
        "\n",
        "To do this you will need to do two things:\n",
        "\n",
        "1. **Important!** Accept the OpenRAIL license for stable diffusion here https://huggingface.co/CompVis/stable-diffusion-v1-4 \n",
        "2. Go to https://huggingface.co/settings/tokens and generate a new token to use.\n",
        "\n",
        "> ***NOTE:*** If you do not complete the second step, or do not do so without generating the token on the same account you accepted the license you will run into errors below."
      ],
      "metadata": {
        "id": "0MSHmEnGuozJ"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!huggingface-cli login # make sure you hit enter after pasting your token"
      ],
      "metadata": {
        "id": "w2hWo_HNNSQC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## The backend\n",
        "\n",
        "Now we can finally setup our backend server. For now we will keep it simple and have one API endpoint. \n",
        "\n",
        "> ⏰ **NOTE**! ⏰\n",
        ">\n",
        "> If you get an error saying the `transformers` library is not installed you probably need to restart the runtime to refresh the environment. \n",
        "> \n",
        "> You can do this by navigating to `Runtime` > `Restart Runtime`. After it refreshes you can try re-running this cell and it should work. *(You do ***not*** need to re-execute the other cells)*\n",
        "\n",
        "---\n",
        "\n",
        "In the future setting up the colab-backend will hopefully be as smple as just doing something like...\n",
        "\n",
        "```python\n",
        "from koi import colab_server\n",
        "\n",
        "run_server()\n",
        "```\n",
        "\n",
        "...but for now we will do it explicitly 🙂\n",
        "\n"
      ],
      "metadata": {
        "id": "5v9LJCRHwqqE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from flask import Flask, Response, request, send_file\n",
        "from PIL import Image\n",
        "from io import BytesIO\n",
        "from torch import autocast\n",
        "from diffusers import StableDiffusionImg2ImgPipeline\n",
        "from click import secho\n",
        "from zipfile import ZipFile\n",
        "\n",
        "# the following line is specific to remote environments (like google colab)\n",
        "from flask_ngrok import run_with_ngrok\n",
        "\n",
        "# Load the model for use (this may take a minute or two...or three)\n",
        "secho(\"Loading Model...\", fg=\"yellow\")\n",
        "\n",
        "pipe = StableDiffusionImg2ImgPipeline.from_pretrained(\n",
        "    \"CompVis/stable-diffusion-v1-4\", \n",
        "    use_auth_token=True,\n",
        "    revision=\"fp16\",\n",
        "    torch_dtype=torch.float16,\n",
        "    safety_checker=None,\n",
        "    requires_safety_checker=False,\n",
        ").to(\"cuda\")\n",
        "\n",
        "secho(\"Finished!\", fg=\"green\")\n",
        "\n",
        "# Start setting up flask\n",
        "\n",
        "app = Flask(__name__)\n",
        "\n",
        "# Define a function to help us \"control the randomness\"\n",
        "\n",
        "def seed_everything(seed: int):\n",
        "    import random, os\n",
        "\n",
        "    random.seed(seed)\n",
        "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
        "    torch.manual_seed(seed)\n",
        "    torch.cuda.manual_seed(seed)\n",
        "    torch.backends.cudnn.deterministic = True\n",
        "    torch.backends.cudnn.benchmark = True\n",
        "\n",
        "def get_name(prompt, seed):\n",
        "  return f'{prompt}-{seed}'\n",
        "\n",
        "# Define one endpoint \"/api/img2img\" for us to communicate with\n",
        "@app.route(\"/api/img2img\", methods=[\"POST\"])\n",
        "def img2img():\n",
        "    global pipe\n",
        "\n",
        "    r = request\n",
        "    headers = r.headers\n",
        "\n",
        "    data = r.data\n",
        "    buff = BytesIO(data)\n",
        "    img = Image.open(buff).convert(\"RGB\")\n",
        "\n",
        "    seed = int(headers[\"seed\"])\n",
        "    prompt = headers['prompt']\n",
        "\n",
        "\n",
        "    print(r.headers)\n",
        "\n",
        "    zip_stream = BytesIO()\n",
        "    with ZipFile(zip_stream, 'w') as zf:\n",
        "\n",
        "        for index in range(int(headers['variations'])):\n",
        "            variation_seed = seed + index\n",
        "            seed_everything(variation_seed)\n",
        "        \n",
        "            with autocast(\"cuda\"):\n",
        "                return_image = pipe(\n",
        "                    init_image=img,\n",
        "                    prompt=prompt,\n",
        "                    strength=float(headers[\"sketch_strength\"]),\n",
        "                    guidance_scale=float(headers[\"prompt_strength\"]),\n",
        "                    num_inference_steps=int(headers[\"steps\"]),\n",
        "                ).images[0]\n",
        "\n",
        "\n",
        "            return_bytes = BytesIO()\n",
        "            return_image.save(return_bytes, format=\"JPEG\")\n",
        "\n",
        "            return_bytes.seek(0)\n",
        "            zf.writestr(get_name(prompt, variation_seed), return_bytes.read())\n",
        "\n",
        "    zip_stream.seek(0)\n",
        "\n",
        "    return send_file(zip_stream, mimetype=\"application/zip\")\n",
        "\n",
        "run_with_ngrok(app)\n",
        "app.run()"
      ],
      "metadata": {
        "id": "1doJCzyLNnJp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Plugging into Krita\n",
        "(last step!)\n",
        "\n",
        "At this point, if everything worked, you should see something like the following!\n",
        "```terminal\n",
        "Finished!\n",
        " * Serving Flask app '__main__'\n",
        " * Debug mode: off\n",
        "INFO:werkzeug:WARNING: This is a development server. Do not use it in a production deployment. Use a production WSGI server instead.\n",
        " * Running on http://127.0.0.1:5000\n",
        "INFO:werkzeug:Press CTRL+C to quit\n",
        " * Running on http://c44b-34-124-187-78.ngrok.io <----❗This is what we need❗\n",
        " * Traffic stats available on http://127.0.0.1:4040 \n",
        " ```\n",
        "\n",
        "Everytime you run this notebook, ngrok will give you a new public-ip to use...in my case this was `http://c44b-34-124-187-78.ngrok.io`.\n",
        "\n",
        "**IMPORTANT: To begin using koi you will paste this public ip into the `endpoint` field, along with our *api route*. Like so:**\n",
        "\n",
        "`http://c44b-34-124-187-78.ngrok.io/api/img2img`\n",
        "\n",
        "---\n",
        "\n",
        "### Have fun! feel free to tweet me your creations--I'd love to see them :)"
      ],
      "metadata": {
        "id": "66d9BQVR014n"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Developers Only\n"
      ],
      "metadata": {
        "id": "uZJdQmUMG5iT"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "If you need to re-start the flask server in the course of development, killing the cell where it is running won't be enough. Instead you will need to use the commands below to kill the current runtime (which will also kill the running flask server). You can then re-run the flask code again and changes will propagate without needing to reinstall dependencies, log back into huggingface or, download the model again."
      ],
      "metadata": {
        "id": "X_ddlOapI0-o"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!sudo netstat -tulnp | grep :5000"
      ],
      "metadata": {
        "id": "bZXIdJkZH_v_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!sudo kill 434"
      ],
      "metadata": {
        "id": "Bycf_WAlH_1H"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}