wasm demo: load models from device files (no download)

Add a multi-file picker so the models can be selected from the device
instead of fetched: each File is read to an ArrayBuffer, handed to the
pipeline (transferred to the worker), and used by name before any network
fetch. This avoids the 380 MB TableFormer download on a phone entirely and
— since reading a File is a single allocation — sidesteps the double-
buffering peak fetchProgress hits on the 225 MB encoder (the likely cause
of the earlier load failure). Layout and TableFormer both resolve
provided-file → local ./models/ → Hugging Face, in that order.

Refs #157

Signed-off-by: artiz <artem.kustikov@gmail.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EY5KAiquN4YpVf2PXEQkVT

Author: artiz

crates/docling-wasm/www/pipeline.js

@@ -16,16 +16,11 @@ ort.env.wasm.numThreads = self.crossOriginIsolated
   : 1;
 export const THREADS = ort.env.wasm.numThreads;
 
-// Layout model candidates, tried in order. Local ./models/ first (populated by
-// download_dependencies.sh for local dev); then a CORS-enabled Hugging Face
-// mirror so the page works when served cross-origin — e.g. the hosted phone
-// demo on raw.githack — since GitHub Release assets carry no CORS header.
+// Model bases, tried in order after any user-provided file: local ./models/
+// first (download_dependencies.sh for local dev), then a CORS-enabled Hugging
+// Face mirror so the page works cross-origin (the hosted phone demo on
+// raw.githack) — GitHub Release assets carry no CORS header.
 const MODEL_BASE = "https://huggingface.co/pivozavrus/docling-rs-models/resolve/main/";
