notebooks/Imagen_pytorch_inference_new.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tXkwem58FrAj"
      },
      "source": [
        "## Installation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 191,
      "metadata": {
        "id": "jFduUePWDiRr"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!git lfs install\n",
        "!git clone https://huggingface.co/Cene655/ImagenT5-3B"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eBuoRWpcY3ph"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip install git+https://github.com/cene555/Imagen-pytorch.git\n",
        "!pip install git+https://github.com/openai/CLIP.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mH-AFZAEfyEJ"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!git clone https://github.com/xinntao/Real-ESRGAN.git"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MvW2RHMHQl9g"
      },
      "outputs": [],
      "source": [
        "%cd Real-ESRGAN"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xIiOlXJzQosB"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip install basicsr\n",
        "# facexlib and gfpgan are for face enhancement\n",
        "!pip install facexlib\n",
        "!pip install gfpgan"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fChuITUYQstj"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip install -r requirements.txt\n",
        "!python setup.py develop\n",
        "!wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth -P experiments/pretrained_models"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h5_rnwU5GTry"
      },
      "source": [
        "## Imports "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HAGoFjvGZJ6s"
      },
      "outputs": [],
      "source": [
        "from PIL import Image\n",
        "from IPython.display import display\n",
        "import torch as th\n",
        "from imagen_pytorch.model_creation import create_model_and_diffusion as create_model_and_diffusion_dalle2\n",
        "from imagen_pytorch.model_creation import model_and_diffusion_defaults as model_and_diffusion_defaults_dalle2\n",
        "from transformers import AutoTokenizer\n",
        "import cv2\n",
        "\n",
        "import glob\n",
        "import os\n",
        "from basicsr.archs.rrdbnet_arch import RRDBNet\n",
        "from realesrgan import RealESRGANer\n",
        "from realesrgan.archs.srvgg_arch import SRVGGNetCompact\n",
        "from gfpgan import GFPGANer\n",
        "\n",
        "has_cuda = th.cuda.is_available()\n",
        "device = th.device('cpu' if not has_cuda else 'cuda')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qhCZMEueGk3Q"
      },
      "source": [
        "## Setting Up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Kp6HRM1vdFZq"
      },
      "outputs": [],
      "source": [
        "def model_fn(x_t, ts, **kwargs):\n",
        "    guidance_scale = 5\n",
        "    half = x_t[: len(x_t) // 2]\n",
        "    combined = th.cat([half, half], dim=0)\n",
        "    model_out = model(combined, ts, **kwargs)\n",
        "    eps, rest = model_out[:, :3], model_out[:, 3:]\n",
        "    cond_eps, uncond_eps = th.split(eps, len(eps) // 2, dim=0)\n",
        "    half_eps = uncond_eps + guidance_scale * (cond_eps - uncond_eps)\n",
        "    eps = th.cat([half_eps, half_eps], dim=0)\n",
        "    return th.cat([eps, rest], dim=1)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RyQgjL0OdMeu"
      },
      "outputs": [],
      "source": [
        "def show_images(batch: th.Tensor):\n",
        "    \"\"\" Display a batch of images inline.\"\"\"\n",
        "    scaled = ((batch + 1)*127.5).round().clamp(0,255).to(th.uint8).cpu()\n",
        "    reshaped = scaled.permute(2, 0, 3, 1).reshape([batch.shape[2], -1, 3])\n",
        "    display(Image.fromarray(reshaped.numpy()))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rVHdkvoGPqJJ"
      },
      "outputs": [],
      "source": [
        "def get_numpy_img(img):\n",
        "    scaled = ((img + 1)*127.5).round().clamp(0,255).to(th.uint8).cpu()\n",
        "    reshaped = scaled.permute(2, 0, 3, 1).reshape([img.shape[2], -1, 3])\n",
        "    return cv2.cvtColor(reshaped.numpy(), cv2.COLOR_BGR2RGB)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qhKnHPtJZCL_"
      },
      "outputs": [],
      "source": [
        "def _fix_path(path):\n",
        "  d = th.load(path)\n",
        "  checkpoint = {}\n",
        "  for key in d.keys():\n",
        "    checkpoint[key.replace('module.','')] = d[key]\n",
        "  return checkpoint"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "veXmk5XTZOuy"
      },
      "outputs": [],
      "source": [
        "options = model_and_diffusion_defaults_dalle2()\n",
        "options['use_fp16'] = False\n",
        "options['diffusion_steps'] = 200\n",
        "options['num_res_blocks'] = 3\n",
        "options['t5_name'] = 't5-3b'\n",
        "options['cache_text_emb'] = True\n",
        "model, diffusion = create_model_and_diffusion_dalle2(**options)\n",
        "\n",
        "model.eval()\n",
        "\n",
        "#if has_cuda:\n",
        "#    model.convert_to_fp16()\n",
        "\n",
        "model.to(device)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 192,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "dtqqsCzzaPzo",
        "outputId": "c8da3d0a-b897-4743-bc48-ab8440f9354b"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "total base parameters 1550556742\n"
          ]
        }
      ],
      "source": [
        "model.load_state_dict(_fix_path('/content/ImagenT5-3B/model.pt'))\n",
        "print('total base parameters', sum(x.numel() for x in model.parameters()))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 193,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_mdhjDRmejr5",
        "outputId": "61049fe6-00fa-4cd3-ef48-c3c8932e15f0"
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "1550556742"
            ]
          },
          "metadata": {},
          "execution_count": 193
        }
      ],
      "source": [
        "num_params = sum(param.numel() for param in model.parameters())\n",
        "num_params"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 194,
      "metadata": {
        "id": "4oDZRKP_NcV0"
      },
      "outputs": [],
      "source": [
        "realesrgan_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64,\n",
        "                           num_block=23, num_grow_ch=32, scale=4)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 195,
      "metadata": {
        "id": "BFl4yR5ONtil"
      },
      "outputs": [],
      "source": [
        "netscale = 4"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 196,
      "metadata": {
        "id": "Vwy6nPleNuUN"
      },
      "outputs": [],
      "source": [
        "upsampler = RealESRGANer(\n",
        "    scale=netscale,\n",
        "    model_path='/content/Real-ESRGAN/experiments/pretrained_models/RealESRGAN_x4plus.pth',\n",
        "    model=realesrgan_model,\n",
        "    tile=0,\n",
        "    tile_pad=10,\n",
        "    pre_pad=0,\n",
        "    half=True\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 197,
      "metadata": {
        "id": "P_PiM5y5PHCe"
      },
      "outputs": [],
      "source": [
        "face_enhancer = GFPGANer(\n",
        "    model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth',\n",
        "    upscale=4,\n",
        "    arch='clean',\n",
        "    channel_multiplier=2,\n",
        "    bg_upsampler=upsampler\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 198,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ilBtcMcGcUSZ",
        "outputId": "063933dc-014d-422b-bc28-f11ffa24ff61"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/transformers/models/t5/tokenization_t5_fast.py:161: FutureWarning: This tokenizer was incorrectly instantiated with a model max length of 512 which will be corrected in Transformers v5.\n",
            "For now, this behavior is kept to avoid breaking backwards compatibility when padding/encoding with `truncation is True`.\n",
            "- Be aware that you SHOULD NOT rely on t5-3b automatically truncating your input to 512 when padding/encoding.\n",
            "- If you want to encode/pad to sequences longer than 512 you can either instantiate this tokenizer with `model_max_length` or pass `max_length` when encoding/padding.\n",
            "- To avoid this warning, please instantiate this tokenizer with `model_max_length` set to your preferred value.\n",
            "  FutureWarning,\n"
          ]
        }
      ],
      "source": [
        "tokenizer = AutoTokenizer.from_pretrained(options['t5_name'])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 312,
      "metadata": {
        "id": "I1E6jjzhvV40"
      },
      "outputs": [],
      "source": [
        "#@title What do you want to generate?\n",
        "\n",
        "prompt = 'A photo of cat'#@param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 313,
      "metadata": {
        "id": "h3FFIeLrcuTX"
      },
      "outputs": [],
      "source": [
        "text_encoding = tokenizer(\n",
        "    prompt,\n",
        "    max_length=128,\n",
        "    padding=\"max_length\",\n",
        "    truncation=True,\n",
        "    return_attention_mask=True,\n",
        "    add_special_tokens=True,\n",
        "    return_tensors=\"pt\"\n",
        ")\n",
        "\n",
        "uncond_text_encoding = tokenizer(\n",
        "    '',\n",
        "    max_length=128,\n",
        "    padding=\"max_length\",\n",
        "    truncation=True,\n",
        "    return_attention_mask=True,\n",
        "    add_special_tokens=True,\n",
        "    return_tensors=\"pt\"\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 314,
      "metadata": {
        "id": "1Imtn9DVZD99"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "batch_size = 4\n",
        "cond_tokens = th.from_numpy(np.array([text_encoding['input_ids'][0].numpy() for i in range(batch_size)]))\n",
        "uncond_tokens = th.from_numpy(np.array([uncond_text_encoding['input_ids'][0].numpy() for i in range(batch_size)]))\n",
        "cond_attention_mask = th.from_numpy(np.array([text_encoding['attention_mask'][0].numpy() for i in range(batch_size)]))\n",
        "uncond_attention_mask = th.from_numpy(np.array([uncond_text_encoding['attention_mask'][0].numpy() for i in range(batch_size)]))\n",
        "model_kwargs = {}\n",
        "model_kwargs[\"tokens\"] = th.cat((cond_tokens,\n",
        "                                 uncond_tokens)).to(device)\n",
        "model_kwargs[\"mask\"] = th.cat((cond_attention_mask,\n",
        "                               uncond_attention_mask)).to(device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y0tzvo1tG9vS"
      },
      "source": [
        "## Generation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D-4GzOBadCbj"
      },
      "outputs": [],
      "source": [
        "model.del_cache()\n",
        "sample = diffusion.p_sample_loop(\n",
        "    model_fn,\n",
        "    (batch_size * 2, 3, 64, 64),\n",
        "    clip_denoised=True,\n",
        "    model_kwargs=model_kwargs,\n",
        "    device='cuda',\n",
        "    progress=True,\n",
        ")[:batch_size]\n",
        "model.del_cache()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hvsr_S-xMwyG"
      },
      "outputs": [],
      "source": [
        "show_images(sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lsYwSFrE6zrJ"
      },
      "outputs": [],
      "source": [
        "for i in sample:\n",
        "    show_images(i.unsqueeze(0))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 318,
      "metadata": {
        "id": "h4hoHPBw2DQh"
      },
      "outputs": [],
      "source": [
        "new_img = get_numpy_img(sample)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VAHtiUiyPZWg"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "for j in range(batch_size):\n",
        "    new_img = get_numpy_img(sample[j].unsqueeze(0))\n",
        "    for i in range(1):\n",
        "        _, _, new_img = face_enhancer.enhance(new_img, has_aligned=False,\n",
        "                                              only_center_face=False, paste_back=True)\n",
        "        cv2.imwrite(f'/content/test_out{j}.jpg', new_img)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 319,
      "metadata": {
        "id": "f_BlQypdaQ92"
      },
      "outputs": [],
      "source": [
        ""
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "machine_shape": "hm",
      "name": "Imagen_pytorch_inference_new.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}