Stable_Diffusion_tokenizer_and_embedding.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YPHN7PJgKOzb"
      },
      "source": [
        "# Stable Diffusion v1 tokenizer and embedding\n",
        "\n",
        "This notebook examines tokens and embeddings used in Stable Diffusion v1.\n",
        "\n",
        "Modified from [Interacting with CLIP notebook](https://colab.research.google.com/github/openai/clip/blob/master/notebooks/Interacting_with_CLIP.ipynb)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "53N4k0pj_9qL"
      },
      "source": [
        "# Setup\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0BpdJkdBssk9",
        "outputId": "d07e78ed-89a8-4f15-fa34-68b4b18b780b"
      },
      "source": [
        "! pip install ftfy regex tqdm\n",
        "! pip install git+https://github.com/openai/CLIP.git\n",
        "\n",
        "import numpy as np\n",
        "import torch\n",
        "from pkg_resources import packaging\n",
        "print(\"Torch version:\", torch.__version__)\n",
        "\n",
        "import clip\n",
        "print('Available models:')\n",
        "print(clip.available_models())\n",
        "\n",
        "model, preprocess = clip.load(\"ViT-L/14\") # used by stable diffusion v1\n",
        "model.cuda().eval()\n",
        "input_resolution = model.visual.input_resolution\n",
        "context_length = model.context_length\n",
        "vocab_size = model.vocab_size\n",
        "\n",
        "print(\"Model parameters:\", f\"{np.sum([int(np.prod(p.shape)) for p in model.parameters()]):,}\")\n",
        "print(\"Input resolution:\", input_resolution)\n",
        "print(\"Context length:\", context_length)\n",
        "print(\"Vocab size:\", vocab_size)\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting ftfy\n",
            "  Downloading ftfy-6.1.1-py3-none-any.whl (53 kB)\n",
            "\u001b[K     |████████████████████████████████| 53 kB 1.9 MB/s \n",
            "\u001b[?25hRequirement already satisfied: regex in /usr/local/lib/python3.8/dist-packages (2022.6.2)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.8/dist-packages (4.64.1)\n",
            "Requirement already satisfied: wcwidth>=0.2.5 in /usr/local/lib/python3.8/dist-packages (from ftfy) (0.2.5)\n",
            "Installing collected packages: ftfy\n",
            "Successfully installed ftfy-6.1.1\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Collecting git+https://github.com/openai/CLIP.git\n",
            "  Cloning https://github.com/openai/CLIP.git to /tmp/pip-req-build-d9d9s9wi\n",
            "  Running command git clone -q https://github.com/openai/CLIP.git /tmp/pip-req-build-d9d9s9wi\n",
            "Requirement already satisfied: ftfy in /usr/local/lib/python3.8/dist-packages (from clip==1.0) (6.1.1)\n",
            "Requirement already satisfied: regex in /usr/local/lib/python3.8/dist-packages (from clip==1.0) (2022.6.2)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.8/dist-packages (from clip==1.0) (4.64.1)\n",
            "Requirement already satisfied: torch in /usr/local/lib/python3.8/dist-packages (from clip==1.0) (1.13.0+cu116)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.8/dist-packages (from clip==1.0) (0.14.0+cu116)\n",
            "Requirement already satisfied: wcwidth>=0.2.5 in /usr/local/lib/python3.8/dist-packages (from ftfy->clip==1.0) (0.2.5)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.8/dist-packages (from torch->clip==1.0) (4.4.0)\n",
            "Requirement already satisfied: pillow!=8.3.*,>=5.3.0 in /usr/local/lib/python3.8/dist-packages (from torchvision->clip==1.0) (7.1.2)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.8/dist-packages (from torchvision->clip==1.0) (2.23.0)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.8/dist-packages (from torchvision->clip==1.0) (1.21.6)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.8/dist-packages (from requests->torchvision->clip==1.0) (2022.12.7)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.8/dist-packages (from requests->torchvision->clip==1.0) (2.10)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.8/dist-packages (from requests->torchvision->clip==1.0) (1.24.3)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.8/dist-packages (from requests->torchvision->clip==1.0) (3.0.4)\n",
            "Building wheels for collected packages: clip\n",
            "  Building wheel for clip (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for clip: filename=clip-1.0-py3-none-any.whl size=1369408 sha256=83353fc7530428ea6d36bc33364ac60757bb68cca0cf06956739efe2ef22ca6f\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-4j8wdad4/wheels/ab/4f/3a/5e51521b55997aa6f0690e095c08824219753128ce8d9969a3\n",
            "Successfully built clip\n",
            "Installing collected packages: clip\n",
            "Successfully installed clip-1.0\n",
            "Torch version: 1.13.0+cu116\n",
            "Available models:\n",
            "['RN50', 'RN101', 'RN50x4', 'RN50x16', 'RN50x64', 'ViT-B/32', 'ViT-B/16', 'ViT-L/14', 'ViT-L/14@336px']\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|███████████████████████████████████████| 890M/890M [00:25<00:00, 37.0MiB/s]\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Model parameters: 427,616,513\n",
            "Input resolution: 224\n",
            "Context length: 77\n",
            "Vocab size: 49408\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xwSB5jZki3Cj"
      },
      "source": [
        "# Token and embedding\n",
        "Modify prompt to see tokens and embeddings"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "qGom156-i2kL",
        "outputId": "4434b86a-536d-4300-873e-943d7de7be79"
      },
      "source": [
        "# modify prompt to check tokens\n",
        "prompt = \"Photo of a cat\"\n",
        "\n",
        "tokens = clip.tokenize(prompt)\n",
        "with torch.no_grad():\n",
        "    embeddings = model.encode_text(tokens.cuda()).float()\n",
        "print(\"text tokens:\")\n",
        "print(tokens)\n",
        "print(\"text tokens size:\", tokens.shape)\n",
        "print(\"Embeddings size:\", embeddings.shape )"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "text tokens:\n",
            "tensor([[49406,  1125,   539,   320,  2368, 49407,     0,     0,     0,     0,\n",
            "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
            "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
            "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
            "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
            "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
            "             0,     0,     0,     0,     0,     0,     0,     0,     0,     0,\n",
            "             0,     0,     0,     0,     0,     0,     0]], dtype=torch.int32)\n",
            "text tokens size: torch.Size([1, 77])\n",
            "Embeddings size: torch.Size([1, 768])\n"
          ]
        }
      ]
    }
  ]
}