Text2Image.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Text2Image",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/tg-bomze/collection-of-notebooks/blob/master/Text2Image.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "toWe1IoH7X35"
      },
      "source": [
        "<b><font color=\"black\" size=\"+4\">Text2Image</font></b>\r\n",
        "\r\n",
        "<b><font color=\"black\" size=\"+2\">Based on:</font></b>\r\n",
        "\r\n",
        "**GitHub repository**: [CLIP](https://github.com/openai/CLIP)\r\n",
        "\r\n",
        "Article: [Learning Transferable Visual Models From Natural Language Supervision](https://cdn.openai.com/papers/Learning_Transferable_Visual_Models_From_Natural_Language_Supervision.pdf)\r\n",
        "\r\n",
        "Creator: **[OpenAI](https://openai.com/)**\r\n",
        "\r\n",
        "Colab based on **[Ryan Murdock's](https://rynmurdock.github.io/)** notebook.\r\n",
        "\r\n",
        "<b><font color=\"black\" size=\"+2\">Colab created by:</font></b>\r\n",
        "\r\n",
        "GitHub: [@tg-bomze](https://github.com/tg-bomze),\r\n",
        "Telegram: [@bomze](https://t.me/bomze),\r\n",
        "Twitter: [@tg_bomze](https://twitter.com/tg_bomze).\r\n",
        "\r\n",
        "---\r\n",
        "\r\n",
        "\r\n",
        "```\r\n",
        "(ENG) To get started, click on the button (where the red arrow indicates). After clicking, wait until the execution is complete.\r\n",
        "```\r\n",
        "```\r\n",
        "(RUS) Чтобы начать, поочередно нажимайте на кнопки (куда указывают красные стрелки), дожидаясь завершения выполнения каждого блока.\r\n",
        "```\r\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "etzxXVZ_r-Nf",
        "cellView": "form"
      },
      "source": [
        "#@title <b><font color=\"red\" size=\"+3\">←</font><font color=\"black\" size=\"+3\"> Install all necessary libraries</font></b>\r\n",
        "#@markdown <b><font color=\"black\" size=\"+1\"> (RUS) Установить все необходимые библиотеки</font></b>\r\n",
        "\r\n",
        "try: \r\n",
        "  !pip3 install googletrans==3.1.0a0\r\n",
        "  from googletrans import Translator, constants\r\n",
        "  from pprint import pprint\r\n",
        "  translator = Translator()\r\n",
        "except: pass\r\n",
        "\r\n",
        "import subprocess\r\n",
        "CUDA_version = [s for s in subprocess.check_output([\"nvcc\", \"--version\"]).decode(\"UTF-8\").split(\", \") if s.startswith(\"release\")][0].split(\" \")[-1]\r\n",
        "print(\"CUDA version:\", CUDA_version)\r\n",
        "\r\n",
        "if CUDA_version == \"10.0\":\r\n",
        "    torch_version_suffix = \"+cu100\"\r\n",
        "elif CUDA_version == \"10.1\":\r\n",
        "    torch_version_suffix = \"+cu101\"\r\n",
        "elif CUDA_version == \"10.2\":\r\n",
        "    torch_version_suffix = \"\"\r\n",
        "else:\r\n",
        "    torch_version_suffix = \"+cu110\"\r\n",
        "\r\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\r\n",
        "import torch\r\n",
        "import numpy as np\r\n",
        "import torchvision\r\n",
        "import torchvision.transforms.functional as TF\r\n",
        "\r\n",
        "import PIL\r\n",
        "import matplotlib.pyplot as plt\r\n",
        "\r\n",
        "import os\r\n",
        "import random\r\n",
        "import imageio\r\n",
        "from IPython import display\r\n",
        "from IPython.core.interactiveshell import InteractiveShell\r\n",
        "InteractiveShell.ast_node_interactivity = \"all\"\r\n",
        "import glob\r\n",
        "from google.colab import output\r\n",
        "\r\n",
        "!git clone https://github.com/openai/CLIP.git\r\n",
        "%cd /content/CLIP/\r\n",
        "!pip install ftfy\r\n",
        "import clip\r\n",
        "import numpy as np\r\n",
        "perceptor, preprocess = clip.load('ViT-B/32')\r\n",
        "\r\n",
        "!mkdir frames\r\n",
        "import moviepy.editor as mpy\r\n",
        "from moviepy.video.io.ffmpeg_writer import FFMPEG_VideoWriter\r\n",
        "from google.colab import files\r\n",
        "import warnings\r\n",
        "from IPython.display import clear_output\r\n",
        "import matplotlib.pyplot as plt\r\n",
        "warnings.filterwarnings(\"ignore\")\r\n",
        "%matplotlib inline\r\n",
        "clear_output()\r\n",
        "!nvidia-smi -L\r\n",
        "print('\\nDone!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Nq0wA-wc-P-s",
        "trusted": true,
        "cellView": "form"
      },
      "source": [
        "#@title <b><font color=\"red\" size=\"+3\">←</font><font color=\"black\" size=\"+3\"> Start Processing</font></b>\n",
        "#@markdown <b><font color=\"black\" size=\"+1\"> (RUS) Начать обработку</font></b>\n",
        "\n",
        "!rm -rf /content/CLIP/frames/*.*\n",
        "#@markdown **(ENG)** Enter a description for the picture and click the \"Play\" button. Text can be written in any language.\n",
        "\n",
        "#@markdown **(RUS)** Введите описание картинки и нажмите кнопку \"Play\". Текст можно писать на любом языке.\n",
        "\n",
        "text = 'Forbidden love' #@param {type:\"string\"}\n",
        "try:\n",
        "  translation = translator.translate(text)\n",
        "  prompt = translation.text\n",
        "except: prompt = text\n",
        "#@markdown **(ENG) Attention!** This process is almost endless, so watch the results and stop this block when you think it's time. Press the \"Play\" button again to stop.\n",
        "\n",
        "#@markdown **(RUS) Внимание!** Этот процесс почти бесконечный, поэтому смотрите за результатами и остановите работу этого блока когда посчитаете, что уже пора. Для остановки снова нажмите кнопку \"Play\".\n",
        "\n",
        "#@markdown ---\n",
        "resolution = \"256\" #@param [128, 256, 512]\n",
        "res = int(resolution)\n",
        "im_shape = [res, res, 3]\n",
        "sideX, sideY, channels = im_shape\n",
        "\n",
        "tx = clip.tokenize(prompt)\n",
        "frames_frequency = \"high\" #@param ['high', 'medium', 'low']\n",
        "if frames_frequency=='low': fr = 150\n",
        "elif frames_frequency=='medium': fr = 100\n",
        "else: fr = 50\n",
        "audio_notification = False #@param {type:\"boolean\"}\n",
        "\n",
        "# Define some helper functions\n",
        "def displ(img, pre_scaled=True):\n",
        "  img = np.array(img)[:,:,:]\n",
        "  img = np.transpose(img, (1, 2, 0))\n",
        "  if not pre_scaled:\n",
        "    img = scale(img, 48*4, 32*4)\n",
        "  imageio.imwrite('result.png', np.array(img))\n",
        "  file_path = '/content/CLIP/frames/{}.png'.format(str(len(os.listdir('/content/CLIP/frames'))).zfill(5))\n",
        "  imageio.imwrite(file_path, np.array(img))\n",
        "  return display.Image('result.png')\n",
        "\n",
        "def card_padded(im, to_pad=3):\n",
        "  return np.pad(np.pad(np.pad(im, [[1,1], [1,1], [0,0]],constant_values=0), [[2,2], [2,2], [0,0]],constant_values=1),\n",
        "            [[to_pad,to_pad], [to_pad,to_pad], [0,0]],constant_values=0)\n",
        "\n",
        "import torch.nn as nn\n",
        "\n",
        "class SineLayer(nn.Module):\n",
        "    def __init__(self, in_features, out_features, bias=True,\n",
        "                 is_first=False, omega_0=30):\n",
        "        super().__init__()\n",
        "        self.omega_0 = omega_0\n",
        "        self.is_first = is_first\n",
        "        \n",
        "        self.in_features = in_features\n",
        "        self.linear = nn.Linear(in_features, out_features, bias=bias)\n",
        "        \n",
        "        self.init_weights()\n",
        "    \n",
        "    def init_weights(self):\n",
        "        with torch.no_grad():\n",
        "            if self.is_first:\n",
        "                self.linear.weight.uniform_(-1 / self.in_features, \n",
        "                                             1 / self.in_features)      \n",
        "            else:\n",
        "                self.linear.weight.uniform_(-np.sqrt(6 / self.in_features) / self.omega_0, \n",
        "                                             np.sqrt(6 / self.in_features) / self.omega_0)\n",
        "        \n",
        "    def forward(self, input):\n",
        "        return torch.sin(self.omega_0 * self.linear(input))\n",
        "    \n",
        "    def forward_with_intermediate(self, input): \n",
        "        # For visualization of activation distributions\n",
        "        intermediate = self.omega_0 * self.linear(input)\n",
        "        return torch.sin(intermediate), intermediate\n",
        "    \n",
        "class Siren(nn.Module):\n",
        "    def __init__(self, in_features, hidden_features, hidden_layers, out_features, outermost_linear=True, \n",
        "                 first_omega_0=30, hidden_omega_0=30.):\n",
        "        super().__init__()\n",
        "        \n",
        "        self.net = []\n",
        "        self.net.append(SineLayer(in_features, hidden_features, \n",
        "                                  is_first=True, omega_0=first_omega_0))\n",
        "\n",
        "        for i in range(hidden_layers):\n",
        "            self.net.append(SineLayer(hidden_features, hidden_features, \n",
        "                                      is_first=False, omega_0=hidden_omega_0))\n",
        "\n",
        "        if outermost_linear:\n",
        "            final_linear = nn.Linear(hidden_features, out_features)\n",
        "            \n",
        "            with torch.no_grad():\n",
        "                final_linear.weight.uniform_(-np.sqrt(6 / hidden_features) / hidden_omega_0, \n",
        "                                              np.sqrt(6 / hidden_features) / hidden_omega_0)\n",
        "                \n",
        "            self.net.append(final_linear)\n",
        "        else:\n",
        "            self.net.append(SineLayer(hidden_features, out_features, \n",
        "                                      is_first=False, omega_0=hidden_omega_0))\n",
        "        \n",
        "        self.net = nn.Sequential(*self.net)\n",
        "    \n",
        "    def forward(self, coords):\n",
        "        coords = coords.clone().detach().requires_grad_(True)\n",
        "        output = self.net(coords.cuda())\n",
        "        return output.view(1, sideX, sideY, 3).permute(0, 3, 1, 2)#.sigmoid_()\n",
        "\n",
        "    def forward_with_activations(self, coords, retain_grad=False):\n",
        "        '''Returns not only model output, but also intermediate activations.\n",
        "        Only used for visualizing activations later!'''\n",
        "        activations = OrderedDict()\n",
        "\n",
        "        activation_count = 0\n",
        "        x = coords.clone().detach().requires_grad_(True)\n",
        "        activations['input'] = x\n",
        "        for i, layer in enumerate(self.net):\n",
        "            if isinstance(layer, SineLayer):\n",
        "                x, intermed = layer.forward_with_intermediate(x)\n",
        "                \n",
        "                if retain_grad:\n",
        "                    x.retain_grad()\n",
        "                    intermed.retain_grad()\n",
        "                    \n",
        "                activations['_'.join((str(layer.__class__), \"%d\" % activation_count))] = intermed\n",
        "                activation_count += 1\n",
        "            else: \n",
        "                x = layer(x)\n",
        "                \n",
        "                if retain_grad:\n",
        "                    x.retain_grad()\n",
        "                    \n",
        "            activations['_'.join((str(layer.__class__), \"%d\" % activation_count))] = x\n",
        "            activation_count += 1\n",
        "\n",
        "        return activations\n",
        "\n",
        "\n",
        "def get_mgrid(sidelen, dim=2):\n",
        "    '''Generates a flattened grid of (x,y,...) coordinates in a range of -1 to 1.\n",
        "    sidelen: int\n",
        "    dim: int'''\n",
        "    tensors = tuple(dim * [torch.linspace(-1, 1, steps=sidelen)])\n",
        "    mgrid = torch.stack(torch.meshgrid(*tensors), dim=-1)\n",
        "    mgrid = mgrid.reshape(-1, dim)\n",
        "    return mgrid\n",
        "\n",
        "\n",
        "model = Siren(in_features=2, out_features=3, hidden_features=256, \n",
        "                  hidden_layers=16, outermost_linear=False)\n",
        "model.cuda()\n",
        "LLL = []\n",
        "eps = 0\n",
        "optimizer = torch.optim.Adam(model.parameters(), .00001)\n",
        "\n",
        "def checkin(loss):\n",
        "  #print(loss)\n",
        "  with torch.no_grad():\n",
        "    al = nom(model(get_mgrid(sideX)).cpu()).numpy()\n",
        "  for allls in al:\n",
        "    displ(allls)\n",
        "    clear_output()\n",
        "    pic_num = str(len(os.listdir('/content/CLIP/frames')))\n",
        "    print(f'Picture number {pic_num}\\n')\n",
        "    if int(pic_num) == 1:\n",
        "      print(\"\\n(ENG) Don't be alarmed. The first image is always bad\")\n",
        "      print('(RUS) Не пугайтесь. Первое изображение всегда плохое\\n')\n",
        "    display.display(display.Image('result.png'))\n",
        "    print('\\n(ENG) Please wait a few seconds until the next picture')\n",
        "    print('(RUS) Пожалуйста подождите несколько секунд, пока не появится новое изображение')\n",
        "  if audio_notification == True: output.eval_js('new Audio(\"https://freesound.org/data/previews/80/80921_1022651-lq.ogg\").play()')\n",
        "\n",
        "\n",
        "def ascend_txt():\n",
        "  out = model(get_mgrid(sideX))\n",
        "\n",
        "  cutn = 64\n",
        "  p_s = []\n",
        "  for ch in range(cutn):\n",
        "    size = torch.randint(int(.5*sideX), int(.98*sideX), ())\n",
        "    offsetx = torch.randint(0, sideX - size, ())\n",
        "    offsety = torch.randint(0, sideX - size, ())\n",
        "    apper = out[:, :, offsetx:offsetx + size, offsety:offsety + size]\n",
        "    apper = torch.nn.functional.interpolate(apper, (224,224), mode='bilinear')\n",
        "    p_s.append(nom(apper))\n",
        "  into = torch.cat(p_s, 0)\n",
        "\n",
        "  iii = perceptor.encode_image(into)\n",
        "  t = perceptor.encode_text(tx.cuda())\n",
        "  return -100*torch.cosine_similarity(t, iii, dim=-1).mean()\n",
        "\n",
        "def train(epoch, i):\n",
        "  loss = ascend_txt()\n",
        "  optimizer.zero_grad()\n",
        "  loss.backward()\n",
        "  optimizer.step()\n",
        "  \n",
        "  if itt % fr == 0:\n",
        "    checkin(loss)\n",
        "\n",
        "nom = torchvision.transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))\n",
        "\n",
        "itt = 0\n",
        "for epochs in range(10000):\n",
        "  for i in range(1000):\n",
        "    train(eps, i)\n",
        "    itt+=1\n",
        "  eps+=1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v68XwQAMv9O6",
        "cellView": "form"
      },
      "source": [
        "#@title <b><font color=\"red\" size=\"+3\">←</font><font color=\"black\" size=\"+3\"> Generate Video</font></b>\r\n",
        "#@markdown <b><font color=\"black\" size=\"+1\"> (RUS) Сгенерировать видео</font></b>\r\n",
        "\r\n",
        "#@markdown **(ENG) Attention!** The duration of the video will depend on how many pictures you generated in the previous block. 10 frames = 1 sec..\r\n",
        "\r\n",
        "#@markdown **(RUS) Внимание!** Длительность видео будет зависеть от того, сколько картинок вы сгенерировали на предыдущем блоке. 10 кадров = 1 сек.\r\n",
        "\r\n",
        "frames = []\r\n",
        "img = os.listdir(\"/content/CLIP/frames\")\r\n",
        "img.sort()\r\n",
        "for i in img:\r\n",
        "  frames.append(imageio.imread(\"/content/CLIP/frames/\"+i))\r\n",
        "frames = np.array(frames)\r\n",
        "imageio.mimsave('/content/CLIP/video.mp4', frames)\r\n",
        "clear_output()\r\n",
        "\r\n",
        "from IPython.display import HTML\r\n",
        "import io\r\n",
        "import base64\r\n",
        "video = io.open('/content/CLIP/video.mp4', 'r+b').read()\r\n",
        "encoded = base64.b64encode(video)\r\n",
        "play_html = ('<video alt=\"test\" controls><source src=\"data:video/mp4;base64,{0}\" type=\"video/mp4\" /> </video>'.format(encoded.decode('ascii')))\r\n",
        "HTML(data=play_html)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7Di3LTW84O1i",
        "cellView": "form"
      },
      "source": [
        "#@title <b><font color=\"red\" size=\"+3\">←</font><font color=\"black\" size=\"+3\"> Download results</font></b>\r\n",
        "#@markdown <b><font color=\"black\" size=\"+1\"> (RUS) Скачать результаты</font></b>\r\n",
        "last_image = True #@param {type:\"boolean\"}\r\n",
        "video = True #@param {type:\"boolean\"}\r\n",
        "if last_image==True: files.download('/content/CLIP/result.png')\r\n",
        "if video==True: files.download('/content/CLIP/video.mp4')"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}