Gate first model download behind start button

static-app.js

@@ -1,14 +1,15 @@
-import { AutoTokenizer, env } from "https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.8.1";
-
-const $ = (id) => document.getElementById(id);
+const $ = (id) => document.getElementById(id);
 
 const statusEl = $("status");
-const imageEl = $("image");
-const canvasEl = $("image-canvas");
-const placeholderEl = $("placeholder");
-const runButton = $("run");
-const warmButton = $("warm");
-const abortButton = $("abort");
+const imageEl = $("image");
+const canvasEl = $("image-canvas");
+const placeholderEl = $("placeholder");
+const downloadGateEl = $("download-gate");
+const startDownloadButton = $("start-download");
+const actionsEl = $("actions");
+const runButton = $("run");
+const warmButton = $("warm");
+const abortButton = $("abort");
 const elapsedEl = $("elapsed");
 const initImageInput = $("init-image");
 const initImageLabelEl = $("init-image-label");
@@ -53,10 +54,15 @@ const THEME_STORAGE_KEY = "flux2-static-theme";
 const GALLERY_LIMIT = 24;
 const DEFAULT_LOCAL_MODEL_BASE_URL = "/models";
 const DEFAULT_LOCAL_RUNTIME_BASE_URL = "/runtime";
-
-let tokenizerPromise = null;
-let engineReadyPromise = null;
-let currentImageUrl = null;
+const DEFAULT_FIRST_LOAD_DOWNLOAD_LABEL = "about 3.8 GB / 3.5 GiB";
+const FULL_MODEL_REPO_DOWNLOAD_LABEL = "12.6 GB";
+
+let tokenizerPromise = null;
+let transformersModulePromise = null;
+let engineReadyPromise = null;
+let engineModulePromise = null;
+let customRuntimeScriptsPromise = null;
+let currentImageUrl = null;
 let currentSourceUrl = null;
 let backgroundWarmPromise = null;
 let backgroundWarmKey = "";
@@ -82,7 +88,9 @@ let backgroundVaePrepareKey = "";
 let backgroundVaeEncoderPreparePromise = null;
 let backgroundVaeEncoderPrepareKey = "";
 let backgroundDeepPrepareTimer = 0;
-let backgroundDeepPrepareToken = 0;
+let backgroundDeepPrepareToken = 0;
+let bootStarted = false;
+let bootReady = false;
 const preparedInitImageCache = new WeakMap();
 const generationResultCache = new Map();
 const tokenizedPromptCache = new Map();
@@ -422,19 +430,72 @@ function setStatus(lines, isError = false) {
   statusEl.classList.toggle("error", Boolean(isError));
 }
 
