arnheim_2.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Copy of arnheim_2.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [
        "e7DCzf-EHzL_",
        "GIXzueO3PB-4"
      ],
      "machine_shape": "hm"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "spHtkST4105t"
      },
      "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.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e7DCzf-EHzL_"
      },
      "source": [
        "\n",
        "# The Arnheim Architecture and its Variants: Colab for Blog Post.\n",
        "\n",
        "**Chrisantha Fernando,  Dylan Banarse, Piotr Mirowski, Jean-Baptiste Alayrac, Sina Samangooei, Mateusz Malinowski, Mikolaj Binkowski, Wojciech Stokowiec, Jordan Hoffmann, El Morrison, Ali Eslami, Charlie Deck, Oriol Vinyals, Simon Osindero**\n",
        "\n",
        "DeepMind, 2021\n",
        "\n",
        "![picture](https://storage.googleapis.com/deepmind-media/Art%20in%20AI/Figure1Arnheim1PhotorealisticChicken.png)\n",
        "\n",
        "\"A Photorealistic Chicken\" by the Arnheim 2 Algorithm\n",
        "\n",
        "**STEPS:**\n",
        "\n",
        "1. Click \"Connect\" in the top right corner\n",
        "2. Runtime -> Change runtime type -> Hardware accelerator -> GPU\n",
        "2. Click the run button on \"Preliminaries\". This will install dependencies, it may take a while.\n",
        "2. **Important:** Runtime -> Restart Runtime\n",
        "2. Select grammar and enter image description prompt\n",
        "3. Run the rest of the cells (tip: simply select menu Runtime -> Run all)\n",
        "\n",
        "**An Exploration of Architectures and Losses for Painting and Drawing**\n",
        "\n",
        "This Colab explores a range of architectures for generating stroke specifications, guided by the CLIPDraw interface produced by Kevin Frans et al recently, and extending our earlier work on Arnheim 1 in [Generative Art Using Neural Visual Grammars and Dual Encoders](https://arxiv.org/abs/2105.00162).\n",
        "\n",
        "For completeness, an evolutionary version of the original Arnheim 1 code is included in a separate Colab. However, its scope is limited with only 1 GPU, but you can modify the generative procedure and genotype to implement your own much more general picture generators that do not need to be differentiable, so you might find it useful for some of your own purposes. Enjoy playing with this stuff, and do post your creations and share your modifications :)\n",
        "\n",
        "The Arnheim 1 architecture was a rather idiosyncratic variant of a seq-to-seq model, more like an L-system in many respects, in which the input sequence embeddings and parameters of an encoder and decoder set of LSTMs were evolved. Unlike standard seq-to-seq models (without attention) however, the outputs of the encoder LSTM were sent to the decoder LSTM directly (not just the final hidden state). Each layer of LSTMs thus re-wrote the sentence produced by the LSTM before, resulting in a hierarchical organization of strokes which leant structure to the drawing.\n",
        "\n",
        "Originally optimized by an evolutionary algorithm using 500 GPUs, the system was rather inefficient (as the Arnheim 1 Colab demonstrates).  In this Colab the Arnheim architecture has been moved into a differentiable form with minor modification and is being used to generate the stroke descriptions for CLIPDraw, released by Kevin Frans, L.B. Soros, Olaf Witkowski https://arxiv.org/abs/2106.14843, see their [original CLIPDraw Colab](https://colab.research.google.com/github/kvfrans/clipdraw/blob/main/clipdraw.ipynb).\n",
        "\n",
        "Thanks to Kevin from CrossLabs, who showed me how to add a simple MLP to the CLIPDraw Colab, I added the Arnheim 2 architecture and a bunch of other architectures I wanted to try out with CLIPDraw. CLIPDraw uses pydiffsvg to allow the gradients from CLIP to pass through the image itself into the stroke descriptions. The original CLIPDraw used direct descriptions of strokes rather than using a neural network to generate the strokes.\n",
        "\n",
        "This Colab contains a variety of different architectures that can be trained by gradients transmitted through the picture itself (using diffsvg). **The style of the painting is highly constrained by the algorithmic procedure for its generation.** Without a generative architecture the painting lacks order to some extent, for example, marks are independent, only given order by feedback from CLIP rather than by any generative priors. The amazing work on photorealistic image generation e.g. with DALL-E https://openai.com/blog/dall-e/ and other VQ-GAN and Diffusion approaches use generators pre-trained on images to constrain this process. However, here we are interested not in imposing photographic priors but in imposing mark making priors explicitly in the architecture of the generator.\n",
        "\n",
        "The second way that style can be encoded is in hand designed **computational aesthetic losses** which measure some property of the painting throughout its production, e.g. the diversity of line thicknesses or some aspect of composition, and reward or punish the generator accordingly. We also explore some of these losses here, but leave them as a template for you to add your own creative new losses.\n",
        "\n",
        "The third kind of style constraint is a **material (motor spec) constraint** determined by the the nature of the marks that are made, e.g. lines (used in our original paper and Arnheim 1) to Beziers (used here in Arnheim 2). We intentionally do not explore that rich domain here, because it could extend to all kinds of processes including collage constructed from a predefined set of primitive cutouts, 3D forms, cabbages, etc...\n",
        "\n",
        "A forth kind of style constraint is more commonly explored, and that is to exploit the knowledge inside CLIP by using a **prompt modifier** such as \"in the style of Van Gough\" etc... We do not explore this here very much as it has been extensively studied elsewhere. Related approaches were tried before which included representing artists and styles in a manifold and exploring intermediate points in that low dimensional space which captures \"the shape of art history\", although this often resulted in rather derivative looking images, rather than genuinely novel styles: https://arxiv.org/abs/1801.07729\n",
        "\n",
        "After running the Colab one time, read the notes at the end of this Colab to see how the various parts of the code can be changed to generate different images.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BfEKdIbK4VZa"
      },
      "source": [
        "# Preliminaries"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "oEbbebio7KaM"
      },
      "source": [
        "#@title Check GPU\n",
        "gpu_spec = !nvidia-smi -L\n",
        "print(gpu_spec[0])\n",
        "if \"Tesla K80\" in gpu_spec[0]:\n",
        "  print(\"\\nDetected a Tesla K80 GPU which results in the following error:\")\n",
        "  print(\"     RuntimeError: radix_sort: failed on 1st step: cudaErrorInvalidDeviceFunction: invalid device function\\n\")\n",
        "  print(\"Either try Factory Resetting the runtime until you get a different GPU or consider Colab Pro which offers higher-spec GPUs by default.\\n\")\n",
        "  raise ValueError(\"Incompatible GPU detected, see message above\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5dyyH781qzIC",
        "cellView": "form"
      },
      "source": [
        "#@title Installation of libraries {vertical-output: true}\n",
        "\n",
        "import subprocess\n",
        "\n",
        "CUDA_version = [s for s in subprocess.check_output([\"nvcc\", \"--version\"]).decode(\"UTF-8\").split(\", \") 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",
        "\n",
        "%cd /content/\n",
        "!pip install svgwrite\n",
        "!pip install svgpathtools\n",
        "!pip install cssutils\n",
        "!pip install numba\n",
        "!pip install torch-tools\n",
        "!pip install visdom\n",
        "\n",
        "!git clone https://github.com/BachiLi/diffvg\n",
        "%cd diffvg\n",
        "!git submodule update --init --recursive\n",
        "!python setup.py install\n",
        "\n",
        "!pip install ftfy regex tqdm\n",
        "!pip install git+https://github.com/openai/CLIP.git --no-deps"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c8bdUyJs4hq3"
      },
      "source": [
        "# Imports and libraries"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hjt9T3ARukAg",
        "cellView": "form"
      },
      "source": [
        "#@title Imports {vertical-output: true}\n",
        "\n",
        "import os\n",
        "import PIL.Image\n",
        "import random\n",
        "import time\n",
        "\n",
        "from matplotlib import pyplot as plt\n",
        "\n",
        "import numpy as np\n",
        "%tensorflow_version 2.x\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "from torch.utils.tensorboard import SummaryWriter\n",
        "import torchvision.transforms as transforms\n",
        "\n",
        "print(\"Torch version:\", torch.__version__)\n",
        "\n",
        "import clip\n",
        "\n",
        "try:\n",
        "  import pydiffvg\n",
        "except:\n",
        "  raise ValueError(\"Please follow Step 4, i.e. restart runtime and re-run.\")\n",
        "  \n",
        "os.environ[\"FFMPEG_BINARY\"] = \"ffmpeg\"\n",
        "import moviepy.editor as mvp\n",
        "from moviepy.video.io.ffmpeg_writer import FFMPEG_VideoWriter\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "Z-Wt7UjTi8Le"
      },
      "source": [
        "#@title Configure pydiffvg and load CLIP model {vertical-output: true}\n",
        "\n",
        "device = torch.device(\"cuda\")\n",
        "pydiffvg.set_print_timing(False)\n",
        "pydiffvg.set_device(device)\n",
        "pydiffvg.set_use_gpu(torch.cuda.is_available())  # Use GPU if available.\n",
        "\n",
        "CLIP_MODEL = \"ViT-B/32\"\n",
        "print(f\"Downloading CLIP model {CLIP_MODEL}...\")\n",
        "clip_model, _ = clip.load(CLIP_MODEL, device, jit=False)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4ckYmVuAAO7x"
      },
      "source": [
        "#Grammar Drawing Network Definitions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "wvWi5M3aQImk"
      },
      "source": [
        "#@title Grammar Specific Drawing Network: Photographic\n",
        "\n",
        "class CurveNetworkPhotographicLSTM(torch.nn.Module):\n",
        "  \"\"\"LSTM-based line-properties with CLIPDraw-like line optimization.\"\"\"\n",
        "\n",
        "  def __init__(self):\n",
        "    \"\"\"Constructor, relying on global parameters.\"\"\"\n",
        "    super().__init__()\n",
        "\n",
        "    # There are 3 LSTMS, for points, widths and colours.\n",
        "    assert NUM_LSTMS == 3, \"Need exactly 3 LSTMs.\"\n",
        "\n",
        "    # Bézier curve parameterisations are added to the generative grammar.\n",
        "    points = []\n",
        "    for _ in range(BATCH_SIZE * NUM_STROKE_TYPES * SEQ_LENGTH):\n",
        "      p0 = (1.5 * (random.random() - 0.5), 1.5 * (random.random() - 0.5))\n",
        "      points.append(p0)\n",
        "      radius = 0.1\n",
        "      p1 = (p0[0] + radius * (\n",
        "          random.random() - 0.5), p0[1] + radius * (random.random() - 0.5))\n",
        "      p2 = (p1[0] + radius * (\n",
        "          random.random() - 0.5), p1[1] + radius * (random.random() - 0.5))\n",
        "      p3 = (p2[0] + radius * (\n",
        "          random.random() - 0.5), p2[1] + radius * (random.random() - 0.5))\n",
        "      points.append(p1)\n",
        "      points.append(p2)\n",
        "      points.append(p3)\n",
        "      p0 = p3\n",
        "    points = np.array(points).flatten()\n",
        "    self.positions = torch.nn.Parameter(torch.Tensor(points))\n",
        "\n",
        "    # Initial sequences.\n",
        "    self._initials = []\n",
        "    for _ in range(NUM_LSTMS):\n",
        "      initial = torch.nn.Parameter(torch.ones(\n",
        "          BATCH_SIZE, SEQ_LENGTH, INPUT_SPEC_SIZE))\n",
        "      self._initials.append(initial)\n",
        "\n",
        "    # Shared input layer to process the initial sequence.\n",
        "    self._input_layer = torch.nn.Sequential(\n",
        "        torch.nn.Linear(INPUT_SPEC_SIZE, NET_LSTM_HIDDENS),\n",
        "        torch.nn.LeakyReLU(0.2, inplace=True))\n",
        "\n",
        "    # Different initial sequences, LSTMs and heads for lines, colorus, width.\n",
        "    lstms = []\n",
        "    heads = []\n",
        "    for _ in range(NUM_LSTMS):\n",
        "      lstm_layer = torch.nn.LSTM(\n",
        "          input_size=NET_LSTM_HIDDENS, hidden_size=NET_LSTM_HIDDENS,\n",
        "          num_layers=1, batch_first=True, bias=True)\n",
        "      head_layer = torch.nn.Sequential(\n",
        "          torch.nn.Linear(NET_LSTM_HIDDENS, NET_MLP_HIDDENS),\n",
        "          torch.nn.LeakyReLU(0.2, inplace=True),\n",
        "          torch.nn.Linear(NET_MLP_HIDDENS, OUTPUT_SIZE))\n",
        "      lstms.append(lstm_layer)\n",
        "      heads.append(head_layer)\n",
        "    self._lstms = torch.nn.ModuleList(lstms)\n",
        "    self._heads = torch.nn.ModuleList(heads)\n",
        "\n",
        "  def forward(self):\n",
        "    \"\"\"Input-less forward function.\"\"\"\n",
        "    pred = []\n",
        "    for i in range(NUM_LSTMS):\n",
        "      x = self._input_layer(self._initials[i])\n",
        "      y, _ = self._lstms[i](x)\n",
        "      y = torch.reshape(self._heads[i](y), (BATCH_SIZE*SEQ_LENGTH*OUTPUT_SIZE,))\n",
        "      if i == 0:\n",
        "        y = y * OUTPUT_COEFF_SYSTEMATICITY + self.positions\n",
        "        points = torch.clamp(y, min=-1, max=1)\n",
        "        pred.append(points)\n",
        "      elif i == 1:\n",
        "        widths = torch.clamp(y * OUTPUT_COEFF_WIDTH, min=1, max=100)\n",
        "        pred.append(widths)\n",
        "      elif i == 2:\n",
        "        colours = torch.clamp(y * OUTPUT_COEFF_COLOUR, min=0, max=1)\n",
        "        pred.append(colours)\n",
        "    # Unused diversity loss term.\n",
        "    pred.append(0)\n",
        "    return pred"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "yqof7yPiPpHS"
      },
      "source": [
        "#@title Grammar Specific Drawing Network: Arnheim 2\n",
        "\n",
        "class CurveNetworkHierarchicalLSTM(torch.nn.Module):\n",
        "  \"\"\"LSTM-based production of drawing sequences.\"\"\"\n",
        "\n",
        "  def __init__(self):\n",
        "    super().__init__()\n",
        "\n",
        "    def init_all(model, init_func, *params, **kwargs):\n",
        "      \"\"\"Init all parameters of model.\n",
        "\n",
        "      Args:\n",
        "        model: TF model top initialise\n",
        "        init_func: TF initialisation function\n",
        "        *params: Params to pass to init_func\n",
        "        **kwargs: kwargs to pass to init_func\n",
        "\n",
        "      Returns:\n",
        "        None\n",
        "      \"\"\"\n",
        "      for p in model.parameters():\n",
        "        init_func(p, *params, **kwargs)\n",
        "\n",
        "    # ********************* Direct encoding*************************************\n",
        "    points_in = []\n",
        "    for _ in range(BATCH_SIZE * LSTM_ITERATIONS * LSTM_ITERATIONS):\n",
        "      p0 = (1.5 * (random.random()-0.5), 1.5 * (random.random() - 0.5))\n",
        "      points_in.append(p0)\n",
        "      radius = 0.1\n",
        "      p1 = (p0[0] + radius * (\n",
        "          random.random() - 0.5), p0[1] + radius * (random.random() - 0.5))\n",
        "      p2 = (p1[0] + radius * (\n",
        "          random.random() - 0.5), p1[1] + radius * (random.random() - 0.5))\n",
        "      p3 = (p2[0] + radius * (\n",
        "          random.random() - 0.5), p2[1] + radius * (random.random() - 0.5))\n",
        "      points_in.append(p1)\n",
        "      points_in.append(p2)\n",
        "      points_in.append(p3)\n",
        "      p0 = p3\n",
        "\n",
        "    self.direct_points = torch.nn.Parameter(\n",
        "        torch.Tensor(np.array(points_in).flatten().reshape(\n",
        "            (BATCH_SIZE*LSTM_ITERATIONS*LSTM_ITERATIONS, 4, 2))))\n",
        "\n",
        "    # Set values to torch.nn.Parameters to make them evolvable.\n",
        "    self.positions_top = torch.nn.Parameter(\n",
        "        2.0 * (torch.rand(BATCH_SIZE, EMBEDDING_SIZE) - 0.5))\n",
        "    # Direct vs indirect position weighting\n",
        "    # One can experiment with replacing the constants below with learnable\n",
        "    # parameters, see suggestions in the comments following them.\n",
        "    self.scale_a = 0.01  # torch.nn.Parameter(1.0*(torch.rand(1)))\n",
        "    # Bottom-level position scaling\n",
        "    self.scale_b = 4  # torch.nn.Parameter(4.0*(torch.rand(1)))\n",
        "    # Softmax logit scaling\n",
        "    self.scale_c = 1  # torch.nn.Parameter(1.0*(torch.rand(1)))\n",
        "    # Point scaling\n",
        "    self.scale_d = 3  # torch.nn.Parameter(10.0*(torch.rand(1)))\n",
        "    # Width scaling\n",
        "    self.scale_e = 2  # torch.nn.Parameter(1.0*(torch.rand(1)))\n",
        "    # Colour scaling\n",
        "    self.scale_f = 1.5  # torch.nn.Parameter(3.0*(torch.rand(1)))\n",
        "    if LEARNABLE_FORCING:\n",
        "      self.scale_g = torch.nn.Parameter(1.0 * (torch.rand(1)))\n",
        "    if VERBOSE:\n",
        "      print(\"self.positions_top\", self.positions_top.shape)\n",
        "\n",
        "    lstms_top = []\n",
        "    heads_top = []\n",
        "    for _ in range(NUM_LSTMS):\n",
        "      lstm_layer_top = torch.nn.LSTM(\n",
        "          input_size=EMBEDDING_SIZE, hidden_size=NET_LSTM_HIDDENS,\n",
        "          num_layers=2, batch_first=True, bias=True)\n",
        "      if WEIGHT_INITIALIZER:\n",
        "        init_all(lstm_layer_top, torch.nn.init.normal_, mean=0., std=WEIGHT_STD)\n",
        "\n",
        "      head_layer_top = torch.nn.Sequential(\n",
        "          torch.nn.Linear(NET_LSTM_HIDDENS, NET_MLP_HIDDENS),\n",
        "          torch.nn.LeakyReLU(0.2, inplace=True),\n",
        "          torch.nn.Linear(NET_MLP_HIDDENS, OUTPUT_SIZE))\n",
        "\n",
        "      lstms_top.append(lstm_layer_top)\n",
        "      heads_top.append(head_layer_top)\n",
        "    self._lstms_top = torch.nn.ModuleList(lstms_top)\n",
        "    self._heads_top = torch.nn.ModuleList(heads_top)\n",
        "\n",
        "    lstms = []\n",
        "    heads = []\n",
        "    for _ in range(NUM_LSTMS):\n",
        "\n",
        "      lstm_layer = torch.nn.LSTM(\n",
        "          input_size=EMBEDDING_SIZE, hidden_size=NET_LSTM_HIDDENS,\n",
        "          num_layers=2, batch_first=True, bias=True)\n",
        "\n",
        "      if WEIGHT_INITIALIZER:\n",
        "        init_all(lstm_layer, torch.nn.init.normal_, mean=0., std=WEIGHT_STD)\n",
        "\n",
        "      head_layer = torch.nn.Sequential(\n",
        "          torch.nn.Linear(NET_LSTM_HIDDENS, NET_MLP_HIDDENS),\n",
        "          torch.nn.LeakyReLU(0.2, inplace=True),\n",
        "          torch.nn.Linear(NET_MLP_HIDDENS, OUTPUT_SIZE))\n",
        "      lstms.append(lstm_layer)\n",
        "      heads.append(head_layer)\n",
        "    self._lstms = torch.nn.ModuleList(lstms)\n",
        "    self._heads = torch.nn.ModuleList(heads)\n",
        "\n",
        "  def forward(self):\n",
        "    # TOP LEVEL\n",
        "    self.interleaved_positions_top = self.positions_top[:, :]\n",
        "    self.interleaved_positions_top = self.interleaved_positions_top.unsqueeze(\n",
        "        1).repeat(1, LSTM_ITERATIONS, 1)\n",
        "    if VERBOSE:\n",
        "      print(\"self.interleaved_positions_top\",\n",
        "            self.interleaved_positions_top.shape)\n",
        "\n",
        "    self.softmax_logits_top = self.positions_top[:, 2:2 + NUM_LSTMS]\n",
        "\n",
        "    if VERBOSE:\n",
        "      print(\"self.softmax_logits_top\", self.softmax_logits_top.shape)\n",
        "\n",
        "    self.non_parametric_positions_top = self.positions_top[:, :2]\n",
        "\n",
        "    if VERBOSE:\n",
        "      print(\"self.non_parametric_positions_top\",\n",
        "            self.non_parametric_positions_top.shape)\n",
        "\n",
        "    self.non_parametric_positions_repeated_top = (\n",
        "        self.non_parametric_positions_top.unsqueeze(1).repeat(\n",
        "            1, LSTM_ITERATIONS, 1).reshape(\n",
        "                BATCH_SIZE*LSTM_ITERATIONS, 2))\n",
        "\n",
        "    if VERBOSE:\n",
        "      print(\"self.non_parametric_positions_repeated_top\",\n",
        "            self.non_parametric_positions_repeated_top.shape)\n",
        "\n",
        "    self.softmax_outputs_top = torch.nn.functional.normalize(\n",
        "        self.softmax_logits_top, p=2, dim=1)\n",
        "\n",
        "    #TOP LEVEL NETWORK FORWARD PASS\n",
        "    pred_top = []\n",
        "    for i in range(NUM_LSTMS):\n",
        "      x_top = self.interleaved_positions_top*10.0\n",
        "      if USE_DROPOUT:\n",
        "        y_top, _ = self._lstms_top[i](\n",
        "            torch.nn.Dropout(DROPOUT_PROP)(x_top))\n",
        "        y_top = self._heads_top[i](\n",
        "            torch.nn.Dropout(DROPOUT_PROP)(y_top))\n",
        "      else:\n",
        "        y_top, _ = self._lstms_top[i](x_top)\n",
        "        y_top = self._heads_top[i](y_top)\n",
        "      pred_top.append(y_top)\n",
        "    preds_top = torch.stack(pred_top, axis=1)\n",
        "    out_top = torch.einsum(\"bijk,bi->bjk\", preds_top, self.softmax_outputs_top)\n",
        "    out_top = out_top.reshape(BATCH_SIZE*LSTM_ITERATIONS, EMBEDDING_SIZE)\n",
        "\n",
        "    # TOGGLE CRASHES TO INCLUDE THE TOP LEVEL NETWORK IN THE GRAPH.\n",
        "    crashes = True\n",
        "    if not crashes:\n",
        "      input_layer = self.input\n",
        "    else:\n",
        "      input_layer = out_top\n",
        "\n",
        "    # BOTTOM LEVEL\n",
        "    if VERBOSE:\n",
        "      print(\"BOTTOM LEVEL <<<<<<<<<<<<\")\n",
        "\n",
        "    self.non_parametric_positions = (self.non_parametric_positions_repeated_top\n",
        "                                     + self.scale_a * input_layer[:, :2])\n",
        "    if VERBOSE:\n",
        "      print(\"self.non_parametric_positions\",\n",
        "            self.non_parametric_positions.shape)\n",
        "\n",
        "    self.non_parametric_positions_repeated = (\n",
        "        self.non_parametric_positions.unsqueeze(1).repeat(\n",
        "            1, LSTM_ITERATIONS, 1).reshape(\n",
        "                (BATCH_SIZE * LSTM_ITERATIONS * LSTM_ITERATIONS, 2)))\n",
        "    if VERBOSE:\n",
        "      print(\"self.non_parametric_positions_repeated\",\n",
        "            self.non_parametric_positions_repeated.shape)\n",
        "    interleaved_positions = self.scale_b * input_layer[:, :]\n",
        "\n",
        "    interleaved_positions = interleaved_positions.unsqueeze(1).repeat(\n",
        "        1, LSTM_ITERATIONS, 1)\n",
        "    if VERBOSE:\n",
        "      print(\"interleaved_positions\", interleaved_positions.shape)\n",
        "    softmax_logits = self.scale_c * input_layer[:, 2:2 + NUM_LSTMS]\n",
        "    softmax_outputs = torch.nn.functional.normalize(softmax_logits, p=2, dim=1)\n",
        "\n",
        "    pred = []\n",
        "    for i in range(NUM_LSTMS):\n",
        "      x = interleaved_positions*10.0\n",
        "      if USE_DROPOUT:\n",
        "        y, _ = self._lstms[i](torch.nn.Dropout(DROPOUT_PROP)(x))\n",
        "        y = self._heads[i](torch.nn.Dropout(DROPOUT_PROP)(y))\n",
        "      else:\n",
        "        y, _ = self._lstms[i](x)\n",
        "        y = self._heads[i](y)\n",
        "\n",
        "      pred.append(y)\n",
        "    preds = torch.stack(pred, axis=1)\n",
        "    out = torch.einsum(\"bijk,bi->bjk\", preds, softmax_outputs)\n",
        "    out = out.reshape(\n",
        "        (BATCH_SIZE * LSTM_ITERATIONS * LSTM_ITERATIONS, EMBEDDING_SIZE))\n",
        "\n",
        "    self.non_parametric_positions_repeated = (\n",
        "        self.non_parametric_positions_repeated.repeat(1, 4))\n",
        "    if VERBOSE:\n",
        "      print(\"self.non_parametric_positions_repeated\",\n",
        "            self.non_parametric_positions_repeated.shape)\n",
        "\n",
        "    if VERBOSE:\n",
        "      print(\"out shape\", out.shape)\n",
        "\n",
        "    if YELLOW_ABSOLUTE_POSITIONS_USED:\n",
        "      points = (out[:, :8] * self.scale_d\n",
        "                + self.non_parametric_positions_repeated)\n",
        "    else:\n",
        "      points = (out[:, :8] * self.scale_d\n",
        "                + self.non_parametric_positions_repeated * 0.0)\n",
        "\n",
        "    points = points.reshape(\n",
        "        (BATCH_SIZE*LSTM_ITERATIONS * LSTM_ITERATIONS, 4, 2))\n",
        "    if LEARNABLE_FORCING:\n",
        "      points = (points * self.scale_g\n",
        "                + torch.nn.self.direct_points * (1-self.scale_g))\n",
        "    else:\n",
        "      points = (points * GRAMMATICAL_FORCING\n",
        "                + self.direct_points * (1-GRAMMATICAL_FORCING))\n",
        "\n",
        "    widths = out[:, 8] * self.scale_e * 20.0\n",
        "    colours = out[:, 9:13] * self.scale_f * 6.0\n",
        "\n",
        "    points = torch.clamp(points, min=-1.1, max=1.1)\n",
        "    widths = torch.clamp(widths, min=1, max=20)\n",
        "    colours = torch.clamp(colours, min=0.0, max=1)\n",
        "\n",
        "    return points, widths, colours, 0"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "ZOiMPnJLwlqA"
      },
      "source": [
        "#@title Grammar Specific Drawing Network: DPPN Generative Grammar\n",
        "\n",
        "class CurveNetworkDPPN(torch.nn.Module):\n",
        "  \"\"\"DPPN-based production of drawing sequences.\"\"\"\n",
        "\n",
        "  def __init__(self):\n",
        "    super().__init__()\n",
        "\n",
        "    def init_all(model, init_func, *params, **kwargs):\n",
        "      for p in model.parameters():\n",
        "        init_func(p, *params, **kwargs)\n",
        "\n",
        "    # ***************************Direct encoding**********************\n",
        "    points_in = []\n",
        "    for _ in range(BATCH_SIZE):\n",
        "      p0 = (1.5*(random.random() - 0.5), 1.5 * (random.random() - 0.5))\n",
        "      points_in.append(p0)\n",
        "      radius = 0.1\n",
        "      p1 = (p0[0] + radius * (random.random() - 0.5),\n",
        "            p0[1] + radius * (random.random() - 0.5))\n",
        "      p2 = (p1[0] + radius * (random.random() - 0.5),\n",
        "            p1[1] + radius * (random.random() - 0.5))\n",
        "      p3 = (p2[0] + radius * (random.random() - 0.5),\n",
        "            p2[1] + radius * (random.random() - 0.5))\n",
        "      points_in.append(p1)\n",
        "      points_in.append(p2)\n",
        "      points_in.append(p3)\n",
        "      p0 = p3\n",
        "\n",
        "    self.direct_points = torch.nn.Parameter(torch.Tensor(\n",
        "        np.array(points_in).flatten().reshape((BATCH_SIZE, 4, 2))))\n",
        "\n",
        "    # ****************************Indirect encoding********************\n",
        "\n",
        "    self.positions_top = torch.nn.Parameter(\n",
        "        2.0 * (torch.rand(BATCH_SIZE, EMBEDDING_SIZE) - 0.5))\n",
        "\n",
        "    self.scale_d = 0.2  #torch.nn.Parameter(10.0*(torch.rand(1))) #Point scaling\n",
        "    self.scale_e = 4  #torch.nn.Parameter(1.0*(torch.rand(1))) #Width scaling\n",
        "    self.scale_f = 1  #torch.nn.Parameter(3.0*(torch.rand(1))) #Colour scaling\n",
        "    if LEARNABLE_FORCING:\n",
        "      self.scale_g = torch.nn.Parameter(1.0*(torch.rand(1)))\n",
        "    if VERBOSE:\n",
        "      print(\"self.positions_top\", self.positions_top.shape)\n",
        "\n",
        "    self.ffnn = torch.nn.Sequential(\n",
        "        torch.nn.Linear(EMBEDDING_SIZE, NET_MLP_HIDDENS),\n",
        "        torch.nn.LeakyReLU(0.2, inplace=True),\n",
        "        torch.nn.Linear(NET_MLP_HIDDENS, NET_MLP_HIDDENS),\n",
        "        torch.nn.LeakyReLU(0.2, inplace=True),\n",
        "        torch.nn.Linear(NET_MLP_HIDDENS, NET_MLP_HIDDENS),\n",
        "        torch.nn.LeakyReLU(0.2, inplace=True),\n",
        "        torch.nn.Linear(NET_MLP_HIDDENS, OUTPUT_SIZE))\n",
        "\n",
        "    if WEIGHT_INITIALIZER:\n",
        "      init_all(self.ffnn, torch.nn.init.normal_, mean=0., std=WEIGHT_STD)\n",
        "\n",
        "  def forward(self):\n",
        "    # TOP LEVEL\n",
        "    self.non_parametric_positions_top = self.positions_top[:, :2]\n",
        "    out = self.ffnn(self.positions_top)\n",
        "    self.non_parametric_positions_repeated = (\n",
        "        self.non_parametric_positions_top.repeat(1, 4))\n",
        "    points = (out[:, :8] * self.scale_d\n",
        "              + self.non_parametric_positions_repeated)\n",
        "\n",
        "    points = points.reshape((BATCH_SIZE, 4, 2))\n",
        "    if LEARNABLE_FORCING:\n",
        "      points = (points * self.scale_g\n",
        "                + torch.nn.self.direct_points * (1 - self.scale_g))\n",
        "    else:\n",
        "      points = (points * GRAMMATICAL_FORCING\n",
        "                + self.direct_points * (1 - GRAMMATICAL_FORCING))\n",
        "\n",
        "    widths = out[:, 8] * self.scale_e\n",
        "    colours = out[:, 9:13] * self.scale_f\n",
        "\n",
        "    points = torch.clamp(points, min=-1.1, max=1.1)\n",
        "    widths = torch.clamp(widths, min=1, max=20)\n",
        "    colours = torch.clamp(colours, min=0.0, max=1)\n",
        "\n",
        "    return points, widths, colours, 0\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "d3AgGcGuASFa"
      },
      "source": [
        "#@title Grammar Specific Drawing Network: SEQ-TO-SEQ\n",
        "class EncoderRNN(torch.nn.Module):\n",
        "  \"\"\"RNN-based sequence encoder.\"\"\"\n",
        "\n",
        "  def __init__(self, input_size, hidden_size):\n",
        "    \"\"\"Constructor.\n",
        "\n",
        "    Args:\n",
        "      input_size: number of input units\n",
        "      hidden_size: number of hidden units\n",
        "    \"\"\"\n",
        "    super(EncoderRNN, self).__init__()\n",
        "    self.hidden_size = hidden_size\n",
        "    self.gru = torch.nn.GRU(\n",
        "        input_size=input_size,\n",
        "        hidden_size=hidden_size,\n",
        "        num_layers=2,\n",
        "        batch_first=True,\n",
        "        bias=True)\n",
        "\n",
        "  def forward(self, input_layer, hidden_layer):\n",
        "\n",
        "    output_layer, hidden_layer = self.gru(input_layer, hidden_layer)\n",
        "    return output_layer, hidden_layer\n",
        "\n",
        "  def init_hidden(self):\n",
        "    return torch.zeros(2, BATCH_SIZE, self.hidden_size)\n",
        "\n",
        "\n",
        "class DecoderRNN(torch.nn.Module):\n",
        "  \"\"\"RNN-based sequence decoder.\"\"\"\n",
        "\n",
        "  def __init__(self, hidden_size, output_size):\n",
        "    \"\"\"Constructor.\n",
        "\n",
        "    Args:\n",
        "      hidden_size: number of hidden units\n",
        "      output_size: number of output units\n",
        "    \"\"\"\n",
        "    super(DecoderRNN, self).__init__()\n",
        "    self.hidden_size = hidden_size\n",
        "\n",
        "    self.embedding = torch.nn.Embedding(output_size, hidden_size)\n",
        "    self.gru = torch.nn.GRU(\n",
        "        input_size=hidden_size,\n",
        "        hidden_size=hidden_size,\n",
        "        num_layers=2,\n",
        "        batch_first=True,\n",
        "        bias=True)\n",
        "    self.out = torch.nn.Sequential(\n",
        "        torch.nn.Linear(hidden_size, hidden_size),\n",
        "        torch.nn.LeakyReLU(0.2, inplace=True),\n",
        "        torch.nn.Linear(hidden_size, output_size))\n",
        "\n",
        "  def forward(self, input_layer, hidden):\n",
        "    output = self.embedding(input_layer)\n",
        "    output = F.relu(output)\n",
        "    output, hidden = self.gru(output, hidden)\n",
        "    output = self.out(output)\n",
        "    return output, hidden\n",
        "\n",
        "\n",
        "class CurveNetworkSEQTOSEQ(torch.nn.Module):\n",
        "  \"\"\"Sequence-to-sequence-based production of drawing sequences.\"\"\"\n",
        "\n",
        "  def __init__(self):\n",
        "    \"\"\"Constructor, relying on global parameters.\"\"\"\n",
        "    super().__init__()\n",
        "\n",
        "    def init_all(model, init_func, *params, **kwargs):\n",
        "      for p in model.parameters():\n",
        "        init_func(p, *params, **kwargs)\n",
        "\n",
        "    self.hidden_size = NET_LSTM_HIDDENS\n",
        "    self.input_length = SEQ_LENGTH\n",
        "    self.output_length = NUM_STROKES\n",
        "    self.input_size = INPUT_SIZE\n",
        "    self.output_size = OUTPUT_SIZE\n",
        "\n",
        "    # *************************************Direct encoding*************************************\n",
        "    points_in = []\n",
        "    for _ in range(self.output_length * BATCH_SIZE):\n",
        "      p0 = (1.5 * (random.random() - 0.5), 1.5 * (random.random() - 0.5))\n",
        "      points_in.append(p0)\n",
        "      radius = 0.1\n",
        "      p1 = (p0[0] + radius * (random.random() - 0.5),\n",
        "            p0[1] + radius * (random.random() - 0.5))\n",
        "      p2 = (p1[0] + radius * (random.random() - 0.5),\n",
        "            p1[1] + radius * (random.random() - 0.5))\n",
        "      p3 = (p2[0] + radius * (random.random() - 0.5),\n",
        "            p2[1] + radius * (random.random() - 0.5))\n",
        "      points_in.append(p1)\n",
        "      points_in.append(p2)\n",
        "      points_in.append(p3)\n",
        "      p0 = p3\n",
        "\n",
        "    self.direct_points = torch.nn.Parameter(\n",
        "        torch.Tensor(\n",
        "            np.array(points_in).flatten().reshape(\n",
        "                (self.output_length * BATCH_SIZE, 4, 2))))\n",
        "\n",
        "    #input_tensor\n",
        "    self.positions_top = torch.nn.Parameter(\n",
        "        2.0 * (torch.rand(BATCH_SIZE, SEQ_LENGTH, self.input_size) - 0.5))\n",
        "    self.positions_top2 = torch.nn.Parameter(2.0 *\n",
        "                                             (torch.rand(BATCH_SIZE, 2) - 0.5))\n",
        "\n",
        "    self.scale_d = 0.4  #torch.nn.Parameter(10.0*(torch.rand(1))) #Point scaling\n",
        "    self.scale_e = 3.0  #torch.nn.Parameter(1.0*(torch.rand(1))) #Width scaling\n",
        "    self.scale_f = 1.0  #torch.nn.Parameter(3.0*(torch.rand(1))) #Colour scaling\n",
        "\n",
        "    if LEARNABLE_FORCING:\n",
        "      self.scale_g = torch.nn.Parameter(1.0 * (torch.rand(1)))\n",
        "    if VERBOSE:\n",
        "      print(\"self.positions_top\", self.positions_top.shape)\n",
        "\n",
        "    self.encoder1 = EncoderRNN(self.input_size, self.hidden_size)\n",
        "    self.decoder1 = DecoderRNN(self.hidden_size, self.output_size)\n",
        "    if WEIGHT_INITIALIZER:\n",
        "      init_all(self.encoder1, torch.nn.init.normal_, mean=0., std=WEIGHT_STD)\n",
        "      init_all(self.decoder1, torch.nn.init.normal_, mean=0., std=WEIGHT_STD)\n",
        "\n",
        "    self.encoder_outputs = torch.zeros(self.output_length,\n",
        "                                       self.encoder1.hidden_size)\n",
        "    self.decoder_input = torch.tensor(np.zeros((BATCH_SIZE, 1), dtype=int))\n",
        "\n",
        "  def forward(self):\n",
        "\n",
        "    self.non_parametric_positions_repeated = (\n",
        "        self.positions_top2.unsqueeze(1).repeat(\n",
        "            1, self.output_length, 1).reshape(BATCH_SIZE, self.output_length,\n",
        "                                              2))\n",
        "    self.non_parametric_positions_repeated = (\n",
        "        self.non_parametric_positions_repeated.repeat(1, 1, 4).reshape(\n",
        "            (BATCH_SIZE * self.output_length, 8)))\n",
        "\n",
        "    self.encoder_hidden = self.encoder1.init_hidden()\n",
        "\n",
        "    # Input-less forward function.\n",
        "    self.encoder_output, self.encoder_hidden = self.encoder1(\n",
        "        self.positions_top * 3.0, self.encoder_hidden)\n",
        "\n",
        "    decoder_hidden = self.encoder_hidden\n",
        "    decoder_input = self.decoder_input\n",
        "    output_history = []\n",
        "\n",
        "    # Without teacher forcing: use its own predictions as the next input\n",
        "    for _ in range(self.output_length):\n",
        "\n",
        "      decoder_output, decoder_hidden = self.decoder1(decoder_input,\n",
        "                                                     decoder_hidden)\n",
        "\n",
        "      _, topi = decoder_output.topk(1)\n",
        "      # detach from history as input\n",
        "      decoder_input = topi.squeeze().detach().reshape((BATCH_SIZE, 1))\n",
        "      output_history.append(decoder_output)\n",
        "\n",
        "    out = torch.stack(output_history, axis=0)\n",
        "    out = torch.reshape(out,\n",
        "                        (BATCH_SIZE * self.output_length, self.output_size))\n",
        "\n",
        "    if VERBOSE:\n",
        "      print(\"out shape\", out.shape)\n",
        "    if YELLOW_ABSOLUTE_POSITIONS_USED:\n",
        "\n",
        "      points = (\n",
        "          out[:, :8] * self.scale_d + self.non_parametric_positions_repeated)\n",
        "    else:\n",
        "      points = out[:, :8] * self.scale_d\n",
        "\n",
        "    points = points.reshape(self.output_length * BATCH_SIZE, 4, 2)\n",
        "    if LEARNABLE_FORCING:\n",
        "      points = (\n",
        "          points * (self.scale_g) + self.direct_points * (1 - self.scale_g))\n",
        "    else:\n",
        "      points = (\n",
        "          points * (GRAMMATICAL_FORCING) + self.direct_points *\n",
        "          (1 - GRAMMATICAL_FORCING))\n",
        "\n",
        "    widths = out[:, 8] * self.scale_e\n",
        "    colours = out[:, 9:13] * self.scale_f\n",
        "\n",
        "    points = torch.clamp(points, min=-1.1, max=1.1)\n",
        "    widths = torch.clamp(widths, min=1, max=20)\n",
        "    colours = torch.clamp(colours, min=0.0, max=1)\n",
        "\n",
        "    return points, widths, colours, 0\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GIXzueO3PB-4"
      },
      "source": [
        "# Function definitions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "urXU9M5hS6H0"
      },
      "source": [
        "#@title Image rendering and display\n",
        "\n",
        "def show_and_save(img, t=None, dpi=75, figsize=(5, 5)):\n",
        "  \"\"\"Display image.\n",
        "\n",
        "  Args:\n",
        "    img: image to display\n",
        "    t: time step\n",
        "    dpi: display resolution\n",
        "    figsize: size of image\n",
        "  \"\"\"\n",
        "  _ = plt.figure(figsize=figsize, dpi=dpi)\n",
        "  img = np.transpose(img.detach().cpu().numpy()[0], (1, 2, 0))\n",
        "  img = np.clip(img, 0.0, 1.0)\n",
        "  plt.imshow(img, interpolation=\"None\")\n",
        "  plt.grid(None)\n",
        "  plt.axis(\"off\")\n",
        "  if img.shape[1] > CANVAS_WIDTH:\n",
        "    type_image = \"highres_image_\"\n",
        "  else:\n",
        "    type_image = \"image_\"\n",
        "  path_fig = DIR_RESULTS + \"/\" + type_image + PROMPT\n",
        "  if t is not None:\n",
        "    path_fig += \"_t_\" + str(t)\n",
        "  path_fig += \".png\"\n",
        "  plt.savefig(path_fig, dpi=dpi)\n",
        "  plt.show()\n",
        "\n",
        "\n",
        "def render(all_points, all_widths, all_colors, multiplier=1):\n",
        "  \"\"\"Render line date to image.\n",
        "\n",
        "  Args:\n",
        "    all_points: points defining the lines\n",
        "    all_widths: the widths of the lines\n",
        "    all_colors: line colours\n",
        "    multiplier: scale factor to enlarge drawing\n",
        "\n",
        "  Returns:\n",
        "    image with lines drawn\n",
        "  \"\"\"\n",
        "  # Store `all_points` definitions as shapes, colours and widths.\n",
        "  shapes = []\n",
        "  shape_groups = []\n",
        "  for p in range(all_points.shape[0]):\n",
        "    points = all_points[p].contiguous().cpu()\n",
        "    width = all_widths[p].cpu()\n",
        "    color = all_colors[p].cpu()\n",
        "    num_ctrl_pts = torch.zeros(NUM_SEGMENTS, dtype=torch.int32) + 2\n",
        "    path = pydiffvg.Path(\n",
        "        num_control_points=num_ctrl_pts, points=points * multiplier,\n",
        "        stroke_width=width * multiplier, is_closed=False)\n",
        "    shapes.append(path)\n",
        "    path_group = pydiffvg.ShapeGroup(\n",
        "        shape_ids=torch.tensor([len(shapes) - 1]),\n",
        "        fill_color=None,\n",
        "        stroke_color=color)\n",
        "    shape_groups.append(path_group)\n",
        "\n",
        "  # Rasterize the image.\n",
        "  scene_args = pydiffvg.RenderFunction.serialize_scene(\n",
        "      CANVAS_WIDTH * multiplier,\n",
        "      CANVAS_HEIGHT * multiplier,\n",
        "      shapes, shape_groups)\n",
        "  img = pydiffvg.RenderFunction.apply(\n",
        "      CANVAS_WIDTH * multiplier,\n",
        "      CANVAS_HEIGHT * multiplier,\n",
        "      2, 2, 0, None, *scene_args)\n",
        "  if DRAW_WHITE_BACKGROUND:\n",
        "    w, h = img.shape[0], img.shape[1]\n",
        "    img = img[:, :, 3:4] * img[:, :, :3] + (\n",
        "        torch.ones(w, h, 3, device=pydiffvg.get_device()) * (1-img[:, :, 3:4]))\n",
        "  else:\n",
        "    img = img[:, :, :3]\n",
        "  img = img.unsqueeze(0)\n",
        "\n",
        "  return img"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "ApgzwCVJcU7v"
      },
      "source": [
        "#@title Image augmentation transformations\n",
        "\n",
        "def augmentation_transforms(canvas_width, use_normalized_clip):\n",
        "  \"\"\"Image transforms to produce distorted crops to augment the evaluation.\n",
        "\n",
        "  Args:\n",
        "    canvas_width: width of the drawing canvas\n",
        "    use_normalized_clip: Normalisation to better suit CLIP's training data\n",
        "\n",
        "  Returns:\n",
        "    transforms\n",
        "  \"\"\"\n",
        "  if use_normalized_clip:\n",
        "    augment_trans = transforms.Compose(\n",
        "        [transforms.RandomPerspective(fill=1, p=1, distortion_scale=0.6),\n",
        "         transforms.RandomResizedCrop(canvas_width, scale=(0.7, 0.9)),\n",
        "         transforms.Normalize((0.48145466, 0.4578275, 0.40821073),\n",
        "                              (0.26862954, 0.26130258, 0.27577711))])\n",
        "  else:\n",
        "    augment_trans = transforms.Compose([\n",
        "        transforms.RandomPerspective(fill=1, p=1, distortion_scale=0.6),\n",
        "        transforms.RandomResizedCrop(canvas_width, scale=(0.7, 0.9)),\n",
        "    ])\n",
        "\n",
        "  return augment_trans"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "uvMiX61xGEOS"
      },
      "source": [
        "#@title Video creator {vertical-output: true}\n",
        "\n",
        "class VideoWriter:\n",
        "  \"\"\"Create a video from image frames.\"\"\"\n",
        "\n",
        "  def __init__(self, filename=\"_autoplay.mp4\", fps=20.0, **kw):\n",
        "    \"\"\"Video creator.\n",
        "\n",
        "    Creates and display a video made from frames. The default\n",
        "    filename causes the video to be displayed on exit.\n",
        "\n",
        "    Args:\n",
        "      filename: name of video file\n",
        "      fps: frames per second for video\n",
        "      **kw: args to be passed to FFMPEG_VideoWriter\n",
        "\n",
        "    Returns:\n",
        "      VideoWriter instance.\n",
        "    \"\"\"\n",
        "\n",
        "    self.writer = None\n",
        "    self.params = dict(filename=filename, fps=fps, **kw)\n",
        "\n",
        "  def add(self, img):\n",
        "    \"\"\"Add image to video.\n",
        "\n",
        "    Add new frame to image file, creating VideoWriter if requried.\n",
        "\n",
        "    Args:\n",
        "      img: array-like frame, shape [X, Y, 3] or [X, Y]\n",
        "\n",
        "    Returns:\n",
        "      None\n",
        "    \"\"\"\n",
        "\n",
        "    img = np.asarray(img)\n",
        "    if self.writer is None:\n",
        "      h, w = img.shape[:2]\n",
        "      self.writer = FFMPEG_VideoWriter(size=(w, h), **self.params)\n",
        "    if img.dtype in [np.float32, np.float64]:\n",
        "      img = np.uint8(img.clip(0, 1)*255)\n",
        "    if len(img.shape) == 2:\n",
        "      img = np.repeat(img[..., None], 3, -1)\n",
        "    self.writer.write_frame(img)\n",
        "\n",
        "  def close(self):\n",
        "    if self.writer:\n",
        "      self.writer.close()\n",
        "\n",
        "  def __enter__(self):\n",
        "    return self\n",
        "\n",
        "  def __exit__(self, *kw):\n",
        "    self.close()\n",
        "    if self.params[\"filename\"] == \"_autoplay.mp4\":\n",
        "      self.show()\n",
        "\n",
        "  def show(self, **kw):\n",
        "    \"\"\"Display video.\n",
        "\n",
        "    Args:\n",
        "      **kw: args to be passed to mvp.ipython_display\n",
        "\n",
        "    Returns:\n",
        "      None\n",
        "    \"\"\"\n",
        "    self.close()\n",
        "    fn = self.params[\"filename\"]\n",
        "    display(mvp.ipython_display(fn, **kw))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "4XIVMSJuWgxG"
      },
      "source": [
        "#@title Training functions\n",
        "\n",
        "def get_features(prompt, negative_prompt_1, negative_prompt_2):\n",
        "  # Tokenize prompts and coompute CLIP features.\n",
        "  text_input = clip.tokenize(prompt).to(device)\n",
        "  text_input_neg1 = clip.tokenize(negative_prompt_1).to(device)\n",
        "  text_input_neg2 = clip.tokenize(negative_prompt_2).to(device)\n",
        "  with torch.no_grad():\n",
        "    features = clip_model.encode_text(text_input)\n",
        "    neg1_features = clip_model.encode_text(text_input_neg1)\n",
        "    neg2_features = clip_model.encode_text(text_input_neg2)\n",
        "  return features, neg1_features, neg2_features\n",
        "\n",
        "# Create writers.\n",
        "def load_torch_img(filename):\n",
        "  img = PIL.Image.open(filename).convert(mode=\"RGB\")\n",
        "  img = img.resize((CANVAS_WIDTH, CANVAS_HEIGHT))\n",
        "  img = np.float32(img)\n",
        "  img = torch.from_numpy(img).to(torch.float32) / 255.0\n",
        "  #img = img.pow(2.0)\n",
        "  img = img.to(pydiffvg.get_device())\n",
        "  img = img.unsqueeze(0)\n",
        "  img = img.permute(0, 3, 1, 2)\n",
        "  return img\n",
        "\n",
        "\n",
        "def create_generator(grammar_type):\n",
        "  \"\"\"Create the drawing generator.\n",
        "\n",
        "  Args:\n",
        "    grammar_type: string defining the class of generator to use\n",
        "\n",
        "  Returns:\n",
        "    stroke generator instance\n",
        "  \"\"\"\n",
        "  if grammar_type == \"Arnheim 2\":\n",
        "    new_generator = CurveNetworkHierarchicalLSTM()\n",
        "  elif grammar_type == \"Photographic\":\n",
        "    new_generator = CurveNetworkPhotographicLSTM()\n",
        "  elif grammar_type == \"DPPN\":\n",
        "    new_generator = CurveNetworkDPPN()\n",
        "  elif grammar_type == \"SEQTOSEQ\":\n",
        "    new_generator = CurveNetworkSEQTOSEQ()\n",
        "  else:\n",
        "    print(\"Unknown drawing function:\", grammar_type)\n",
        "\n",
        "  if LOAD_MODEL:\n",
        "    state_dict = torch.load(DIR_RESULTS + \"/generator.pt\")\n",
        "    new_generator.load_state_dict(state_dict)\n",
        "    with torch.no_grad():\n",
        "      for name, param in new_generator.named_parameters():\n",
        "        if \"positions_top\" in name:\n",
        "          print(\"resetting positions TOP \", flush=True)\n",
        "          param.copy_(2.0*(torch.rand(BATCH_SIZE, EMBEDDING_SIZE)-0.5))\n",
        "        if \"direct_points\" in name:\n",
        "          print(\"resetting direct points\")\n",
        "          points_in = []\n",
        "          for _ in range(BATCH_SIZE * LSTM_ITERATIONS * LSTM_ITERATIONS):\n",
        "            p0 = (1.5*(random.random()-0.5), 1.5*(random.random()-0.5))\n",
        "            points_in.append(p0)\n",
        "            for _ in range(1):\n",
        "              radius = 0.1\n",
        "              p1 = (p0[0] + radius * (random.random() - 0.5),\n",
        "                    p0[1] + radius * (random.random() - 0.5))\n",
        "              p2 = (p1[0] + radius * (random.random() - 0.5),\n",
        "                    p1[1] + radius * (random.random() - 0.5))\n",
        "              p3 = (p2[0] + radius * (random.random() - 0.5),\n",
        "                    p2[1] + radius * (random.random() - 0.5))\n",
        "              points_in.append(p1)\n",
        "              points_in.append(p2)\n",
        "              points_in.append(p3)\n",
        "              p0 = p3\n",
        "\n",
        "          param.copy_(torch.Tensor(\n",
        "              np.array(points_in).flatten().reshape(\n",
        "                  (BATCH_SIZE*LSTM_ITERATIONS*LSTM_ITERATIONS, 4, 2))))\n",
        "  return new_generator\n",
        "\n",
        "def make_optimizer(generator, learning_rate, input_learing_rate, decay_rate):\n",
        "  \"\"\"Make optimizer for generator's parameters.\n",
        "\n",
        "  Args:\n",
        "    generator: generator model\n",
        "    learning_rate: learning rate\n",
        "    input_learing_rate: learning rate for input\n",
        "    decay_rate: optional learning ratge decay rate\n",
        "\n",
        "  Returns:\n",
        "    optimizer\n",
        "  \"\"\"\n",
        "  if decay_rate is not None:\n",
        "    optim = torch.optim.SGD(generator.parameters(), lr=learning_rate)\n",
        "    lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(\n",
        "        optimizer=optim, gamma=decay_rate)\n",
        "  else:\n",
        "    my_list = [\"positions_top\", \"positions_top2\"]\n",
        "    params = list(map(\n",
        "        lambda x: x[1],\n",
        "        list(filter(lambda kv: kv[0] in my_list,\n",
        "                    generator.named_parameters()))))\n",
        "    base_params = list(map(\n",
        "        lambda x: x[1],\n",
        "        list(filter(\n",
        "            lambda kv: kv[0] not in my_list, generator.named_parameters()\n",
        "            ))))\n",
        "    lr_scheduler = torch.optim.SGD(\n",
        "        [{\"params\": base_params}, {\"params\": params, \"lr\": input_learing_rate}],\n",
        "        lr=learning_rate)\n",
        "  return lr_scheduler\n",
        "\n",
        "def step_optimization(t, clip_enc, lr_scheduler, generator, augment_trans,\n",
        "                      text_features, final_step=False):\n",
        "  \"\"\"Do a step of optimization.\n",
        "\n",
        "  Args:\n",
        "    t: step count\n",
        "    clip_enc: model for CLIP encoding\n",
        "    lr_scheduler: optimizer\n",
        "    generator: drawing generator to optimise\n",
        "    augment_trans: transforms for image augmentation\n",
        "    text_features: tuple with the prompt two negative prompts\n",
        "    final_step: if True does extras such as saving the model\n",
        "  \"\"\"\n",
        "  # Anneal learning rate (NOTE THIS REDUCES the learning rate for the LSTM\n",
        "  # whether USE_DECAY is set or not!!!)\n",
        "  if t == int(OPTIM_STEPS * 0.5):\n",
        "    for g in lr_scheduler.param_groups:\n",
        "      g[\"lr\"] = g[\"lr\"] / 2.0\n",
        "  if t == int(OPTIM_STEPS * 0.75):\n",
        "    for g in lr_scheduler.param_groups:\n",
        "      g[\"lr\"] = g[\"lr\"] / 2.0\n",
        "\n",
        "  # Rebuild the generator.\n",
        "  t0 = time.time()\n",
        "  lr_scheduler.zero_grad()\n",
        "\n",
        "  all_points, all_widths, all_colors, _ = generator()\n",
        "\n",
        "  if isinstance(generator, CurveNetworkPhotographicLSTM):\n",
        "    all_widths = all_widths[0:NUM_PATHS]\n",
        "    all_points = all_points[0:(2 * NUM_PATHS * (NUM_SEGMENTS * 3 + 1))]\n",
        "    all_points = all_points.view(NUM_PATHS, -1, 2)\n",
        "    all_points = all_points * (CANVAS_HEIGHT // 2 - 2) + CANVAS_HEIGHT // 2\n",
        "    all_colors = all_colors[:(NUM_PATHS * 4)].view(NUM_PATHS, 4)\n",
        "  else:\n",
        "    all_points = all_points * (CANVAS_HEIGHT // 2 - 2) + CANVAS_HEIGHT // 2\n",
        "\n",
        "  # Convert points to Bézier curves, widths and colours, and rasterize to img.\n",
        "  t1 = time.time()\n",
        "  img = render(all_points, all_widths, all_colors)\n",
        "\n",
        "  if t % VIDEO_STEPS == 0:\n",
        "    # Write image to video.\n",
        "    t2 = time.time()\n",
        "    video_writer.add(img.cpu().detach().numpy()[0])\n",
        "\n",
        "  # Compute and add losses after augmenting the image with transforms.\n",
        "  t3 = time.time()\n",
        "  img = img.permute(0, 3, 1, 2)  # NHWC -> NCHW\n",
        "  loss = 0\n",
        "  if not MSE_LOSS:\n",
        "    img_augs = []\n",
        "    for n in range(NUM_AUGS):\n",
        "      img_augs.append(augment_trans(img))\n",
        "    img_batch = torch.cat(img_augs)\n",
        "    image_features = clip_enc.encode_image(img_batch)\n",
        "    for n in range(NUM_AUGS):\n",
        "      loss -= torch.cosine_similarity(text_features[0],\n",
        "                                      image_features[n:n+1], dim=1)\n",
        "      if USE_NEG_PROMPTS:\n",
        "        loss += torch.cosine_similarity(text_features[1],\n",
        "                                        image_features[n:n+1], dim=1) * 0.3\n",
        "        loss += torch.cosine_similarity(text_features[2],\n",
        "                                        image_features[n:n+1], dim=1) * 0.3\n",
        "  else:\n",
        "    loss += (img - target_img).pow(2).mean()\n",
        "\n",
        "  if WIDTH_DIVERSITY_LOSS:\n",
        "    width_diversity_loss = torch.std(all_widths)\n",
        "    print(\"clip loss = \", loss)\n",
        "    print(\"WDL = \", -WDL_COEFFICIENT * width_diversity_loss)\n",
        "    loss = loss - WDL_COEFFICIENT*width_diversity_loss\n",
        "\n",
        "  writer.add_scalar(\"Loss/train\", loss, t)\n",
        "  writer.flush()\n",
        "\n",
        "  # Backpropagate the gradients.\n",
        "  t4 = time.time()\n",
        "  loss.backward()\n",
        "  # Decay the learning rate.\n",
        "  lr_scheduler.step()\n",
        "\n",
        "  # Render the big version.\n",
        "  t5 = time.time()\n",
        "  if final_step:\n",
        "    img_big = render(all_points, all_widths, all_colors,\n",
        "                     multiplier=MULTIPLIER_BIG_IMAGE)\n",
        "    img_big = img_big.permute(0, 3, 1, 2)  # NHWC -> NCHW\n",
        "    show_and_save(img_big, t=t, dpi=300)\n",
        "    print(\"Saving model...\")\n",
        "    torch.save(generator.state_dict(), DIR_RESULTS + \"/generator.pt\")\n",
        "\n",
        "  # Trace the learning error and images.\n",
        "  if t % TRACE_EVERY == 0:\n",
        "    # Send gradients to tensorboard.\n",
        "    grads = []\n",
        "    for name, param in generator.named_parameters():\n",
        "      grads.append(param.grad.view(-1))\n",
        "      writer.add_histogram(name, param.grad.view(-1), t)\n",
        "    # Show and trace.\n",
        "    show_and_save(img, dpi=75)\n",
        "    if USE_DECAY:\n",
        "      lr = lr_scheduler.get_last_lr()[0]\n",
        "    else:\n",
        "      lr = LEARNING_RATE\n",
        "    print(\"Iteration {:3d}, lr {}, rendering loss {:.6f}\".format(\n",
        "        t, lr, loss.item()))\n",
        "\n",
        "  t6 = time.time()\n",
        "  if PLOT_DURATIONS:\n",
        "    print(f\"gen_fwd: {t1-t0:.4f}s render: {t2-t1:.4f}s video: {t3-t2:.4f}s clip_loss: {t4-t3:.4f}s bprop: {t5-t4:.4f}s big+trace: {t6-t5:.4f}s\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NvgEJY4XGfGI"
      },
      "source": [
        "# Hyperparameters"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "9dqM1FUdELjn"
      },
      "source": [
        "\n",
        "#@markdown Drawing size for evaluation\n",
        "CANVAS_WIDTH = 224  #@param {type:\"integer\"}\n",
        "CANVAS_HEIGHT = 224  #@param {type:\"integer\"}\n",
        "\n",
        "#@markdown Relative size of final large image\n",
        "MULTIPLIER_BIG_IMAGE = 10  #@param {type:\"integer\"}\n",
        "\n",
        "#@markdown Number of augmentations to use in evaluation\n",
        "NUM_AUGS = 4  #@param {type:\"integer\"}\n",
        "\n",
        "#@markdown Extra loss to increase diversity of stroke widths\n",
        "WIDTH_DIVERSITY_LOSS = False  #@param {type:\"boolean\"}\n",
        "WDL_COEFFICIENT = 0.01  #@param {type:\"number\"}\n",
        "\n",
        "#@markdown Ought to be True but results are better when False\n",
        "USE_NORMALIZED_CLIP = False  #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown MSE_LOSS attempts to match an image loaded from Drive\n",
        "MSE_LOSS = False  #@param {type:\"boolean\"}\n",
        "DRIVE_IMAGE_FOR_MSE_LOSS = \"circle.png\"  #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Debugging and monitoring\n",
        "VERBOSE = False  #@param {type:\"boolean\"}\n",
        "PLOT_DURATIONS = False  #@param {type:\"boolean\"}\n",
        "VIDEO_STEPS = 10  #@param {type:\"integer\"}\n",
        "TRACE_EVERY = 10\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "37P41H1vGu-0"
      },
      "source": [
        "# @title Images can be saved on Drive\n",
        "STORE_ON_GOOGLE_DRIVE = False  #@param {type:\"boolean\"}\n",
        "DIR_RESULTS = \"/content\"  #@param {type:\"string\"}\n",
        "\n",
        "if STORE_ON_GOOGLE_DRIVE:\n",
        "  from google.colab import drive\n",
        "  mount_dir = \"/content/drive\"\n",
        "  drive.mount(mount_dir)\n",
        "  DIR_RESULTS = mount_dir + \"/My Drive\"\n",
        "print(f\"Storing results in {DIR_RESULTS}\")\n",
        "\n",
        "\n",
        "target_img = None\n",
        "if MSE_LOSS:\n",
        "  if not STORE_ON_GOOGLE_DRIVE:\n",
        "    raise ValueError(\n",
        "        \"Need access to Google drive to load target image\")\n",
        "  target_img = load_torch_img(\n",
        "      f\"{DIR_RESULTS}/{DRIVE_IMAGE_FOR_MSE_LOSS}\")\n",
        "# show_img(target_img)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tyyw6Glris_e"
      },
      "source": [
        "#Make Drawings"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "Bl9BHoyRHjmf"
      },
      "source": [
        "#@title Create Tensorboard\n",
        "\n",
        "!cd /content/diffvg/apps/\n",
        "\n",
        "LOGS_BASE_DIR = \"runs\"\n",
        "os.makedirs(LOGS_BASE_DIR, exist_ok=True)\n",
        "%load_ext tensorboard\n",
        "%tensorboard --logdir logs/tensorboard"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "p6drbSbshNN0"
      },
      "source": [
        "# @title Configure Image\n",
        "#@markdown Select drawing grammar\n",
        "USE_GRAMMAR = \"Arnheim 2\"  #@param [\"Arnheim 2\", \"Photographic\", \"DPPN\", \"SEQTOSEQ\"]\n",
        "#@markdown Enter a description of the image, e.g. 'a photorealistic chicken'\n",
        "PROMPT = \"a chicken\"  #@param {type:\"string\"}\n",
        "#@markdown Optional negative prompts to reduce certain image elements or characteristics\n",
        "USE_NEG_PROMPTS = False  #@param {type:\"boolean\"}\n",
        "NEG_PROMPT_1 = \"messy\"  #@param {type:\"string\"}\n",
        "NEG_PROMPT_2 = \"cluttered\"  #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Use a white background instead of black\n",
        "DRAW_WHITE_BACKGROUND = False  #@param {type:\"boolean\"}\n",
        "\n",
        "if USE_GRAMMAR == \"Photographic\":  # Produces more photorealistic paintings.\n",
        "  BATCH_SIZE = 1\n",
        "  LEARNING_RATE = 0.004\n",
        "  INPUT_LEARNING_RATE = 0.1\n",
        "  NUM_LSTMS = 3\n",
        "  INPUT_SPEC_SIZE = 50\n",
        "  NET_LSTM_HIDDENS = 100\n",
        "  NET_MLP_HIDDENS = 100\n",
        "  NUM_STROKE_TYPES = 10\n",
        "  OUTPUT_SIZE = 8 * NUM_STROKE_TYPES\n",
        "  NUM_PATHS = 2000\n",
        "  SEQ_LENGTH = int(NUM_PATHS/NUM_STROKE_TYPES)\n",
        "  OPTIM_STEPS = 1000\n",
        "  NUM_SEGMENTS = 1\n",
        "  USE_DECAY = False\n",
        "  OUTPUT_COEFF_SYSTEMATICITY = 0.1\n",
        "  OUTPUT_COEFF_WIDTH = 25.0\n",
        "  OUTPUT_COEFF_COLOUR = 10.0\n",
        "  LOAD_MODEL = False\n",
        "  DECAY_RATE = None\n",
        "\n",
        "if USE_GRAMMAR == \"Arnheim 2\":  # Produces hierarchically structured paintings.\n",
        "  BATCH_SIZE = 200  #30  # 100\n",
        "  LEARNING_RATE = 0.0004  #0.0008  # 0.0004\n",
        "  INPUT_LEARNING_RATE = 0.01  #0.05  # 0.01\n",
        "  NUM_LSTMS = 10  # 10\n",
        "  EMBEDDING_SIZE = 100\n",
        "  LSTM_ITERATIONS = 5  #10  # 5\n",
        "  NET_LSTM_HIDDENS = 250\n",
        "  NET_MLP_HIDDENS = 250\n",
        "  OUTPUT_SIZE = 100\n",
        "  OPTIM_STEPS = 1000\n",
        "  NUM_SEGMENTS = 1\n",
        "  USE_DECAY = False\n",
        "  LEARNABLE_FORCING = False\n",
        "  GRAMMATICAL_FORCING = 1.0\n",
        "  USE_DROPOUT = False\n",
        "  DROPOUT_PROP = 0.2\n",
        "  YELLOW_ABSOLUTE_POSITIONS_USED = True\n",
        "  WEIGHT_INITIALIZER = True\n",
        "  WEIGHT_STD = 0.1\n",
        "  LOAD_MODEL = False\n",
        "  DECAY_RATE = None\n",
        "\n",
        "if USE_GRAMMAR == \"DPPN\":  # Paintings with a feed forward neural network.\n",
        "  BATCH_SIZE = 2000\n",
        "  LEARNING_RATE = 0.004\n",
        "  INPUT_LEARNING_RATE = 0.02\n",
        "  EMBEDDING_SIZE = 30\n",
        "  MLP_LAYERS = 3\n",
        "  NET_MLP_HIDDENS = 250\n",
        "  OUTPUT_SIZE = 20\n",
        "  OPTIM_STEPS = 2000\n",
        "  NUM_SEGMENTS = 1\n",
        "  USE_DECAY = False\n",
        "  LEARNABLE_FORCING = False\n",
        "  GRAMMATICAL_FORCING = 1.00\n",
        "  USE_DROPOUT = False\n",
        "  DROPOUT_PROP = 0.2\n",
        "  YELLOW_ABSOLUTE_POSITIONS_USED = True\n",
        "  WEIGHT_INITIALIZER = True\n",
        "  WEIGHT_STD = 0.1\n",
        "  LOAD_MODEL = False\n",
        "  DECAY_RATE = None\n",
        "\n",
        "if USE_GRAMMAR == \"SEQTOSEQ\":  # Paintings with a simple language model\n",
        "  BATCH_SIZE = 30\n",
        "  SEQ_LENGTH = 20\n",
        "  LEARNING_RATE = 0.0008\n",
        "  INPUT_LEARNING_RATE = 0.03\n",
        "  NUM_STROKES = 20\n",
        "  NET_LSTM_HIDDENS = 250\n",
        "  EMBEDDING_SIZE = 13\n",
        "  INPUT_SIZE = OUTPUT_SIZE = EMBEDDING_SIZE\n",
        "  OPTIM_STEPS = 2000\n",
        "  NUM_SEGMENTS = 1\n",
        "  USE_DECAY = False\n",
        "  LEARNABLE_FORCING = False\n",
        "  GRAMMATICAL_FORCING = 1.0\n",
        "  YELLOW_ABSOLUTE_POSITIONS_USED = True\n",
        "  WEIGHT_INITIALIZER = True\n",
        "  WEIGHT_STD = 0.1\n",
        "  LOAD_MODEL = False\n",
        "  DECAY_RATE = None\n",
        "\n",
        "if USE_DECAY:\n",
        "  DECAY_RATE = 0.999\n",
        "  LEARNING_RATE = 0.1\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "09akCs1QeRMz"
      },
      "source": [
        "#@title Create Drawing!\n",
        "#@markdown Run this cell to create the drawing. A series of improving images will appear ending with a final large render.\n",
        "\n",
        "#@markdown Performance time varies depending on assigned GPU but expect to wait in the order of ~1 hour.\n",
        "writer = SummaryWriter(\"logs/tensorboard\")\n",
        "video_writer = VideoWriter()\n",
        "\n",
        "prompt_features = get_features(PROMPT, NEG_PROMPT_1, NEG_PROMPT_2)\n",
        "augmentations = augmentation_transforms(CANVAS_WIDTH, USE_NORMALIZED_CLIP)\n",
        "\n",
        "stroke_generator = create_generator(USE_GRAMMAR)\n",
        "optimizer = make_optimizer(\n",
        "    stroke_generator, LEARNING_RATE, INPUT_LEARNING_RATE, DECAY_RATE)\n",
        "clipping_value = 0.1  # arbitrary value of your choosing\n",
        "torch.nn.utils.clip_grad_norm(stroke_generator.parameters(), clipping_value)\n",
        "\n",
        "for step in range(OPTIM_STEPS):\n",
        "  last_step = step == (OPTIM_STEPS - 1)\n",
        "  step_optimization(step, clip_model, optimizer, stroke_generator,\n",
        "                    augmentations, prompt_features, final_step=last_step)\n",
        "\n",
        "video_writer.close()\n",
        "video_writer.show()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2RR09pp3ggv0"
      },
      "source": [
        "# Investigations\n",
        "\n",
        "## Generative Architectures\n",
        "\n",
        "### Arnheim 2\n",
        "\n",
        "Running the standard \"Arnheim 2\" architecture produces images of the kinds below. They are characterised by an ordered production of marks which show smooth local variations in colour, thickness and shape. This is imposed by the re-write nature of the architecture. The videos show the intermediate stages in the optimization of the picture, and the kinds of variation that arise. As with CPPNs (Compositional Pattern Producing Networks) there are correlated changes in the marks throughout learning.\n",
        "\n",
        "\n",
        "![picture](https://storage.googleapis.com/deepmind-media/Art%20in%20AI/Figure1Arnheim1PhotorealisticChicken.png)\n",
        "\n",
        "\"A Photorealistic chicken\"\n",
        "\n",
        "Video of [\"Flemish still life with steak and tulips\" https://youtu.be/EKzmqa9Ol-g\n",
        "\n",
        "Video of \"A red coral\" https://youtu.be/wB3iE-1H4js\n",
        "\n",
        "Video of \"A chicken\" https://youtu.be/U7guaMdeF4g\n",
        "\n",
        "Video of \"An apple\" with only two input sequence elements https://youtu.be/ZIAjyS5bU0g\n",
        "\n",
        "Video of \"clouds\" https://youtu.be/y7a8rldcJps\n",
        "\n",
        "Let us investigate a few other architectures to get a feeling for how the architecture constrains the style of the painting that is learned.\n",
        "\n",
        "### Photographic\n",
        "\n",
        "If you choose \"Photographic\" instead of \"Arnheim 2\" then the direct encoding of stroke positions used in CLIPDraw is used, but the colour, width, and some small positional modifications of the strokes are produced by 3 LSTMs (separate ones for stroke, width and colour), which adds some element of coherence to the overall image, whilst not containing any of the hierarchical stroke encoding that Arnheim 2 has. \n",
        "Watching and comparing the videos of images generated in this way, with those produced by Arnheim 2 reveals that structured global variation in the image does not occur at the level of stroke complexes.\n",
        "\n",
        "![picture](https://storage.googleapis.com/deepmind-media/Art%20in%20AI/PhotographicHen.png)\n",
        "\n",
        "Photographic \"A hen\"\n",
        "\n",
        "Video of \"A hen\" https://youtu.be/T_MXpv_4NJM\n",
        "\n",
        "Another video of \"A hen\" https://youtu.be/zh0goLbS-l0\n",
        "\n",
        "### MLP and DPPN\n",
        "\n",
        "Running multiple levels of LSTM can be rather slow. The original motivation for the CPPN (compositional pattern producing network https://en.wikipedia.org/wiki/Compositional_pattern-producing_network) was to remove the need for recurrance in generating an image. In the original CPPN the x,y coordinates are given to a network and at each specified pixel position a single pixel colour is outputted.\n",
        "A while ago we published the DPPN (differentiable pattern producing network (https://arxiv.org/abs/1606.02580)) which was extended to HyperNets (https://arxiv.org/abs/1609.09106). We generalize this principle to optimizing a set of x,y coordinates input into a network and interpreting the outputs as a stroke to be drawn at that x,y position.\n",
        "This is equivalent to having a big residual connection directly from the x,y input encoding to the output, and allows gradients to have much greater control over the x,y position of marks than if gradients are needed to adjust all the parameters in the network to achieve a change in the \"where\" of the stroke rather than the \"what\" of a stroke.\n",
        "The videos and images below compare the result of allowing this residual connection and not doing so, and examine the kinds of variation in the image that are produced as the depth of the network is modified. With deeper networks more correlated changes in the image become possible, e.g a global translation of the whole image is more likely with deeper networks. If you choose \"DPPN\" you can generate these kinds of images yourself.\n",
        "\n",
        "![picture](https://storage.googleapis.com/deepmind-media/Art%20in%20AI/SimpleDPPNChicken.png)\n",
        "\n",
        "Simple MLP \"Chicken\"\n",
        "\n",
        "Video of \"Simple MLP chicken\" formation https://youtu.be/mCgtvVBGA90\n",
        "\n",
        "![picture](https://storage.googleapis.com/deepmind-media/Art%20in%20AI/DPPNChickens.png)\n",
        "\n",
        "DPPN Chickens of 2, 4 and 7 layers.\n",
        "\n",
        "Video of 2 layer DPPN chicken:\n",
        "https://youtu.be/splxpSxnIsc\n",
        "\n",
        "Video of 4 layer DPPN chicken:\n",
        "https://youtu.be/U_RKupUf2uU\n",
        "\n",
        "Video of 7 layer DPPN chicken:\n",
        "https://youtu.be/jzyTnuY2tKc\n",
        "\n",
        "### Seq-to-Seq\n",
        "\n",
        "Next we use a standard practically off the shelf simple seq-to-seq model to see how existing architectures compare with those above. You can run this by choosing SEQTOSEQ.\n",
        "\n",
        "![picture](https://storage.googleapis.com/deepmind-media/Art%20in%20AI/SeqtoSeqChicken.png)\n",
        "\n",
        "A variety of chickens produced by the SeqtoSeq with the prompt \"A chicken\"\n",
        "\n",
        "Video of SeqtoSeq \"A chicken\" with only two input vectors\n",
        "https://youtu.be/mksVKPUHIiA\n",
        "\n",
        "Video of SeqtoSeq \"A chicken\" with 10 input vectors.\n",
        "https://youtu.be/VZ9a5ECVOFo\n",
        "\n",
        "Video of SeqtoSeq \"A chicken\" with 10 input vectors and greater scale of output strokes\n",
        "https://youtu.be/F-HS8n7RLLo\n",
        "\n",
        "Video of SeqtoSeq \"A chicken\" with many input vectors. The chicken is constructed from negative space.\n",
        "https://youtu.be/RcvXp85CsvU\n",
        "\n",
        "Video of SeqtoSeq \"A chicken with 200 input vectors.\n",
        "https://youtu.be/SYJGNt7yu6M\n",
        "\n",
        "Video of SeqtoSeq \"A chicken\" with only 1 input vector but a longer sequence length.\n",
        "https://youtu.be/fEp4jP7p3i0\n",
        "\n",
        "## Additional Computational Aesthetic Losses\n",
        "\n",
        "Here we make an additional computational aesthetic loss which promotes pictures which possess a diversity of thick and thin marks. This can be easily calculated by taking the standard deviation of the width specificiation vector, and then combining this with some co-efficient with the CLIP loss.\n",
        "If you set WIDTH_DIVERSITY_LOSS = True then you see the modified results. Feel free to invent your own computational aesthetic losses, e.g. to promote certain compositions, etc...\n",
        "\n",
        "![picture](https://storage.googleapis.com/deepmind-media/Art%20in%20AI/WidthDiversityLoss.png)\n",
        "\n",
        "Pictures showing the influence of the WIDTH_DIVERSITY_LOSS coefficient when producing a \"Photorealistic Chicken\".\n",
        "\n",
        "## Prompt engineering to exploit CLIP\n",
        "\n",
        "The beautiful images produced by RiversHaveWings are partly due to discoveries made about the prompts that result in interesting images, by exploiting the knowledge that CLIP has about styles. For example adding \"A photorealistic\" or \"X\" before the prompt has drastic effects, e.g. see below for some examples.\n",
        "\n",
        "## Replacing CLIP with a simple MSE loss\n",
        "\n",
        "It is not even necessary to use CLIP, instead we can use the MSE loss between a target image and the painted image. You can load an image into your google Drive and use that as a template for a painting directly. Set MSE_LOSS = True to make this happen and modify the name of the target_img to the image you desire. This allows us much more direct access to the generative process, uncomplicated by representational forces from CLIP.\n",
        "\n",
        "Video trying to fit a red circle with Arnheim 2 using various batch sizes and sequence lengths:\n",
        "\n",
        "https://youtu.be/IMJD4vgoGkA\n",
        "\n",
        "https://youtu.be/4ITeqY6PIcg\n",
        "\n",
        "https://youtu.be/MxkYKa0x5AU\n",
        " \n"
      ]
    }
  ]
}