stable-diffusion-notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8e908f76",
   "metadata": {},
   "source": [
    "# Paperspace Stable Diffusion Notebook\n",
    "\n",
    "This Notebook contains instructions and easy steps for generating images with Stable Diffusion on any of Paperspace's GPUs with little set up required. Follow the instructions in the markdown cells in this Notebook to quickly launch the Stable Diffusion Web UI or begin generating images using the Diffusers Stable Diffusion Pipeline.\n",
    "\n",
    "Contents:\n",
    "\n",
    "- Use the Stable Diffusion Web UI\n",
    "    - Loading in the Stable Diffusion models for the Web UI with support for v1-5 and v2 models\n",
    "- The Diffusers Stable Diffusion Pipeline\n",
    "    - Loading in the Stable Diffusion models for the Diffusers StableDiffusionPipeline with support for v1-4, v1-5 and v2 models\n",
    "    - Access additional models: log in to HuggingFace for online access model files\n",
    "    - Inference with the Diffusers pipeline\n",
    "- CLIP reranking\n",
    "- GFPGAN face restoration and image upscaling\n",
    "- CodeFormer face restoration and image upscaling"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9a384a1",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "First, run the cell below to get everything we need for this Notebook to run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "52004bf6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-08T16:15:05.358599Z",
     "iopub.status.busy": "2023-11-08T16:15:05.357686Z",
     "iopub.status.idle": "2023-11-08T16:18:02.681550Z",
     "shell.execute_reply": "2023-11-08T16:18:02.680130Z",
     "shell.execute_reply.started": "2023-11-08T16:15:05.358558Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting git+https://github.com/huggingface/diffusers.git\n",
      "  Cloning https://github.com/huggingface/diffusers.git to /tmp/pip-req-build-exwgb9g8\n",
      "  Running command git clone --filter=blob:none --quiet https://github.com/huggingface/diffusers.git /tmp/pip-req-build-exwgb9g8\n",
      "  Resolved https://github.com/huggingface/diffusers.git to commit 17528afcba583fa0e49504208af33b1a62ff1294\n",
      "  Installing build dependencies ... \u001b[?25ldone\n",
      "\u001b[?25h  Getting requirements to build wheel ... \u001b[?25ldone\n",
      "\u001b[?25h  Preparing metadata (pyproject.toml) ... \u001b[?25ldone\n",
      "\u001b[?25hRequirement already satisfied: numpy in /usr/local/lib/python3.9/dist-packages (from diffusers==0.23.0.dev0) (1.23.4)\n",
      "Requirement already satisfied: filelock in /usr/local/lib/python3.9/dist-packages (from diffusers==0.23.0.dev0) (3.9.0)\n",
      "Requirement already satisfied: requests in /usr/local/lib/python3.9/dist-packages (from diffusers==0.23.0.dev0) (2.28.2)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.9/dist-packages (from diffusers==0.23.0.dev0) (2022.10.31)\n",
      "Requirement already satisfied: Pillow in /usr/local/lib/python3.9/dist-packages (from diffusers==0.23.0.dev0) (9.2.0)\n",
      "Collecting huggingface-hub>=0.13.2\n",
      "  Downloading huggingface_hub-0.19.0-py3-none-any.whl (311 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m311.2/311.2 kB\u001b[0m \u001b[31m20.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: importlib-metadata in /usr/local/lib/python3.9/dist-packages (from diffusers==0.23.0.dev0) (6.0.0)\n",
      "Collecting safetensors>=0.3.1\n",
      "  Downloading safetensors-0.4.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (1.3 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m52.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: tqdm>=4.42.1 in /usr/local/lib/python3.9/dist-packages (from huggingface-hub>=0.13.2->diffusers==0.23.0.dev0) (4.64.1)\n",
      "Collecting fsspec>=2023.5.0\n",
      "  Downloading fsspec-2023.10.0-py3-none-any.whl (166 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m166.4/166.4 kB\u001b[0m \u001b[31m31.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.9/dist-packages (from huggingface-hub>=0.13.2->diffusers==0.23.0.dev0) (5.4.1)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.9/dist-packages (from huggingface-hub>=0.13.2->diffusers==0.23.0.dev0) (4.4.0)\n",
      "Requirement already satisfied: packaging>=20.9 in /usr/local/lib/python3.9/dist-packages (from huggingface-hub>=0.13.2->diffusers==0.23.0.dev0) (23.0)\n",
      "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.9/dist-packages (from importlib-metadata->diffusers==0.23.0.dev0) (3.11.0)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /usr/lib/python3/dist-packages (from requests->diffusers==0.23.0.dev0) (2.8)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/lib/python3/dist-packages (from requests->diffusers==0.23.0.dev0) (2019.11.28)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.9/dist-packages (from requests->diffusers==0.23.0.dev0) (2.1.1)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.9/dist-packages (from requests->diffusers==0.23.0.dev0) (1.26.14)\n",
      "Building wheels for collected packages: diffusers\n",
      "  Building wheel for diffusers (pyproject.toml) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for diffusers: filename=diffusers-0.23.0.dev0-py3-none-any.whl size=1660842 sha256=aae00ae6b7ad37230a26ba2f36df737e2c41758902de4252f626f8460fba691d\n",
      "  Stored in directory: /tmp/pip-ephem-wheel-cache-gspd5u6v/wheels/3a/75/ac/2724a3c688984d6b999006bb9f6b10275f59cf50377b74c585\n",
      "Successfully built diffusers\n",
      "Installing collected packages: safetensors, fsspec, huggingface-hub, diffusers\n",
      "  Attempting uninstall: fsspec\n",
      "    Found existing installation: fsspec 2023.1.0\n",
      "    Uninstalling fsspec-2023.1.0:\n",
      "      Successfully uninstalled fsspec-2023.1.0\n",
      "  Attempting uninstall: huggingface-hub\n",
      "    Found existing installation: huggingface-hub 0.12.0\n",
      "    Uninstalling huggingface-hub-0.12.0:\n",
      "      Successfully uninstalled huggingface-hub-0.12.0\n",
      "Successfully installed diffusers-0.23.0.dev0 fsspec-2023.10.0 huggingface-hub-0.19.0 safetensors-0.4.0\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mRequirement already satisfied: transformers in /usr/local/lib/python3.9/dist-packages (4.21.3)\n",
      "Collecting transformers\n",
      "  Downloading transformers-4.35.0-py3-none-any.whl (7.9 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m7.9/7.9 MB\u001b[0m \u001b[31m57.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting accelerate\n",
      "  Downloading accelerate-0.24.1-py3-none-any.whl (261 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m261.4/261.4 kB\u001b[0m \u001b[31m39.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: scipy in /usr/local/lib/python3.9/dist-packages (1.9.2)\n",
      "Collecting scipy\n",
      "  Downloading scipy-1.11.3-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (36.6 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m36.6/36.6 MB\u001b[0m \u001b[31m29.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.9/dist-packages (from transformers) (2022.10.31)\n",
      "Requirement already satisfied: requests in /usr/local/lib/python3.9/dist-packages (from transformers) (2.28.2)\n",
      "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.9/dist-packages (from transformers) (1.23.4)\n",
      "Requirement already satisfied: huggingface-hub<1.0,>=0.16.4 in /usr/local/lib/python3.9/dist-packages (from transformers) (0.19.0)\n",
      "Requirement already satisfied: filelock in /usr/local/lib/python3.9/dist-packages (from transformers) (3.9.0)\n",
      "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.9/dist-packages (from transformers) (23.0)\n",
      "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.9/dist-packages (from transformers) (4.64.1)\n",
      "Requirement already satisfied: safetensors>=0.3.1 in /usr/local/lib/python3.9/dist-packages (from transformers) (0.4.0)\n",
      "Collecting tokenizers<0.15,>=0.14\n",
      "  Downloading tokenizers-0.14.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (3.8 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.8/3.8 MB\u001b[0m \u001b[31m75.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m:00:01\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.9/dist-packages (from transformers) (5.4.1)\n",
      "Requirement already satisfied: psutil in /usr/local/lib/python3.9/dist-packages (from accelerate) (5.9.4)\n",
      "Requirement already satisfied: torch>=1.10.0 in /usr/local/lib/python3.9/dist-packages (from accelerate) (1.12.1+cu116)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.9/dist-packages (from huggingface-hub<1.0,>=0.16.4->transformers) (4.4.0)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /usr/local/lib/python3.9/dist-packages (from huggingface-hub<1.0,>=0.16.4->transformers) (2023.10.0)\n",
      "Collecting huggingface-hub<1.0,>=0.16.4\n",
      "  Downloading huggingface_hub-0.17.3-py3-none-any.whl (295 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m295.0/295.0 kB\u001b[0m \u001b[31m41.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.9/dist-packages (from requests->transformers) (1.26.14)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.9/dist-packages (from requests->transformers) (2.1.1)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/lib/python3/dist-packages (from requests->transformers) (2019.11.28)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /usr/lib/python3/dist-packages (from requests->transformers) (2.8)\n",
      "Installing collected packages: scipy, huggingface-hub, tokenizers, accelerate, transformers\n",
      "  Attempting uninstall: scipy\n",
      "    Found existing installation: scipy 1.9.2\n",
      "    Uninstalling scipy-1.9.2:\n",
      "      Successfully uninstalled scipy-1.9.2\n",
      "  Attempting uninstall: huggingface-hub\n",
      "    Found existing installation: huggingface-hub 0.19.0\n",
      "    Uninstalling huggingface-hub-0.19.0:\n",
      "      Successfully uninstalled huggingface-hub-0.19.0\n",
      "  Attempting uninstall: tokenizers\n",
      "    Found existing installation: tokenizers 0.12.1\n",
      "    Uninstalling tokenizers-0.12.1:\n",
      "      Successfully uninstalled tokenizers-0.12.1\n",
      "  Attempting uninstall: transformers\n",
      "    Found existing installation: transformers 4.21.3\n",
      "    Uninstalling transformers-4.21.3:\n",
      "      Successfully uninstalled transformers-4.21.3\n",
      "Successfully installed accelerate-0.24.1 huggingface-hub-0.17.3 scipy-1.11.3 tokenizers-0.14.1 transformers-4.35.0\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mCollecting flax==0.5.0\n",
      "  Downloading flax-0.5.0-py3-none-any.whl (190 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m190.2/190.2 kB\u001b[0m \u001b[31m22.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: flax\n",
      "  Attempting uninstall: flax\n",
      "    Found existing installation: flax 0.6.3\n",
      "    Uninstalling flax-0.6.3:\n",
      "      Successfully uninstalled flax-0.6.3\n",
      "Successfully installed flax-0.5.0\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mRequirement already satisfied: msgpack in /usr/local/lib/python3.9/dist-packages (1.0.4)\n",
      "Requirement already satisfied: rich in /usr/local/lib/python3.9/dist-packages (13.2.0)\n",
      "Requirement already satisfied: optax in /usr/local/lib/python3.9/dist-packages (0.1.4)\n",
      "Requirement already satisfied: accelerate in /usr/local/lib/python3.9/dist-packages (0.24.1)\n",
      "Collecting ftfy\n",
      "  Downloading ftfy-6.1.1-py3-none-any.whl (53 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m53.1/53.1 kB\u001b[0m \u001b[31m7.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: scipy in /usr/local/lib/python3.9/dist-packages (1.11.3)\n",
      "Requirement already satisfied: markdown-it-py<3.0.0,>=2.1.0 in /usr/local/lib/python3.9/dist-packages (from rich) (2.1.0)\n",
      "Requirement already satisfied: pygments<3.0.0,>=2.6.0 in /usr/local/lib/python3.9/dist-packages (from rich) (2.14.0)\n",
      "Requirement already satisfied: absl-py>=0.7.1 in /usr/local/lib/python3.9/dist-packages (from optax) (1.4.0)\n",
      "Requirement already satisfied: jax>=0.1.55 in /usr/local/lib/python3.9/dist-packages (from optax) (0.4.1)\n",
      "Requirement already satisfied: chex>=0.1.5 in /usr/local/lib/python3.9/dist-packages (from optax) (0.1.5)\n",
      "Requirement already satisfied: jaxlib>=0.1.37 in /usr/local/lib/python3.9/dist-packages (from optax) (0.4.1+cuda11.cudnn82)\n",
      "Requirement already satisfied: numpy>=1.18.0 in /usr/local/lib/python3.9/dist-packages (from optax) (1.23.4)\n",
      "Requirement already satisfied: typing-extensions>=3.10.0 in /usr/local/lib/python3.9/dist-packages (from optax) (4.4.0)\n",
      "Requirement already satisfied: huggingface-hub in /usr/local/lib/python3.9/dist-packages (from accelerate) (0.17.3)\n",
      "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.9/dist-packages (from accelerate) (23.0)\n",
      "Requirement already satisfied: torch>=1.10.0 in /usr/local/lib/python3.9/dist-packages (from accelerate) (1.12.1+cu116)\n",
      "Requirement already satisfied: pyyaml in /usr/local/lib/python3.9/dist-packages (from accelerate) (5.4.1)\n",
      "Requirement already satisfied: psutil in /usr/local/lib/python3.9/dist-packages (from accelerate) (5.9.4)\n",
      "Requirement already satisfied: wcwidth>=0.2.5 in /usr/local/lib/python3.9/dist-packages (from ftfy) (0.2.6)\n",
      "Requirement already satisfied: dm-tree>=0.1.5 in /usr/local/lib/python3.9/dist-packages (from chex>=0.1.5->optax) (0.1.8)\n",
      "Requirement already satisfied: toolz>=0.9.0 in /usr/local/lib/python3.9/dist-packages (from chex>=0.1.5->optax) (0.12.0)\n",
      "Requirement already satisfied: opt-einsum in /usr/local/lib/python3.9/dist-packages (from jax>=0.1.55->optax) (3.3.0)\n",
      "Requirement already satisfied: mdurl~=0.1 in /usr/local/lib/python3.9/dist-packages (from markdown-it-py<3.0.0,>=2.1.0->rich) (0.1.2)\n",
      "Requirement already satisfied: fsspec in /usr/local/lib/python3.9/dist-packages (from huggingface-hub->accelerate) (2023.10.0)\n",
      "Requirement already satisfied: tqdm>=4.42.1 in /usr/local/lib/python3.9/dist-packages (from huggingface-hub->accelerate) (4.64.1)\n",
      "Requirement already satisfied: filelock in /usr/local/lib/python3.9/dist-packages (from huggingface-hub->accelerate) (3.9.0)\n",
      "Requirement already satisfied: requests in /usr/local/lib/python3.9/dist-packages (from huggingface-hub->accelerate) (2.28.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /usr/lib/python3/dist-packages (from requests->huggingface-hub->accelerate) (2.8)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.9/dist-packages (from requests->huggingface-hub->accelerate) (2.1.1)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/lib/python3/dist-packages (from requests->huggingface-hub->accelerate) (2019.11.28)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.9/dist-packages (from requests->huggingface-hub->accelerate) (1.26.14)\n",
      "Installing collected packages: ftfy\n",
      "Successfully installed ftfy-6.1.1\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mRequirement already satisfied: ipywidgets in /usr/local/lib/python3.9/dist-packages (8.0.2)\n",
      "Requirement already satisfied: ipython>=6.1.0 in /usr/local/lib/python3.9/dist-packages (from ipywidgets) (8.5.0)\n",
      "Requirement already satisfied: ipykernel>=4.5.1 in /usr/local/lib/python3.9/dist-packages (from ipywidgets) (6.16.0)\n",
      "Requirement already satisfied: traitlets>=4.3.1 in /usr/local/lib/python3.9/dist-packages (from ipywidgets) (5.8.1)\n",
      "Requirement already satisfied: widgetsnbextension~=4.0 in /usr/local/lib/python3.9/dist-packages (from ipywidgets) (4.0.5)\n",
      "Requirement already satisfied: jupyterlab-widgets~=3.0 in /usr/local/lib/python3.9/dist-packages (from ipywidgets) (3.0.5)\n",
      "Requirement already satisfied: tornado>=6.1 in /usr/local/lib/python3.9/dist-packages (from ipykernel>=4.5.1->ipywidgets) (6.1)\n",
      "Requirement already satisfied: debugpy>=1.0 in /usr/local/lib/python3.9/dist-packages (from ipykernel>=4.5.1->ipywidgets) (1.6.6)\n",
      "Requirement already satisfied: matplotlib-inline>=0.1 in /usr/local/lib/python3.9/dist-packages (from ipykernel>=4.5.1->ipywidgets) (0.1.6)\n",
      "Requirement already satisfied: jupyter-client>=6.1.12 in /usr/local/lib/python3.9/dist-packages (from ipykernel>=4.5.1->ipywidgets) (7.3.4)\n",
      "Requirement already satisfied: pyzmq>=17 in /usr/local/lib/python3.9/dist-packages (from ipykernel>=4.5.1->ipywidgets) (25.0.0)\n",
      "Requirement already satisfied: psutil in /usr/local/lib/python3.9/dist-packages (from ipykernel>=4.5.1->ipywidgets) (5.9.4)\n",
      "Requirement already satisfied: packaging in /usr/local/lib/python3.9/dist-packages (from ipykernel>=4.5.1->ipywidgets) (23.0)\n",
      "Requirement already satisfied: nest-asyncio in /usr/local/lib/python3.9/dist-packages (from ipykernel>=4.5.1->ipywidgets) (1.5.6)\n",
      "Requirement already satisfied: backcall in /usr/local/lib/python3.9/dist-packages (from ipython>=6.1.0->ipywidgets) (0.2.0)\n",
      "Requirement already satisfied: decorator in /usr/local/lib/python3.9/dist-packages (from ipython>=6.1.0->ipywidgets) (5.1.1)\n",
      "Requirement already satisfied: stack-data in /usr/local/lib/python3.9/dist-packages (from ipython>=6.1.0->ipywidgets) (0.6.2)\n",
      "Requirement already satisfied: pexpect>4.3 in /usr/local/lib/python3.9/dist-packages (from ipython>=6.1.0->ipywidgets) (4.8.0)\n",
      "Requirement already satisfied: prompt-toolkit<3.1.0,>3.0.1 in /usr/local/lib/python3.9/dist-packages (from ipython>=6.1.0->ipywidgets) (3.0.36)\n",
      "Requirement already satisfied: jedi>=0.16 in /usr/local/lib/python3.9/dist-packages (from ipython>=6.1.0->ipywidgets) (0.18.2)\n",
      "Requirement already satisfied: pickleshare in /usr/local/lib/python3.9/dist-packages (from ipython>=6.1.0->ipywidgets) (0.7.5)\n",
      "Requirement already satisfied: pygments>=2.4.0 in /usr/local/lib/python3.9/dist-packages (from ipython>=6.1.0->ipywidgets) (2.14.0)\n",
      "Requirement already satisfied: parso<0.9.0,>=0.8.0 in /usr/local/lib/python3.9/dist-packages (from jedi>=0.16->ipython>=6.1.0->ipywidgets) (0.8.3)\n",
      "Requirement already satisfied: jupyter-core>=4.9.2 in /usr/local/lib/python3.9/dist-packages (from jupyter-client>=6.1.12->ipykernel>=4.5.1->ipywidgets) (5.1.5)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.9/dist-packages (from jupyter-client>=6.1.12->ipykernel>=4.5.1->ipywidgets) (2.8.2)\n",
      "Requirement already satisfied: entrypoints in /usr/local/lib/python3.9/dist-packages (from jupyter-client>=6.1.12->ipykernel>=4.5.1->ipywidgets) (0.4)\n",
      "Requirement already satisfied: ptyprocess>=0.5 in /usr/local/lib/python3.9/dist-packages (from pexpect>4.3->ipython>=6.1.0->ipywidgets) (0.7.0)\n",
      "Requirement already satisfied: wcwidth in /usr/local/lib/python3.9/dist-packages (from prompt-toolkit<3.1.0,>3.0.1->ipython>=6.1.0->ipywidgets) (0.2.6)\n",
      "Requirement already satisfied: pure-eval in /usr/local/lib/python3.9/dist-packages (from stack-data->ipython>=6.1.0->ipywidgets) (0.2.2)\n",
      "Requirement already satisfied: asttokens>=2.1.0 in /usr/local/lib/python3.9/dist-packages (from stack-data->ipython>=6.1.0->ipywidgets) (2.2.1)\n",
      "Requirement already satisfied: executing>=1.2.0 in /usr/local/lib/python3.9/dist-packages (from stack-data->ipython>=6.1.0->ipywidgets) (1.2.0)\n",
      "Requirement already satisfied: six in /usr/lib/python3/dist-packages (from asttokens>=2.1.0->stack-data->ipython>=6.1.0->ipywidgets) (1.14.0)\n",
      "Requirement already satisfied: platformdirs>=2.5 in /usr/local/lib/python3.9/dist-packages (from jupyter-core>=4.9.2->jupyter-client>=6.1.12->ipykernel>=4.5.1->ipywidgets) (2.6.2)\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mEnabling notebook extension jupyter-js-widgets/extension...\n",
      "      - Validating: \u001b[32mOK\u001b[0m\n",
      "fatal: destination path 'stable-diffusion-webui' already exists and is not an empty directory.\n",
      "mkdir: cannot create directory ‘outputs’: File exists\n",
      "Found existing installation: torch 1.12.1+cu116\n",
      "Uninstalling torch-1.12.1+cu116:\n",
      "  Successfully uninstalled torch-1.12.1+cu116\n",
      "Found existing installation: torchvision 0.13.1+cu116\n",
      "Uninstalling torchvision-0.13.1+cu116:\n",
      "  Successfully uninstalled torchvision-0.13.1+cu116\n",
      "Found existing installation: torchaudio 0.12.1+cu116\n",
      "Uninstalling torchaudio-0.12.1+cu116:\n",
      "  Successfully uninstalled torchaudio-0.12.1+cu116\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mLooking in indexes: https://download.pytorch.org/whl/cu118\n",
      "Collecting torch\n",
      "  Downloading https://download.pytorch.org/whl/cu118/torch-2.1.0%2Bcu118-cp39-cp39-linux_x86_64.whl (2325.9 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.3/2.3 GB\u001b[0m \u001b[31m137.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting torchvision\n",
      "  Downloading https://download.pytorch.org/whl/cu118/torchvision-0.16.0%2Bcu118-cp39-cp39-linux_x86_64.whl (6.2 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m6.2/6.2 MB\u001b[0m \u001b[31m34.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hCollecting torchaudio\n",
      "  Downloading https://download.pytorch.org/whl/cu118/torchaudio-2.1.0%2Bcu118-cp39-cp39-linux_x86_64.whl (3.2 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.2/3.2 MB\u001b[0m \u001b[31m128.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting triton==2.1.0\n",
      "  Downloading https://download.pytorch.org/whl/triton-2.1.0-0-cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.whl (89.3 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m89.3/89.3 MB\u001b[0m \u001b[31m159.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: networkx in /usr/local/lib/python3.9/dist-packages (from torch) (3.0)\n",
      "Requirement already satisfied: fsspec in /usr/local/lib/python3.9/dist-packages (from torch) (2023.10.0)\n",
      "Requirement already satisfied: filelock in /usr/local/lib/python3.9/dist-packages (from torch) (3.9.0)\n",
      "Collecting sympy\n",
      "  Downloading https://download.pytorch.org/whl/sympy-1.12-py3-none-any.whl (5.7 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m5.7/5.7 MB\u001b[0m \u001b[31m145.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: typing-extensions in /usr/local/lib/python3.9/dist-packages (from torch) (4.4.0)\n",
      "Requirement already satisfied: jinja2 in /usr/local/lib/python3.9/dist-packages (from torch) (3.1.2)\n",
      "Requirement already satisfied: numpy in /usr/local/lib/python3.9/dist-packages (from torchvision) (1.23.4)\n",
      "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /usr/local/lib/python3.9/dist-packages (from torchvision) (9.2.0)\n",
      "Requirement already satisfied: requests in /usr/local/lib/python3.9/dist-packages (from torchvision) (2.28.2)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.9/dist-packages (from jinja2->torch) (2.1.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /usr/lib/python3/dist-packages (from requests->torchvision) (2.8)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.9/dist-packages (from requests->torchvision) (1.26.14)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/lib/python3/dist-packages (from requests->torchvision) (2019.11.28)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.9/dist-packages (from requests->torchvision) (2.1.1)\n",
      "Collecting mpmath>=0.19\n",
      "  Downloading https://download.pytorch.org/whl/mpmath-1.3.0-py3-none-any.whl (536 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m536.2/536.2 kB\u001b[0m \u001b[31m186.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: mpmath, triton, sympy, torch, torchvision, torchaudio\n",
      "Successfully installed mpmath-1.3.0 sympy-1.12 torch-2.1.0+cu118 torchaudio-2.1.0+cu118 torchvision-0.16.0+cu118 triton-2.1.0\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!pip install --upgrade git+https://github.com/huggingface/diffusers.git \n",
    "!pip install -U transformers accelerate scipy\n",
    "!pip install flax==0.5.0 --no-deps\n",
    "!pip install msgpack rich optax accelerate ftfy scipy\n",
    "!pip install ipywidgets\n",
    "!jupyter nbextension enable --py widgetsnbextension\n",
    "!git clone https://github.com/AUTOMATIC1111/stable-diffusion-webui\n",
    "!mkdir outputs\n",
    "!pip uninstall -y torch torchvision torchaudio\n",
    "!pip3 install torch torchvision torchaudio --no-cache-dir --index-url https://download.pytorch.org/whl/cu118\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1d63370",
   "metadata": {},
   "source": [
    "## Use the Stable Diffusion Web UI\n",
    "\n",
    "Run the cell below to setup, and launch the Automatic1111 Stable Diffusion Web UI from this notebook. Click the Gradio link after the setup completes to access the Web UI in your local browser. \n",
    "\n",
    "Skip this step to use the Diffusers library methodology.\n",
    "\n",
    "### Loading in the Stable Diffusion models for the Web UI \n",
    "\n",
    "To make accessing the Stable Diffusion models easy and not take up any storage, we have added the Stable Diffusion v1-5 and v2 models as mountable public datasets!\n",
    "\n",
    "> These files will not count towards the storage limit for free notebooks.\n",
    "\n",
    "To access the model files:\n",
    "\n",
    "First, navigate to the \"Data Sources\" tab using the navigator on the far left of the page.\n",
    "\n",
    "Next, click \"Public\" to switch into the Gradient Public Datasets, and scroll down until you find the Stable Diffusion datasets near the bottom of the list. You can mount `stable-diffusion-classic` for access to the v1-5 model checkpoints for the web UI, `stable-diffusion-diffusers` for the v1-5 models as Diffusers StableDiffusionPipeline binaries, `stable-diffusion-classic-v2` for access to the v2 model checkpoints and config file for use with the web UI, and `stable-diffusion-diffusers-v2` for access to the v2 models as Diffusers StableDiffusionPipeline binaries.\n",
    "\n",
    "Finally, click \"mount\" to make any of these Public Datasets accessible from the \"datasets\" directory. This directory is in the root folder, so access it with the path `~/../datasets/stable-diffusion-classic/SDv1.5.ckpt`. For example, after mounting, the v1-5 checkpoint can be found at `~/../datasets/stable-diffusion-classic/SDv1.5.ckpt`\n",
    "\n",
    "See the video below for a video guide to for help mounting these files. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f29ed4b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2023-11-08T16:18:02.686029Z",
     "iopub.status.busy": "2023-11-08T16:18:02.684755Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/notebooks/stable-diffusion-webui\n",
      "Python 3.9.16 (main, Dec  7 2022, 01:11:51) \n",
      "[GCC 9.4.0]\n",
      "Version: v1.6.0-2-g4afaaf8a\n",
      "Commit hash: 4afaaf8a020c1df457bcf7250cb1c7f609699fa7\n",
      "Installing clip\n",
      "Installing open_clip\n",
      "Installing requirements for CodeFormer\n",
      "Installing requirements\n",
      "Launching Web UI with arguments: --share --enable-insecure-extension-access --config /root/../datasets/stable-diffusion-classic-v2/768-v-ema.yaml --ckpt /root/../datasets/stable-diffusion-classic-v2/768-v-ema.ckpt\n",
      "no module 'xformers'. Processing without...\n",
      "no module 'xformers'. Processing without...\n",
      "No module 'xformers'. Proceeding without it.\n",
      "Loading weights [bfcaf07557] from /root/../datasets/stable-diffusion-classic-v2/768-v-ema.ckpt\n",
      "Running on local URL:  http://127.0.0.1:7860\n",
      "Running on public URL: https://2efdbd2acf81f81e1c.gradio.live\n",
      "\n",
      "This share link expires in 72 hours. For free permanent hosting and GPU upgrades, run `gradio deploy` from Terminal to deploy to Spaces (https://huggingface.co/spaces)\n",
      "[Lobe]: Initializing Lobe\n",
      "Startup time: 109.8s (prepare environment: 79.9s, import torch: 13.8s, import gradio: 1.8s, setup paths: 1.2s, initialize shared: 0.6s, other imports: 1.6s, setup codeformer: 0.4s, load scripts: 2.2s, create ui: 1.3s, gradio launch: 6.6s).\n",
      "Creating model from config: /root/../datasets/stable-diffusion-classic-v2/768-v-ema.yaml\n",
      "Applying attention optimization: Doggettx... done.\n",
      "Model loaded in 58.9s (load weights from disk: 52.5s, create model: 0.2s, apply weights to model: 5.4s, calculate empty prompt: 0.6s).\n"
     ]
    }
   ],
   "source": [
    "# Launch the Stable Diffusion Web UI. \n",
    "# Be sure to hash out whichever you are not using.\n",
    "%cd stable-diffusion-webui\n",
    "\n",
    "## Launch Web UI for Stable Diffusion v2 models. Note: this prevents users from using checkpoints and embeddings trained on v1-X models.\n",
    "!python launch.py --share --enable-insecure-extension-access --config ~/../datasets/stable-diffusion-classic-v2/768-v-ema.yaml --ckpt ~/../datasets/stable-diffusion-classic-v2/768-v-ema.ckpt\n",
    "\n",
    "## Launch Web UI for Stable Diffusion v1-x models.\n",
    "!python launch.py --share --enable-insecure-extension-access --ckpt ../../datasets/stable-diffusion-classic/SDv1.5.ckpt\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cccad2f7",
   "metadata": {},
   "source": [
    "# The Diffusers Stable Diffusion Pipeline\n",
    "\n",
    "## Loading in the Stable Diffusion v1-5 and v2 models for the Diffusers StableDiffusionPipeline\n",
    "\n",
    "To make accessing the Stable Diffusion models easy and not take up any storage, we have added the Stable Diffusion models as mountable public datasets! To access the models this way, simply navigate to the \"Data Sources\" tab using the navigator on the far left of the page. \n",
    "\n",
    "Then click \"Public\" to switch into the Gradient Public Datasets, and scroll down until you find the datasets starting with \"stable-diffusion\" near the bottom of the list. We recommend mounting each of these public datasets: `stable-diffusion-classic` (v1-5 checkpoint), `stable-diffusion-diffusers` (v1-5 diffusers), `stable-diffusion-classic-v2` (v2 checkpoint & config), and `stable-diffusion-diffusers-v2` (v2 diffusers). \n",
    "\n",
    "Click \"mount\" to make each of them accessible from the \"datasets\" directory. This directory is in the root folder, so access the two directories with the paths `~/../datasets/` and `../datasets/stable-diffusion-diffusers-v2/stable-diffusion-2/`\n",
    "\n",
    "<video controls src=\"assets/stab-upload.mp4\" />\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c7d9971",
   "metadata": {},
   "source": [
    "## Access additional models: log in to HuggingFace for online access to v1-4 and v2 model files\n",
    "\n",
    "In order to access the models from CompVis for Stable Diffusion, you must follow three steps:\n",
    "\n",
    "1. You must acknowledge and agree to their user requirements and license for their models. You can do so by reading the instructions found on these pages: \n",
    "- [CompVis/stable-diffusion-v1-4](https://huggingface.co/CompVis/stable-diffusion-v1-4)\n",
    "- [RunwayML/stable-diffusion-v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5)\n",
    "- [StabilityAI/stable-diffusion-2](https://huggingface.co/stabilityai/stable-diffusion-2) \n",
    "\n",
    "2. You must login to Huggingface, and then create and retrieve an access token (found here: https://huggingface.co/settings/tokens)\n",
    "\n",
    "3. Finally, replace the segment of the cell below `<your_huggingface_token>` with your own token, and run the cell. \n",
    "\n",
    "If you follow these steps, you will be able to access the model for free!\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f34b957",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python login.py --token <your_huggingface_token>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0b17152",
   "metadata": {},
   "source": [
    "## Inference with the Diffusers pipeline\n",
    "\n",
    "Use the following cells to begin synthesizing images. Use the radio buttons in the cell below to select which model type to use. \n",
    "\n",
    "### Setup\n",
    "\n",
    "Use the following cell to set up your inference by\n",
    "\n",
    "- Choose which Stable Diffusion checkpoint to use\n",
    "- Creating the outputs directory\n",
    "- Choose whether to use FP16 or FP32 (use FP16 for Free-GPU Notebooks)\n",
    "- import packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c867ed86",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import autocast\n",
    "from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler\n",
    "import ipywidgets as widgets\n",
    "from ipywidgets import interact\n",
    "\n",
    "%cd ~/../notebooks\n",
    "\n",
    "## Use the lists below to select our model type, image output resolution, and the computer number format for inference\n",
    "device = 'cuda'\n",
    "## We recommend either v1-5 (index 1) or v2 (index 2)\n",
    "model_id = ['CompVis/stable-diffusion-v1-4',\"../datasets/stable-diffusion-diffusers/stable-diffusion-v1-5/\",'../datasets/stable-diffusion-diffusers-v2/stable-diffusion-2/']\n",
    "## We strongly recommend using the v1-5 and v1-4 models with size 512, and the v2 models with size 768.\n",
    "size = [512, 768]\n",
    "## If you are using a GPU with ~ 8 GBs of RAM, like the Free-GPU Notebooks, then use torch.float16\n",
    "## otherwise, either selection works\n",
    "dtype = [torch.float16, torch.float32]\n",
    "\n",
    "\n",
    "# Set your values here - defaults are for v2, size 768 x 768 with FP16 computer number format\n",
    "model_ = model_id[2]\n",
    "size_ = size[1] \n",
    "precision = dtype[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdc25595",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Depending on the selections made above, the type of model loaded will change using this short if/else statement\n",
    "\n",
    "print(f'Now loading Stable Diffusion {model_} model ({precision})...')\n",
    "scheduler = EulerDiscreteScheduler.from_pretrained('stabilityai/stable-diffusion-2', subfolder=\"scheduler\")\n",
    "if precision == torch.float16:\n",
    "    if model_ == '../datasets/stable-diffusion-diffusers-v2/stable-diffusion-2/':\n",
    "        pipe = StableDiffusionPipeline.from_pretrained(model_, torch_dtype=torch.float16, revision=\"fp16\", scheduler = scheduler)\n",
    "        pipe = pipe.to(device)\n",
    "    else:\n",
    "        pipe = StableDiffusionPipeline.from_pretrained(model_, torch_dtype=torch.float16, revision=\"fp16\")\n",
    "        pipe = pipe.to(device)\n",
    "else:\n",
    "    if model_ == '../datasets/stable-diffusion-diffusers-v2/stable-diffusion-2/':\n",
    "        pipe = StableDiffusionPipeline.from_pretrained(model_, scheduler = scheduler)\n",
    "        pipe = pipe.to(device)\n",
    "    else:\n",
    "        pipe = StableDiffusionPipeline.from_pretrained(model_)\n",
    "        pipe = pipe.to(device)\n",
    "print('Model now loaded. Head to next cell to generate images.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa7fe95f",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_num = 5\n",
    "lst = []\n",
    "\n",
    "print(f'Generating samples from Stable Diffusion {model_} checkpoint ({precision})')\n",
    "\n",
    "prompt = 'a corgi astronaut on mars'\n",
    "for j in range(sample_num):\n",
    "    a = pipe(prompt,\n",
    "             guidance_scale=7.5,\n",
    "             height = size_,\n",
    "             width = size_,\n",
    "             num_inference_steps=50)['images'][0]\n",
    "    lst.append(a)\n",
    "    display(a)\n",
    "    a.save(f'outputs/gen-image-{j}.png')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c39f33d",
   "metadata": {},
   "source": [
    "# CLIP reranking\n",
    "\n",
    "Find the best of your images from the generated samples by ranking them with CLIP by their accuracy to the prompt. \n",
    "\n",
    "Code adapted from Boris Dayma's [DALL-E Mini Inference Pipeline](https://github.com/borisdayma/dalle-mini/blob/main/tools/inference/inference_pipeline.ipynb) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a932bf4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import CLIPProcessor, FlaxCLIPModel\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from flax.jax_utils import replicate\n",
    "from functools import partial\n",
    "\n",
    "\n",
    "# CLIP model\n",
    "CLIP_REPO = \"openai/clip-vit-base-patch32\"\n",
    "CLIP_COMMIT_ID = None\n",
    "\n",
    "# Load CLIP\n",
    "clip, clip_params = FlaxCLIPModel.from_pretrained(\n",
    "    CLIP_REPO, revision=CLIP_COMMIT_ID, dtype=jnp.float16, _do_init=False\n",
    ")\n",
    "clip_processor = CLIPProcessor.from_pretrained(CLIP_REPO, revision=CLIP_COMMIT_ID)\n",
    "clip_params = replicate(clip_params)\n",
    "\n",
    "# score images\n",
    "@partial(jax.pmap, axis_name=\"batch\")\n",
    "def p_clip(inputs, params):\n",
    "    logits = clip(params=params, **inputs).logits_per_image\n",
    "    return logits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cdb0685",
   "metadata": {},
   "outputs": [],
   "source": [
    "from flax.training.common_utils import shard\n",
    "import numpy as np\n",
    "# get clip scores\n",
    "clip_inputs = clip_processor(\n",
    "    text=prompt * jax.device_count(),\n",
    "    images=lst,\n",
    "    return_tensors=\"np\",\n",
    "    padding=\"max_length\",\n",
    "    max_length=77,\n",
    "    truncation=True,\n",
    ").data\n",
    "logits = p_clip(shard(clip_inputs), clip_params)\n",
    "\n",
    "out = list(reversed(sorted(zip(logits[0], lst))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e3b5efc",
   "metadata": {},
   "outputs": [],
   "source": [
    "for idx, v in enumerate(out):\n",
    "    display(v[1])\n",
    "    print(f\"Score: {v[0][0]:.2f}\\n\")\n",
    "    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d2f1439",
   "metadata": {},
   "source": [
    "# GFPGAN\n",
    "\n",
    "GFPGAN aims at developing a Practical Algorithm for Real-world Face Restoration. It leverages rich and diverse priors encapsulated in a pretrained face GAN (e.g., StyleGAN2) for blind face restoration. GFPGAN is combined with a suite of other restoration and upscaling models to help make your pictures much clearer and cleaner. \n",
    "\n",
    "credit to: https://github.com/TencentARC/GFPGAN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38eb726b",
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone https://github.com/TencentARC/GFPGAN.git\n",
    "\n",
    "%cd GFPGAN\n",
    "# Install basicsr - https://github.com/xinntao/BasicSR\n",
    "# We use BasicSR for both training and inference\n",
    "# Install facexlib - https://github.com/xinntao/facexlib\n",
    "# We use face detection and face restoration helper in the facexlib package\n",
    "!pip install facexlib\n",
    "\n",
    "!pip install - r requirements.txt\n",
    "!python setup.py develop\n",
    "!pip install realesrgan basicsr optax\n",
    "\n",
    "# If you want to enhance the background (non-face) regions with Real-ESRGAN,\n",
    "# you also need to install the realesrgan package\n",
    "!pip install realesrgan\n",
    "\n",
    "# get model\n",
    "%cd GFPGAN/experiments/pretrained_models\n",
    "!wget https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth\n",
    "%cd ~/../notebooks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a31a6bde",
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd GFPGAN\n",
    "!python inference_gfpgan.py -i ../outputs -o results -v 1.3 -s 2\n",
    "%cd ~/../notebooks\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb4bc651",
   "metadata": {},
   "source": [
    "# CodeFormer\n",
    "\n",
    "'Transformer-based prediction network, named CodeFormer, to model global composition and context of the low-quality faces for code prediction, enabling the discovery of natural faces that closely approximate the target faces even when the inputs are severely degraded.' - Project credit to [Source](https://shangchenzhou.com/projects/CodeFormer/)\n",
    "\n",
    "This alternative to GFPGAN also does a great job of upscaling images. Much like above, you can run the next two cells to upscale the full images from our outputs directory.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e81f5a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install and setup\n",
    "# git clone this repository\n",
    "!git clone https://github.com/sczhou/CodeFormer\n",
    "%cd CodeFormer\n",
    "# install python dependencies\n",
    "!pip install realesrgan basicsr optax lpips\n",
    "!python scripts/download_pretrained_models.py facelib\n",
    "!python scripts/download_pretrained_models.py CodeFormer\n",
    "%cd .."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afc0d67e",
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd CodeFormer\n",
    "!python inference_codeformer.py -w 0.7 --input_path ../outputs\n",
    "%cd .."
   ]
  }
 ],
 "metadata": {
  "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}