Story2Hallucination_preset.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Story2Hallucination_preset.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3hElmF-malqM"
      },
      "source": [
        "## Restart after running this cell!\n",
        "\n",
        "You must run this cell and then restart and rerun everything for the PyTorch version to be correct. Otherwise the model will run but not produce any meaningful output."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "c33IJYCRafX0"
      },
      "source": [
        "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",
        "! pip install torch==1.7.1{torch_version_suffix} torchvision==0.8.2{torch_version_suffix} -f https://download.pytorch.org/whl/torch_stable.html ftfy regex"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xzl4Zrd7SEvG"
      },
      "source": [
        "Install Big Sleep from pip. (Note that the latest major upgrade to 0.5.x breaks this script. I'm working on a fix to make it compatible soon.)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "D2jUsCZXaqcw"
      },
      "source": [
        "!pip install big-sleep==0.4.11"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "El_t0JgjSIye"
      },
      "source": [
        "Mount your Google Drive. You will need a folder called story_halluc in drive to store the images there."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xjQaqenajK83"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m99xCwxSSRwt"
      },
      "source": [
        "This block takes the text and generates the images.  Update the all_text variable to change the story. You can adjust the parameters to change how many images are run before changing phrases and when to restart with a fresh image."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NfJ0RMCAauV3"
      },
      "source": [
        "from IPython.display import Image, display\n",
        "import string\n",
        "import torch\n",
        "from torchvision.utils import save_image\n",
        "import numpy as np\n",
        "\n",
        "from big_sleep import Imagine\n",
        "from big_sleep.clip import tokenize\n",
        "\n",
        "from nltk.corpus import stopwords\n",
        "\n",
        "from skimage.measure import compare_ssim\n",
        "\n",
        "import cv2\n",
        "from pathlib import Path\n",
        "\n",
        "import PIL\n",
        "from PIL import ImageFont, ImageDraw\n",
        "\n",
        "TEXT = 'story_hallucinator' \n",
        "SAVE_EVERY = 1\n",
        "SAVE_PROGRESS = True\n",
        "LEARNING_RATE = 0.1\n",
        "ITERATIONS =  1\n",
        "\n",
        "def train_step(self, epoch, i, rand=0):\n",
        "  total_loss = 0\n",
        "\n",
        "  for _ in range(self.gradient_accumulate_every):\n",
        "      losses = self.model(self.encoded_text) \n",
        "      loss = (sum(losses) / self.gradient_accumulate_every) + rand*np.random.randn()\n",
        "      total_loss += loss\n",
        "      loss.backward()\n",
        "\n",
        "  self.optimizer.step()\n",
        "  self.optimizer.zero_grad()\n",
        "\n",
        "  if (i + 1) % self.save_every == 0:\n",
        "      with torch.no_grad():\n",
        "          # best = torch.topk(losses[2], k = 1, largest = False)[1]\n",
        "          mres = self.model.model()\n",
        "          image = mres[len(mres)-1].cpu()\n",
        "          num = i // self.save_every\n",
        "          save_image(image, Path(f'./{self.textpath}.{num}.png'))\n",
        "\n",
        "model = Imagine(\n",
        "    text = TEXT,\n",
        "    save_every = SAVE_EVERY,\n",
        "    lr = LEARNING_RATE,\n",
        "    iterations = ITERATIONS,\n",
        "    save_progress = SAVE_PROGRESS\n",
        ")\n",
        "filename = TEXT.replace(' ', '_')\n",
        "\n",
        "\n",
        "burnin=20 #\n",
        "checkin_gap = 10\n",
        "long_sim_gap = 10\n",
        "span = 6\n",
        "iterations = 100\n",
        "display_gap = 50\n",
        "similarity = 0.9\n",
        "\n",
        "all_text_list = [\"The moon has a face like the clock in the hall,\",\n",
        "                 \"She shines on thieves on the garden wall,\",\n",
        "                 \"On streets and fields and harbour quays,\",\n",
        "                 \"And birdies asleep in the forks of the trees,\"]\n",
        "\n",
        "iter_num = 0\n",
        "last_one = 0\n",
        "rand = 0\n",
        "model.text = \" \".join(words[:span])\n",
        "model.encoded_text = tokenize(model.text).cuda()\n",
        "for j in range(burnin):\n",
        "    train_step(model, 0, 0, rand)\n",
        "for epoch in range(0, len(all_text_list)):\n",
        "    restart_point = iter_num\n",
        "    i = 0\n",
        "    while i < iterations:\n",
        "        phrase =all_text_list[epoch]\n",
        "        model.text = phrase.translate(str.maketrans('', '', string.punctuation))\n",
        "        model.encoded_text = tokenize(model.text).cuda()\n",
        "        train_step(model, epoch, iter_num, rand)\n",
        "        \n",
        "        if iter_num % display_gap == 0:\n",
        "          print(f'iter: {iter_num} text={phrase}')\n",
        "          image_cur = Image(f'./{filename}.{iter_num}.png')\n",
        "          display(image_cur)\n",
        "        \n",
        "        if i % checkin_gap == 0 and i > 0:\n",
        "          imageA = cv2.imread(f'./{filename}.{iter_num}.png')\n",
        "          imageB = cv2.imread(f'./{filename}.{restart_point}.png')\n",
        "          # convert the images to grayscale\n",
        "          grayA = cv2.cvtColor(imageA, cv2.COLOR_BGR2GRAY)\n",
        "          grayB = cv2.cvtColor(imageB, cv2.COLOR_BGR2GRAY)\n",
        "          (score, diff) = compare_ssim(grayA, grayB, full=True)\n",
        "          toinc = checkin_gap\n",
        "          print(f'iter{iter_num}: rand={rand} sim={score} smooth={grayB.std()}, ext={((grayB < 50) | (grayB > 205)).mean()}')\n",
        "          if score>similarity or grayB.std()<15 or ((grayB < 50) | (grayB > 205)).mean()>0.9:\n",
        "              print(f'restart!')\n",
        "              model = Imagine(\n",
        "                  text = TEXT,\n",
        "                  save_every = SAVE_EVERY,\n",
        "                  lr = LEARNING_RATE,\n",
        "                  iterations = ITERATIONS,\n",
        "                  save_progress = SAVE_PROGRESS\n",
        "              )\n",
        "              model.text = \" \".join(words[epoch:epoch+span]).translate(str.maketrans('', '', string.punctuation))\n",
        "              model.encoded_text = tokenize(model.text).cuda()\n",
        "              for j in range(burnin):\n",
        "                train_step(model, epoch, iter_num, rand) \n",
        "              iter_num = restart_point\n",
        "              i = 0\n",
        "              rand = 0\n",
        "              continue\n",
        "        i += 1\n",
        "        iter_num += 1\n",
        "        \n",
        "        \n",
        "    for i in range(last_one,iter_num):\n",
        "      msg_orig =all_text_list[epoch]\n",
        "      img = PIL.Image.open(f'./{filename}.{i}.png')\n",
        "      W, H = img.size\n",
        "      draw = ImageDraw.Draw(img)\n",
        "      font = ImageFont.truetype(\"/usr/share/fonts/truetype/liberation/LiberationMono-Bold.ttf\", 18)\n",
        "      msgs = [msg_orig]\n",
        "      w, h = draw.textsize(msg_orig, font=font)\n",
        "      if w>W:\n",
        "        split = span // 2\n",
        "        msgs = [\" \".join(words[epoch:epoch+split]), \" \".join(words[epoch+split:epoch+span])]\n",
        "      for shift, msg in enumerate(msgs): \n",
        "        w, h = draw.textsize(msg, font=font)\n",
        "        x, y = (W-w)/2, 7*(H-h)/8 + shift*h\n",
        "        adj = 1\n",
        "        #move right\n",
        "        shadowColor = \"black\"\n",
        "        draw.text((x-adj, y), msg, fill=shadowColor, font=font)\n",
        "        #move left\n",
        "        draw.text((x+adj, y), msg, fill=shadowColor, font=font)\n",
        "        #move up\n",
        "        draw.text((x, y+adj), msg, fill=shadowColor, font=font)\n",
        "        #move down\n",
        "        draw.text((x, y-adj), msg, fill=shadowColor, font=font)\n",
        "        #diagnal left up\n",
        "        draw.text((x-adj, y+adj), msg, fill=shadowColor, font=font)\n",
        "        #diagnal right up\n",
        "        draw.text((x+adj, y+adj), msg, fill=shadowColor, font=font)\n",
        "        #diagnal left down\n",
        "        draw.text((x-adj, y-adj), msg, fill=shadowColor, font=font)\n",
        "        #diagnal right down\n",
        "        draw.text((x+adj, y-adj), msg, fill=shadowColor, font=font)\n",
        "        draw.text((x, y), msg, fill=\"white\", font=font)\n",
        "      img.save(f'./{filename}.{i}.png')\n",
        "    last_one = iter_num\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e922of2WR3tB"
      },
      "source": [
        "This section zips your image files and stores the zip file in a folder called story_halluc on your google drive."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gGpSuIxjFSps"
      },
      "source": [
        "!zip s2h.zip *.png"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kWfYQ7yvE6qr"
      },
      "source": [
        "!cp s2h.zip /content/drive/MyDrive/story_halluc"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pcNMH32_TPe7"
      },
      "source": [
        "Next download the zip file from Drive, unzip it in a folder and run the following:\n",
        "\n",
        "```\n",
        "ffmpeg -framerate 10 -i story_hallucinator.%d.png -c:v libx264 -crf 0 story_hallucination.mp4\n",
        "```\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "W6TgvSSQN3zS"
      },
      "source": [
        "!rm story_*"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bIg29_14jwNx"
      },
      "source": [
        "!rm s2h.zip"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iVuvkIWJnjWY"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}