-const LAYOUT_PATHS = [
-  "./models/layout_heron_int8.onnx",
-  MODEL_BASE + "layout_heron_int8.onnx",
-  "./models/layout_heron.onnx",
-];
 
 const REC_MODELS = {
   en: {
@@ -104,6 +99,29 @@ class JsTfSession {
 export function createOcr({ onStatus }) {
   const status = (msg, spinning = true) => onStatus && onStatus(msg, spinning);
 
+  // Model files the user picked from the device (basename → ArrayBuffer),
+  // used instead of any network fetch — see scan.html's file picker. Reading a
+  // File to an ArrayBuffer is a single allocation, so it also sidesteps the
+  // double-buffering peak fetchProgress hits on the 225 MB encoder.
+  let provided = {};
+  function setProvidedModels(map) {
+    provided = map || {};
+  }
+  // Resolve a model by name: a user-provided file wins, else the first base
+  // that serves it. Returns { buf, base } (base is null for a provided file).
+  async function fetchModel(name, bases, label) {
+    if (provided[name]) return { buf: provided[name], base: null };
+    let lastErr = null;
+    for (const b of bases) {
+      try {
+        return { buf: await fetchProgress(b + name, label), base: b };
+      } catch (e) {
+        lastErr = e;
+      }
+    }
+    throw new Error(`${name} failed: ${(lastErr && lastErr.message) || lastErr}`);
+  }
+
   // fetch() with a live "x / y MB" progress line.
   async function fetchProgress(url, label) {
     const resp = await fetch(url, { cache: "force-cache" });
@@ -159,28 +177,32 @@ export function createOcr({ onStatus }) {
   let layout = null;
   let layoutKind = null;
   async function loadLayout() {
-    for (const path of LAYOUT_PATHS) {
+    for (const [name, kind] of [
+      ["layout_heron_int8.onnx", "int8"],
+      ["layout_heron.onnx", "fp32"],
+    ]) {
+      let buf;
       try {
-        const buf = await fetchProgress(path, "layout model (first load only)");
-        status("starting layout session …", true);
-        const session = await ort.InferenceSession.create(buf, {
-          executionProviders: ["wasm"],
-          logSeverityLevel: 3,
-        });
-        layoutKind = path.includes("int8") ? "int8" : "fp32";
-        layout = {
-          run: async (data) => {
-            const results = await session.run({
-              pixel_values: new ort.Tensor("float32", data, [1, 3, 640, 640]),
-            });
-            const t = (n) => ({ data: results[n].data, dims: Array.from(results[n].dims) });
-            return { logits: t("logits"), boxes: t("pred_boxes") };
-          },
-        };
-        return layoutKind;
+        ({ buf } = await fetchModel(name, ["./models/", MODEL_BASE], "layout model (first load only)"));
       } catch (e) {
-        // 404 / decode failure → try the next candidate.
+        continue; // not provided and not fetchable → try the next candidate
       }
+      status("starting layout session …", true);
+      const session = await ort.InferenceSession.create(buf, {
+        executionProviders: ["wasm"],
+        logSeverityLevel: 3,
+      });
+      layoutKind = kind;
+      layout = {
+        run: async (data) => {
+          const results = await session.run({
+            pixel_values: new ort.Tensor("float32", data, [1, 3, 640, 640]),
+          });
+          const t = (n) => ({ data: results[n].data, dims: Array.from(results[n].dims) });
+          return { logits: t("logits"), boxes: t("pred_boxes") };
+        },
+      };
+      return layoutKind;
     }
     return null;
   }
@@ -210,24 +232,16 @@ export function createOcr({ onStatus }) {
   async function ensureTf() {
     if (tf) return tf;
     const load = async (name, external) => {
-      // Find the first base (local, then HF) that serves the .onnx; take its
-      // external data from the same base.
-      let base = null;
-      let model = null;
-      let lastErr = null;
-      for (const b of TF_DIRS) {
-        try {
-          model = await fetchProgress(b + name + ".onnx", `tableformer ${name}`);
-          base = b;
-          break;
-        } catch (e) {
-          lastErr = e; // keep the real cause (HTTP status, out-of-memory, …)
-        }
-      }
-      if (!model) throw new Error(`tableformer ${name}.onnx failed: ${(lastErr && lastErr.message) || lastErr}`);
+      // A provided file wins; else the first base (local, then HF). External
+      // data comes from a provided file or the same base the .onnx came from.
+      const { buf: model, base } = await fetchModel(name + ".onnx", TF_DIRS, `tableformer ${name}`);
       const opts = { executionProviders: ["wasm"], logSeverityLevel: 3 };
       if (external) {
-        const data = await fetchProgress(base + name + ".onnx.data", `tableformer ${name} data`);
+        const { buf: data } = await fetchModel(
+          name + ".onnx.data",
+          base ? [base] : TF_DIRS,
+          `tableformer ${name} data`,
+        );
         opts.externalData = [{ path: name + ".onnx.data", data: new Uint8Array(data) }];
       }
       return ort.InferenceSession.create(model, opts);
@@ -268,7 +282,7 @@ export function createOcr({ onStatus }) {
   }
 
   return {
-    boot, warmup, recFor, startDoc, addPage, finishDoc, convertImage,
+    boot, warmup, recFor, startDoc, addPage, finishDoc, convertImage, setProvidedModels,
     get layoutKind() { return layoutKind; },
   };
 }

crates/docling-wasm/www/scan.html

@@ -84,12 +84,17 @@ <h1>docling.rs → wasm: scanned PDF / image, full lite profile</h1>
       <option value="cyrillic">Cyrillic + Latin (cyrillic_PP-OCRv5)</option>
       <option value="ch">Chinese + Latin (ch_PP-OCRv3)</option>
     </select>
-    <label style="margin-left:1rem" title="Stage 3: real table structure via the TableFormer ONNX graphs, instead of the geometric reconstruction. Fetches ~380 MB of models (encoder/decoder/bbox) on first use, same-origin from ./models/tableformer/.">
+    <label style="margin-left:1rem" title="Stage 3: real table structure via the TableFormer ONNX graphs, instead of the geometric reconstruction.">
       <input type="checkbox" id="tf" disabled> TableFormer tables
-      <span style="color:#888">(+~380 MB models, PDFs)</span>
     </label>
     <span id="meta"></span>
   </p>
+  <p style="font-size:0.85rem; color:#888">
+    Models from device (optional — skips the download):
+    <input type="file" id="models" multiple accept=".onnx,.data,.txt" disabled>
+    <br>pick <code>layout_heron_int8.onnx</code> and, for TableFormer, <code>encoder.onnx</code>,
+    <code>decoder_kv.onnx</code>(+<code>.data</code>), <code>bbox.onnx</code>(+<code>.data</code>).
+  </p>
 
   <div id="diag" style="font: 0.8rem ui-monospace, monospace; color: #888; white-space: pre-wrap; margin: 0.3rem 0;"></div>
   <div id="out"></div>
@@ -125,6 +130,7 @@ <h1>docling.rs → wasm: scanned PDF / image, full lite profile</h1>
     const diagEl = document.getElementById("diag");
     const lang = document.getElementById("lang");
     const tf = document.getElementById("tf");
+    const modelsInput = document.getElementById("models");
     const drop = document.getElementById("drop");
     const pick = document.getElementById("pick");
 
@@ -143,12 +149,12 @@ <h1>docling.rs → wasm: scanned PDF / image, full lite profile</h1>
     }
     function busy(msg) {
       document.body.classList.add("busy");
-      pick.disabled = lang.disabled = tf.disabled = true;
+      pick.disabled = lang.disabled = tf.disabled = modelsInput.disabled = true;
       status(msg, true);
     }
     function ready(msg) {
       document.body.classList.remove("busy");
-      pick.disabled = lang.disabled = tf.disabled = false;
+      pick.disabled = lang.disabled = tf.disabled = modelsInput.disabled = false;
       status(msg, false);
     }
     const onStatus = (msg, spinning) => {
@@ -192,6 +198,7 @@ <h1>docling.rs → wasm: scanned PDF / image, full lite profile</h1>
         addPage: (rgba, w, h, scale) => rpc("doc-page", { rgba, w, h, scale }, [rgba]),
         docFinish: (name) => rpc("doc-finish", { name }).then((r) => r.md),
         convertImage: (bytes, name, l) => rpc("convert-image", { bytes, name, lang: l }, [bytes]).then((r) => r.md),
+        setModels: (models) => rpc("set-models", { models }, Object.values(models)),
       };
     }
     async function mainDriver() {
@@ -210,6 +217,7 @@ <h1>docling.rs → wasm: scanned PDF / image, full lite profile</h1>
         addPage: (rgba, w, h, scale) => ocr.addPage(new Uint8Array(rgba), w, h, scale),
         docFinish: (name) => Promise.resolve(ocr.finishDoc(name)),
         convertImage: (bytes, name, l) => ocr.convertImage(bytes, name, l),
+        setModels: (models) => ocr.setProvidedModels(models),
       };
     }
 
@@ -280,6 +288,23 @@ <h1>docling.rs → wasm: scanned PDF / image, full lite profile</h1>
       }
     });
 
+    // Pick model files from the device: read each to an ArrayBuffer and hand
+    // them to the pipeline (transferred to the worker), so layout/TableFormer
+    // load from local files instead of fetching — no 380 MB download, and a
+    // single allocation per file (no fetch double-buffering).
+    modelsInput.addEventListener("change", async () => {
+      if (!drv || !modelsInput.files.length) return;
+      busy(`reading ${modelsInput.files.length} model file(s) …`);
+      try {
+        const map = {};
+        for (const f of modelsInput.files) map[f.name] = await f.arrayBuffer();
+        await drv.setModels(map);
+        ready(`models ready: ${Object.keys(map).join(", ")} — drop a scanned PDF or image`);
+      } catch (e) {
+        ready(`model file read failed: ${(e && e.message) || e}`);
+      }
+    });
+
     drop.addEventListener("click", () => pick.click());
     pick.addEventListener("change", () => pick.files[0] && handle(pick.files[0]));
     for (const ev of ["dragover", "dragleave", "drop"]) {

crates/docling-wasm/www/worker.js

@@ -24,6 +24,9 @@ const ocr = createOcr({
 
 async function handle(m) {
   switch (m.type) {
+    case "set-models":
+      ocr.setProvidedModels(m.models);
+      return {};
     case "boot": {
       const kind = await ocr.boot();
       if (!kind) return { noLayout: true };