Latent_Diffusion_Upscale.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "mount_file_id": "1xxXnX6GMlUifTA6mYq3uRmbDNKv9D-bH",
      "authorship_tag": "ABX9TyNFIQVWf/70yls85G3+65PG",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/olaviinha/NeuralImageSuperResolution/blob/master/Latent_Diffusion_Upscale.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wCskRiNym_LW"
      },
      "source": [
        "#<font face=\"Trebuchet MS\" size=\"6\">Neural Image Super-Resolution<font color=\"#999\" size=\"4\">&nbsp;&nbsp;&nbsp;|&nbsp;&nbsp;&nbsp;</font><font color=\"#999\" size=\"4\">Latent Diffusion upscale</font><font color=\"#999\" size=\"4\">&nbsp;&nbsp;&nbsp;|&nbsp;&nbsp;&nbsp;</font><a href=\"https://github.com/olaviinha/NeuralImageSuperResolution\" target=\"_blank\"><font color=\"#999\" size=\"4\">Github</font></a>\n",
        "\n",
        "This notebook implements Superresolution Upscale from [Latent Diffusion](https://github.com/CompVis/latent-diffusion) in an attempt to improve and enhance image quality.\n",
        "\n",
        "`input` may be a file path or a directory path. All paths should berelative to your Google Drive root. I.e. if your Google Drive has a directory called _images_ and under that directory you have a file _face.jpg_, then `input` value should be `images/face.jpg`"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A6-8KXU6m8TM",
        "cellView": "form"
      },
      "source": [
        "\n",
        "#@title #Setup\n",
        "#@markdown This cell needs to be run only once. It will mount your Google Drive and setup prerequisites.<br>\n",
        "#@markdown <small>Mounting Drive will enable this notebook to save outputs directly to your Drive. Otherwise you will need to copy/download them manually from this notebook.</small>\n",
        "\n",
        "force_setup = False\n",
        "repositories = [\n",
        "  'https://github.com/olaviinha/latent-diffusion.git',\n",
        "  'https://github.com/olaviinha/taming-transformers'\n",
        "]\n",
        "pip_packages = 'ipywidgets omegaconf>=2.0.0 pytorch-lightning>=1.0.8 torch-fidelity einops'\n",
        "apt_packages = ''\n",
        "mount_drive = True #@param {type:\"boolean\"}\n",
        "skip_setup = False #@ param {type:\"boolean\"}\n",
        "\n",
        "# Download the repo from Github\n",
        "import os\n",
        "from google.colab import output, files\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "%cd /content/\n",
        "\n",
        "# inhagcutils\n",
        "if not os.path.isfile('/content/inhagcutils.ipynb') and force_setup == False:\n",
        "  !pip -q install import-ipynb {pip_packages}\n",
        "  if apt_packages != '':\n",
        "    !apt-get update && apt-get install {apt_packages}\n",
        "  !curl -s -O https://raw.githubusercontent.com/olaviinha/inhagcutils/master/inhagcutils.ipynb\n",
        "import import_ipynb\n",
        "from inhagcutils import *\n",
        "\n",
        "# Mount Drive\n",
        "if mount_drive == True:\n",
        "  if not os.path.isdir('/content/drive'):\n",
        "    from google.colab import drive\n",
        "    drive.mount('/content/drive')\n",
        "    drive_root = '/content/drive/My Drive'\n",
        "  if not os.path.isdir('/content/mydrive'):\n",
        "    os.symlink('/content/drive/My Drive', '/content/mydrive')\n",
        "    drive_root = '/content/mydrive/'\n",
        "  drive_root_set = True\n",
        "else:\n",
        "  create_dirs(['/content/faux_drive'])\n",
        "  drive_root = '/content/faux_drive/'\n",
        "\n",
        "if len(repositories) > 0 and skip_setup == False:\n",
        "  for repo in repositories:\n",
        "    %cd /content/\n",
        "    install_dir = fix_path('/content/'+path_leaf(repo).replace('.git', ''))\n",
        "    repo = repo if '.git' in repo else repo+'.git'\n",
        "    !git clone {repo}\n",
        "    if os.path.isfile(install_dir+'setup.py') or os.path.isfile(install_dir+'setup.cfg'):\n",
        "      !pip install -e {install_dir}\n",
        "    if os.path.isfile(install_dir+'requirements.txt'):\n",
        "      !pip install -r {install_dir}/requirements.txt\n",
        "\n",
        "if len(repositories) == 1:\n",
        "  %cd {install_dir}\n",
        "\n",
        "dir_tmp = '/content/tmp/'\n",
        "dir_tmp_in = '/content/tmp_in/'\n",
        "dir_tmp_out = '/content/tmp_out/'\n",
        "create_dirs([dir_tmp, dir_tmp_in, dir_tmp_out])\n",
        "\n",
        "import time, sys\n",
        "from datetime import timedelta\n",
        "import math\n",
        "\n",
        "\n",
        "\n",
        "#-- start --\n",
        "\n",
        "\n",
        "sys.path.append(\".\")\n",
        "sys.path.append('./taming-transformers')\n",
        "from taming.models import vqgan # checking correct import from taming\n",
        "\n",
        "%cd latent-diffusion\n",
        "import ipywidgets as widgets\n",
        "from IPython.display import display\n",
        "\n",
        "mode = widgets.Select(options=['superresolution'], value='superresolution', description='Task:')\n",
        "#display(mode)\n",
        "\n",
        "from notebook_helpers import get_model\n",
        "model = get_model(mode.value)\n",
        "\n",
        "from notebook_helpers import run\n",
        "import torch\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "\n",
        "#-- end --\n",
        "output.clear()\n",
        "# !nvidia-smi\n",
        "op(c.ok, 'Setup finished.')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "H2z1os4Bocdy",
        "cellView": "form"
      },
      "source": [
        "#@title # Do stuff\n",
        "input = \"\" #@param {type:\"string\"}\n",
        "output_dir = \"\" #@param {type:\"string\"}\n",
        "steps = 100 #@param {type:\"integer\"}\n",
        "\n",
        "uniq_id = gen_id()\n",
        "\n",
        "if os.path.isfile(drive_root+input):\n",
        "  inputs = [drive_root+input]\n",
        "  dir_in = path_dir(drive_root+input)\n",
        "elif os.path.isdir(drive_root+input):\n",
        "  dir_in = drive_root+fix_path(input)\n",
        "  # What to do if input is directory path\n",
        "  inputs = list_images(dir_in)\n",
        "elif os.path.isdir(drive_root+input) and '*' in input:\n",
        "  dir_in = path_dir(drive_root+input)\n",
        "  inputs = glob(drive_root+input)\n",
        "else:\n",
        "  op(c.fail, 'FAIL!', 'Input should be a path to an image file or a directory of image files.')\n",
        "  sys.exit('Input not understood.')\n",
        "\n",
        "# Output\n",
        "if output_dir == '':\n",
        "  dir_out = dir_in\n",
        "else:\n",
        "  if not os.path.isdir(drive_root+output_dir):\n",
        "    os.mkdir(drive_root+output_dir)\n",
        "  dir_out = drive_root+fix_path(output_dir)\n",
        "\n",
        "total = len(inputs)\n",
        "timer_start = time.time()\n",
        "\n",
        "# -- DO THINGS --\n",
        "for i, input in enumerate(inputs, 1):\n",
        "  ndx = str(i)+'/'str(total)+' '\n",
        "  op(c.title, ndx+'Processing', input.replace(drive_root, ''), time=True)\n",
        "\n",
        "  img_out = dir_out+uniq_id+'_'+str(steps)+'steps_'+path_leaf(input)\n",
        "  logs = run(model[\"model\"], input, mode.value, steps)\n",
        "\n",
        "  sample = logs[\"sample\"]\n",
        "  sample = sample.detach().cpu()\n",
        "  sample = torch.clamp(sample, -1., 1.)\n",
        "  sample = (sample + 1.) / 2. * 255\n",
        "  sample = sample.numpy().astype(np.uint8)\n",
        "  sample = np.transpose(sample, (0, 2, 3, 1))\n",
        "  a = Image.fromarray(sample[0])\n",
        "  a.save(img_out)\n",
        "\n",
        "  if os.path.isfile(img_out):\n",
        "    # display(a)\n",
        "    op(c.ok, 'Upscaled image saved as', img_out.replace(drive_root, ''))\n",
        "  else:\n",
        "    op(c.fail, 'Error occurred: ', input.replace(drive_root, ''))\n",
        "\n",
        "# -- END THINGS --\n",
        "\n",
        "timer_end = time.time()\n",
        "\n",
        "print('\\nElapsed', timedelta(seconds=timer_end-timer_start))\n",
        "op(c.ok, 'FIN.')\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}