-function setButtonDisabled(button, disabled) {
-  if (button) button.disabled = disabled;
-}
-
-function setElapsed(text, active = false) {
-  if (!elapsedEl) return;
-  elapsedEl.textContent = text;
-  elapsedEl.classList.toggle("active", active);
-}
-
-function updateRunTimer(prefix = "Generating") {
-  if (!runTimerStarted) {
-    setElapsed("0.0s");
+function setButtonDisabled(button, disabled) {
+  if (button) button.disabled = disabled;
+}
+
+function setElapsed(text, active = false) {
+  if (!elapsedEl) return;
+  elapsedEl.textContent = text;
+  elapsedEl.classList.toggle("active", active);
+}
+
+function setActionsVisible(visible) {
+  if (actionsEl) actionsEl.hidden = !visible;
+}
+
+function setDownloadGateVisible(visible) {
+  if (downloadGateEl) downloadGateEl.hidden = !visible;
+}
+
+function canRunRuntimeWork() {
+  return bootStarted;
+}
+
+function showDownloadGate({ retry = false, preserveStatus = false } = {}) {
+  bootReady = false;
+  setDownloadGateVisible(true);
+  setActionsVisible(false);
+  if (startDownloadButton) startDownloadButton.disabled = false;
+  runButton.disabled = true;
+  setButtonDisabled(warmButton, true);
+  abortButton.disabled = true;
+  setElapsed("0.0s");
+  placeholderEl.hidden = false;
+  placeholderEl.textContent = "Model not loaded";
+  if (!preserveStatus) {
+    setStatus([
+      retry ? "Download failed" : "Ready to download",
+      `Default fast path: ${DEFAULT_FIRST_LOAD_DOWNLOAD_LABEL}`,
+      `Full model repo: ${FULL_MODEL_REPO_DOWNLOAD_LABEL}`,
+    ], retry);
+  }
+}
+
+function showRuntimeLoading() {
+  setDownloadGateVisible(false);
+  setActionsVisible(false);
+  if (startDownloadButton) startDownloadButton.disabled = true;
+  runButton.disabled = true;
+  setButtonDisabled(warmButton, true);
+  abortButton.disabled = true;
+  setElapsed("0.0s");
+  placeholderEl.hidden = false;
+  placeholderEl.textContent = "Loading";
+}
+
+function showRuntimeReady() {
+  bootReady = true;
+  setDownloadGateVisible(false);
+  setActionsVisible(true);
+  runButton.disabled = false;
+  setButtonDisabled(warmButton, false);
+  abortButton.disabled = true;
+}
+
+function updateRunTimer(prefix = "Generating") {
+  if (!runTimerStarted) {
+    setElapsed("0.0s");
     return 0;
   }
   const seconds = (performance.now() - runTimerStarted) / 1000;
@@ -641,12 +702,60 @@ function initTheme() {
   applyTheme(requestedTheme || "dark", { persist: false });
 }
 
-function isAbortError(err) {
-  return String(err?.stack || err?.message || err || "").includes("ABORTED_BY_CLIENT");
-}
-
-function waitForEngine() {
-  return new Promise((resolve, reject) => {
+function isAbortError(err) {
+  return String(err?.stack || err?.message || err || "").includes("ABORTED_BY_CLIENT");
+}
+
+function loadClassicScript(src) {
+  return new Promise((resolve, reject) => {
+    const existing = Array.from(document.scripts).find((script) => script.dataset.flux2Src === src);
+    if (existing?.dataset.loaded === "true") {
+      resolve();
+      return;
+    }
+    if (existing) existing.remove();
+    const script = document.createElement("script");
+    script.src = src;
+    script.dataset.flux2Src = src;
+    script.async = false;
+    script.onload = () => {
+      script.dataset.loaded = "true";
+      resolve();
+    };
+    script.onerror = () => reject(new Error(`Failed to load ${src}`));
+    if (!existing) document.body.appendChild(script);
+  });
+}
+
+async function loadCustomRuntimeScripts() {
+  if (!customRuntimeScriptsPromise) {
+    customRuntimeScriptsPromise = (async () => {
+      setLoadPhase("Loading custom WebGPU kernels");
+      await loadClassicScript("tools/custom_lowbit_webgpu/custom_lowbit_linear.js");
+      await loadClassicScript("tools/custom_lowbit_webgpu/custom_transformer_runtime.js");
+    })();
+  }
+  return customRuntimeScriptsPromise;
+}
+
+async function loadEngineModule() {
+  if (!engineModulePromise) {
+    setLoadPhase("Loading browser engine");
+    engineModulePromise = import("./flux2-engine.js");
+  }
+  return engineModulePromise;
+}
+
+function resetRuntimeLoadPromises() {
+  engineReadyPromise = null;
+  tokenizerPromise = null;
+  if (!window.flux2Engine) engineModulePromise = null;
+  if (typeof globalThis.createCustomFluxTransformerRuntime !== "function") customRuntimeScriptsPromise = null;
+  transformersModulePromise = null;
+}
+
+function waitForEngine() {
+  return new Promise((resolve, reject) => {
     const deadline = performance.now() + 30000;
     const tick = () => {
       if (window.flux2Engine && typeof window.flux2Engine.init === "function") {
@@ -724,6 +833,7 @@ function maskFromLength(length, seqLen) {
 async function loadTokenizer() {
   if (tokenizerPromise) return tokenizerPromise;
   setLoadPhase("Loading tokenizer");
+  const { AutoTokenizer, env } = await loadTransformersModule();
   const modelBaseUrl = configuredModelBaseUrl();
   const tokenizerSpec = tokenizerLoadSpec(modelBaseUrl);
   env.allowLocalModels = !tokenizerSpec.remote;
@@ -740,6 +850,14 @@ async function loadTokenizer() {
   return tokenizerPromise;
 }
 
+async function loadTransformersModule() {
+  if (!transformersModulePromise) {
+    setLoadPhase("Loading tokenizer runtime");
+    transformersModulePromise = import("https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.8.1");
+  }
+  return transformersModulePromise;
+}
+
 async function tokenizePrompt(prompt, seqLen) {
   const cachePayload = JSON.stringify({ model: "flux.2-klein-4b", text_seq_len: seqLen, prompt });
   const promptCacheKey = await sha256Hex(cachePayload);
@@ -779,6 +897,7 @@ async function tokenizePrompt(prompt, seqLen) {
 async function initEngine() {
   if (engineReadyPromise) return engineReadyPromise;
   engineReadyPromise = (async () => {
+    await Promise.all([loadCustomRuntimeScripts(), loadEngineModule()]);
     const engine = await waitForEngine();
     const modelBaseUrl = configuredModelBaseUrl();
     const runtimeBaseUrl = configuredRuntimeBaseUrl(modelBaseUrl);
@@ -1418,9 +1537,10 @@ function shouldTextContextPrepareRequest(request) {
   return request.transformerBackend === "custom-lowbit-webgpu";
 }
 
-function startBackgroundTextContextPrepareForCurrentSettings(request = readRequest()) {
-  const prepareRequest = backgroundPrepareRequestForRequest(request);
-  if (!shouldTextContextPrepareRequest(prepareRequest)) return false;
+function startBackgroundTextContextPrepareForCurrentSettings(request = readRequest()) {
+  if (!canRunRuntimeWork()) return false;
+  const prepareRequest = backgroundPrepareRequestForRequest(request);
+  if (!shouldTextContextPrepareRequest(prepareRequest)) return false;
   const key = textContextPrepareKeyForRequest(prepareRequest);
   if (backgroundTextContextPreparePromise && backgroundTextContextPrepareKey === key) return true;
   backgroundTextContextPrepareKey = key;
@@ -1483,9 +1603,10 @@ function shouldTransformerDispatchWarmupRequest(request) {
   return plannedMaxDimension <= 512;
 }
 
-function startBackgroundTransformerSetupPrepareForCurrentSettings(request = readRequest()) {
-  const prepareRequest = backgroundPrepareRequestForRequest(request);
-  if (!shouldTransformerSetupPrepareRequest(prepareRequest)) return false;
+function startBackgroundTransformerSetupPrepareForCurrentSettings(request = readRequest()) {
+  if (!canRunRuntimeWork()) return false;
+  const prepareRequest = backgroundPrepareRequestForRequest(request);
+  if (!shouldTransformerSetupPrepareRequest(prepareRequest)) return false;
   const dispatchWarmup = shouldTransformerDispatchWarmupRequest(prepareRequest);
   const key = `${transformerSetupPrepareKeyForRequest(prepareRequest)}|dispatch=${dispatchWarmup ? 1 : 0}`;
   if (backgroundTransformerSetupPreparePromise && backgroundTransformerSetupPrepareKey === key) return true;
@@ -1569,8 +1690,9 @@ function shouldAutoWarmRequest(request) {
     maxDimension <= 256;
 }
 
-function startBackgroundAssetPrepareForCurrentSettings(request = readRequest()) {
-  if (!shouldAssetPrepareRequest(request)) return false;
+function startBackgroundAssetPrepareForCurrentSettings(request = readRequest()) {
+  if (!canRunRuntimeWork()) return false;
+  if (!shouldAssetPrepareRequest(request)) return false;
   const key = assetPrepareKeyForRequest(request);
   if (backgroundAssetPreparePromise && backgroundAssetPrepareKey === key) return true;
   backgroundAssetPrepareKey = key;
@@ -1631,8 +1753,9 @@ function shouldVaeEncoderPrepareRequest(request) {
   return true;
 }
 
-function startBackgroundVaeEncoderPrepareForCurrentSettings(request = readRequest()) {
-  if (!shouldVaeEncoderPrepareRequest(request)) return false;
+function startBackgroundVaeEncoderPrepareForCurrentSettings(request = readRequest()) {
+  if (!canRunRuntimeWork()) return false;
+  if (!shouldVaeEncoderPrepareRequest(request)) return false;
   const key = vaeEncoderPrepareKeyForRequest(request);
   if (backgroundVaeEncoderPreparePromise && backgroundVaeEncoderPrepareKey === key) return true;
   backgroundVaeEncoderPrepareKey = key;
@@ -1695,8 +1818,9 @@ async function waitForStartedPreparePromises() {
   }));
 }
 
-function startStageAndVaePrepareForRequest(request) {
-  const startedText = startBackgroundTextContextPrepareForCurrentSettings(request);
+function startStageAndVaePrepareForRequest(request) {
+  if (!canRunRuntimeWork()) return false;
+  const startedText = startBackgroundTextContextPrepareForCurrentSettings(request);
   const startedTransformer = startBackgroundTransformerSetupPrepareForCurrentSettings(request);
   if (request.hasInitImage && shouldVaeEncoderPrepareRequest(request)) {
     const startedVaeEncoder = startBackgroundVaeEncoderPrepareForCurrentSettings(request);
@@ -1706,8 +1830,9 @@ function startStageAndVaePrepareForRequest(request) {
   return startedText || startedTransformer || startedVae;
 }
 
-function startBackgroundVaePrepareForCurrentSettings(request = readRequest()) {
-  const prepareRequest = backgroundPrepareRequestForRequest(request);
+function startBackgroundVaePrepareForCurrentSettings(request = readRequest()) {
+  if (!canRunRuntimeWork()) return false;
+  const prepareRequest = backgroundPrepareRequestForRequest(request);
   if (!shouldVaePrepareRequest(prepareRequest)) return false;
   const dispatchWarmup = shouldVaeDispatchWarmupRequest(prepareRequest);
   const key = `${vaePrepareKeyForRequest(prepareRequest)}|dispatch=${dispatchWarmup ? 1 : 0}`;
@@ -1770,8 +1895,9 @@ function cancelBackgroundWarm(message = "") {
   return true;
 }
 
-function startBackgroundWarmForCurrentSettings() {
-  if (foregroundBusy) return false;
+function startBackgroundWarmForCurrentSettings() {
+  if (!canRunRuntimeWork()) return false;
+  if (foregroundBusy) return false;
   const request = readRequest();
   if (!shouldAutoWarmRequest(request)) return false;
   const key = warmKeyForRequest(request);
@@ -1809,8 +1935,9 @@ function startBackgroundWarmForCurrentSettings() {
   return true;
 }
 
-function scheduleBackgroundStageAndVaePrepareForCurrentSettings(request = readRequest(), delayMs = 750) {
-  if (foregroundBusy || urlFlag("noAutowarm")) return false;
+function scheduleBackgroundStageAndVaePrepareForCurrentSettings(request = readRequest(), delayMs = 750) {
+  if (!canRunRuntimeWork()) return false;
+  if (foregroundBusy || urlFlag("noAutowarm")) return false;
   if (backgroundDeepPrepareTimer) {
     clearTimeout(backgroundDeepPrepareTimer);
     backgroundDeepPrepareTimer = 0;
@@ -1847,9 +1974,10 @@ function scheduleBackgroundStageAndVaePrepareForCurrentSettings(request = readRe
   return true;
 }
 
-function scheduleAutoWarmForCurrentSettings(delayMs = 350) {
-  clearWarmSchedule();
-  if (foregroundBusy || urlFlag("noAutowarm")) return;
+function scheduleAutoWarmForCurrentSettings(delayMs = 350) {
+  clearWarmSchedule();
+  if (!canRunRuntimeWork()) return;
+  if (foregroundBusy || urlFlag("noAutowarm")) return;
   const request = readRequest();
   startBackgroundAssetPrepareForCurrentSettings(request);
   scheduleBackgroundStageAndVaePrepareForCurrentSettings(request, request.hasInitImage ? delayMs : 750);
@@ -2344,9 +2472,10 @@ function displayGenerationResult(result, elapsed, seqLen, heading = "Elapsed", e
   }
 }
 
-async function generate() {
-  const requestedStarted = performance.now();
-  const request = readRequest();
+async function generate() {
+  if (!bootReady) throw new Error("Model assets are not loaded yet.");
+  const requestedStarted = performance.now();
+  const request = readRequest();
   runButton.disabled = true;
   setButtonDisabled(warmButton, true);
   abortButton.disabled = false;
@@ -2621,8 +2750,9 @@ async function waitForBackgroundVaePrepare(request) {
   return (performance.now() - started) / 1000;
 }
 
-async function warmCurrentSettings({ background = false } = {}) {
-  const request = readRequest();
+async function warmCurrentSettings({ background = false } = {}) {
+  if (!canRunRuntimeWork()) throw new Error("Model assets are not loaded yet.");
+  const request = readRequest();
   const { warmRequest, warmProfile } = warmPlanForRequest(request);
   const { engine, init } = await initEngine();
   const started = performance.now();
@@ -2730,19 +2860,20 @@ async function prepareInitialFastPath(request) {
 }
 
 async function boot() {
-  beginLoadProgress("First load");
+  if (!bootStarted) bootStarted = true;
+  showRuntimeLoading();
+  beginLoadProgress("Downloading assets");
   try {
     installLoadProgressFetch();
     setStatus("Loading browser runtime");
     setLoadPhase("Loading browser runtime");
     const [{ init }, tokenizer] = await Promise.all([initEngine(), loadTokenizer()]);
     const bootRequest = readRequest();
-    const fastPathPrepared = await prepareInitialFastPath(bootRequest);
-    runButton.disabled = false;
-    setButtonDisabled(warmButton, false);
-    setStatus([
-      "Ready",
-      fastPathPrepared ? "Fast path prepared" : "",
+    const fastPathPrepared = await prepareInitialFastPath(bootRequest);
+    showRuntimeReady();
+    setStatus([
+      "Ready",
+      fastPathPrepared ? "Fast path prepared" : "",
       `Model: ${init.config.model_variant || init.config.model_name}`,
       seedInputLine(),
       `Tokenizer: ${tokenizer.constructor?.name || "loaded"}`,
@@ -2754,18 +2885,31 @@ async function boot() {
     }
     if (shouldAutoWarmRequest(bootRequest)) startBackgroundWarmForCurrentSettings();
     finishLoadProgress(fastPathPrepared ? "Fast path prepared" : "Ready");
+    return true;
   } catch (err) {
+    bootStarted = false;
+    bootReady = false;
+    resetRuntimeLoadPromises();
     runButton.disabled = true;
     setButtonDisabled(warmButton, true);
+    abortButton.disabled = true;
     setStatus(err?.stack || err?.message || String(err), true);
     finishLoadProgress("Load failed", false);
+    showDownloadGate({ retry: true, preserveStatus: true });
+    return false;
   }
 }
-
-warmButton?.addEventListener("click", () => {
-  if (foregroundBusy) return;
-  clearWarmSchedule();
-  foregroundBusy = true;
+
+startDownloadButton?.addEventListener("click", () => {
+  if (bootStarted) return;
+  bootStarted = true;
+  void boot();
+});
+
+warmButton?.addEventListener("click", () => {
+  if (!bootReady || foregroundBusy) return;
+  clearWarmSchedule();
+  foregroundBusy = true;
   warmCurrentSettings().catch((err) => {
     if (isAbortError(err)) return;
     placeholderEl.hidden = false;
@@ -2779,9 +2923,9 @@ warmButton?.addEventListener("click", () => {
   });
 });
 
-runButton.addEventListener("click", () => {
-  if (foregroundBusy) return;
-  clearWarmSchedule();
+runButton.addEventListener("click", () => {
+  if (!bootReady || foregroundBusy) return;
+  clearWarmSchedule();
   foregroundBusy = true;
   startRunTimer();
   generate().catch((err) => {
@@ -2961,7 +3105,7 @@ galleryGridEl?.addEventListener("click", (event) => {
   });
 });
 
-initTheme();
-applyUrlDefaults();
-updateStrengthDisplay();
-boot();
+initTheme();
+applyUrlDefaults();
+updateStrengthDisplay();
+showDownloadGate();

static.html

@@ -337,17 +337,44 @@
       gap: 10px;
     }
     .wide { grid-column: 1 / -1; }
-    .actions {
-      display: flex;
-      gap: 10px;
-      margin-top: 18px;
-      align-items: center;
-      flex-wrap: wrap;
-    }
-    button {
-      min-height: 42px;
-      border: 0;
-      border-radius: var(--radius);
+    .actions {
+      display: flex;
+      gap: 10px;
+      margin-top: 18px;
+      align-items: center;
+      flex-wrap: wrap;
+    }
+    .actions[hidden] {
+      display: none;
+    }
+    .download-gate {
+      margin-top: 18px;
+      padding: 14px;
+      border: 1px solid color-mix(in srgb, var(--accent) 42%, var(--line));
+      border-radius: var(--radius);
+      background: color-mix(in srgb, var(--accent-soft) 74%, var(--panel-2));
+    }
+    .download-gate[hidden] {
+      display: none;
+    }
+    .download-gate-title {
+      font-size: 13px;
+      font-weight: 850;
+      color: var(--text);
+    }
+    .download-gate-copy {
+      margin: 7px 0 12px;
+      color: var(--muted);
+      font-size: 12px;
+      line-height: 1.45;
+    }
+    .download-gate button {
+      width: 100%;
+    }
+    button {
+      min-height: 42px;
+      border: 0;
+      border-radius: var(--radius);
       padding: 0 14px;
       background: var(--accent);
       color: #061715;
@@ -698,7 +725,12 @@
         <input id="decode-max" type="hidden" value="0">
         <input id="render-max" type="hidden" value="0">
       </div>
-      <div class="actions">
+      <div id="download-gate" class="download-gate">
+        <div class="download-gate-title">Download model assets</div>
+        <div class="download-gate-copy">Default fast path downloads about 3.8 GB / 3.5 GiB before the first generation. The full public model repo is 12.6 GB; ONNX fallback modes can fetch more later.</div>
+        <button id="start-download" type="button"><span class="icon">download</span>Start download</button>
+      </div>
+      <div id="actions" class="actions" hidden>
         <button id="run" type="button" disabled><span class="icon">auto_awesome</span>Generate</button>
         <button id="abort" class="secondary" type="button" disabled><span class="icon">stop</span>Stop</button>
         <div id="elapsed" class="elapsed-meter" aria-live="polite">0.0s</div>
@@ -732,11 +764,8 @@
         <div id="gallery-grid" class="gallery-grid"></div>
       </div>
     </section>
-  </main>
+  </main>
   <script src="space-config.js"></script>
-  <script src="tools/custom_lowbit_webgpu/custom_lowbit_linear.js"></script>
-  <script src="tools/custom_lowbit_webgpu/custom_transformer_runtime.js"></script>
-  <script type="module" src="flux2-engine.js"></script>
   <script type="module" src="static-app.js"></script>
 </body>
 </html>