src/lib/engine.js

// The engine is the flux-klein.js package (npm).
// It picks one of two pipelines behind one entry point:
//   desktop — int8 transformer resident on the GPU (~4 GB), fastest
//   mobile  — int4 shards streamed through a small GPU ring (~0.6 GB), built to survive the
//             iPhone's 2 GB tab cap; also fine on any low-memory device
// This module is the page's view of it: which engine to pick, the cache probe that runs before
// anything else is loaded, and one status line per progress event.
import { isIOS, isMobile } from 'flux-klein.js/platform'
import { formatProgress, progressFraction } from 'flux-klein.js/progress'
import { defaultCache } from 'flux-klein.js/cache'
import { MAX_AREA, MAX_AREA_ONE_THING, MAX_SIDE, maxAreaFor } from 'flux-klein.js/limits'
// onnxruntime-web's wasm runtime as Vite assets, so it is served from here rather than from the
// package's default (the jsDelivr CDN). Both files are subpath exports of onnxruntime-web.
import ortMjs from 'onnxruntime-web/ort-wasm-simd-threaded.mjs?url'
import ortWasm from 'onnxruntime-web/ort-wasm-simd-threaded.wasm?url'
import { HF_BASE, SPRINT_LORA_DIR, SPRINT_LORA_LOCAL_BASE } from './hf.js'

export const IOS = isIOS()
export const IS_MOBILE = isMobile()

// Dev only: send just the LoRA's requests to a local server (SPRINT_LORA_LOCAL_BASE), leaving the
// base transformer / text encoder / VAE on the hub. Keyed on the directory name, so the cache key
// stays the hub URL and the browser's existing OPFS cache still hits for everything else.
function loraFetchImpl() {
  if (!SPRINT_LORA_LOCAL_BASE) return undefined
  const localRoot = SPRINT_LORA_LOCAL_BASE.replace(/\/+$/, '')
  return (url, opts) => {
    const s = typeof url === 'string' ? url : url.url
    const i = s.indexOf(`/${SPRINT_LORA_DIR}/`)
    return fetch(i >= 0 ? localRoot + s.slice(i) : url, opts)
  }
}

// ?engine=mobile|desktop overrides the detection (e.g. to try the streaming engine on a laptop).
export function pickEngine() {
  const q = new URLSearchParams(location.search).get('engine')
  return q === 'mobile' || q === 'desktop' ? q : IS_MOBILE ? 'mobile' : 'desktop'
}

const ROOT = new URL(import.meta.env.BASE_URL, location.href).href // where public/ is served

// The limits the engine will enforce, known before it loads: each side up to MAX_SIDE, and an
// area budget — the phone build on an iPhone (one heavy thing in memory at a time) caps it at
// MAX_AREA_ONE_THING, everything else at what the GPU can bind in one attention buffer
// (flux-klein.js/limits maxAreaFor), MAX_AREA at most. Knowing them up front keeps the size menu
// from offering presets the first run would refuse; the loaded engine confirms them (loadEngine).
const F16 = () => new URLSearchParams(location.search).get('f16') !== '0'
export const limitsFor = (kind, gpuLimits = null) => ({
  maxSide: MAX_SIDE,
  maxArea: maxAreaFor(gpuLimits, {
    es: F16() ? 2 : 4,
    cap: kind === 'mobile' && IOS ? MAX_AREA_ONE_THING : MAX_AREA,
  }),
})
// The same, refined with the adapter's limits. An adapter is a handle, not a device: no GPU
// memory is taken. Falls back to the static guess where WebGPU is missing.
export async function probeLimits(kind) {
  try {
    const adapter = await navigator.gpu?.requestAdapter()
    return limitsFor(kind, adapter?.limits ?? null)
  } catch {
    return limitsFor(kind)
  }
}

// The weights already in this browser, and wiping them — both without loading the engine
// (flux-klein.js/cache reads OPFS directory metadata only: no weights, no GPU device).
export const cacheInfo = () => defaultCache.cacheInfo()
export const clearCache = () => defaultCache.clearCache()

// `onStatus(text, fraction)`: one status line at a time, with a 0…1 fraction while the line
// carries one (download, upload, denoise) or null. The package reports { stage, detail } events
// (flux-klein.js/progress); the wording and the fraction come from there, not from regexes.
export async function loadEngine(kind, onStatus) {
  // The package and the tokenizer class come in only now: nothing loads before the first Generate.
  const [{ createFluxKlein }, { PreTrainedTokenizer }] = await Promise.all([
    import('flux-klein.js'),
    import('@huggingface/transformers'),
  ])
  const fetchImpl = loraFetchImpl()
  const klein = await createFluxKlein({
    mode: kind,
    base: HF_BASE,
    tokenizerUrl: ROOT + 'tokenizer/',
    f16: F16(),
    // Absolute: the mobile build hands these to its ONNX workers, whose base URL is not the page.
    wasmPaths: {
      mjs: new URL(ortMjs, location.href).href,
      wasm: new URL(ortWasm, location.href).href,
    },
    PreTrainedTokenizer,
    // The SANA-Sprint LoRA (the few-step distillation), applied at any step count; the engine
    // downloads it (once, cached) the first time a run asks for it. Each generate() says whether
    // it runs (`lora: true`) or the stock weights do (`lora: false`) — the switch in the prompt bar.
    lora: SPRINT_LORA_DIR,
    ...(fetchImpl ? { fetchImpl } : {}),
    // Threads need cross-origin isolation; inside the hub's iframe that is never granted, and the
    // package falls back to one thread on its own.
    onEvent: (ev) => onStatus(formatProgress(ev), progressFraction(ev)),
  })
  addEventListener('pagehide', () => klein.destroy(), { once: true })
  return {
    kind: klein.mode,
    klein,
    f16: klein.f16,
    maxSide: klein.limits.maxSide, // per side 128…maxSide, in steps of 16
    maxArea: klein.limits.maxArea, // the phone build caps the area below maxSide²
    cacheInfo: () => klein.cacheInfo(),
    clearCache: () => klein.clearCache(),
  }
}

// Generate one image. `ref` (ImageBitmap) turns generation into an edit of that image: the package
// cover-fits it to a square, encodes it with the VAE (at 256² unless the output is a 128² or 256²
// square — enc512 overflows ORT-web's shape math) and the prompt describes the change.
// `decoder` picks the VAE decoder per call: 'full' (default) or 'tiny' — fal's 2.9M
// FLUX.2 Tiny AutoEncoder, ~20× less decode work, slightly softer textures. `lora` picks the
// model: true = the SANA-Sprint LoRA on the base, false = the stock distilled weights.
export function generate(engine, { prompt, width, height, steps, seed, ref, decoder, lora }) {
  return engine.klein.generate({
    prompt,
    width,
    height,
    steps,
    seed,
    reference: ref ?? null,
    ...(decoder ? { decoder } : {}),
    ...(lora !== undefined ? { lora } : {}),
  })
}