CLIP_VQGAN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "Text2Image_VQGAN.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",
        "Based on [CLIP](https://github.com/openai/CLIP) + VQGAN from [Taming Transformers](https://github.com/CompVis/taming-transformers) // made by Vadim Epstein [[eps696](https://github.com/eps696)]  \n",
        "thanks to [Ryan Murdock](https://twitter.com/advadnoun), [Jonathan Fly](https://twitter.com/jonathanfly), [Hannu Toyryla](https://twitter.com/htoyryla) for ideas\n",
        "\n",
        "## Features \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",
        "\n",
        "* various VQGAN models (incl. newest Gumbel-F8)\n",
        "* various CLIP models (incl. multi-language from [SBERT](https://sbert.net))\n",
        "* saving/loading VQGAN snapshots to resume processing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QytcEMSKBtN-"
      },
      "source": [
        "**Run the cell below after each session restart**\n",
        "\n",
        "First select `VQGAN_model` for generation.  \n",
        "`Gumbel` is probably the best, but eats more RAM (max resolution on Colab ~900x500). `F16-1024` can go up to ~1000x600.  \n",
        "`resume` if you want to start from the saved snapshot."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "etzxXVZ_r-Nf",
        "cellView": "form"
      },
      "source": [
        "#@title General setup\n",
        "\n",
        "VQGAN_model = \"gumbel_f8-8192\" #@param ['gumbel_f8-8192', 'imagenet_f16-1024', 'imagenet_f16-16384']\n",
        "resume = False #@param {type:\"boolean\"}\n",
        "\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 collections import OrderedDict\n",
        "from base64 import b64encode\n",
        "# from moviepy.editor import ImageSequenceClip\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",
        "%cd /content\n",
        "!rm -rf aphantasia\n",
        "!git clone https://github.com/eps696/aphantasia\n",
        "%cd aphantasia/\n",
        "from clip_fft import slice_imgs\n",
        "from utils import pad_up_to, basename, img_list, img_read, plot_text, txt_clean\n",
        "import transforms\n",
        "from progress_bar import ProgressIPy as ProgressBar\n",
        "\n",
        "!pip install omegaconf>=2.0.0 pytorch-lightning>=1.0.8 einops>=0.3.0\n",
        "import pytorch_lightning as pl\n",
        "!git clone https://github.com/CompVis/taming-transformers\n",
        "!mv taming-transformers/* ./\n",
        "import yaml\n",
        "from omegaconf import OmegaConf\n",
        "from taming.modules.diffusionmodules.model import Decoder\n",
        "from taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer\n",
        "from taming.modules.vqvae.quantize import GumbelQuantize\n",
        "\n",
        "class VQModel(pl.LightningModule):\n",
        "  def __init__(self, ddconfig, n_embed, embed_dim, remap=None, sane_index_shape=False, **kwargs_ignore):  # tell vector quantizer to return indices as bhw\n",
        "    super().__init__()\n",
        "    self.decoder = Decoder(**ddconfig)\n",
        "    self.quantize = VectorQuantizer(n_embed, embed_dim, beta=0.25, remap=remap, sane_index_shape=sane_index_shape)\n",
        "  def decode(self, quant):\n",
        "    return self.decoder(quant)\n",
        "\n",
        "class GumbelVQ(VQModel):\n",
        "  def __init__(self, ddconfig, n_embed, embed_dim, kl_weight=1e-8, remap=None, **kwargs_ignore):\n",
        "    z_channels = ddconfig[\"z_channels\"]\n",
        "    super().__init__(ddconfig, n_embed, embed_dim)\n",
        "    self.quantize = GumbelQuantize(z_channels, embed_dim, n_embed=n_embed, kl_weight=kl_weight, temp_init=1.0, remap=remap)\n",
        "\n",
        "if not os.path.isdir('/content/models_TT'):\n",
        "  !mkdir -p /content/models_TT\n",
        "def getm(url, path):\n",
        "  if os.path.isfile(path) and os.stat(path).st_size > 0: \n",
        "    print(' already exists', path, os.stat(path).st_size)\n",
        "  else:\n",
        "    !wget $url -O $path\n",
        "\n",
        "if VQGAN_model == \"gumbel_f8-8192\" and not os.path.isfile('/content/models_TT/gumbel_f8-8192.ckpt'):\n",
        "  getm('https://heibox.uni-heidelberg.de/f/34a747d5765840b5a99d/?dl=1', '/content/models_TT/gumbel_f8-8192.ckpt')\n",
        "  getm('https://heibox.uni-heidelberg.de/f/b24d14998a8d4f19a34f/?dl=1', '/content/models_TT/gumbel_f8-8192.yaml')\n",
        "elif VQGAN_model == \"imagenet_f16-1024\" and not os.path.isfile('/content/models_TT/imagenet_f16-1024.ckpt'):\n",
        "  getm('https://heibox.uni-heidelberg.de/f/140747ba53464f49b476/?dl=1', '/content/models_TT/imagenet_f16-1024.ckpt')\n",
        "  getm('https://heibox.uni-heidelberg.de/f/6ecf2af6c658432c8298/?dl=1', '/content/models_TT/imagenet_f16-1024.yaml')\n",
        "elif VQGAN_model == \"imagenet_f16-16384\" and not os.path.isfile('/content/models_TT/imagenet_f16-16384.ckpt'):\n",
        "  getm('https://heibox.uni-heidelberg.de/f/867b05fc8c4841768640/?dl=1', '/content/models_TT/imagenet_f16-16384.ckpt')\n",
        "  getm('https://heibox.uni-heidelberg.de/f/274fb24ed38341bfa753/?dl=1', '/content/models_TT/imagenet_f16-16384.yaml')\n",
        "\n",
        "clear_output()\n",
        "\n",
        "if resume:\n",
        "  resumed = files.upload()\n",
        "  params_pt = list(resumed.values())[0]\n",
        "  params_pt = torch.load(io.BytesIO(params_pt))\n",
        "\n",
        "if VQGAN_model == \"gumbel_f8-8192\":\n",
        "  scale_res = 8\n",
        "else:\n",
        "  scale_res = 16\n",
        "\n",
        "def load_config(config_path):\n",
        "  config = OmegaConf.load(config_path)\n",
        "  return config\n",
        "\n",
        "def load_vqgan(config, ckpt_path=None):\n",
        "  if VQGAN_model == \"gumbel_f8-8192\":\n",
        "    model = GumbelVQ(**config.model.params)\n",
        "  else:\n",
        "    model = VQModel(**config.model.params)\n",
        "  if ckpt_path is not None:\n",
        "    sd = torch.load(ckpt_path, map_location=\"cpu\")[\"state_dict\"]\n",
        "    missing, unexpected = model.load_state_dict(sd, strict=False)\n",
        "  return model.eval()\n",
        "\n",
        "def vqgan_image(model, z):\n",
        "  x = model.decode(z)\n",
        "  x = (x+1.)/2.\n",
        "  return x\n",
        "\n",
        "class latents(torch.nn.Module):\n",
        "  def __init__(self, shape):\n",
        "    super(latents, self).__init__()\n",
        "    init_rnd = torch.zeros(shape).normal_(0.,4.)\n",
        "    self.lats = torch.nn.Parameter(init_rnd.cuda())\n",
        "  def forward(self):\n",
        "    return self.lats\n",
        "\n",
        "config_vqgan = load_config(\"/content/models_TT/%s.yaml\" % VQGAN_model)\n",
        "model_vqgan  = load_vqgan(config_vqgan, ckpt_path=\"/content/models_TT/%s.ckpt\" % VQGAN_model).cuda()\n",
        "\n",
        "def makevid(seq_dir, size=None):\n",
        "  out_sequence = seq_dir + '/%04d.jpg'\n",
        "  out_video = seq_dir + '.mp4'\n",
        "  !ffpb -y -i $out_sequence -codec nvenc $out_video\n",
        "  # 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] # XXX: only one GPU on Colab and isn’t guaranteed\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 CLIP 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",
        "`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` or resolution if you face OOM.  \n",
        "\n",
        "Generation video and final parameters snapshot are saved automatically.  \n",
        "NB: Requests are cumulative (start near the end of the previous run). To start generation from scratch, re-run General setup."
      ]
    },
    {
      "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 = 900 #@param {type:\"integer\"}\n",
        "sideY =  500#@param {type:\"integer\"}\n",
        "#@markdown > Config\n",
        "model = 'ViT-B/16' #@param ['ViT-B/16', 'ViT-B/32', 'RN101', 'RN50x16', 'RN50x4', 'RN50']\n",
        "align = 'uniform' #@param ['central', 'uniform', 'overscan']\n",
        "aug_transform = True #@param {type:\"boolean\"}\n",
        "sync =  0.4 #@param {type:\"number\"}\n",
        "#@markdown > Training\n",
        "steps = 200 #@param {type:\"integer\"}\n",
        "samples = 60 #@param {type:\"integer\"}\n",
        "learning_rate = 0.1 #@param {type:\"number\"}\n",
        "save_freq = 1 #@param {type:\"integer\"}\n",
        "\n",
        "if resume:\n",
        "  if not isinstance(params_pt, dict):\n",
        "    params_pt = OrderedDict({'lats': params_pt})\n",
        "  ps = params_pt['lats'].shape\n",
        "  size = [s*scale_res for s in ps[2:]]\n",
        "  lats = latents(ps).cuda()\n",
        "  _ = lats.load_state_dict(params_pt)\n",
        "  print(' resumed with size', size)\n",
        "else:\n",
        "  lats = latents([1, 256, sideY//scale_res, sideX//scale_res]).cuda()\n",
        "\n",
        "if multilang: model = 'ViT-B/32' # sbert model is trained with ViT\n",
        "\n",
        "if len(subtract) > 0:\n",
        "  samples = int(samples * 0.75)\n",
        "if sync > 0 and upload_image:\n",
        "  samples = int(samples * 0.5)\n",
        "print(' using %d samples' % samples)\n",
        "\n",
        "use_jit = True if float(torch.__version__[:3]) < 1.8 else False\n",
        "model_clip, _ = clip.load(model, jit=use_jit)\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:\n",
        "  model_lang = SentenceTransformer('clip-ViT-B-32-multilingual-v1').cuda()\n",
        "\n",
        "def enc_text(txt):\n",
        "  if multilang:\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",
        "sign = 1. if invert else -1.\n",
        "if aug_transform:\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(' text:', 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",
        "\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: del model_lang\n",
        "\n",
        "optimizer = torch.optim.AdamW(lats.parameters(), learning_rate, weight_decay=0.01, amsgrad=True)\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 = vqgan_image(model_vqgan, lats()).cpu().numpy()[0]\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",
        "def train(i):\n",
        "  loss = 0\n",
        "  img_out = vqgan_image(model_vqgan, lats())\n",
        "  img_sliced = slice_imgs([img_out], samples, modsize, trform_f, align)[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 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 * 0.5 * 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",
        "  del img_out, img_sliced, out_enc; torch.cuda.empty_cache()\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(lats.lats, tempdir + '.pt')\n",
        "files.download(tempdir + '.pt')\n",
        "files.download(tempdir + '.mp4')\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}