arnheim_1.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0kaSrjMifU0e"
      },
      "source": [
        "Copyright 2021 DeepMind Technologies Limited\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "    https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lFoZUnRBb1ag"
      },
      "source": [
        "#Generative Art Using Neural Visual Grammars and Dual Encoders\n",
        "\n",
        "**Chrisantha Fernando, Piotr Mirowski, Dylan Banarse, S. M. Ali Eslami, Jean-Baptiste Alayrac, Simon Osindero**\n",
        "\n",
        "DeepMind, 2021\n",
        "\n",
        "##Arnheim 1\n",
        "###Generate paintings from text prompts.\n",
        "\n",
        " Whilst there are perhaps only a few scientific methods, there seem to be almost as many artistic methods as there are artists. Artistic processes appear to inhabit the highest order of open-endedness. To begin to understand some of the processes of art making it is helpful to try to automate them even partially.\n",
        "In this paper, a novel algorithm for producing generative art is described which allows a user to input a text string, and which in a creative response to this string, outputs an image which interprets that string. It does so by evolving images using a hierarchical neural [Lindenmeyer system](https://en.wikipedia.org/wiki/L-system), and evaluating these images along the way using an image text dual encoder trained on billions of images and their associated text from the internet.\n",
        "In doing so we have access to and control over an instance of an artistic process, allowing analysis of which aspects of the artistic process become the task of the algorithm, and which elements remain the responsibility of the artist.\n",
        "\n",
        "This colab accompanies the paper [Generative Art Using Neural Visual Grammars and Dual Encoders](https://arxiv.org/abs/2105.00162)\n",
        "\n",
        "##Instructions\n",
        "\n",
        "1. Click \"Connect\" button in the top right corner of this Colab\n",
        "1. Select Runtime -\u003e Change runtime type -\u003e Hardware accelerator -\u003e GPU\n",
        "1. Select High-RAM for \"Runtime shape\" option\n",
        "1. Navigate to \"Get text input\"\n",
        "1. Enter text for IMAGE_NAME\n",
        "1. Select \"Run All\" from Runtime menu\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZysZip8gE24t"
      },
      "source": [
        "# Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Qcig_xWLipmI"
      },
      "outputs": [],
      "source": [
        "#@title Set CUDA version for PyTorch\n",
        "\n",
        "import subprocess\n",
        "\n",
        "CUDA_version = [s for s in subprocess.check_output([\"nvcc\", \"--version\"]\n",
        "                                                   ).decode(\"UTF-8\").split(\", \")\n",
        "                if s.startswith(\"release\")][0].split(\" \")[-1]\n",
        "print(\"CUDA version:\", CUDA_version)\n",
        "\n",
        "if CUDA_version == \"10.0\":\n",
        "  torch_version_suffix = \"+cu100\"\n",
        "elif CUDA_version == \"10.1\":\n",
        "  torch_version_suffix = \"+cu101\"\n",
        "elif CUDA_version == \"10.2\":\n",
        "  torch_version_suffix = \"\"\n",
        "else:\n",
        "  torch_version_suffix = \"+cu110\"\n",
        "! nvidia-smi"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "blPNV7eKjWNr"
      },
      "outputs": [],
      "source": [
        "#@title Install and import PyTorch and Clip\n",
        "\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\n",
        "! pip install git+https://github.com/openai/CLIP.git --no-deps\n",
        "! pip install ftfy regex\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import clip\n",
        "print(\"Torch version:\", torch.__version__)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "7JPKvjxH902m"
      },
      "outputs": [],
      "source": [
        "#@title Install and import ray multiprocessing\n",
        "\n",
        "! pip install -q -U ray[default]\n",
        "import ray"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "A-dXbcBgi-59"
      },
      "outputs": [],
      "source": [
        "#@title Import all other needed libraries\n",
        "\n",
        "import collections\n",
        "import copy\n",
        "import cloudpickle\n",
        "import time\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import math\n",
        "from PIL import Image\n",
        "from PIL import ImageDraw\n",
        "from skimage import transform"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "PHTI5tQekl4S"
      },
      "outputs": [],
      "source": [
        "#@title Load CLIP {vertical-output: true}\n",
        "\n",
        "CLIP_MODEL = \"ViT-B/32\"\n",
        "device = torch.device(\"cuda\")\n",
        "print(f\"Downloading CLIP model {CLIP_MODEL}...\")\n",
        "model, _ = clip.load(CLIP_MODEL, device, jit=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mNH-R_bs6zFZ"
      },
      "source": [
        "# Neural Visual Grammar"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rvo1d2z_Hv1B"
      },
      "source": [
        "### Drawing primitives"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "93HEhxx-7L5S"
      },
      "outputs": [],
      "source": [
        "def to_homogeneous(p):\n",
        "  r, c = p\n",
        "  return np.stack((r, c, np.ones_like(p[0])), axis=0)\n",
        "\n",
        "def from_homogeneous(p):\n",
        "  p = p / p.T[:, 2]\n",
        "  return p[0].astype(\"int32\"), p[1].astype(\"int32\")\n",
        "\n",
        "def apply_scale(scale, lineh):\n",
        "  return np.stack([lineh[0, :] * scale,\n",
        "                   lineh[1, :] * scale,\n",
        "                   lineh[2, :]])\n",
        "\n",
        "def apply_translation(translation, lineh, offset_r=0, offset_c=0):\n",
        "  r, c = translation\n",
        "  return np.stack([lineh[0, :] + c + offset_c,\n",
        "                   lineh[1, :] + r + offset_r,\n",
        "                   lineh[2, :]])\n",
        "\n",
        "def apply_rotation(translation, rad, lineh):\n",
        "  r, c = translation\n",
        "  cos_rad = np.cos(rad)\n",
        "  sin_rad = np.sin(rad)\n",
        "  return np.stack(\n",
        "      [(lineh[0, :] - c) * cos_rad - (lineh[1, :] - r) * sin_rad + c,\n",
        "       (lineh[0, :] - c) * sin_rad + (lineh[1, :] - r) * cos_rad + r,\n",
        "       lineh[2, :]])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vJkqCvbHmkpe"
      },
      "outputs": [],
      "source": [
        "def transform_lines(line_from, line_to, translation, angle, scale,\n",
        "                    translation2, angle2, scale2, img_siz2):\n",
        "  \"\"\"Transform lines by translation, angle and scale, twice.\n",
        "\n",
        "  Args:\n",
        "    line_from: Line start point.\n",
        "    line_to: Line end point.\n",
        "    translation: 1st translation to line.\n",
        "    angle: 1st angle of rotation for line.\n",
        "    scale: 1st scale for line.\n",
        "    translation2: 2nd translation to line.\n",
        "    angle2: 2nd angle of rotation for line.\n",
        "    scale2: 2nd scale for line.\n",
        "    img_siz2: Offset for 2nd translation.\n",
        "\n",
        "  Returns:\n",
        "    Transformed lines.\n",
        "  \"\"\"\n",
        "  if len(line_from.shape) == 1:\n",
        "    line_from = np.expand_dims(line_from, 0)\n",
        "  if len(line_to.shape) == 1:\n",
        "    line_to = np.expand_dims(line_to, 0)\n",
        "\n",
        "  # First transform.\n",
        "  line_from_h = to_homogeneous(line_from.T)\n",
        "  line_to_h = to_homogeneous(line_to.T)\n",
        "  line_from_h = apply_scale(scale, line_from_h)\n",
        "  line_to_h = apply_scale(scale, line_to_h)\n",
        "  translated_line_from = apply_translation(translation, line_from_h)\n",
        "  translated_line_to = apply_translation(translation, line_to_h)\n",
        "  translated_mid_point = (translated_line_from + translated_line_to) / 2.0\n",
        "  translated_mid_point = translated_mid_point[[1, 0]]\n",
        "  line_from_transformed = apply_rotation(translated_mid_point,\n",
        "                                         np.pi * angle,\n",
        "                                         translated_line_from)\n",
        "  line_to_transformed = apply_rotation(translated_mid_point,\n",
        "                                       np.pi * angle,\n",
        "                                       translated_line_to)\n",
        "  line_from_transformed = np.array(from_homogeneous(line_from_transformed))\n",
        "  line_to_transformed = np.array(from_homogeneous(line_to_transformed))\n",
        "\n",
        "  # Second transform.\n",
        "  line_from_h = to_homogeneous(line_from_transformed)\n",
        "  line_to_h = to_homogeneous(line_to_transformed)\n",
        "  line_from_h = apply_scale(scale2, line_from_h)\n",
        "  line_to_h = apply_scale(scale2, line_to_h)\n",
        "  translated_line_from = apply_translation(\n",
        "      translation2, line_from_h, offset_r=img_siz2, offset_c=img_siz2)\n",
        "  translated_line_to = apply_translation(\n",
        "      translation2, line_to_h, offset_r=img_siz2, offset_c=img_siz2)\n",
        "  translated_mid_point = (translated_line_from + translated_line_to) / 2.0\n",
        "  translated_mid_point = translated_mid_point[[1, 0]]\n",
        "  line_from_transformed = apply_rotation(translated_mid_point,\n",
        "                                         np.pi * angle2,\n",
        "                                         translated_line_from)\n",
        "  line_to_transformed = apply_rotation(translated_mid_point,\n",
        "                                       np.pi * angle2,\n",
        "                                       translated_line_to)\n",
        "  return np.concatenate([from_homogeneous(line_from_transformed),\n",
        "                         from_homogeneous(line_to_transformed)],\n",
        "                        axis=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ElxgpBwkJuFj"
      },
      "source": [
        "### Hierarchical stroke painting functions"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sUtr0-pAJ6y2"
      },
      "outputs": [],
      "source": [
        "# PaintingCommand\n",
        "#   origin_top: Origin of line defined by top level LSTM\n",
        "#   angle_top: Angle of line defined by top level LSTM\n",
        "#   scale_top: Scale for line defined by top level LSTM\n",
        "#   origin_bottom: Origin of line defined by bottom level LSTM\n",
        "#   angle_bottom: Angle of line defined by bottom level LSTM\n",
        "#   scale_bottom: Scale for line defined by bottom level LSTM\n",
        "#   position_choice: Selects between use of:\n",
        "#     Origin, angle and scale from both LSTM levels\n",
        "#     Origin, angle and scale just from top level LSTM\n",
        "#     Origin, angle and scale just from bottom level LSTM\n",
        "#   transparency: Line transparency determined by bottom level LSTM\n",
        "PaintingCommand = collections.namedtuple(\"PaintingCommand\",\n",
        "                                         [\"origin_top\",\n",
        "                                          \"angle_top\",\n",
        "                                          \"scale_top\",\n",
        "                                          \"origin_bottom\",\n",
        "                                          \"angle_bottom\",\n",
        "                                          \"scale_bottom\",\n",
        "                                          \"position_choice\",\n",
        "                                          \"transparency\"])\n",
        "\n",
        "def paint_over_image(img, strokes, painting_commands,\n",
        "                     allow_strokes_beyond_image_edges, coeff_size=1):\n",
        "  \"\"\"Make marks over an existing image.\n",
        "\n",
        "  Args:\n",
        "    img: Image to draw on.\n",
        "    strokes: Stroke descriptions.\n",
        "    painting_commands: Top-level painting commands with transforms for the i\n",
        "      sets of strokes.\n",
        "    allow_strokes_beyond_image_edges: Allow strokes beyond image boundary.\n",
        "    coeff_size: Determines low res (1) or high res (10) image will be drawn.\n",
        "\n",
        "  Returns:\n",
        "    num_strokes: The number of strokes made.\n",
        "  \"\"\"\n",
        "  img_center = 112. * coeff_size\n",
        "  # a, b and c: determines the stroke width distribution (see 'weights' below)\n",
        "  a = 10. * coeff_size\n",
        "  b = 2. * coeff_size\n",
        "  c = 300. * coeff_size\n",
        "  # d: extent that the strokes are allowed to go beyond the edge of the canvas\n",
        "  d = 223 * coeff_size\n",
        "\n",
        "  def _clip_colour(col):\n",
        "    return np.clip((np.round(col * 255. + 128.)).astype(np.int32), 0, 255)\n",
        "\n",
        "  # Loop over all the top level...\n",
        "  t0_over = time.time()\n",
        "  num_strokes = sum(len(s) for s in strokes)\n",
        "  translations = np.zeros((2, num_strokes,), np.float32)\n",
        "  translations2 = np.zeros((2, num_strokes,), np.float32)\n",
        "  angles = np.zeros((num_strokes,), np.float32)\n",
        "  angles2 = np.zeros((num_strokes,), np.float32)\n",
        "  scales = np.zeros((num_strokes,), np.float32)\n",
        "  scales2 = np.zeros((num_strokes,), np.float32)\n",
        "  weights = np.zeros((num_strokes,), np.float32)\n",
        "  lines_from = np.zeros((num_strokes, 2), np.float32)\n",
        "  lines_to = np.zeros((num_strokes, 2), np.float32)\n",
        "  rgbas = np.zeros((num_strokes, 4), np.float32)\n",
        "  k = 0\n",
        "  for i in range(len(strokes)):\n",
        "\n",
        "    # Get the top-level transforms for the i-th bunch of strokes\n",
        "    painting_comand = painting_commands[i]\n",
        "    translation_a = painting_comand.origin_top\n",
        "    angle_a = (painting_comand.angle_top + 1) / 5.0\n",
        "    scale_a = 0.5 + (painting_comand.scale_top + 1) / 3.0\n",
        "    translation_b = painting_comand.origin_bottom\n",
        "    angle_b = (painting_comand.angle_bottom + 1) / 5.0\n",
        "    scale_b = 0.5 + (painting_comand.scale_bottom + 1) / 3.0\n",
        "    position_choice = painting_comand.position_choice\n",
        "    solid_colour = painting_comand.transparency\n",
        "\n",
        "    # Do we use origin, angle and scale from both, top or bottom LSTM levels?\n",
        "    if position_choice \u003e 0.33:\n",
        "      translation = translation_a\n",
        "      angle = angle_a\n",
        "      scale = scale_a\n",
        "      translation2 = translation_b\n",
        "      angle2 = angle_b\n",
        "      scale2 = scale_b\n",
        "    elif position_choice \u003e -0.33:\n",
        "      translation = translation_a\n",
        "      angle = angle_a\n",
        "      scale = scale_a\n",
        "      translation2 = [-img_center, -img_center]\n",
        "      angle2 = 0.\n",
        "      scale2 = 1.\n",
        "    else:\n",
        "      translation = translation_b\n",
        "      angle = angle_b\n",
        "      scale = scale_b\n",
        "      translation2 = [-img_center, -img_center]\n",
        "      angle2 = 0.\n",
        "      scale2 = 1.\n",
        "\n",
        "    # Store top-level transforms\n",
        "    strokes_i = strokes[i]\n",
        "    n_i = len(strokes_i)\n",
        "    angles[k:(k+n_i)] = angle\n",
        "    angles2[k:(k+n_i)] = angle2\n",
        "    scales[k:(k+n_i)] = scale\n",
        "    scales2[k:(k+n_i)] = scale2\n",
        "    translations[0, k:(k+n_i)] = translation[0]\n",
        "    translations[1, k:(k+n_i)] = translation[1]\n",
        "    translations2[0, k:(k+n_i)] = translation2[0]\n",
        "    translations2[1, k:(k+n_i)] = translation2[1]\n",
        "\n",
        "    # ... and the bottom level stroke definitions.\n",
        "    for j in range(n_i):\n",
        "      z_ij = strokes_i[j]\n",
        "\n",
        "      # Store line weight (we will process micro-strokes later)\n",
        "      weights[k] = z_ij[4]\n",
        "      # Store line endpoints\n",
        "      lines_from[k, :] = (z_ij[0], z_ij[1])\n",
        "      lines_to[k, :] = (z_ij[2], z_ij[3])\n",
        "\n",
        "      # Store colour and alpha\n",
        "      rgbas[k, 0] = z_ij[7]\n",
        "      rgbas[k, 1] = z_ij[8]\n",
        "      rgbas[k, 2] = z_ij[9]\n",
        "      if solid_colour \u003e -0.5:\n",
        "        rgbas[k, 3] = 25.5\n",
        "      else:\n",
        "        rgbas[k, 3] = z_ij[11]\n",
        "      k += 1\n",
        "\n",
        "  # Draw all the strokes in a batch as sequence of length 2 * num_strokes\n",
        "  t1_over = time.time()\n",
        "  lines_from *= img_center/2.0\n",
        "  lines_to *= img_center/2.0\n",
        "  rr, cc = transform_lines(lines_from, lines_to, translations, angles, scales,\n",
        "                           translations2, angles2, scales2, img_center)\n",
        "  if not allow_strokes_beyond_image_edges:\n",
        "    rrm = np.round(np.clip(rr, 1, d-1)).astype(int)\n",
        "    ccm = np.round(np.clip(cc, 1, d-1)).astype(int)\n",
        "  else:\n",
        "    rrm = np.round(rr).astype(int)\n",
        "    ccm = np.round(cc).astype(int)\n",
        "\n",
        "  # Plot all the strokes\n",
        "  t2_over = time.time()\n",
        "  img_pil = Image.fromarray(img)\n",
        "  canvas = ImageDraw.Draw(img_pil, \"RGBA\")\n",
        "  rgbas[:, :3] = _clip_colour(rgbas[:, :3])\n",
        "  rgbas[:, 3] = (np.clip(5.0 * np.abs(rgbas[:, 3]), 0, 255)).astype(np.int32)\n",
        "  weights = (np.clip(np.round(weights * b + a), 2, c)).astype(np.int32)\n",
        "  for k in range(num_strokes):\n",
        "    canvas.line((rrm[k], ccm[k], rrm[k+num_strokes], ccm[k+num_strokes]),\n",
        "                fill=tuple(rgbas[k]), width=weights[k])\n",
        "  img[:] = np.asarray(img_pil)[:]\n",
        "  t3_over = time.time()\n",
        "  if VERBOSE_CODE:\n",
        "    print(\"{:.2f}s to store {} stroke defs, {:.4f}s to \"\n",
        "          \"compute them, {:.4f}s to plot them\".format(\n",
        "              t1_over - t0_over, num_strokes, t2_over - t1_over,\n",
        "              t3_over - t2_over))\n",
        "  return num_strokes\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D6Yvwl6f6r0i"
      },
      "source": [
        "### Recurrent Neural Network Layer Generator"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4zECYk7v7oZw"
      },
      "outputs": [],
      "source": [
        "# DrawingLSTMSpec - parameters defining the LSTM architecture\n",
        "#   input_spec_size: Size if sequence elements\n",
        "#   num_lstms: Number of LSTMs at each layer\n",
        "#   net_lstm_hiddens: Number of hidden LSTM units\n",
        "#   net_mlp_hiddens: Number of hidden units in MLP layer\n",
        "DrawingLSTMSpec = collections.namedtuple(\"DrawingLSTMSpec\",\n",
        "                                         [\"input_spec_size\",\n",
        "                                          \"num_lstms\",\n",
        "                                          \"net_lstm_hiddens\",\n",
        "                                          \"net_mlp_hiddens\"])\n",
        "\n",
        "\n",
        "class MakeGeneratorLstm(nn.Module):\n",
        "  \"\"\"Block of parallel LSTMs with MLP output heads.\"\"\"\n",
        "\n",
        "  def __init__(self, drawing_lstm_spec, output_size):\n",
        "    \"\"\"Build drawing LSTM architecture using spec.\n",
        "\n",
        "    Args:\n",
        "      drawing_lstm_spec: DrawingLSTMSpec with architecture parameters\n",
        "      output_size: Number of outputs for the MLP head layer\n",
        "    \"\"\"\n",
        "    super(MakeGeneratorLstm, self).__init__()\n",
        "    self._num_lstms = drawing_lstm_spec.num_lstms\n",
        "    self._input_layer = nn.Sequential(\n",
        "        nn.Linear(drawing_lstm_spec.input_spec_size,\n",
        "                  drawing_lstm_spec.net_lstm_hiddens),\n",
        "        torch.nn.LeakyReLU(0.2, inplace=True))\n",
        "    lstms = []\n",
        "    heads = []\n",
        "    for _ in range(self._num_lstms):\n",
        "      lstm_layer = nn.LSTM(\n",
        "          input_size=drawing_lstm_spec.net_lstm_hiddens,\n",
        "          hidden_size=drawing_lstm_spec.net_lstm_hiddens,\n",
        "          num_layers=2, batch_first=True, bias=True)\n",
        "      head_layer = nn.Sequential(\n",
        "          nn.Linear(drawing_lstm_spec.net_lstm_hiddens,\n",
        "                    drawing_lstm_spec.net_mlp_hiddens),\n",
        "          torch.nn.LeakyReLU(0.2, inplace=True),\n",
        "          nn.Linear(drawing_lstm_spec.net_mlp_hiddens, output_size))\n",
        "      lstms.append(lstm_layer)\n",
        "      heads.append(head_layer)\n",
        "    self._lstms = nn.ModuleList(lstms)\n",
        "    self._heads = nn.ModuleList(heads)\n",
        "\n",
        "  def forward(self, x):\n",
        "    pred = []\n",
        "    x = self._input_layer(x)*10.0\n",
        "    for i in range(self._num_lstms):\n",
        "      y, _ = self._lstms[i](x)\n",
        "      y = self._heads[i](y)\n",
        "      pred.append(y)\n",
        "    return pred"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H0DsS6D87jGM"
      },
      "source": [
        "### DrawingLSTM - A Drawing Recurrent Neural Network"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RzEGVkiFjV-T"
      },
      "outputs": [],
      "source": [
        "Genotype = collections.namedtuple(\"Genotype\",\n",
        "                                  [\"top_lstm\",\n",
        "                                   \"bottom_lstm\",\n",
        "                                   \"input_sequence\",\n",
        "                                   \"initial_img\"])\n",
        "\n",
        "class DrawingLSTM:\n",
        "  \"\"\"LSTM for processing input sequences and generating resultant drawings.\n",
        "\n",
        "  Comprised of two LSTM layers.\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self, drawing_lstm_spec, allow_strokes_beyond_image_edges):\n",
        "    \"\"\"Create DrawingLSTM to interpret input sequences and paint an image.\n",
        "\n",
        "    Args:\n",
        "      drawing_lstm_spec: DrawingLSTMSpec with LSTM architecture parameters\n",
        "      allow_strokes_beyond_image_edges: Draw lines outside image boundary\n",
        "    \"\"\"\n",
        "    self._input_spec_size = drawing_lstm_spec.input_spec_size\n",
        "    self._num_lstms = drawing_lstm_spec.num_lstms\n",
        "    self._allow_strokes_beyond_image_edges = allow_strokes_beyond_image_edges\n",
        "    with torch.no_grad():\n",
        "      self.top_lstm = MakeGeneratorLstm(drawing_lstm_spec,\n",
        "                                        self._input_spec_size)\n",
        "      self.bottom_lstm = MakeGeneratorLstm(drawing_lstm_spec, 12)\n",
        "    self._init_all(self.top_lstm, torch.nn.init.normal_, mean=0., std=0.2)\n",
        "    self._init_all(self.bottom_lstm, torch.nn.init.normal_, mean=0., std=0.2)\n",
        "\n",
        "  def _init_all(self, a_model, init_func, *params, **kwargs):\n",
        "    \"\"\"Method for initialising model with given init_func, params and kwargs.\"\"\"\n",
        "    for p in a_model.parameters():\n",
        "      init_func(p, *params, **kwargs)\n",
        "\n",
        "  def _feed_top_lstm(self, input_seq):\n",
        "    \"\"\"Feed all input sequences input_seq into the LSTM models.\"\"\"\n",
        "\n",
        "    x_in = input_seq.reshape((len(input_seq), 1, self._input_spec_size))\n",
        "    x_in = np.tile(x_in, (SEQ_LENGTH, 1))\n",
        "    x_torch = torch.from_numpy(x_in).type(torch.FloatTensor)\n",
        "    y_torch = self.top_lstm(x_torch)\n",
        "    y_torch = [y_torch_k.detach().numpy() for y_torch_k in y_torch]\n",
        "    del x_in\n",
        "    del x_torch\n",
        "\n",
        "    # There are multiple LSTM heads. For each sequence, read out the head and\n",
        "    # length of intermediary output to keep and return intermediary outputs.\n",
        "    readouts_top = np.clip(\n",
        "        np.round(self._num_lstms/2.0 * (1 + input_seq[:, 1])).astype(np.int32),\n",
        "        0, self._num_lstms-1)\n",
        "    lengths_top = np.clip(\n",
        "        np.round(10.0 * (1 + input_seq[:, 0])).astype(np.int32),\n",
        "        0, SEQ_LENGTH) + 1\n",
        "    intermediate_strings = []\n",
        "    for i in range(len(readouts_top)):\n",
        "      y_torch_i = y_torch[readouts_top[i]][i]\n",
        "      intermediate_strings.append(y_torch_i[0:lengths_top[i], :])\n",
        "    return intermediate_strings\n",
        "\n",
        "  def _feed_bottom_lstm(self, intermediate_strings, input_seq, coeff_size=1):\n",
        "    \"\"\"Feed all input sequences into the LSTM models.\n",
        "\n",
        "    Args:\n",
        "      intermediate_strings: top level strings\n",
        "      input_seq: input sequences fed to the top LSTM\n",
        "      coeff_size: sets centre origin\n",
        "\n",
        "    Returns:\n",
        "      strokes: Painting strokes.\n",
        "      painting_commands: Top-level painting commands with origin, angle and scale\n",
        "        information, as well as transparency.\n",
        "    \"\"\"\n",
        "    img_center = 112. * coeff_size\n",
        "    coeff_origin = 100. * coeff_size\n",
        "    top_lengths = []\n",
        "    for i in range(len(intermediate_strings)):\n",
        "      top_lengths.append(len(intermediate_strings[i]))\n",
        "    y_flat = np.concatenate(intermediate_strings, axis=0)\n",
        "    tiled_y_flat = y_flat.reshape((len(y_flat), 1, self._input_spec_size))\n",
        "    tiled_y_flat = np.tile(tiled_y_flat, (SEQ_LENGTH, 1))\n",
        "    y_torch = torch.from_numpy(tiled_y_flat).type(torch.FloatTensor)\n",
        "    z_torch = self.bottom_lstm(y_torch)\n",
        "    z_torch = [z_torch_k.detach().numpy() for z_torch_k in z_torch]\n",
        "    del tiled_y_flat\n",
        "    del y_torch\n",
        "\n",
        "    # There are multiple LSTM heads. For each sequence, read out the head and\n",
        "    # length of intermediary output to keep and return intermediary outputs.\n",
        "    readouts = np.clip(np.round(\n",
        "        NUM_LSTMS/2.0 * (1 + y_flat[:, 0])).astype(np.int32), 0, NUM_LSTMS-1)\n",
        "    lengths_bottom = np.clip(\n",
        "        np.round(10.0 * (1 + y_flat[:, 1])).astype(np.int32), 0, SEQ_LENGTH) + 1\n",
        "    strokes = []\n",
        "    painting_commands = []\n",
        "    offset = 0\n",
        "    for i in range(len(intermediate_strings)):\n",
        "      origin_top = [(1+input_seq[i, 2]) * img_center,\n",
        "                    (1+input_seq[i, 3]) * img_center]\n",
        "      angle_top = input_seq[i, 4]\n",
        "      scale_top = input_seq[i, 5]\n",
        "      for j in range(len(intermediate_strings[i])):\n",
        "        k = j + offset\n",
        "        z_torch_ij = z_torch[readouts[k]][k]\n",
        "        strokes.append(z_torch_ij[0:lengths_bottom[k], :])\n",
        "        y_ij = y_flat[k]\n",
        "        origin_bottom = [y_ij[2] * coeff_origin, y_ij[3] * coeff_origin]\n",
        "        angle_bottom = y_ij[4]\n",
        "        scale_bottom = y_ij[5]\n",
        "        position_choice = y_ij[6]\n",
        "        transparency = y_ij[7]\n",
        "        painting_command = PaintingCommand(\n",
        "            origin_top, angle_top, scale_top, origin_bottom, angle_bottom,\n",
        "            scale_bottom, position_choice, transparency)\n",
        "        painting_commands.append(painting_command)\n",
        "      offset += top_lengths[i]\n",
        "    del y_flat\n",
        "    return strokes, painting_commands\n",
        "\n",
        "  def make_initial_genotype(self, initial_img, sequence_length,\n",
        "                            input_spec_size):\n",
        "    \"\"\"Make and return initial DNA weights for LSTMs, input sequence, and image.\n",
        "\n",
        "    Args:\n",
        "      initial_img: Image (to be appended to the genotype)\n",
        "      sequence_length: Length of the input sequence (i.e. number of strokes)\n",
        "      input_spec_size: Number of inputs for each element in the input sequences\n",
        "    Returns:\n",
        "      Genotype NamedTuple with fields: [parameters of network 0,\n",
        "                                        parameters of network 1,\n",
        "                                        input sequence,\n",
        "                                        initial_img]\n",
        "    \"\"\"\n",
        "    dna_top = []\n",
        "    with torch.no_grad():\n",
        "      for _, params in self.top_lstm.named_parameters():\n",
        "        dna_top.append(params.clone())\n",
        "        param_size = params.numpy().shape\n",
        "        dna_top[-1] = np.random.uniform(\n",
        "            0.1 * DNA_SCALE, 0.3\n",
        "            * DNA_SCALE) * np.random.normal(size=param_size)\n",
        "    dna_bottom = []\n",
        "    with torch.no_grad():\n",
        "      for _, params in self.bottom_lstm.named_parameters():\n",
        "        dna_bottom.append(params.clone())\n",
        "        param_size = params.numpy().shape\n",
        "        dna_bottom[-1] = np.random.uniform(\n",
        "            0.1 * DNA_SCALE, 0.3\n",
        "            * DNA_SCALE) * np.random.normal(size=param_size)\n",
        "    input_sequence = np.random.uniform(\n",
        "        -1, 1, size=(sequence_length, input_spec_size))\n",
        "    return Genotype(dna_top, dna_bottom, input_sequence, initial_img)\n",
        "\n",
        "  def draw(self, img, genotype):\n",
        "    \"\"\"Add to the image using the latest genotype and get latest input sequence.\n",
        "\n",
        "    Args:\n",
        "      img: image to add to.\n",
        "      genotype: as created by make_initial_genotype.\n",
        "\n",
        "    Returns:\n",
        "      image with new strokes added.\n",
        "    \"\"\"\n",
        "    t0_draw = time.time()\n",
        "    img = img + genotype.initial_img\n",
        "    input_sequence = genotype.input_sequence\n",
        "\n",
        "    # Generate the strokes for drawing in batch mode.\n",
        "    # input_sequence is between 10 and 20 but is evolved, can go to 200.\n",
        "    intermediate_strings = self._feed_top_lstm(input_sequence)\n",
        "    strokes, painting_commands = self._feed_bottom_lstm(\n",
        "        intermediate_strings, input_sequence)\n",
        "    del intermediate_strings\n",
        "\n",
        "    # Now we can go through the output strings producing the strokes.\n",
        "    t1_draw = time.time()\n",
        "    num_strokes = paint_over_image(\n",
        "        img, strokes, painting_commands, self._allow_strokes_beyond_image_edges,\n",
        "        coeff_size=1)\n",
        "\n",
        "    t2_draw = time.time()\n",
        "    if VERBOSE_CODE:\n",
        "      print(\n",
        "          \"Draw {:.2f}s (net {:.2f}s plot {:.2f}s {:.1f}ms/strk {}\".format(\n",
        "              t2_draw - t0_draw, t1_draw - t0_draw, t2_draw - t1_draw,\n",
        "              (t2_draw - t1_draw) / num_strokes * 1000, num_strokes))\n",
        "    return img"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g_G5xd574hBO"
      },
      "source": [
        "## DrawingGenerator"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YN3zl5D84fCt"
      },
      "outputs": [],
      "source": [
        "class DrawingGenerator:\n",
        "  \"\"\"Creates a drawing using a DrawingLSTM.\"\"\"\n",
        "\n",
        "  def __init__(self, image_size, drawing_lstm_spec,\n",
        "               allow_strokes_beyond_image_edges):\n",
        "    self.primitives = [\"c\", \"r\", \"l\", \"b\", \"p\", \"j\"]\n",
        "    self.pop = []\n",
        "    self.size = image_size\n",
        "    self.fitnesses = np.zeros(1)\n",
        "    self.noise = 2\n",
        "    self.mutation_std = 0.0004\n",
        "    # input_spec_size, num_lstms, net_lstm_hiddens,\n",
        "    # net_mlp_hiddens, output_size, allow_strokes_beyond_image_edges\n",
        "    self.drawing_lstm = DrawingLSTM(drawing_lstm_spec,\n",
        "                                    allow_strokes_beyond_image_edges)\n",
        "\n",
        "  def make_initial_genotype(self, initial_img, sequence_length, input_spec_size):\n",
        "    \"\"\"Use drawing_lstm to create initial genotypye.\"\"\"\n",
        "\n",
        "    self.genotype = self.drawing_lstm.make_initial_genotype(\n",
        "        initial_img, sequence_length, input_spec_size)\n",
        "    return self.genotype\n",
        "\n",
        "\n",
        "  def _copy_genotype_to_generator(self, genotype):\n",
        "    \"\"\"Copy genotype's data into generator's parameters.\n",
        "\n",
        "    Copies the parameters in genotype (genotype.top_lstm[:] and\n",
        "    genotype.bottom_lstm[:]) into the parmaters for the drawing network so it\n",
        "    can be used to evaluate the genotype.\n",
        "\n",
        "    Args:\n",
        "      genotype: as created by make_initial_genotype.\n",
        "\n",
        "    Returns:\n",
        "      None\n",
        "    \"\"\"\n",
        "    self.genotype = copy.deepcopy(genotype)\n",
        "    i = 0\n",
        "    with torch.no_grad():\n",
        "      for _, param in self.drawing_lstm.top_lstm.named_parameters():\n",
        "        param.copy_(torch.tensor(self.genotype.top_lstm[i]))\n",
        "        i = i + 1\n",
        "    i = 0\n",
        "    with torch.no_grad():\n",
        "      for _, param in self.drawing_lstm.bottom_lstm.named_parameters():\n",
        "        param.copy_(torch.tensor(self.genotype.bottom_lstm[i]))\n",
        "        i = i + 1\n",
        "\n",
        "  def _interpret_genotype(self, genotype):\n",
        "    img = np.zeros((self.size, self.size, 3), dtype=np.uint8)\n",
        "    img = self.drawing_lstm.draw(img, genotype)\n",
        "    return img\n",
        "\n",
        "  def draw_from_genotype(self, genotype):\n",
        "    \"\"\"Copy input sequence and LSTM weights from `genotype`, run and draw.\"\"\"\n",
        "    self._copy_genotype_to_generator(genotype)\n",
        "    return self._interpret_genotype(self.genotype)\n",
        "\n",
        "  def visualize_genotype(self, genotype):\n",
        "    \"\"\"Plot histograms of genotype\"s data.\"\"\"\n",
        "\n",
        "    plt.show()\n",
        "    inp_seq = np.array(genotype.input_sequence).flatten()\n",
        "    plt.title(\"input seq\")\n",
        "    plt.hist(inp_seq)\n",
        "    plt.show()\n",
        "\n",
        "    inp_seq = np.array(genotype.top_lstm).flatten()\n",
        "    plt.title(\"LSTM top\")\n",
        "    plt.hist(inp_seq)\n",
        "    plt.show()\n",
        "\n",
        "    inp_seq = np.array(genotype.bottom_lstm).flatten()\n",
        "    plt.title(\"LSTM bottom\")\n",
        "    plt.hist(inp_seq)\n",
        "\n",
        "    plt.show()\n",
        "\n",
        "  def mutate(self, genotype):\n",
        "    \"\"\"Mutates `genotype`. This function is static.\n",
        "\n",
        "    Args:\n",
        "      genotype: genotype structure to mutate parameters of.\n",
        "\n",
        "    Returns:\n",
        "      new_genotype: Mutated copy of supplied genotype.\n",
        "    \"\"\"\n",
        "\n",
        "    new_genotype = copy.deepcopy(genotype)\n",
        "    new_input_seq = new_genotype.input_sequence\n",
        "    n = len(new_input_seq)\n",
        "\n",
        "    if np.random.uniform() \u003c 1.0:\n",
        "\n",
        "      # Standard gaussian small mutation of input sequence.\n",
        "      if np.random.uniform() \u003e 0.5:\n",
        "        new_input_seq += (\n",
        "            np.random.uniform(0.001, 0.2) * np.random.normal(\n",
        "                size=new_input_seq.shape))\n",
        "\n",
        "      # Low frequency large mutation of individual parts of the input sequence.\n",
        "      for i in range(n):\n",
        "        if np.random.uniform() \u003c 2.0/n:\n",
        "          for j in range(len(new_input_seq[i])):\n",
        "            if np.random.uniform() \u003c 2.0/len(new_input_seq[i]):\n",
        "              new_input_seq[i][j] = new_input_seq[i][j] + 0.5*np.random.normal()\n",
        "\n",
        "      # Adding and deleting elements from the input sequence.\n",
        "      if np.random.uniform() \u003c 0.01:\n",
        "        if VERBOSE_MUTATION:\n",
        "          print(\"Mutation: adding\")\n",
        "        a = np.random.uniform(-1, 1, size=(1, INPUT_SPEC_SIZE))\n",
        "        pos = np.random.randint(1, len(new_input_seq))\n",
        "        new_input_seq = np.insert(new_input_seq, pos, a, axis=0)\n",
        "      if np.random.uniform() \u003c 0.02:\n",
        "        if VERBOSE_MUTATION:\n",
        "          print(\"Mutation: deleting\")\n",
        "        pos = np.random.randint(1, len(new_input_seq))\n",
        "        new_input_seq = np.delete(new_input_seq, pos, axis=0)\n",
        "      n = len(new_input_seq)\n",
        "\n",
        "      # Swapping two elements in the input sequence.\n",
        "      if np.random.uniform() \u003c 0.01:\n",
        "        element1 = np.random.randint(0, n)\n",
        "        element2 = np.random.randint(0, n)\n",
        "        while element1 == element2:\n",
        "          element2 = np.random.randint(0, n)\n",
        "        temp = copy.deepcopy(new_input_seq[element1])\n",
        "        new_input_seq[element1] = copy.deepcopy(new_input_seq[element2])\n",
        "        new_input_seq[element2] = temp\n",
        "\n",
        "      # Duplicate an element in the input sequence (with some mutation).\n",
        "      if np.random.uniform() \u003c 0.01:\n",
        "        if VERBOSE_MUTATION:\n",
        "          print(\"Mutation: duplicating\")\n",
        "        element1 = np.random.randint(0, n)\n",
        "        element2 = np.random.randint(0, n)\n",
        "        while element1 == element2:\n",
        "          element2 = np.random.randint(0, n)\n",
        "        new_input_seq[element1] = copy.deepcopy(new_input_seq[element2])\n",
        "        noise = 0.05 * np.random.normal(size=new_input_seq[element1].shape)\n",
        "        new_input_seq[element1] += noise\n",
        "\n",
        "      # Ensure that the input sequence is always between -1 and 1\n",
        "      # so that positions make sense.\n",
        "      new_genotype = new_genotype._replace(\n",
        "          input_sequence=np.clip(new_input_seq, -1.0, 1.0))\n",
        "\n",
        "    # Mutates dna of networks.\n",
        "    if np.random.uniform() \u003c 1.0:\n",
        "      for net in range(2):\n",
        "        for layer in range(len(new_genotype[net])):\n",
        "          weights = new_genotype[net][layer]\n",
        "          if np.random.uniform() \u003c 0.5:\n",
        "            noise = 0.00001 * np.random.standard_cauchy(size=weights.shape)\n",
        "            weights += noise\n",
        "          else:\n",
        "            noise = np.random.normal(size=weights.shape)\n",
        "            noise *= np.random.uniform(0.0001, 0.006)\n",
        "            weights += noise\n",
        "\n",
        "          if np.random.uniform() \u003c 0.01:\n",
        "            noise = np.random.normal(size=weights.shape)\n",
        "            noise *= np.random.uniform(0.1, 0.3)\n",
        "            weights = noise\n",
        "\n",
        "          # Ensure weights are between -10 and 10.\n",
        "          weights = np.clip(weights, -1.0, 1.0)\n",
        "          new_genotype[net][layer] = weights\n",
        "\n",
        "    return new_genotype"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "osQ1Zp40-VWT"
      },
      "source": [
        "## Evaluator"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "code",
        "id": "9YOsJQkPbPlA"
      },
      "outputs": [],
      "source": [
        "class Evaluator:\n",
        "  \"\"\"Evaluator for a drawing.\"\"\"\n",
        "\n",
        "  def __init__(self, image_size, drawing_lstm_spec,\n",
        "               allow_strokes_beyond_image_edges):\n",
        "    self.drawing_generator = DrawingGenerator(image_size, drawing_lstm_spec,\n",
        "                                              allow_strokes_beyond_image_edges)\n",
        "    self.calls = 0\n",
        "\n",
        "  def make_initial_genotype(self, img, sequence_length, input_spec_size):\n",
        "    return self.drawing_generator.make_initial_genotype(img, sequence_length,\n",
        "                                                        input_spec_size)\n",
        "\n",
        "  def evaluate_genotype(self, pickled_genotype, id_num):\n",
        "    \"\"\"Evaluate genotype and return genotype's image.\n",
        "\n",
        "    Args:\n",
        "      pickled_genotype: pickled genotype to be evaluated.\n",
        "      id_num: ID number of genotype.\n",
        "\n",
        "    Returns:\n",
        "      dict: drawing and id_num.\n",
        "    \"\"\"\n",
        "\n",
        "    genotype = cloudpickle.loads(pickled_genotype)\n",
        "    drawing = self.drawing_generator.draw_from_genotype(genotype)\n",
        "    self.calls += 1\n",
        "    return {\"drawing\": drawing, \"id\": id_num}\n",
        "\n",
        "  def mutate(self, genotype):\n",
        "    \"\"\"Create a mutated version of genotype.\"\"\"\n",
        "    return self.drawing_generator.mutate(genotype)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k3jrlAWcy9On"
      },
      "source": [
        "# Evolution"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iViNYPmHJ_Mb"
      },
      "source": [
        "## Fitness calculation, tournament, and crossover"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DZx0l-f-KI42"
      },
      "outputs": [],
      "source": [
        "IMAGE_MEAN = torch.tensor([0.48145466, 0.4578275, 0.40821073]).cuda()\n",
        "IMAGE_STD = torch.tensor([0.26862954, 0.26130258, 0.27577711]).cuda()\n",
        "\n",
        "def get_fitness(pictures, use_projective_transform,\n",
        "                projective_transform_coefficient):\n",
        "  \"\"\"Run CLIP on a batch of `pictures` and return `fitnesses`.\n",
        "\n",
        "  Args:\n",
        "    pictures: batch if images to evaluate\n",
        "    use_projective_transform: Add transformed versions of the image\n",
        "    projective_transform_coefficient: Degree of transform\n",
        "\n",
        "  Returns:\n",
        "    Similarities between images and the text\n",
        "  \"\"\"\n",
        "\n",
        "  # Do we use projective transforms of images before CLIP eval?\n",
        "  t0 = time.time()\n",
        "  pictures_trans = np.swapaxes(np.array(pictures), 1, 3) / 244.0\n",
        "  if use_projective_transform:\n",
        "    for i in range(len(pictures_trans)):\n",
        "      matrix = np.eye(3) + (\n",
        "          projective_transform_coefficient * np.random.normal(size=(3, 3)))\n",
        "      tform = transform.ProjectiveTransform(matrix=matrix)\n",
        "      pictures_trans[i] = transform.warp(pictures_trans[i], tform.inverse)\n",
        "\n",
        "  # Run the CLIP evaluator.\n",
        "  t1 = time.time()\n",
        "  image_input = torch.tensor(np.stack(pictures_trans)).cuda()\n",
        "  image_input -= IMAGE_MEAN[:, None, None]\n",
        "  image_input /= IMAGE_STD[:, None, None]\n",
        "  with torch.no_grad():\n",
        "    image_features = model.encode_image(image_input).float()\n",
        "  t2 = time.time()\n",
        "  similarity = torch.cosine_similarity(\n",
        "      text_features, image_features, dim=1).cpu().numpy()\n",
        "  t3 = time.time()\n",
        "  if VERBOSE_CODE:\n",
        "    print(f\"get_fitness init {t1-t0:.4f}s, CLIP {t2-t1:.4f}s, sim {t3-t2:.4f}s\")\n",
        "  return similarity\n",
        "\n",
        "\n",
        "def crossover(dna_winner, dna_loser, crossover_prob):\n",
        "  \"\"\"Create new genotype by combining two genotypes.\n",
        "\n",
        "  Randomly replaces parts of the genotype 'dna_winner' with parts of dna_loser\n",
        "  to create a new genotype based mostly on on both 'parents'.\n",
        "\n",
        "  Args:\n",
        "    dna_winner: The high-fitness parent genotype - gets replaced with child.\n",
        "    dna_loser: The lower-fitness parent genotype.\n",
        "    crossover_prob: Probability of crossover between winner and loser.\n",
        "\n",
        "  Returns:\n",
        "    dna_winner: The result of crossover from parents.\n",
        "  \"\"\"\n",
        "\n",
        "  # Copy single input signals\n",
        "  for i in range(len(dna_winner[2])):\n",
        "    if i \u003c len(dna_loser[2]):\n",
        "      if np.random.uniform() \u003c crossover_prob:\n",
        "        dna_winner[2][i] = copy.deepcopy(dna_loser[2][i])\n",
        "\n",
        "  # Copy whole modules\n",
        "  for i in range(len(dna_winner[0])):\n",
        "    if i \u003c len(dna_loser[0]):\n",
        "      if np.random.uniform() \u003c crossover_prob:\n",
        "        dna_winner[0][i] = copy.deepcopy(dna_loser[0][i])\n",
        "\n",
        "  # Copy whole modules\n",
        "  for i in range(len(dna_winner[1])):\n",
        "    if i \u003c len(dna_loser[1]):\n",
        "      if np.random.uniform() \u003c crossover_prob:\n",
        "        dna_winner[1][i] = copy.deepcopy(dna_loser[1][i])\n",
        "\n",
        "  return dna_winner\n",
        "\n",
        "\n",
        "def truncation_selection(population, fitnesses, evaluator, use_crossover,\n",
        "                         crossover_prob):\n",
        "  \"\"\"Create new population using truncation selection.\n",
        "\n",
        "  Creates a new population by copying across the best 50% genotypes and\n",
        "  filling the rest with (for use_crossover==False) a mutated copy of each\n",
        "  genotype or (for use_crossover==True) with children created through crossover\n",
        "  between each winner and a genotype in the bottom 50%.\n",
        "\n",
        "  Args:\n",
        "    population: list of current population genotypes.\n",
        "    fitnesses: list of evaluated fitnesses.\n",
        "    evaluator: class that evaluates a draw generator.\n",
        "    use_crossover: Whether to use crossover between winner and loser.\n",
        "    crossover_prob: Probability of crossover between winner and loser.\n",
        "\n",
        "  Returns:\n",
        "    new_pop: the new population.\n",
        "    best: genotype.\n",
        "  \"\"\"\n",
        "\n",
        "  fitnesses = np.array(-fitnesses)\n",
        "  ordered_fitness_ids = fitnesses.argsort()\n",
        "  best = copy.deepcopy(population[ordered_fitness_ids[0]])\n",
        "  pop_size = len(population)\n",
        "\n",
        "  if not use_crossover:\n",
        "    new_pop = []\n",
        "    for i in range(int(pop_size/2)):\n",
        "      new_pop.append(copy.deepcopy(population[ordered_fitness_ids[i]]))\n",
        "    for i in range(int(pop_size/2)):\n",
        "      new_pop.append(evaluator.mutate(\n",
        "          copy.deepcopy(population[ordered_fitness_ids[i]])))\n",
        "  else:\n",
        "    new_pop = []\n",
        "    for i in range(int(pop_size/2)):\n",
        "      new_pop.append(copy.deepcopy(population[ordered_fitness_ids[i]]))\n",
        "    for i in range(int(pop_size/2)):\n",
        "      new_pop.append(evaluator.mutate(crossover(\n",
        "          copy.deepcopy(population[ordered_fitness_ids[i]]),\n",
        "          population[ordered_fitness_ids[int(pop_size/2) + i]], crossover_prob\n",
        "          )))\n",
        "\n",
        "  return new_pop, best"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kos2j3qALqf3"
      },
      "source": [
        "##Remote workers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HvMKhZ4LK5Iu"
      },
      "outputs": [],
      "source": [
        "VERBOSE_DURATION = False\n",
        "\n",
        "@ray.remote\n",
        "class Worker(object):\n",
        "  \"\"\"Takes a pickled dna and evaluates it, returning result.\"\"\"\n",
        "\n",
        "  def __init__(self, image_size, drawing_lstm_spec,\n",
        "               allow_strokes_beyond_image_edges):\n",
        "    self.evaluator = Evaluator(image_size, drawing_lstm_spec,\n",
        "                               allow_strokes_beyond_image_edges)\n",
        "\n",
        "  def compute(self, dna_pickle, genotype_id):\n",
        "    if VERBOSE_DURATION:\n",
        "      t0 = time.time()\n",
        "    res = self.evaluator.evaluate_genotype(dna_pickle, genotype_id)\n",
        "    if VERBOSE_DURATION:\n",
        "      duration = time.time() - t0\n",
        "      print(f\"Worker {genotype_id} evaluated params in {duration:.1f}sec\")\n",
        "    return res\n",
        "\n",
        "\n",
        "def create_workers(num_workers, image_size, drawing_lstm_spec,\n",
        "                   allow_strokes_beyond_image_edges):\n",
        "  \"\"\"Create the workers.\n",
        "\n",
        "  Args:\n",
        "    num_workers: Number of parallel workers for evaluation.\n",
        "    image_size: Length of side of (square) image\n",
        "    drawing_lstm_spec: DrawingLSTMSpec for LSTM network\n",
        "    allow_strokes_beyond_image_edges: Whether to draw outside the edges\n",
        "  Returns:\n",
        "    List of workers.\n",
        "  \"\"\"\n",
        "  worker_pool = []\n",
        "  for w_i in range(num_workers):\n",
        "    print(\"Creating worker\", w_i, flush=True)\n",
        "    worker_pool.append(Worker.remote(image_size, drawing_lstm_spec,\n",
        "                                     allow_strokes_beyond_image_edges))\n",
        "  return worker_pool\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZA3BNTXtMkyp"
      },
      "source": [
        "##Plotting"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IkdAw6h2rSHL"
      },
      "outputs": [],
      "source": [
        "def plot_training_res(batch_drawings, fitness_history, idx=None):\n",
        "  \"\"\"Plot fitnesses and timings.\n",
        "\n",
        "  Args:\n",
        "    batch_drawings: Drawings\n",
        "    fitness_history: History of fitnesses\n",
        "    idx: Index of drawing to show, default is highest fitness\n",
        "  \"\"\"\n",
        "  _, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))\n",
        "  if idx is None:\n",
        "    idx = np.argmax(fitness_history[-1])\n",
        "  ax1.plot(fitness_history, \".\")\n",
        "  ax1.set_title(\"Fitnesses\")\n",
        "  ax2.imshow(batch_drawings[idx])\n",
        "  ax2.set_title(f\"{PROMPT} (fit: {fitness_history[-1][idx]:.3f})\")\n",
        "  plt.show()\n",
        "\n",
        "def plot_samples(batch_drawings, num_samples=16):\n",
        "  \"\"\"Plot sample of drawings.\n",
        "\n",
        "  Args:\n",
        "    batch_drawings: Batch of drawings to sample from\n",
        "    num_samples: Number to displa\n",
        "  \"\"\"\n",
        "  num_samples = min(len(batch_drawings), num_samples)\n",
        "  num_rows = int(math.floor(np.sqrt(num_samples)))\n",
        "  num_cols = int(math.ceil(num_samples / num_rows))\n",
        "  row_images = []\n",
        "  for c in range(0, num_samples, num_cols):\n",
        "    if c + num_cols \u003c= num_samples:\n",
        "      row_images.append(np.concatenate(batch_drawings[c:(c+num_cols)], axis=1))\n",
        "  composite_image = np.concatenate(row_images, axis=0)\n",
        "  _, ax = plt.subplots(1, 1, figsize=(20, 20))\n",
        "  ax.imshow(composite_image)\n",
        "  ax.set_title(PROMPT)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fM_2F-K9MHRO"
      },
      "source": [
        "## Population and evolution main loop"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "x52rxKlyKWiE"
      },
      "outputs": [],
      "source": [
        "def make_population(pop_size, evaluator, image_size, input_spec_size,\n",
        "                    sequence_length):\n",
        "  \"\"\"Make initial population.\n",
        "\n",
        "  Args:\n",
        "    pop_size: number of genotypes in population.\n",
        "    evaluator: An Evaluator class instance for generating initial genotype.\n",
        "    image_size: Size of initial image for genotype to draw on.\n",
        "    input_spec_size: Sequence element size\n",
        "    sequence_length: Initial length of sequences\n",
        "\n",
        "  Returns:\n",
        "    Initialised population.\n",
        "  \"\"\"\n",
        "  print(f\"Creating initial population of size {pop_size}\")\n",
        "  pop = []\n",
        "  for _ in range(pop_size):\n",
        "    a_genotype = evaluator.make_initial_genotype(\n",
        "        img=np.zeros((image_size, image_size, 3), dtype=np.uint8),\n",
        "        sequence_length=sequence_length,\n",
        "        input_spec_size=input_spec_size)\n",
        "    pop.append(a_genotype)\n",
        "  return pop\n",
        "\n",
        "def evolution_loop(population, worker_pool, evaluator, num_generations,\n",
        "                   use_crossover, crossover_prob,\n",
        "                   use_projective_transform, projective_transform_coefficient,\n",
        "                   plot_every, plot_batch):\n",
        "  \"\"\"Create population and run evolution.\n",
        "\n",
        "  Args:\n",
        "    population: Initial population of genotypes\n",
        "    worker_pool: List of workers of parallel evaluations\n",
        "    evaluator: image evaluator to calculate fitnesses\n",
        "    num_generations: number of generations to run\n",
        "    use_crossover: Whether crossover is used for offspring\n",
        "    crossover_prob: Probability that crossover takes place\n",
        "    use_projective_transform: Use projective transforms in evaluation\n",
        "    projective_transform_coefficient: Degree of projective transform\n",
        "    plot_every: number of generations between new plots\n",
        "    plot_batch: whether to show all samples in the batch then plotting\n",
        "  \"\"\"\n",
        "  population_size = len(population)\n",
        "  num_workers = len(worker_pool)\n",
        "  print(\"Population of {} genotypes being evaluated by {} workers\".format(\n",
        "      population_size, num_workers))\n",
        "  drawings = {}\n",
        "  fitness_history = []\n",
        "  init_gen = len(fitness_history)\n",
        "  print(f\"(Re)starting evolution at generation {init_gen}\")\n",
        "  for gen in range(init_gen, num_generations):\n",
        "\n",
        "    # Drawing\n",
        "    t0_loop = time.time()\n",
        "    futures = []\n",
        "    for j in range(0, population_size, num_workers):\n",
        "      for i in range(num_workers):\n",
        "        futures.append(worker_pool[i].compute.remote(\n",
        "            cloudpickle.dumps(population[i+j]), i+j))\n",
        "      data = ray.get(futures)\n",
        "      for i in range(num_workers):\n",
        "        drawings[data[i+j][\"id\"]] = data[j+i][\"drawing\"]\n",
        "    batch_drawings = []\n",
        "    for i in range(population_size):\n",
        "      batch_drawings.append(drawings[i])\n",
        "\n",
        "    # Fitness evaluation using CLIP\n",
        "    t1_loop = time.time()\n",
        "    fitnesses = get_fitness(batch_drawings, use_projective_transform,\n",
        "                            projective_transform_coefficient)\n",
        "    fitness_history.append(copy.deepcopy(fitnesses))\n",
        "\n",
        "    # Tournament\n",
        "    t2_loop = time.time()\n",
        "    population, best_genotype = truncation_selection(\n",
        "        population, fitnesses, evaluator, use_crossover, crossover_prob)\n",
        "    t3_loop = time.time()\n",
        "    duration_draw = t1_loop - t0_loop\n",
        "    duration_fit = t2_loop - t1_loop\n",
        "    duration_tournament = t3_loop - t2_loop\n",
        "    duration_total = t3_loop - t0_loop\n",
        "    if gen % plot_every == 0:\n",
        "      if VISUALIZE_GENOTYPE:\n",
        "        evaluator.drawing_generator.visualize_genotype(best_genotype)\n",
        "      print(\"Draw: {:.2f}s fit: {:.2f}s evol: {:.2f}s total: {:.2f}s\".format(\n",
        "          duration_draw, duration_fit, duration_tournament, duration_total))\n",
        "      plot_training_res(batch_drawings, fitness_history)\n",
        "      if plot_batch:\n",
        "        num_samples_to_plot = int(math.pow(\n",
        "            math.floor(np.sqrt(population_size)), 2))\n",
        "        plot_samples(batch_drawings, num_samples=num_samples_to_plot)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GNGqpXS1919a"
      },
      "source": [
        "# Configure and Generate"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "kSGgjwTM8uFn"
      },
      "outputs": [],
      "source": [
        "#@title Hyperparameters\n",
        "\n",
        "#@markdown Evolution parameters: population size and number of generations.\n",
        "POPULATION_SIZE = 10  #@param {type:\"slider\", min:4, max:100, step:2}\n",
        "NUM_GENERATIONS = 5000  #@param {type:\"integer\", min:100}\n",
        "#@markdown Number of workers working in parallel (should be equal to or smaller than the population size).\n",
        "NUM_WORKERS = 10  #@param {type:\"slider\", min:4, max:100, step:2}\n",
        "#@markdown Crossover in evolution.\n",
        "USE_CROSSOVER = True  #@param {type:\"boolean\"}\n",
        "CROSSOVER_PROB = 0.01  #@param {type:\"number\"}\n",
        "#@markdown Number of LSTMs, each one encoding a group of strokes.\n",
        "NUM_LSTMS = 5  #@param {type:\"integer\", min:1, max:5}\n",
        "#@markdown Number of inputs for each element in the input sequences.\n",
        "INPUT_SPEC_SIZE = 10  #@param {type:\"integer\"}\n",
        "#@markdown Length of the input sequence fed to the LSTMs (determines number of strokes).\n",
        "SEQ_LENGTH = 20  #@param {type:\"integer\", min:20, max:200}\n",
        "#@markdown Rendering parameter.\n",
        "ALLOW_STROKES_BEYOND_IMAGE_EDGES = True  #@param {type:\"boolean\"}\n",
        "#@markdown CLIP evaluation: do we use projective transforms of images?\n",
        "USE_PROJECTIVE_TRANSFORM = True  #@param {type:\"boolean\"}\n",
        "PROJECTIVE_TRANSFORM_COEFFICIENT = 0.000001  #@param {type:\"number\"}\n",
        "#@markdown These parameters should be edited mostly only for debugging reasons.\n",
        "NET_LSTM_HIDDENS = 40  #@param {type:\"integer\"}\n",
        "NET_MLP_HIDDENS = 20  #@param {type:\"integer\"}\n",
        "# Scales the values used in genotype's initialisation.\n",
        "DNA_SCALE = 1.0  #@param {type:\"number\"}\n",
        "IMAGE_SIZE = 224  #@param {type:\"integer\"}\n",
        "VERBOSE_CODE = False  #@param {type:\"boolean\"}\n",
        "VISUALIZE_GENOTYPE = False  #@param {type:\"boolean\"}\n",
        "VERBOSE_MUTATION = False  #@param {type:\"boolean\"}\n",
        "#@markdown Number of generations between new plots.\n",
        "PLOT_EVERY_NUM_GENS = 5  #@param {type:\"integer\"}\n",
        "#@markdown Whether to show all samples in the batch when plotting.\n",
        "PLOT_BATCH = True  # @param {type:\"boolean\"}\n",
        "\n",
        "assert POPULATION_SIZE % NUM_WORKERS == 0, \"POPULATION_SIZE not multiple of NUM_WORKERS\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U3mjWOijW5m_"
      },
      "source": [
        "#Running the original evolutionary algorithm\n",
        "This is the original inefficient version of Arnheim which uses a genetic algorithm to optimize the picture. It takes at least 12 hours to produce an image using 50 workers. In our paper we used 500-1000 GPUs which speeded things up considerably. Refer to Arnheim 2 for a far more efficient way to generate images with a similar architecture.\n",
        "\n",
        "Try prompts like “A photorealistic chicken”. Feel free to modify this colab to include your own way of generating and evolving images like we did in figure 2 here https://arxiv.org/pdf/2105.00162.pdf."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "kONtUlRK9Rcu"
      },
      "outputs": [],
      "source": [
        "# @title Get text input and run evolution\n",
        "PROMPT = \"an apple\"  #@param {type:\"string\"}\n",
        "\n",
        "# Tokenize prompts and coompute CLIP features.\n",
        "text_input = clip.tokenize(PROMPT).to(device)\n",
        "with torch.no_grad():\n",
        "  text_features = model.encode_text(text_input)\n",
        "\n",
        "ray.shutdown()\n",
        "ray.init()\n",
        "\n",
        "drawing_lstm_arch = DrawingLSTMSpec(INPUT_SPEC_SIZE,\n",
        "                                    NUM_LSTMS,\n",
        "                                    NET_LSTM_HIDDENS,\n",
        "                                    NET_MLP_HIDDENS)\n",
        "\n",
        "workers = create_workers(NUM_WORKERS, IMAGE_SIZE, drawing_lstm_arch,\n",
        "                         ALLOW_STROKES_BEYOND_IMAGE_EDGES)\n",
        "\n",
        "\n",
        "drawing_evaluator = Evaluator(IMAGE_SIZE, drawing_lstm_arch,\n",
        "                              ALLOW_STROKES_BEYOND_IMAGE_EDGES)\n",
        "\n",
        "drawing_population = make_population(POPULATION_SIZE, drawing_evaluator,\n",
        "                                     IMAGE_SIZE, INPUT_SPEC_SIZE, SEQ_LENGTH)\n",
        "\n",
        "evolution_loop(drawing_population, workers, drawing_evaluator, NUM_GENERATIONS,\n",
        "               USE_CROSSOVER, CROSSOVER_PROB,\n",
        "               USE_PROJECTIVE_TRANSFORM, PROJECTIVE_TRANSFORM_COEFFICIENT,\n",
        "               PLOT_EVERY_NUM_GENS, PLOT_BATCH)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "ZysZip8gE24t",
        "mNH-R_bs6zFZ",
        "rvo1d2z_Hv1B",
        "ElxgpBwkJuFj",
        "D6Yvwl6f6r0i",
        "H0DsS6D87jGM",
        "g_G5xd574hBO",
        "osQ1Zp40-VWT",
        "k3jrlAWcy9On",
        "iViNYPmHJ_Mb",
        "Kos2j3qALqf3",
        "ZA3BNTXtMkyp",
        "fM_2F-K9MHRO"
      ],
      "machine_shape": "hm",
      "name": "arnheim_1.ipynb",
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}