rq_vae_transformer_demo.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Uh3sl4iPV84A"
      },
      "source": [
        "論文  \n",
        "https://arxiv.org/abs/2203.01941<br>\n",
        "<br>  \n",
        "GitHub  \n",
        "https://github.com/kakaobrain/rq-vae-transformer<br>\n",
        "<br>\n",
        "<a href=\"https://colab.research.google.com/github/kaz12tech/ai_demos/blob/master/rq_vae_transformer_demo.ipynb\" target=\"_blank\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J_pJP01IV84D"
      },
      "source": [
        "## 環境セットアップ"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P84G1Pi_V84D"
      },
      "source": [
        "## GPU確認"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "2PcII5fLdHnJ",
        "outputId": "95dfb1e9-5acd-4cc3-8ef0-92e1fb9c2e8f"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H8jM3WJyV84F"
      },
      "source": [
        "## GitHubからコード取得"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "R-Y5auuvnnnT",
        "outputId": "a489dc4f-67f3-486a-8cce-5183812195dc"
      },
      "outputs": [],
      "source": [
        "%cd /content\n",
        "\n",
        "!git clone https://github.com/kakaobrain/rq-vae-transformer.git"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N-oImepVV84F"
      },
      "source": [
        "## ライブラリのインストール"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "LScsvO3UppqT",
        "outputId": "a07a296e-9d97-4513-8586-1db25eab9983"
      },
      "outputs": [],
      "source": [
        "%cd /content/rq-vae-transformer\n",
        "\n",
        "# もし RESTART RUNTIMEが表示されたら「ランタイム」→「ランタイムを再起動」\n",
        "!pip install -r requirements.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DuZozvw4V84G"
      },
      "source": [
        "## ライブラリのインポート"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "o2lxptJeo6kr",
        "outputId": "41409c09-a249-4446-c9da-30f9e9412d02"
      },
      "outputs": [],
      "source": [
        "%cd /content/rq-vae-transformer\n",
        "\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "%matplotlib inline\n",
        "\n",
        "from PIL import Image\n",
        "import yaml\n",
        "import torch\n",
        "import torchvision\n",
        "import clip\n",
        "import torch.nn.functional as F\n",
        "\n",
        "from notebooks.notebook_utils import TextEncoder, load_model, get_generated_images_by_texts"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V-ZF2lL_V84H"
      },
      "source": [
        "## 学習済みモデルのセットアップ"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "QYmO4couoYhz",
        "outputId": "ab25c870-ad64-4116-e388-ab0b661e1c36"
      },
      "outputs": [],
      "source": [
        "%cd /content/rq-vae-transformer\n",
        "!mkdir pretrained\n",
        "%cd pretrained\n",
        "\n",
        "!wget https://arena.kakaocdn.net/brainrepo/models/RQVAE/dcd95e8f08408e113aab6451fae895f5/cc3m.tar.gz\n",
        "!tar -xvf cc3m.tar.gz"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dzQvCUYKZLKo"
      },
      "source": [
        "## モデルのロード"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9yTv_4FtZPKT",
        "outputId": "ae0fc7d0-8095-470f-8048-2836aa1158fd"
      },
      "outputs": [],
      "source": [
        "vqvae_path = '/content/rq-vae-transformer/pretrained/cc3m/stage1/model.pt'\n",
        "model_path = '/content/rq-vae-transformer/pretrained/cc3m/stage2/model.pt'\n",
        "\n",
        "# load stage 1 model: RQ-VAE\n",
        "model_vqvae, _ = load_model(vqvae_path)\n",
        "\n",
        "# load stage 2 model: RQ-Transformer\n",
        "model_ar, config = load_model(model_path, ema=False)\n",
        "\n",
        "# GPUにセット\n",
        "model_ar = model_ar.cuda().eval()\n",
        "model_vqvae = model_vqvae.cuda().eval()\n",
        "\n",
        "# CLIPモデルのダウンロード\n",
        "model_clip, preprocess_clip = clip.load(\"ViT-B/32\", device='cpu')\n",
        "model_clip = model_clip.cuda().eval()\n",
        "\n",
        "# prepare text encoder to tokenize natual languages\n",
        "text_encoder = TextEncoder(tokenizer_name=config.dataset.txt_tok_name, \n",
        "                           context_length=config.dataset.context_length)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NRO5D0XdZ7QW"
      },
      "source": [
        "# Text to Image"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8ztqzt5dZ8rE"
      },
      "outputs": [],
      "source": [
        "#@markdown 入力テキストを設定してください。\n",
        "text_prompts = \"'a photo of \\\"Cherry blossoms in the snow\\\"'\" #@param {type:\"string\"}\n",
        "\n",
        "num_samples = 16\n",
        "temperature= 1.0\n",
        "top_k=1024\n",
        "top_p=0.95"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lm82lWyNamqY",
        "outputId": "e85ca770-4488-4c0c-d00a-c4f361ede73f"
      },
      "outputs": [],
      "source": [
        "pixels = get_generated_images_by_texts(model_ar,\n",
        "                                       model_vqvae,\n",
        "                                       text_encoder,\n",
        "                                       model_clip,\n",
        "                                       preprocess_clip,\n",
        "                                       text_prompts,\n",
        "                                       num_samples,\n",
        "                                       temperature,\n",
        "                                       top_k,\n",
        "                                       top_p,\n",
        "                                      )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3FQdzmrSaxAc"
      },
      "source": [
        "## 結果の表示"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "bgqTfIRpazLA",
        "outputId": "2ea6c9d3-97b8-4aea-8e7f-25833cad721a"
      },
      "outputs": [],
      "source": [
        "num_visualize_samples = 16\n",
        "images = [pixel.cpu().numpy() * 0.5 + 0.5 for pixel in pixels]\n",
        "images = torch.from_numpy(np.array(images[:num_visualize_samples]))\n",
        "images = torch.clamp(images, 0, 1)\n",
        "grid = torchvision.utils.make_grid(images, nrow=4)\n",
        "\n",
        "img = Image.fromarray(np.uint8(grid.numpy().transpose([1,2,0])*255))\n",
        "display(img)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "rq_vae_transformer_demo.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "orig_nbformat": 4
  },
  "nbformat": 4,
  "nbformat_minor": 0
}