CLIP_&_gradient_ascent_for_text_to_image_(Deep_Daze ).ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "CLIP & gradient ascent for text-to-image (Deep Daze?).ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Zfn2m8QCTUO"
      },
      "source": [
        "# Description"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1rQVANPk-2n-"
      },
      "source": [
        "A Colab notebook for generating images using OpenAI's CLIP model.\r\n",
        "\r\n",
        "Heavily influenced by Alexander Mordvintsev's Deep Dream, this work uses CLIP to match an image learned by a SIREN network with a given textual description.\r\n",
        "\r\n",
        "\r\n",
        "As a good launching point for future directions and to find more related work, see https://distill.pub/2017/feature-visualization/\r\n",
        "\r\n",
        "\r\n",
        "If you have questions, please see my twitter at https://twitter.com/advadnoun\r\n",
        "\r\n",
        "This is all free! But if you're feeling generous, you can donate to my venmo @rynnn while your \"a beautiful Waluigi\" loads ;)\r\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KAcixx9Z3XYH"
      },
      "source": [
        "# Top\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VAbC90KCAUVR"
      },
      "source": [
        "Import the usual libraries"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "piJOg9MY7khd",
        "trusted": true
      },
      "source": [
        "import torch\n",
        "import numpy as np\n",
        "import torchvision\n",
        "import torchvision.transforms.functional as TF\n",
        "\n",
        "import PIL\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import os\n",
        "import random\n",
        "import imageio\n",
        "from IPython import display\n",
        "from IPython.core.interactiveshell import InteractiveShell\n",
        "InteractiveShell.ast_node_interactivity = \"all\"\n",
        "\n",
        "import glob\n",
        "\n",
        "from google.colab import output\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4pWYOXvI_AtL"
      },
      "source": [
        "Check which GPU Colab has assigned to you"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EUOAL4UcytFb",
        "outputId": "03019bd4-ce5e-4a9b-daee-d9d9d24e1266"
      },
      "source": [
        "!nvidia-smi -L\r\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "GPU 0: Tesla V100-SXM2-16GB (UUID: GPU-45599a17-8938-faab-73b4-cbdcbf7a005b)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7b4_ZrdU_FdV"
      },
      "source": [
        "# Restart after running this cell!"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ftkN5c3f_MuZ"
      },
      "source": [
        "You must run this cell and then restart and rerun everything for the PyTorch version to be correct. Otherwise the model will run but not produce any meaningful output."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "etzxXVZ_r-Nf",
        "outputId": "9b90e547-4726-4afe-ab87-7e4a22e88701"
      },
      "source": [
        "import subprocess\r\n",
        "\r\n",
        "CUDA_version = [s for s in subprocess.check_output([\"nvcc\", \"--version\"]).decode(\"UTF-8\").split(\", \") if s.startswith(\"release\")][0].split(\" \")[-1]\r\n",
        "print(\"CUDA version:\", CUDA_version)\r\n",
        "\r\n",
        "if CUDA_version == \"10.0\":\r\n",
        "    torch_version_suffix = \"+cu100\"\r\n",
        "elif CUDA_version == \"10.1\":\r\n",
        "    torch_version_suffix = \"+cu101\"\r\n",
        "elif CUDA_version == \"10.2\":\r\n",
        "    torch_version_suffix = \"\"\r\n",
        "else:\r\n",
        "    torch_version_suffix = \"+cu110\"\r\n",
        "\r\n",
        "! pip install torch==1.7.1{torch_version_suffix} torchvision==0.8.2{torch_version_suffix} -f https://download.pytorch.org/whl/torch_stable.html ftfy regex"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "CUDA version: 10.1\n",
            "Looking in links: https://download.pytorch.org/whl/torch_stable.html\n",
            "Requirement already satisfied: torch==1.7.1+cu101 in /usr/local/lib/python3.6/dist-packages (1.7.1+cu101)\n",
            "Requirement already satisfied: torchvision==0.8.2+cu101 in /usr/local/lib/python3.6/dist-packages (0.8.2+cu101)\n",
            "Requirement already satisfied: ftfy in /usr/local/lib/python3.6/dist-packages (5.8)\n",
            "Requirement already satisfied: regex in /usr/local/lib/python3.6/dist-packages (2019.12.20)\n",
            "Requirement already satisfied: dataclasses; python_version < \"3.7\" in /usr/local/lib/python3.6/dist-packages (from torch==1.7.1+cu101) (0.8)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from torch==1.7.1+cu101) (1.19.5)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.6/dist-packages (from torch==1.7.1+cu101) (3.7.4.3)\n",
            "Requirement already satisfied: pillow>=4.1.1 in /usr/local/lib/python3.6/dist-packages (from torchvision==0.8.2+cu101) (7.0.0)\n",
            "Requirement already satisfied: wcwidth in /usr/local/lib/python3.6/dist-packages (from ftfy) (0.2.5)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JeMCHwDdHIcu"
      },
      "source": [
        "# CLIP"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-irecVJr_Zc9"
      },
      "source": [
        "Setup CLIP and set it to be the \"perceptor\" used to determine the loss for the SIREN network.\r\n",
        "\r\n",
        "Thanks to the authors below & OpenAI for sharing! https://github.com/openai/CLIP\r\n",
        "\r\n",
        "Alec Radford \\* Jong Wook Kim \\* Chris Hallacy Aditya Ramesh Gabriel Goh Sandhini Agarwal\r\n",
        "Girish Sastry Amanda Askell Pamela Mishkin Jack Clark Gretchen Krueger\r\n",
        "Ilya Sutskever\r\n",
        "\r\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Tm3_VmxpAiB1",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "9e4bd94b-c601-4716-81e4-d1e2c7550a2a"
      },
      "source": [
        "\n",
        "\n",
        "%cd /content/\n",
        "\n",
        "!git clone https://github.com/openai/CLIP.git\n",
        "\n",
        "\n",
        "%cd /content/CLIP/\n",
        "\n",
        "!pip install ftfy\n",
        "\n",
        "import clip\n",
        "\n",
        "\n",
        "import numpy as np\n",
        "# Load the model\n",
        "perceptor, preprocess = clip.load('ViT-B/32')\n"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "/content\n",
            "Cloning into 'CLIP'...\n",
            "remote: Enumerating objects: 24, done.\u001b[K\n",
            "remote: Total 24 (delta 0), reused 0 (delta 0), pack-reused 24\u001b[K\n",
            "Unpacking objects: 100% (24/24), done.\n",
            "/content/CLIP\n",
            "Requirement already satisfied: ftfy in /usr/local/lib/python3.6/dist-packages (5.8)\n",
            "Requirement already satisfied: wcwidth in /usr/local/lib/python3.6/dist-packages (from ftfy) (0.2.5)\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|███████████████████████| 353976522/353976522 [00:04<00:00, 75406040.59it/s]\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mw0KBLebMywW"
      },
      "source": [
        "# Params"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0WMEM-RK_ra0"
      },
      "source": [
        "Determine the output dimensions of the image and the number of channels.\r\n",
        "\r\n",
        "Set the text to be matched"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Nq0wA-wc-P-s",
        "trusted": true
      },
      "source": [
        "im_shape = [512, 512, 3]\n",
        "sideX, sideY, channels = im_shape\n",
        "\n",
        "tx = clip.tokenize(\"a beautiful Waluigi\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UlIUu0jK3S19"
      },
      "source": [
        "# Define"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XsYea-43_xdZ"
      },
      "source": [
        "Define some helper functions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gCj1wr8pOqOe",
        "trusted": true
      },
      "source": [
        "def displ(img, pre_scaled=True):\n",
        "  img = np.array(img)[:,:,:]\n",
        "  img = np.transpose(img, (1, 2, 0))\n",
        "  if not pre_scaled:\n",
        "    img = scale(img, 48*4, 32*4)\n",
        "  imageio.imwrite(str(3) + '.png', np.array(img))\n",
        "  return display.Image(str(3)+'.png')\n",
        "\n",
        "def card_padded(im, to_pad=3):\n",
        "  return np.pad(np.pad(np.pad(im, [[1,1], [1,1], [0,0]],constant_values=0), [[2,2], [2,2], [0,0]],constant_values=1),\n",
        "            [[to_pad,to_pad], [to_pad,to_pad], [0,0]],constant_values=0)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5odD-rH8t6oz"
      },
      "source": [
        "# SIREN"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tcGF8UUG_5Zi"
      },
      "source": [
        "Thanks to the authors of SIREN! https://github.com/vsitzmann/siren\r\n",
        "\r\n",
        "@inproceedings{sitzmann2019siren,\r\n",
        "    author = {Sitzmann, Vincent\r\n",
        "              and Martel, Julien N.P.\r\n",
        "              and Bergman, Alexander W.\r\n",
        "              and Lindell, David B.\r\n",
        "              and Wetzstein, Gordon},\r\n",
        "    title = {Implicit Neural Representations\r\n",
        "              with Periodic Activation Functions},\r\n",
        "    booktitle = {arXiv},\r\n",
        "    year={2020}\r\n",
        "}\r\n",
        "\r\n",
        "\r\n",
        "The number of layers is 8 right now, but if the machine OOMs (runs out of RAM), it can naturally be tweaked. I've found that 16 layers for the SIREN works best, but I'm not always blessed with a V100 GPU."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GdCh2D8Dt8Xd"
      },
      "source": [
        "\r\n",
        "import torch.nn as nn\r\n",
        "\r\n",
        "class SineLayer(nn.Module):\r\n",
        "    # See paper sec. 3.2, final paragraph, and supplement Sec. 1.5 for discussion of omega_0.\r\n",
        "    \r\n",
        "    # If is_first=True, omega_0 is a frequency factor which simply multiplies the activations before the \r\n",
        "    # nonlinearity. Different signals may require different omega_0 in the first layer - this is a \r\n",
        "    # hyperparameter.\r\n",
        "    \r\n",
        "    # If is_first=False, then the weights will be divided by omega_0 so as to keep the magnitude of \r\n",
        "    # activations constant, but boost gradients to the weight matrix (see supplement Sec. 1.5)\r\n",
        "    \r\n",
        "    def __init__(self, in_features, out_features, bias=True,\r\n",
        "                 is_first=False, omega_0=30):\r\n",
        "        super().__init__()\r\n",
        "        self.omega_0 = omega_0\r\n",
        "        self.is_first = is_first\r\n",
        "        \r\n",
        "        self.in_features = in_features\r\n",
        "        self.linear = nn.Linear(in_features, out_features, bias=bias)\r\n",
        "        \r\n",
        "        self.init_weights()\r\n",
        "    \r\n",
        "    def init_weights(self):\r\n",
        "        with torch.no_grad():\r\n",
        "            if self.is_first:\r\n",
        "                self.linear.weight.uniform_(-1 / self.in_features, \r\n",
        "                                             1 / self.in_features)      \r\n",
        "            else:\r\n",
        "                self.linear.weight.uniform_(-np.sqrt(6 / self.in_features) / self.omega_0, \r\n",
        "                                             np.sqrt(6 / self.in_features) / self.omega_0)\r\n",
        "        \r\n",
        "    def forward(self, input):\r\n",
        "        return torch.sin(self.omega_0 * self.linear(input))\r\n",
        "    \r\n",
        "    def forward_with_intermediate(self, input): \r\n",
        "        # For visualization of activation distributions\r\n",
        "        intermediate = self.omega_0 * self.linear(input)\r\n",
        "        return torch.sin(intermediate), intermediate\r\n",
        "    \r\n",
        "    \r\n",
        "class Siren(nn.Module):\r\n",
        "    def __init__(self, in_features, hidden_features, hidden_layers, out_features, outermost_linear=True, \r\n",
        "                 first_omega_0=30, hidden_omega_0=30.):\r\n",
        "        super().__init__()\r\n",
        "        \r\n",
        "        self.net = []\r\n",
        "        self.net.append(SineLayer(in_features, hidden_features, \r\n",
        "                                  is_first=True, omega_0=first_omega_0))\r\n",
        "\r\n",
        "        for i in range(hidden_layers):\r\n",
        "            self.net.append(SineLayer(hidden_features, hidden_features, \r\n",
        "                                      is_first=False, omega_0=hidden_omega_0))\r\n",
        "\r\n",
        "        if outermost_linear:\r\n",
        "            final_linear = nn.Linear(hidden_features, out_features)\r\n",
        "            \r\n",
        "            with torch.no_grad():\r\n",
        "                final_linear.weight.uniform_(-np.sqrt(6 / hidden_features) / hidden_omega_0, \r\n",
        "                                              np.sqrt(6 / hidden_features) / hidden_omega_0)\r\n",
        "                \r\n",
        "            self.net.append(final_linear)\r\n",
        "        else:\r\n",
        "            self.net.append(SineLayer(hidden_features, out_features, \r\n",
        "                                      is_first=False, omega_0=hidden_omega_0))\r\n",
        "        \r\n",
        "        self.net = nn.Sequential(*self.net)\r\n",
        "    \r\n",
        "    def forward(self, coords):\r\n",
        "        coords = coords.clone().detach().requires_grad_(True)\r\n",
        "        output = self.net(coords.cuda())\r\n",
        "        return output.view(1, sideX, sideY, 3).permute(0, 3, 1, 2)#.sigmoid_()\r\n",
        "\r\n",
        "    def forward_with_activations(self, coords, retain_grad=False):\r\n",
        "        '''Returns not only model output, but also intermediate activations.\r\n",
        "        Only used for visualizing activations later!'''\r\n",
        "        activations = OrderedDict()\r\n",
        "\r\n",
        "        activation_count = 0\r\n",
        "        x = coords.clone().detach().requires_grad_(True)\r\n",
        "        activations['input'] = x\r\n",
        "        for i, layer in enumerate(self.net):\r\n",
        "            if isinstance(layer, SineLayer):\r\n",
        "                x, intermed = layer.forward_with_intermediate(x)\r\n",
        "                \r\n",
        "                if retain_grad:\r\n",
        "                    x.retain_grad()\r\n",
        "                    intermed.retain_grad()\r\n",
        "                    \r\n",
        "                activations['_'.join((str(layer.__class__), \"%d\" % activation_count))] = intermed\r\n",
        "                activation_count += 1\r\n",
        "            else: \r\n",
        "                x = layer(x)\r\n",
        "                \r\n",
        "                if retain_grad:\r\n",
        "                    x.retain_grad()\r\n",
        "                    \r\n",
        "            activations['_'.join((str(layer.__class__), \"%d\" % activation_count))] = x\r\n",
        "            activation_count += 1\r\n",
        "\r\n",
        "        return activations\r\n",
        "\r\n",
        "\r\n",
        "def get_mgrid(sidelen, dim=2):\r\n",
        "    '''Generates a flattened grid of (x,y,...) coordinates in a range of -1 to 1.\r\n",
        "    sidelen: int\r\n",
        "    dim: int'''\r\n",
        "    tensors = tuple(dim * [torch.linspace(-1, 1, steps=sidelen)])\r\n",
        "    mgrid = torch.stack(torch.meshgrid(*tensors), dim=-1)\r\n",
        "    mgrid = mgrid.reshape(-1, dim)\r\n",
        "    return mgrid\r\n",
        "\r\n",
        "\r\n",
        "model = Siren(2, 256, 16, 3).cuda()\r\n",
        "LLL = []\r\n",
        "eps = 0\r\n",
        "\r\n",
        "optimizer = torch.optim.Adam(model.parameters(), .00001)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Q9lHgtYLrUn9"
      },
      "source": [
        "make sure you're on this page when train first starts running, and it dings\r\n",
        "otherwise the thing'll never start"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WztSrRF23Rqg"
      },
      "source": [
        "# Train"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nbeZVdA_AGK7"
      },
      "source": [
        "Train and output samples every 150 iterations\r\n",
        "\r\n",
        "We create batches of images at different resolutions in different parts of the SIREN image and resize them with bilinear upsampling. This seems to work very, very well as regularization for visualizing networks with larger images than their usual input resolution."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7EuUz-ICNKUr",
        "cellView": "both",
        "trusted": true
      },
      "source": [
        "\n",
        "\n",
        "\n",
        "def checkin(loss):\n",
        "  print(loss)\n",
        "  with torch.no_grad():\n",
        "    al = nom(model(get_mgrid(sideX)).cpu()).numpy()\n",
        "  for allls in al:\n",
        "    displ(allls)\n",
        "    display.display(display.Image(str(3)+'.png'))\n",
        "    print('\\n')\n",
        "  output.eval_js('new Audio(\"https://freesound.org/data/previews/80/80921_1022651-lq.ogg\").play()')\n",
        "\n",
        "\n",
        "def ascend_txt():\n",
        "  out = model(get_mgrid(sideX))\n",
        "\n",
        "  cutn = 64\n",
        "  p_s = []\n",
        "  for ch in range(cutn):\n",
        "    size = torch.randint(int(.5*sideX), int(.98*sideX), ())\n",
        "    offsetx = torch.randint(0, sideX - size, ())\n",
        "    offsety = torch.randint(0, sideX - size, ())\n",
        "    apper = out[:, :, offsetx:offsetx + size, offsety:offsety + size]\n",
        "    apper = torch.nn.functional.interpolate(apper, (224,224), mode='bilinear')\n",
        "    p_s.append(nom(apper))\n",
        "  into = torch.cat(p_s, 0)\n",
        "\n",
        "\n",
        "  \n",
        "  iii = perceptor.encode_image(into)\n",
        "  t = perceptor.encode_text(tx.cuda())\n",
        "  return -100*torch.cosine_similarity(t, iii, dim=-1).mean()\n",
        "\n",
        "def train(epoch, i):\n",
        "  loss = ascend_txt()\n",
        "  optimizer.zero_grad()\n",
        "  loss.backward()\n",
        "  optimizer.step()\n",
        "  \n",
        "\n",
        "  \n",
        "  \n",
        "  if itt % 150 == 0:\n",
        "    checkin(loss)\n",
        "\n",
        "nom = torchvision.transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))\n",
        "\n",
        "\n",
        "itt = 0\n",
        "for epochs in range(10000):\n",
        "  for i in range(1000):\n",
        "    train(eps, i)\n",
        "    itt+=1\n",
        "  eps+=1\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YQ3AK_iIjwaw"
      },
      "source": [
        "# Bot"
      ]
    }
  ]
}