From 59d4e26509d65bb6210992d423f64ab3205d5cd8 Mon Sep 17 00:00:00 2001
From: Sebastian Itokazu <sebitokazu@gmail.com>
Date: Thu, 27 Aug 2026 12:53:50 -0300
Subject: [PATCH] Enable graph optimization for the tract backend
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

`ort-tract` runs tract's `into_optimized()` only when the session requests a
graph optimization level, and `Session::builder()` sets none by default, so
every target except aarch64 Windows ran the unoptimized graph.

Requesting `Level3` — ONNX Runtime's own default, leaving the native backend
unaffected — makes `predict()` over a 2 s window 7.4x faster on an M-series
release build (534.5 ms to 72.5 ms median, n=25). Fixture scores are unchanged.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
---
 .../wakeword_enable_graph_optimization.md     | 12 ++++++++++++
 livekit-wakeword/src/lib.rs                   | 19 ++++++++++++++++---
 2 files changed, 28 insertions(+), 3 deletions(-)
 create mode 100644 .changeset/wakeword_enable_graph_optimization.md

diff --git a/.changeset/wakeword_enable_graph_optimization.md b/.changeset/wakeword_enable_graph_optimization.md
new file mode 100644
index 000000000..feb4716ac
--- /dev/null
+++ b/.changeset/wakeword_enable_graph_optimization.md
@@ -0,0 +1,12 @@
+---
+livekit-wakeword: patch
+---
+
+# Enable graph optimization for the tract backend
+
+Sessions were built with `Session::builder()` without setting a graph optimization
+level. The `ort-tract` backend used on every target except aarch64 Windows runs
+tract's `into_optimized()` only when the session requests an optimization level, so
+wake word inference ran the unoptimized graph. Requesting `Level3` — ONNX Runtime's
+own default, so the native backend is unaffected — made `predict()` over a 2 s window
+7.4x faster (534.5 ms to 72.5 ms median, Apple M-series release build).
diff --git a/livekit-wakeword/src/lib.rs b/livekit-wakeword/src/lib.rs
index 619869de9..dc0ef3cd4 100644
--- a/livekit-wakeword/src/lib.rs
+++ b/livekit-wakeword/src/lib.rs
@@ -16,7 +16,10 @@ use std::path::Path;
 #[cfg(use_tract)]
 use std::sync::Once;
 
-use ort::session::Session;
+use ort::session::{
+    builder::{GraphOptimizationLevel, SessionBuilder},
+    Session,
+};
 
 #[cfg(use_tract)]
 static INIT_TRACT: Once = Once::new();
@@ -74,15 +77,25 @@ pub(crate) fn to_resampler_rate(hz: u32) -> Result<resampler::SampleRate, WakeWo
     }
 }
 
+// Graph optimization is opt-in for the alternative backend: `ort-tract` only runs
+// tract's `into_optimized()` when the session asks for an optimization level, and
+// `Session::builder()` does not set one by default. Without this call the tract
+// backend executes the unoptimized graph, which measured several times slower on
+// the wake word models. `Level3` matches ONNX Runtime's own default, so the native
+// backend is unaffected.
+fn session_builder() -> Result<SessionBuilder, WakeWordError> {
+    Ok(Session::builder()?.with_optimization_level(GraphOptimizationLevel::Level3)?)
+}
+
 pub(crate) fn build_session_from_memory(bytes: &[u8]) -> Result<Session, WakeWordError> {
     #[cfg(use_tract)]
     ensure_tract_backend();
-    Ok(Session::builder()?.commit_from_memory(bytes)?)
+    Ok(session_builder()?.commit_from_memory(bytes)?)
 }
 
 pub(crate) fn build_session_from_file(path: impl AsRef<Path>) -> Result<Session, WakeWordError> {
     #[cfg(use_tract)]
     ensure_tract_backend();
     let bytes = std::fs::read(path)?;
-    Ok(Session::builder()?.commit_from_memory(&bytes)?)
+    Ok(session_builder()?.commit_from_memory(&bytes)?)
 }