Dallemini_seq_upscaler.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Dallemini_seq_upscaler.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "UHSFfKPt9TV7"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# setup"
      ],
      "metadata": {
        "id": "UHSFfKPt9TV7"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hyzu1Pn37ygR"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from threading import Thread\n",
        "\n",
        "\n",
        "def bkgw():\n",
        "  !apt-get install -y libvulkan-dev libomp5\n",
        "  !pip install ncnn-vulkan\n",
        "  !mv /usr/local/lib/python3.7/dist-packages/ncnn_vulkan/*.so /usr/local/lib/python3.7/dist-packages/ncnn/\n",
        "  os.link('/content/vqdec18.param','/content/setupfin.txt')\n",
        "\n",
        "\n",
        "t3 = Thread(target = bkgw)\n",
        "a3 = t3.start()\n",
        "!wget https://github.com/TabuaTambalam/vqqncnn/releases/download/0.2/rcup.zip\n",
        "!7z e rcup.zip\n",
        "!wget https://github.com/TabuaTambalam/vqqncnn/releases/download/0.2/rcub.7z\n",
        "!wget https://github.com/TabuaTambalam/vqqncnn/releases/download/0.0/vq_anime.bin\n",
        "!mkdir rcubw\n",
        "!7z e -o/content/rcubw rcub.7z\n",
        "!rm rcub.7z\n",
        "!rm rcup.zip"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# o"
      ],
      "metadata": {
        "id": "zPlDVKc59WAa"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "UseVulkan=False #@param {type:'boolean'}\n",
        "\n",
        "import gc\n",
        "import os\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "\n",
        "\n",
        "class ExtractorGPU:\n",
        "  def __init__(self,c):\n",
        "    self.cnet=c\n",
        "  def __enter__(self):\n",
        "    self.ex = self.cnet.create_extractor()\n",
        "    self.ex.set_blob_vkallocator(blob_vkallocator)\n",
        "    self.ex.set_workspace_vkallocator(blob_vkallocator)\n",
        "    self.ex.set_staging_vkallocator(staging_vkallocator)\n",
        "    return self.ex\n",
        "\n",
        "  def __exit__(self, exc_type, exc_val, exc_tb):\n",
        "    blob_vkallocator.clear()\n",
        "    staging_vkallocator.clear()\n",
        "    self.ex.clear()\n",
        "\n",
        "class ExtractorCPU:\n",
        "  def __init__(self,c):\n",
        "    self.cnet=c\n",
        "  def __enter__(self):\n",
        "    self.ex = self.cnet.create_extractor()\n",
        "    return self.ex\n",
        "\n",
        "  def __exit__(self, exc_type, exc_val, exc_tb):\n",
        "    self.ex.clear()\n",
        "\n",
        "def loaddumps(fna):\n",
        "  return np.fromfile(fna,dtype=np.uint16).astype(np.uint32).reshape((-1,16,16))\n",
        "\n",
        "def showp(n):\n",
        "  uwa=np.pad(dumped_seqs[n],((1,1),(1,1)), 'reflect' )\n",
        "  with NCNNex(net) as ex:\n",
        "    ex.input('in0', ncnn.Mat(uwa.reshape(324)).clone())\n",
        "    hrr, out0 = ex.extract('out0')\n",
        "  with NCNNex(netrcu) as ex:\n",
        "    ex.input('in0', out0)\n",
        "    hrr, out1 = ex.extract('out0')\n",
        "\n",
        "  ymg=np.array(out1).transpose(1,2,0)\n",
        "  if upscaleNx == 4:\n",
        "    krop=np.array(out0).transpose(1,2,0)[16:-16,16:-16,:]\n",
        "    h,w=ymg.shape[:-1]\n",
        "    for y in range(h):\n",
        "      for x in range(w):\n",
        "        ymg[y,x,:]+=krop[y>>2,x>>2,:]\n",
        "    \n",
        "\n",
        "  uz=Image.fromarray( ( ymg*255 ).clip(0,255).astype(np.uint8) )\n",
        "  uz.save('/content/sample_data/%d.png'%n)\n",
        "  return uz\n",
        "\n",
        "\n",
        "ntai='-n%d.bin'\n",
        "def rculoader():\n",
        "  global typo\n",
        "  prm=typo+'.param'\n",
        "  if not os.path.isfile(prm):\n",
        "    print('try again, no '+prm)\n",
        "    return\n",
        "  if typo.startswith('upn'):\n",
        "    typo='upn2x'\n",
        "    nstr='-n0.bin'\n",
        "  elif typo == 'up2x':\n",
        "    if Denoise < 5:\n",
        "      nstr=ntai%Denoise\n",
        "    else:\n",
        "      n=Denoise//21\n",
        "      nstr=ntai%n\n",
        "  else:\n",
        "    if Denoise < 3:\n",
        "      nstr=ntai%Denoise\n",
        "    else:\n",
        "      n=Denoise//34\n",
        "      nstr=ntai%n\n",
        "\n",
        "  netrcu.load_param(  prm   ) \n",
        "  netrcu.load_model('rcubw/'+typo+nstr)\n",
        "\n",
        "\n",
        "try:\n",
        "  del net\n",
        "  gc.collect()\n",
        "except:\n",
        "  pass\n",
        "\n",
        "if os.path.isfile('/content/setupfin.txt'):\n",
        "  import ncnn\n",
        "  net = ncnn.Net()\n",
        "  net.opt.use_vulkan_compute = UseVulkan\n",
        "\n",
        "  if UseVulkan:\n",
        "    vkdev = ncnn.get_gpu_device(ncnn.get_default_gpu_index())\n",
        "    blob_vkallocator = ncnn.VkBlobAllocator(vkdev)\n",
        "    staging_vkallocator = ncnn.VkStagingAllocator(vkdev)\n",
        "    NCNNex = ExtractorGPU\n",
        "\n",
        "  else:\n",
        "    NCNNex=ExtractorCPU\n",
        "\n",
        "\n",
        "  net.opt.use_fp16_packed = False\n",
        "  net.opt.use_fp16_storage = False\n",
        "  net.opt.use_fp16_arithmetic = False\n",
        "\n",
        "  net.load_param( 'vqdec18.param'  )\n",
        "  net.load_model('vq_anime.bin')\n",
        "else:\n",
        "  print('wait ncnn setup')"
      ],
      "metadata": {
        "cellView": "form",
        "id": "FYNJSNc28Xar"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "typo='pro' #@param ['up', 'upn', 'pro']\n",
        "upscaleNx=3 #@param {type:'integer'}\n",
        "Denoise=2 #@param {type:'integer'}\n",
        "\n",
        "if upscaleNx <2:\n",
        "  if typo=='upn':\n",
        "    upscaleNx=1\n",
        "  else:\n",
        "    upscaleNx=2\n",
        "elif upscaleNx >4:\n",
        "  upscaleNx=4\n",
        "typo+='%dx'%upscaleNx\n",
        "\n",
        "\n",
        "try:\n",
        "  del netrcu\n",
        "  gc.collect()\n",
        "except:\n",
        "  pass\n",
        "\n",
        "netrcu = ncnn.Net()\n",
        "netrcu.opt.use_vulkan_compute = UseVulkan\n",
        "rculoader()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "KJXanYP48aZK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "dumped_seqs=loaddumps('/content/Dallesav14.bin')"
      ],
      "metadata": {
        "id": "3zVRtzD58cBy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "showp(10)"
      ],
      "metadata": {
        "id": "DSX3MdkW8dRi"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}