Fix static tokenizer loading from Hugging Face model repo

README.md

@@ -14,4 +14,4 @@ Static browser-only FLUX.2 Klein 4B inference through WebGPU.
 
 The Space serves a small HTML/JS app. Model files are fetched from
 [ryanhlewis/flux2-klein-4b-webgpu-lowbit](https://huggingface.co/ryanhlewis/flux2-klein-4b-webgpu-lowbit)
-and generation runs in the user's browser.
+and generation runs in the user's browser.

space-config.js

@@ -1,3 +1,3 @@
 globalThis.FLUX2_MODEL_BASE_URL = "https://huggingface.co/ryanhlewis/flux2-klein-4b-webgpu-lowbit/resolve/main";
 globalThis.FLUX2_RUNTIME_BASE_URL = globalThis.FLUX2_MODEL_BASE_URL;
-globalThis.FLUX2_CUSTOM_KERNEL_BASE_URL = `${globalThis.FLUX2_MODEL_BASE_URL}/custom_lowbit`;
+globalThis.FLUX2_CUSTOM_KERNEL_BASE_URL = `${globalThis.FLUX2_MODEL_BASE_URL}/custom_lowbit`;

static-app.js

@@ -129,12 +129,52 @@ function configuredCustomKernelBaseUrl(modelBaseUrl = configuredModelBaseUrl(),
 function resolveAssetUrl(baseUrl, path) {
   return new URL(path, new URL(withTrailingSlash(baseUrl), window.location.href)).toString();
 }
-
-function setStatus(lines, isError = false) {
-  const text = Array.isArray(lines) ? lines.filter(Boolean).join("\n") : String(lines);
-  statusEl.textContent = text;
-  statusEl.classList.toggle("error", Boolean(isError));
-}
+
+function huggingFaceResolveSpec(baseUrl) {
+  let url;
+  try {
+    url = new URL(baseUrl);
+  } catch {
+    return null;
+  }
+  if (!/^https?:$/i.test(url.protocol)) return null;
+  if (!/^(huggingface\.co|hf\.co)$/i.test(url.hostname)) return null;
+  const parts = url.pathname.split("/").filter(Boolean);
+  const resolveIndex = parts.indexOf("resolve");
+  if (resolveIndex < 2 || resolveIndex + 1 >= parts.length) return null;
+  return {
+    repoId: `${parts[0]}/${parts[1]}`,
+    revision: decodeURIComponent(parts[resolveIndex + 1]),
+  };
+}
+
+function tokenizerLoadSpec(modelBaseUrl) {
+  const hfSpec = huggingFaceResolveSpec(modelBaseUrl);
+  if (hfSpec) {
+    return {
+      source: hfSpec.repoId,
+      remote: true,
+      options: {
+        local_files_only: false,
+        revision: hfSpec.revision || "main",
+      },
+    };
+  }
+  const remote = /^https?:\/\//i.test(modelBaseUrl);
+  return {
+    source: remote ? resolveAssetUrl(modelBaseUrl, "tokenizer") : "models/tokenizer",
+    remote,
+    options: {
+      local_files_only: !remote,
+    },
+  };
+}
+
+function setStatus(lines, isError = false) {
+  const text = Array.isArray(lines) ? lines.filter(Boolean).join("\n") : String(lines);
+  statusEl.textContent = text;
+  statusEl.classList.toggle("error", Boolean(isError));
+}
 
 function setButtonDisabled(button, disabled) {
   if (button) button.disabled = disabled;
@@ -438,16 +478,11 @@ function maskFromLength(length, seqLen) {
 async function loadTokenizer() {
   if (tokenizerPromise) return tokenizerPromise;
   const modelBaseUrl = configuredModelBaseUrl();
-  const remoteTokenizer = /^https?:\/\//i.test(modelBaseUrl);
-  env.allowLocalModels = true;
-  env.allowRemoteModels = remoteTokenizer;
+  const tokenizerSpec = tokenizerLoadSpec(modelBaseUrl);
+  env.allowLocalModels = !tokenizerSpec.remote;
+  env.allowRemoteModels = tokenizerSpec.remote;
   env.localModelPath = "/";
-  const tokenizerPath = remoteTokenizer
-    ? resolveAssetUrl(modelBaseUrl, "tokenizer")
-    : "models/tokenizer";
-  tokenizerPromise = AutoTokenizer.from_pretrained(tokenizerPath, {
-    local_files_only: !remoteTokenizer,
-  });
+  tokenizerPromise = AutoTokenizer.from_pretrained(tokenizerSpec.source, tokenizerSpec.options);
   return tokenizerPromise;
 }