Support TensorFlow 2.20 inference compatibility

README.md

@@ -98,6 +98,9 @@ curl -LsSf https://astral.sh/uv/install.sh | sh
 git clone https://github.com/google-research/inksight.git
 cd inksight
 uv sync
+
+# Start the configured notebook
+uv run jupyter notebook colab.ipynb
 ```
 
 ### Using Conda
@@ -109,7 +112,15 @@ conda env create -f environment.yml
 conda activate inksight
 ```
 
-> **Important**: Use TensorFlow 2.15.0-2.17.0. Later versions may cause unexpected behavior.
+The uv environment uses TensorFlow 2.20 and TensorFlow Text 2.20. The
+notebook configures TensorFlow before its first import to preserve the pre-2.18
+XLA GEMM lowering and deterministic autotuning used by the released checkpoint.
+For the same configuration in a Python script, import TensorFlow through the
+InkSight utility:
+
+```python
+from utils.tensorflow import tf
+```
 
 ## Local Playground Setup
 

colab.ipynb

@@ -56,33 +56,25 @@
       },
       "outputs": [],
       "source": [
-        "# @title Dependencies\n",
-        "from IPython.display import Markdown\n",
-        "import time\n",
+        "# @title Install dependencies\n",
         "import os\n",
+        "import time\n",
         "\n",
-        "display(Markdown(\"## 📦 Installing required packages...\\nThis may take a minute. Please wait...\"))\n",
-        "\n",
-        "!sudo apt -qq install tesseract-ocr\n",
-        "!uv pip install -q --system \"tensorflow[and-cuda]==2.17.0\" \"tensorflow-text==2.17.0\" pytesseract \"tf-keras==2.17.0\" \"python-doctr[tf,viz]==0.10.0\"\n",
-        "display(Markdown(\"✅ **Installation complete!**\"))\n",
+        "!apt-get -qq update\n",
+        "!apt-get -qq install --yes tesseract-ocr\n",
+        "%pip install --quiet \\\n",
+        "    \"tensorflow[and-cuda]==2.20.0\" \\\n",
+        "    \"tensorflow-text==2.20.1\" \\\n",
+        "    \"tf-keras==2.20.1\" \\\n",
+        "    \"python-doctr[tf,viz]==0.10.0\" \\\n",
+        "    \"pytesseract>=0.3.13\"\n",
+        "\n",
+        "print(\n",
+        "    \"Dependencies installed. Restarting the runtime... \"\n",
+        "    \"Once it reconnects, continue with the next cell.\",\n",
+        "    flush=True,\n",
+        ")\n",
         "time.sleep(1)\n",
-        "\n",
-        "display(Markdown(\"\"\"\n",
-        "---\n",
-        "### 🔄 Restarting Runtime\n",
-        "To finalize the installation, we need to restart the Colab runtime.\n",
-        "\n",
-        "> **Why?** TensorFlow and system-level packages need a restart to properly initialize with new dependencies.\n",
-        "\n",
-        "⏳ Restarting in 3 seconds...\n",
-        "\n",
-        "✅ Restartng done! **Please continue to the following cells**\n",
-        "\"\"\"))\n",
-        "\n",
-        "time.sleep(3)\n",
-        "\n",
-        "# Kill the current process to force a runtime restart\n",
         "os.kill(os.getpid(), 9)"
       ]
     },
@@ -95,10 +87,20 @@
       },
       "outputs": [],
       "source": [
-        "# @title docTR Preparation\n",
+        "# @title Configure TensorFlow and docTR\n",
+        "\n",
+        "import os\n",
+        "\n",
+        "_inksight_xla_flags = (\n",
+        "    \"--xla_gpu_autotune_level=0 \"\n",
+        "    \"--xla_disable_hlo_passes=custom-kernel-fusion-rewriter,custom_kernel-fusion-autotuner\"\n",
+        ")\n",
+        "if _inksight_xla_flags not in os.environ.get(\"XLA_FLAGS\", \"\"):\n",
+        "    os.environ[\"XLA_FLAGS\"] = f'{os.environ.get(\"XLA_FLAGS\", \"\")} {_inksight_xla_flags}'.strip()\n",
         "\n",
         "from doctr.io import DocumentFile\n",
         "from doctr.models import ocr_predictor\n",
+        "\n",
         "predictor = ocr_predictor(pretrained=True)\n",
         "print(\"doctr predictor loaded.\")"
       ]

pyproject.toml

@@ -3,7 +3,7 @@ name = "inksight"
 version = "0.1.0"
 description = "Offline-to-Online Handwriting Conversion by Teaching Vision-Language Models to Read and Write"
 readme = "README.md"
-requires-python = ">=3.11"
+requires-python = ">=3.11,<3.14"
 keywords = ["handwriting", "computer vision", "machine learning", "OCR", "digital ink", "transformer", "vision-language"]
 authors = [
     {name = "Google"},
@@ -26,8 +26,8 @@ dependencies = [
     "jupyter>=1.1.1",
     "pytesseract>=0.3.13",
     "python-doctr[tf,viz]==0.10.0",
-    "tensorflow-text==2.17.0",
-    "tensorflow[and-cuda]==2.17.0",
+    "tensorflow-text==2.20.1",
+    "tensorflow[and-cuda]==2.20.0",
 ]
 
 [project.urls]

utils/tensorflow.py

@@ -0,0 +1,45 @@
+"""TensorFlow import configured for stable InkSight inference."""
+
+from __future__ import annotations
+
+import importlib
+import os
+import shlex
+
+_INCOMPATIBLE_XLA_PASSES = (
+    "custom-kernel-fusion-rewriter",
+    "custom_kernel-fusion-autotuner",
+)
+
+
+def _configure_xla_flags() -> None:
+    tokens = shlex.split(os.environ.get("XLA_FLAGS", ""))
+    disabled_passes: list[str] = []
+    preserved_flags: list[str] = []
+
+    for token in tokens:
+        if token.startswith("--xla_disable_hlo_passes="):
+            disabled_passes.extend(token.partition("=")[2].split(","))
+        elif not token.startswith("--xla_gpu_autotune_level="):
+            preserved_flags.append(token)
+
+    for pass_name in _INCOMPATIBLE_XLA_PASSES:
+        if pass_name not in disabled_passes:
+            disabled_passes.append(pass_name)
+
+    preserved_flags.extend(
+        (
+            "--xla_gpu_autotune_level=0",
+            "--xla_disable_hlo_passes=" + ",".join(filter(None, disabled_passes)),
+        )
+    )
+    os.environ["XLA_FLAGS"] = shlex.join(preserved_flags)
+
+
+_configure_xla_flags()
+
+# These modules must load after the XLA configuration.
+tf = importlib.import_module("tensorflow")
+tf_text = importlib.import_module("tensorflow_text")
+
+__all__ = ["tf", "tf_text"]