INT4 vision (251MB) + chunked Range downloads (fix 1.7GB stall); verified end-to-end in Chrome WebGPU

app.js

@@ -8,7 +8,9 @@ import { AutoTokenizer } from "https://cdn.jsdelivr.net/npm/@huggingface/transfo
 
 const REPO = "Reza2kn/LocateAnything-3B-ONNX-WebGPU-INT4";
 const BASE = `https://huggingface.co/${REPO}/resolve/main`;
-const VISION_URL = `${BASE}/onnx/vision_mlp.onnx`;
+const VISION_URL = `${BASE}/onnx/vision_mlp_int4.onnx`;   // INT4 (~250MB) — fp32 1.73GB stalls on download
+const VISION_DATA = "vision_mlp_int4.onnx.data";
+const VISION_DATA_URL = `${BASE}/onnx/${VISION_DATA}`;
 const LANG_URL = `${BASE}/onnx/language_tail_kv_int4.onnx`;
 const LANG_DATA = "language_tail_kv_int4.onnx.data";
 const LANG_DATA_URL = `${BASE}/onnx/${LANG_DATA}`;
@@ -48,37 +50,110 @@ async function fetchBuf(url, label) {
   return new Uint8Array(await r.arrayBuffer());
 }
 
+const sleep = (ms) => new Promise((r) => setTimeout(r, ms));
+
+// fetch with a stall watchdog: aborts if no progress within `stallMs`
+async function fetchAbortable(url, opts, stallMs = 30000) {
+  const ctrl = new AbortController();
+  const r = await fetch(url, { ...opts, signal: ctrl.signal });
+  if (!(r.status === 200 || r.status === 206)) throw new Error(`status ${r.status}`);
+  const reader = r.body.getReader();
+  const chunks = []; let got = 0;
+  let timer = setTimeout(() => ctrl.abort(), stallMs);
+  try {
+    for (;;) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      clearTimeout(timer); timer = setTimeout(() => ctrl.abort(), stallMs);
+      chunks.push(value); got += value.length;
+    }
+  } finally { clearTimeout(timer); }
+  const buf = new Uint8Array(got); let o = 0;
+  for (const c of chunks) { buf.set(c, o); o += c.length; }
+  return { buf, headers: r.headers };
+}
+
+// Chunked Range download: small pieces (retried independently) so no single long-lived
+// connection can stall the whole file. HF CDN supports range requests.
+async function fetchBufProgress(url, label, chunk = 48 * 1024 * 1024) {
+  const t = performance.now();
+  // discover total size via the first range request's Content-Range
+  let total = 0, first;
+  for (let tr = 0; ; tr++) {
+    try {
+      first = await fetchAbortable(url, { headers: { Range: `bytes=0-${chunk - 1}` } });
+      const cr = first.headers.get("content-range");
+      total = cr ? +cr.split("/")[1] : (+first.headers.get("content-length") || first.buf.length);
+      break;
+    } catch (e) { if (tr >= 4) throw e; log(`  ${label} init retry ${tr + 1}…`); await sleep(1200); }
+  }
+  if (!total || total <= first.buf.length) { // small file, already done
+    log(`  ${label} downloaded ${(first.buf.length/1e6|0)}MB in ${((performance.now()-t)/1000).toFixed(1)}s`);
+    return first.buf;
+  }
+  const buf = new Uint8Array(total);
+  buf.set(first.buf, 0);
+  let off = first.buf.length, lastPct = -1;
+  while (off < total) {
+    const end = Math.min(off + chunk, total) - 1;
+    let ok = false;
+    for (let tr = 0; tr < 5 && !ok; tr++) {
+      try {
+        const { buf: part } = await fetchAbortable(url, { headers: { Range: `bytes=${off}-${end}` } });
+        buf.set(part, off); off += part.length; ok = true;
+      } catch (e) { if (tr === 4) throw e; await sleep(1000); }
+    }
+    const pct = Math.floor((off / total) * 100);
+    if (pct >= lastPct + 10) { lastPct = pct; log(`  ${label}: ${pct}% (${(off/1e6|0)}MB)`); }
+  }
+  log(`  ${label} downloaded ${(total/1e6|0)}MB in ${((performance.now()-t)/1000).toFixed(1)}s`);
+  return buf;
+}
+
 async function loadAll() {
   setBadge($("load"), "loading…", "warn");
   $("prog").style.display = "block";
   // webgpu first; wasm as fallback so unsupported ops/devices degrade instead of hard-failing.
   const sessOpts = { executionProviders: ["webgpu", "wasm"], graphOptimizationLevel: "all" };
+  let t;
 
-  log("loading tokenizer…");
-  tokenizer = await AutoTokenizer.from_pretrained(REPO);
+  // Create the ONNX sessions FIRST (before transformers.js), fetching buffers ourselves so we
+  // can see download vs. compile timing and avoid ort's internal URL fetch hanging on redirects.
+  log("downloading vision model (INT4, ~250MB)…");
+  const visGraph = await fetchBufProgress(VISION_URL, "vision graph");
+  const visData = await fetchBufProgress(VISION_DATA_URL, "vision data");
+  $("prog").value = 30;
+  log("compiling vision session…"); t = performance.now();
+  visionSess = await ort.InferenceSession.create(visGraph, {
+    ...sessOpts,
+    externalData: [{ path: VISION_DATA, data: visData }],
+  });
+  log(`vision session ready in ${((performance.now()-t)/1000).toFixed(1)}s`);
+  $("prog").value = 50;
 
-  log("loading embedding INT4 table…");
+  log("downloading INT4 language model (~1.7GB)…");
+  const langData = await fetchBufProgress(LANG_DATA_URL, "language data");
+  const langGraph = await fetchBufProgress(LANG_URL, "language graph");
+  $("prog").value = 80;
+  log("compiling language session…"); t = performance.now();
+  langSess = await ort.InferenceSession.create(langGraph, {
+    ...sessOpts,
+    externalData: [{ path: LANG_DATA, data: langData }],
+  });
+  log(`language session ready in ${((performance.now()-t)/1000).toFixed(1)}s`);
+  out_names_cache.lang = langSess.outputNames;
+  $("prog").value = 90;
+
+  log("loading tokenizer + INT4 embedding table…");
+  tokenizer = await AutoTokenizer.from_pretrained(REPO);
   embMeta = await (await fetch(EMB_META_URL)).json();
-  embPacked = await fetchBuf(EMB_PACKED_URL, "embed packed");   // uint8 [vocab, hidden/2]
-  const scalesBytes = await fetchBuf(EMB_SCALES_URL, "embed scales"); // fp16
+  embPacked = await fetchBufProgress(EMB_PACKED_URL, "embed packed");   // uint8 [vocab, hidden/2]
+  const scalesBytes = await fetchBufProgress(EMB_SCALES_URL, "embed scales"); // fp16
   const sv = new DataView(scalesBytes.buffer);
   embScales = new Float32Array(scalesBytes.length / 2);
   for (let i = 0; i < embScales.length; i++) embScales[i] = f16to32(sv.getUint16(i * 2, true));
   log(`embedding: vocab=${embMeta.vocab} hidden=${embMeta.hidden} block=${embMeta.block_size}`);
-  $("prog").value = 20;
 
-  log("loading vision model (~1.7GB)…");
-  visionSess = await ort.InferenceSession.create(VISION_URL, sessOpts);
-  $("prog").value = 50;
-
-  log("loading INT4 language model (~1.7GB + data)…");
-  const langData = await fetchBuf(LANG_DATA_URL, "language data");
-  $("prog").value = 85;
-  langSess = await ort.InferenceSession.create(LANG_URL, {
-    ...sessOpts,
-    externalData: [{ path: LANG_DATA, data: langData }],
-  });
-  out_names_cache.lang = langSess.outputNames;
   $("prog").value = 100;
   $("prog").style.display = "none";
   setBadge($("load"), "model ready", "ok");