Aphantasia.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Aphantasia.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.9"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "toWe1IoH7X35"
      },
      "source": [
        "# Text to Image tool\n",
        "\n",
        "Part of [Aphantasia](https://github.com/eps696/aphantasia) suite, made by Vadim Epstein [[eps696](https://github.com/eps696)]  \n",
        "Based on [CLIP](https://github.com/openai/CLIP) + FFT from [Lucent](https://github.com/greentfrapp/lucent).  \n",
        "Thanks to [Ryan Murdock](https://twitter.com/advadnoun), [Jonathan Fly](https://twitter.com/jonathanfly), [Hannu Toyryla](https://twitter.com/htoyryla), [@eduwatch2](https://twitter.com/eduwatch2) for ideas.\n",
        "\n",
        "## Features \n",
        "* generates massive detailed imagery, a la deepdream \n",
        "* high resolution (up to 12K on RTX 3090)\n",
        "* directly parameterized with [FFT](https://github.com/greentfrapp/lucent/blob/master/lucent/optvis/param/spatial.py) [Fourier] or DWT [wavelets] (no pretrained GANs)\n",
        "* various CLIP models (including multi-language from [SBERT](https://sbert.net))\n",
        "* starting/resuming process from saved FFT parameters or from an image\n",
        "* complex requests:\n",
        "  * image and/or text as main prompts  \n",
        "   (composition similarity controlled with [LPIPS](https://github.com/richzhang/PerceptualSimilarity) loss)\n",
        "  * separate text prompts for image style and to subtract (suppress) topics\n",
        "  * criteria inversion (show \"the opposite\")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QytcEMSKBtN-"
      },
      "source": [
        "**Run the cell below after each session restart**\n",
        "\n",
        "Mark `resume` and upload `.pt` file, if you're resuming from the saved snapshot. Or you can simply upload any image to start from.  \n",
        "Resolution settings below will be overwritten in this case. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "etzxXVZ_r-Nf",
        "cellView": "form"
      },
      "source": [
        "#@title General setup\n",
        "\n",
        "# !pip install torchtext==0.8.0 torch==1.7.1 pytorch-lightning==1.2.2 torchvision==0.8.2 ftfy==5.8 regex\n",
        "!pip install ftfy==5.8 transformers==4.6.0\n",
        "!pip install gputil ffpb \n",
        "\n",
        "try: \n",
        "  !pip3 install googletrans==3.1.0a0\n",
        "  from googletrans import Translator, constants\n",
        "  translator = Translator()\n",
        "except: pass\n",
        "\n",
        "import os\n",
        "import io\n",
        "import time\n",
        "from math import exp\n",
        "import random\n",
        "import imageio\n",
        "import numpy as np\n",
        "import PIL\n",
        "from base64 import b64encode\n",
        "# import moviepy, moviepy.editor\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torchvision\n",
        "from torch.autograd import Variable\n",
        "\n",
        "from IPython.display import HTML, Image, display, clear_output\n",
        "from IPython.core.interactiveshell import InteractiveShell\n",
        "InteractiveShell.ast_node_interactivity = \"all\"\n",
        "import ipywidgets as ipy\n",
        "from google.colab import output, files\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")\n",
        "\n",
        "!pip install git+https://github.com/openai/CLIP.git --no-deps\n",
        "import clip\n",
        "!pip install sentence_transformers\n",
        "from sentence_transformers import SentenceTransformer\n",
        "!pip install kornia\n",
        "import kornia\n",
        "!pip install lpips\n",
        "import lpips\n",
        "\n",
        "!pip install PyWavelets==1.1.1\n",
        "!pip install git+https://github.com/fbcotter/pytorch_wavelets\n",
        "import pywt\n",
        "from pytorch_wavelets import DWTForward, DWTInverse\n",
        "# from pytorch_wavelets import DTCWTForward, DTCWTInverse\n",
        "\n",
        "%cd /content\n",
        "!rm -rf aphantasia\n",
        "!git clone https://github.com/eps696/aphantasia\n",
        "%cd aphantasia/\n",
        "from clip_fft import to_valid_rgb, fft_image, img2fft, dwt_image, img2dwt, init_dwt, dwt_scale\n",
        "from utils import slice_imgs, derivat, pad_up_to, basename, img_list, img_read, plot_text, txt_clean, checkout, old_torch\n",
        "import transforms\n",
        "from progress_bar import ProgressIPy as ProgressBar\n",
        "\n",
        "clear_output()\n",
        "\n",
        "resume = False #@param {type:\"boolean\"}\n",
        "if resume:\n",
        "  resumed = files.upload()\n",
        "  resumed_filename = list(resumed)[0]\n",
        "  resumed_bytes = list(resumed.values())[0]\n",
        "\n",
        "def makevid(seq_dir, size=None):\n",
        "  out_sequence = seq_dir + '/%04d.jpg'\n",
        "  out_video = seq_dir + '.mp4'\n",
        "  print('.. generating video ..')\n",
        "  !ffmpeg -y -v warning -i $out_sequence -crf 20 $out_video\n",
        "  # moviepy.editor.ImageSequenceClip(img_list(seq_dir), fps=25).write_videofile(out_video, verbose=False)\n",
        "  data_url = \"data:video/mp4;base64,\" + b64encode(open(out_video,'rb').read()).decode()\n",
        "  wh = '' if size is None else 'width=%d height=%d' % (size, size)\n",
        "  return \"\"\"<video %s controls><source src=\"%s\" type=\"video/mp4\"></video>\"\"\" % (wh, data_url)\n",
        "\n",
        "# Hardware check\n",
        "!ln -sf /opt/bin/nvidia-smi /usr/bin/nvidia-smi\n",
        "import GPUtil as GPU\n",
        "gpu = GPU.getGPUs()[0]\n",
        "!nvidia-smi -L\n",
        "print(\"GPU RAM {0:.0f}MB | Free {1:.0f}MB)\".format(gpu.memoryTotal, gpu.memoryFree))\n",
        "print('\\nDone!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CbJ9K4Cq8MtB"
      },
      "source": [
        "Type some `text` and/or upload some image to start.  \n",
        "Describe `style`, which you'd like to apply to the imagery.  \n",
        "Put to `subtract` the topics, which you would like to avoid in the result.  \n",
        "`invert` the whole criteria, if you want to see \"the totally opposite\".\n",
        "\n",
        "Options for non-English languages (use only one of them!):  \n",
        "`multilang` = use multi-language model, trained with ViT  \n",
        "`translate` = use Google translate (works with any visual model)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JUvpdy8BWGuM",
        "cellView": "form"
      },
      "source": [
        "#@title Input\n",
        "\n",
        "text = \"\" #@param {type:\"string\"}\n",
        "style = \"\" #@param {type:\"string\"}\n",
        "subtract = \"\" #@param {type:\"string\"}\n",
        "multilang = False #@param {type:\"boolean\"}\n",
        "translate = False #@param {type:\"boolean\"}\n",
        "invert = False #@param {type:\"boolean\"}\n",
        "upload_image = False #@param {type:\"boolean\"}\n",
        "\n",
        "if translate:\n",
        "  text = translator.translate(text, dest='en').text\n",
        "if upload_image:\n",
        "  uploaded = files.upload()\n",
        "\n",
        "workdir = '_out'\n",
        "tempdir = os.path.join(workdir, '%s-%s' % (txt_clean(text)[:50], txt_clean(style)[:50]))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f3Sj0fxmtw6K"
      },
      "source": [
        "### Settings\n",
        "\n",
        "Select visual `model` (results do vary!). I prefer ViT for consistency (and it's the only native multi-language option).  \n",
        "`align` option is about composition. `uniform` looks most adequate, `overscan` can make semi-seamless tileable texture.  \n",
        "`use_wavelets` for DWT encoding instead of FFT. Select `wave` method if needed.  \n",
        "`aug_transform` applies some augmentations, inhibiting image fragmentation & \"graffiti\" printing (slower, yet recommended).  \n",
        "`sync` value adds LPIPS loss between the output and input image (if there's one), allowing to \"redraw\" it with controlled similarity.  \n",
        "Decrease `samples` if you face OOM (it's the main RAM eater).  \n",
        "\n",
        "Setting `steps` much higher (1000-..) will elaborate details and make tones smoother, but may start throwing texts like graffiti.  \n",
        "Tune `decay` (compositional softness) and `sharpness`, `colors` (saturation) and `contrast` as needed.  \n",
        "\n",
        "Experimental tricks:  \n",
        "`aug_noise` augmentation, `macro` (from 0 to 1) and `progressive_grow` (read more [here](https://github.com/eps696/aphantasia/issues/2)) may boost bigger forms, making composition less disperse.  \n",
        "`no_text` tries to remove \"graffiti\" by subtracting plotted text prompt  \n",
        "`enhance` boosts training consistency (of simultaneous samples) and steps progress. good start is 0.1~0.2.  "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Nq0wA-wc-P-s",
        "cellView": "form"
      },
      "source": [
        "#@title Generate\n",
        "\n",
        "!rm -rf $tempdir\n",
        "os.makedirs(tempdir, exist_ok=True)\n",
        "\n",
        "sideX = 1280 #@param {type:\"integer\"}\n",
        "sideY = 720 #@param {type:\"integer\"}\n",
        "#@markdown > Config\n",
        "model = 'ViT-B/32' #@param ['ViT-B/16', 'ViT-B/32', 'RN101', 'RN50x16', 'RN50x4', 'RN50']\n",
        "align = 'uniform' #@param ['central', 'uniform', 'overscan']\n",
        "use_wavelets = False #@param {type:\"boolean\"}\n",
        "wave = 'coif1' #@param ['db2', 'db3', 'db4', 'coif1', 'coif2', 'coif3', 'coif4']\n",
        "aug_transform = True #@param {type:\"boolean\"}\n",
        "sync =  0.4 #@param {type:\"number\"}\n",
        "#@markdown > Look\n",
        "decay = 1.5 #@param {type:\"number\"}\n",
        "colors = 1.5 #@param {type:\"number\"}\n",
        "contrast = 0.9 #@param {type:\"number\"}\n",
        "sharpness = 0.3 #@param {type:\"number\"}\n",
        "#@markdown > Training\n",
        "steps = 300 #@param {type:\"integer\"}\n",
        "samples = 200 #@param {type:\"integer\"}\n",
        "learning_rate = .05 #@param {type:\"number\"}\n",
        "save_freq = 1 #@param {type:\"integer\"}\n",
        "#@markdown > Tricks\n",
        "aug_noise = 0.2 #@param {type:\"number\"}\n",
        "no_text = 0.07 #@param {type:\"number\"}\n",
        "enhance = 0. #@param {type:\"number\"}\n",
        "macro = 0.4 #@param {type:\"number\"}\n",
        "progressive_grow = False #@param {type:\"boolean\"}\n",
        "if multilang: model = 'ViT-B/32' # sbert model is trained with ViT\n",
        "diverse = -enhance\n",
        "expand = abs(enhance)\n",
        "\n",
        "shape = [1, 3, sideY, sideX]\n",
        "\n",
        "if use_wavelets:\n",
        "  if resume is True:\n",
        "    if os.path.splitext(resumed_filename)[1].lower()[1:] in ['jpg','png','tif','bmp']:\n",
        "      img_in = imageio.imread(resumed_bytes)\n",
        "      init_pt = img2dwt(img_in, wave, 1.5)\n",
        "      scale = dwt_scale(init_pt, sharpness)\n",
        "      for i in range(len(init_pt)-1):\n",
        "        init_pt[i+1] /= (scale[i] * 10)\n",
        "      sideY, sideX = img_in.shape[0], img_in.shape[1]\n",
        "    else:\n",
        "      init_pt = torch.load(io.BytesIO(resumed_bytes))\n",
        "      # init_pt = [y.detach().cuda() for y in init_pt]\n",
        "  else:\n",
        "    init_pt, _, _, _ = init_dwt(None, shape, wave, colors)\n",
        "\n",
        "  shape = [1, 3, sideY, sideX]\n",
        "  params, image_f, _ = dwt_image(shape, wave, sharpness, colors, init_pt)\n",
        "\n",
        "else:\n",
        "  if resume is True:\n",
        "    if os.path.splitext(resumed_filename)[1].lower()[1:] in ['jpg','png','tif','bmp']:\n",
        "      img_in = imageio.imread(resumed_bytes)\n",
        "      init_pt = img2fft(img_in, 1.5, 1.5) * 0.1\n",
        "    else:\n",
        "      init_pt = torch.load(io.BytesIO(resumed_bytes))\n",
        "      if isinstance(init_pt, list): init_pt = init_pt[0]\n",
        "  #  init_pt = init_pt.cuda()\n",
        "    sideY, sideX = init_pt.shape[2], (init_pt.shape[3]-1)*2\n",
        "  else:\n",
        "    params_shape = [1, 3, sideY, sideX//2+1, 2]\n",
        "    init_pt = torch.randn(*params_shape) * 0.01\n",
        "\n",
        "  shape = [1, 3, sideY, sideX]\n",
        "  params, image_f, _ = fft_image(shape, 1, decay, init_pt)\n",
        "\n",
        "image_f = to_valid_rgb(image_f, colors = colors)\n",
        "\n",
        "if progressive_grow is True:\n",
        "  lr1 = learning_rate * 2\n",
        "  lr0 = lr1 * 0.01\n",
        "else:\n",
        "  lr0 = learning_rate\n",
        "optimizer = torch.optim.AdamW(params, lr0, weight_decay=0.01, amsgrad=True)\n",
        "\n",
        "if len(subtract) > 0:\n",
        "  samples = int(samples * 0.75)\n",
        "print(' using %d samples' % samples)\n",
        "\n",
        "model_clip, _ = clip.load(model, jit=old_torch())\n",
        "modsize = model_clip.visual.input_resolution\n",
        "xmem = {'ViT-B/16':0.25, 'RN50':0.5, 'RN50x4':0.16, 'RN50x16':0.06, 'RN101':0.33}\n",
        "if model in xmem.keys():\n",
        "  samples = int(samples * xmem[model])\n",
        "\n",
        "if multilang is True:\n",
        "    model_lang = SentenceTransformer('clip-ViT-B-32-multilingual-v1').cuda()\n",
        "\n",
        "def enc_text(txt):\n",
        "  if multilang is True:\n",
        "    emb = model_lang.encode([txt], convert_to_tensor=True, show_progress_bar=False)\n",
        "  else:\n",
        "    emb = model_clip.encode_text(clip.tokenize(txt).cuda())\n",
        "  return emb.detach().clone()\n",
        "\n",
        "if diverse != 0:\n",
        "  samples = int(samples * 0.5)\n",
        "if sync > 0 and upload_image:\n",
        "  samples = int(samples * 0.6)\n",
        "        \n",
        "sign = 1. if invert is True else -1.\n",
        "if aug_transform is True:\n",
        "  trform_f = transforms.transforms_elastic\n",
        "  samples = int(samples * 0.95)\n",
        "else:\n",
        "  trform_f = transforms.normalize()\n",
        "\n",
        "if upload_image:\n",
        "  in_img = list(uploaded.values())[0]\n",
        "  print(' image:', list(uploaded)[0])\n",
        "  img_in = torch.from_numpy(imageio.imread(in_img).astype(np.float32)/255.).unsqueeze(0).permute(0,3,1,2).cuda()[:,:3,:,:]\n",
        "  in_sliced = slice_imgs([img_in], samples, modsize, transforms.normalize(), align)[0]\n",
        "  img_enc = model_clip.encode_image(in_sliced).detach().clone()\n",
        "  if sync > 0:\n",
        "    align = 'overscan'\n",
        "    sim_loss = lpips.LPIPS(net='vgg', verbose=False).cuda()\n",
        "    sim_size = [sideY//4, sideX//4]\n",
        "    img_in = F.interpolate(img_in, sim_size).float()\n",
        "    # img_in = F.interpolate(img_in, (sideY, sideX)).float()\n",
        "  else:\n",
        "    del img_in\n",
        "  del in_sliced; torch.cuda.empty_cache()\n",
        "\n",
        "if len(text) > 0:\n",
        "  print(' main topic:', text)\n",
        "  if translate:\n",
        "    translator = Translator()\n",
        "    text = translator.translate(text, dest='en').text\n",
        "    print(' translated to:', text) \n",
        "  txt_enc = enc_text(text)\n",
        "  if no_text > 0:\n",
        "    txt_plot = torch.from_numpy(plot_text(text, modsize)/255.).unsqueeze(0).permute(0,3,1,2).cuda()\n",
        "    txt_plot_enc = model_clip.encode_image(txt_plot).detach().clone()\n",
        "\n",
        "if len(style) > 0:\n",
        "  print(' style:', style)\n",
        "  if translate:\n",
        "    translator = Translator()\n",
        "    style = translator.translate(style, dest='en').text\n",
        "    print(' translated to:', style) \n",
        "  txt_enc2 = enc_text(style)\n",
        "\n",
        "if len(subtract) > 0:\n",
        "  print(' without:', subtract)\n",
        "  if translate:\n",
        "    translator = Translator()\n",
        "    subtract = translator.translate(subtract, dest='en').text\n",
        "    print(' translated to:', subtract) \n",
        "  txt_enc0 = enc_text(subtract)\n",
        "\n",
        "if multilang is True: del model_lang\n",
        "\n",
        "def save_img(img, fname=None):\n",
        "  img = np.array(img)[:,:,:]\n",
        "  img = np.transpose(img, (1,2,0))  \n",
        "  img = np.clip(img*255, 0, 255).astype(np.uint8)\n",
        "  if fname is not None:\n",
        "    imageio.imsave(fname, np.array(img))\n",
        "    imageio.imsave('result.jpg', np.array(img))\n",
        "\n",
        "def checkout(num):\n",
        "  with torch.no_grad():\n",
        "    img = image_f(contrast=contrast).cpu().numpy()[0]\n",
        "    # empirical tone mapping\n",
        "    if sync > 0 and upload_image:\n",
        "      img = img **1.3 \n",
        "    if sharpness != 0:\n",
        "      img = img ** (1 + sharpness/2.)\n",
        "  save_img(img, os.path.join(tempdir, '%04d.jpg' % num))\n",
        "  outpic.clear_output()\n",
        "  with outpic:\n",
        "    display(Image('result.jpg'))\n",
        "\n",
        "prev_enc = 0\n",
        "def train(i):\n",
        "  loss = 0\n",
        "  noise = aug_noise * torch.randn(1, 1, *params[0].shape[2:4], 1).cuda() if aug_noise > 0 else None\n",
        "  img_out = image_f(noise)\n",
        "  img_sliced = slice_imgs([img_out], samples, modsize, trform_f, align, macro=macro)[0]\n",
        "  out_enc = model_clip.encode_image(img_sliced)\n",
        "\n",
        "  if len(text) > 0: # input text\n",
        "    loss += sign * torch.cosine_similarity(txt_enc, out_enc, dim=-1).mean()\n",
        "    if no_text > 0:\n",
        "      loss -= sign * no_text * torch.cosine_similarity(txt_plot_enc, out_enc, dim=-1).mean()\n",
        "  if len(style) > 0: # input text - style\n",
        "    loss += sign * 0.5 * torch.cosine_similarity(txt_enc2, out_enc, dim=-1).mean()\n",
        "  if len(subtract) > 0: # subtract text\n",
        "    loss -= sign * torch.cosine_similarity(txt_enc0, out_enc, dim=-1).mean()\n",
        "  if upload_image:\n",
        "    loss += sign * 0.5 * torch.cosine_similarity(img_enc, out_enc, dim=-1).mean()\n",
        "  if sync > 0 and upload_image: # image composition sync\n",
        "    prog_sync = (steps - i) / steps \n",
        "    loss += prog_sync * sync * sim_loss(F.interpolate(img_out, sim_size).float(), img_in, normalize=True).squeeze()\n",
        "  if sharpness != 0 and not use_wavelets: # mode = scharr|sobel|default\n",
        "    loss -= sharpness * derivat(img_out, mode='sobel')\n",
        "    # loss -= sharpness * derivat(img_sliced, mode='scharr')\n",
        "  if diverse != 0:\n",
        "    img_sliced = slice_imgs([image_f(noise)], samples, modsize, trform_f, align, macro=macro)[0]\n",
        "    out_enc2 = model_clip.encode_image(img_sliced)\n",
        "    loss += diverse * torch.cosine_similarity(out_enc, out_enc2, dim=-1).mean()\n",
        "    del out_enc2; torch.cuda.empty_cache()\n",
        "  if expand > 0:\n",
        "    global prev_enc\n",
        "    if i > 0:\n",
        "      loss += expand * torch.cosine_similarity(out_enc, prev_enc, dim=-1).mean()\n",
        "    prev_enc = out_enc.detach()\n",
        "  del img_out, img_sliced, out_enc; torch.cuda.empty_cache()\n",
        "\n",
        "  if progressive_grow is True:\n",
        "    lr_cur = lr0 + (i / steps) * (lr1 - lr0)\n",
        "    for g in optimizer.param_groups: \n",
        "      g['lr'] = lr_cur\n",
        "\n",
        "  optimizer.zero_grad()\n",
        "  loss.backward()\n",
        "  optimizer.step()\n",
        "  \n",
        "  if i % save_freq == 0:\n",
        "    checkout(i // save_freq)\n",
        "\n",
        "outpic = ipy.Output()\n",
        "outpic\n",
        "\n",
        "pbar = ProgressBar(steps)\n",
        "for i in range(steps):\n",
        "  train(i)\n",
        "  _ = pbar.upd()\n",
        "\n",
        "HTML(makevid(tempdir))\n",
        "torch.save(params, tempdir + '.pt')\n",
        "files.download(tempdir + '.pt')\n",
        "files.download(tempdir + '.mp4')\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}