Update hific colab to work with TFC 2.0.

PiperOrigin-RevId: 363445673
Change-Id: Ifba60b483cd8f9f3660e42c289881b8db693ff79

Author: Fabian Mentzer

models/hific/colab.ipynb

@@ -1,19 +1,4 @@
 {
-  "nbformat": 4,
-  "nbformat_minor": 0,
-  "metadata": {
-    "colab": {
-      "name": "HiFiC_Interactive_Colab.ipynb",
-      "provenance": [],
-      "collapsed_sections": [],
-      "toc_visible": true
-    },
-    "kernelspec": {
-      "display_name": "Python 3",
-      "name": "python3"
-    },
-    "accelerator": "GPU"
-  },
   "cells": [
     {
       "cell_type": "markdown",
@@ -40,18 +25,17 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {
         "id": "LO_MNEQ7Bhbw"
       },
+      "outputs": [],
       "source": [
-        "%tensorflow_version 1.x\n",
-        "!pip install tensorflow-compression\\<2.0\n",
+        "!pip install tensorflow-compression\n",
         "![[ -e /tfc ]] || git clone https://github.com/tensorflow/compression /tfc\n",
         "%cd /tfc/models\n",
-        "import tfci  # Check if tfci.py is available.\n"
-      ],
-      "execution_count": null,
-      "outputs": []
+        "import tfci  # Check if tfci.py is available."
+      ]
     },
     {
       "cell_type": "markdown",
@@ -68,19 +52,19 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {
         "id": "x-yLUG_tmo3M"
       },
+      "outputs": [],
       "source": [
         "import tensorflow as tf\n",
         "\n",
-        "if not tf.test.is_gpu_available():\n",
+        "if not tf.config.list_physical_devices('GPU'):\n",
         "  print('WARNING: No GPU found. Might be slow!')\n",
         "else:\n",
         "  print('Found GPU.')"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "markdown",
@@ -93,9 +77,11 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {
         "id": "vtd1l70Pf95V"
       },
+      "outputs": [],
       "source": [
         "import os\n",
         "import zipfile\n",
@@ -121,20 +107,18 @@
         "File = collections.namedtuple('File', ['full_path', 'num_bytes', 'bpp'])\n",
         "\n",
         "def print_html(html):\n",
-        "  display(HTML(html + '<br/>'))\n",
+        "  display(HTML(html + '\u003cbr/\u003e'))\n",
         "\n",
         "def make_cell_large():\n",
         "  display(Javascript(\n",
         "      '''google.colab.output.setIframeHeight(0, true, {maxHeight: 5000})'''))\n",
         "\n",
         "def get_default_image(output_dir):\n",
         "  output_path = os.path.join(output_dir, os.path.basename(DEFAULT_IMAGE_URL))\n",
-        "  print('Downloading', DEFAULT_IMAGE_URL, '\\n->', output_path)\n",
+        "  print('Downloading', DEFAULT_IMAGE_URL, '\\n-\u003e', output_path)\n",
         "  urllib.request.urlretrieve(DEFAULT_IMAGE_URL, output_path)\n",
         "\n"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "markdown",
@@ -147,9 +131,11 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {
         "id": "NgtIlL2ADCI2"
       },
+      "outputs": [],
       "source": [
         "#@title Setup { vertical-output: false, run: \"auto\", display-mode: \"form\" }\n",
         "#@markdown #### Custom Images\n",
@@ -173,30 +159,30 @@
         "\n",
         "model = 'hific-lo' #@param [\"hific-lo\", \"hific-mi\", \"hific-hi\"]\n",
         "\n"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {
         "id": "GYcbc2HupTRD"
       },
+      "outputs": [],
       "source": [
         "if 'upload_custom_images' not in locals():\n",
         "  print('ERROR: Please run the previous cell!')\n",
         "  # Setting defaults anyway.\n",
         "  upload_custom_images = False\n",
         "  model = 'hific-lo'"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {
         "id": "e0C4vMqZsnqA"
       },
+      "outputs": [],
       "source": [
         "all_files = os.listdir(FILES_DIR)\n",
         "if not upload_custom_images or not all_files:\n",
@@ -213,22 +199,7 @@
         "  img = img.resize((w // 15, h // 15))\n",
         "  print('- ' + file_name + ':')\n",
         "  display(img)"
-      ],
-      "execution_count": null,
-      "outputs": []
-    },
-    {
-      "cell_type": "code",
-      "metadata": {
-        "id": "46eq2rCRpXzC"
-      },
-      "source": [
-        "print(f'Caching model \"{model}\"...')\n",
-        "tfci.import_metagraph(model)\n",
-        "print('Done')"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "markdown",
@@ -241,9 +212,11 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {
         "id": "kd02HOhLBj6e"
       },
+      "outputs": [],
       "source": [
         "SUPPORTED_EXT = {'.png', '.jpg'}\n",
         "\n",
@@ -296,9 +269,7 @@
         "           get_bpp(Image.open(full_path).size, num_bytes)))\n",
         "\n",
         "print('All done!')"
-      ],
-      "execution_count": null,
-      "outputs": []
+      ]
     },
     {
       "cell_type": "markdown",
@@ -311,20 +282,20 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {
         "id": "3nVCPeDnskD8"
       },
+      "outputs": [],
       "source": [
         "make_cell_large()  # Larger output window.\n",
         "\n",
         "for file in all_outputs:\n",
-        "  print_html('<hr/>')\n",
+        "  print_html('\u003chr/\u003e')\n",
         "  print(f'Showing {file.full_path} | {file.num_bytes//1000}kB | {file.bpp:.4f}bpp')\n",
         "  display(Image.open(file.full_path))\n",
-        "  print_html('<hr/>')"
-      ],
-      "execution_count": null,
-      "outputs": []
+        "  print_html('\u003chr/\u003e')"
+      ]
     },
     {
       "cell_type": "markdown",
@@ -345,9 +316,11 @@
     },
     {
       "cell_type": "code",
+      "execution_count": null,
       "metadata": {
         "id": "9BKccvcTpj1k"
       },
+      "outputs": [],
       "source": [
         "ZIP = '/content/images.zip'\n",
         "\n",
@@ -357,9 +330,31 @@
         "    zf.write(f.full_path, os.path.basename(path_with_bpp))\n",
         "\n",
         "files.download(ZIP) "
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "collapsed_sections": [],
+      "last_runtime": {
+        "build_target": "//research/colab/notebook:notebook_backend_py3",
+        "kind": "private"
+      },
+      "name": "HiFiC_Interactive_Colab.ipynb",
+      "provenance": [
+        {
+
+          "timestamp": 1615973093906
+        }
       ],
-      "execution_count": null,
-      "outputs": []
+      "toc_visible": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
     }
-  ]
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
 }