const urlParams = new URLSearchParams(globalThis.location?.search || ''); const ortVersion = urlParams.get('ort') || '1.26.0'; const ortBaseUrl = `https://cdn.jsdelivr.net/npm/onnxruntime-web@${ortVersion}/dist/`; const ort = await import(`${ortBaseUrl}ort.all.mjs`); ort.env.wasm.wasmPaths = ortBaseUrl; ort.env.wasm.numThreads = globalThis.crossOriginIsolated ? (navigator.hardwareConcurrency || 4) : 1; ort.env.wasm.simd = true; ort.env.webgpu.powerPreference = 'high-performance'; let config = null; let modelBaseUrl = null; let runtimeOptions = {}; let abortRequested = false; const customKernelState = { available: false, checked: false, lastStatus: null, }; const sessionCache = new Map(); const textContextCache = new Map(); const vaeEncoderCache = new Map(); const staticGpuTensorCache = new Map(); const customTransformerRuntimeCache = new Map(); let latentUpdatePipeline = null; let latentUpdateBindGroupLayout = null; let webGpuProfileEvents = []; let activeProfileLabel = null; let textContextDbPromise = null; let lastTextEncoderDetails = null; let lastVaeEncoderDetails = null; function optionList(value) { if (Array.isArray(value)) return value.map((item) => String(item).trim()).filter(Boolean); if (typeof value === 'string') return value.split(',').map((item) => item.trim()).filter(Boolean); return []; } function stageMatches(stage, patterns) { if (!patterns.length) return false; return patterns.some((pattern) => { if (pattern === '*' || pattern === stage) return true; return stage.startsWith(pattern); }); } function shouldUseGraphCapture(stage) { if (!runtimeOptions.enableGraphCapture) return false; if (String(stage || '').startsWith('vae-decoder-split')) return false; const allowedStages = optionList(runtimeOptions.graphCaptureStages); return !allowedStages.length || stageMatches(stage, allowedStages); } function profileDurationMs(event) { const raw = Number(event.endTime || 0) - Number(event.startTime || 0); if (!Number.isFinite(raw) || raw < 0) return 0; return raw > 100000 ? raw / 1000000 : raw; } function log(stage, detail) { console.log(`[${stage}] ${detail}`); } function summarizeProfileEvents(stage, label, events) { if (!runtimeOptions.profileWebGpu || !events.length) return; const grouped = new Map(); let total = 0; for (const event of events) { const duration = profileDurationMs(event); if (!Number.isFinite(duration) || duration < 0) continue; total += duration; const name = event.programName || event.kernelName || event.kernelType || 'unknown'; const entry = grouped.get(name) || { name, duration: 0, count: 0, kernelTypes: new Set() }; entry.duration += duration; entry.count += 1; if (event.kernelType) entry.kernelTypes.add(event.kernelType); grouped.set(name, entry); } const limit = Number(runtimeOptions.profileTopK || 12); const rows = Array.from(grouped.values()) .sort((a, b) => b.duration - a.duration) .slice(0, Number.isFinite(limit) && limit > 0 ? limit : 12) .map((entry) => `${entry.name}:${entry.duration.toFixed(3)}ms/${entry.count}${entry.kernelTypes.size ? ` types=${Array.from(entry.kernelTypes).join('|')}` : ''}`); log(stage, `webgpu profile ${label}: events=${events.length} total=${total.toFixed(3)}ms top=${rows.join(' ; ')}`); } async function profiledRun(session, feeds, fetches, stage, label) { if (!runtimeOptions.profileWebGpu) { return fetches ? session.run(feeds, fetches) : session.run(feeds); } const startIndex = webGpuProfileEvents.length; const previousLabel = activeProfileLabel; activeProfileLabel = `${stage}:${label}`; try { return fetches ? await session.run(feeds, fetches) : await session.run(feeds); } finally { activeProfileLabel = previousLabel; summarizeProfileEvents(stage, label, webGpuProfileEvents.slice(startIndex)); } } function openTextContextDb() { if (typeof indexedDB === 'undefined') return Promise.resolve(null); if (textContextDbPromise) return textContextDbPromise; textContextDbPromise = new Promise((resolve) => { const request = indexedDB.open('flux2-browser-cache', 3); request.onupgradeneeded = () => { const db = request.result; const textStore = db.objectStoreNames.contains('text-contexts') ? request.transaction.objectStore('text-contexts') : db.createObjectStore('text-contexts', { keyPath: 'key' }); if (!textStore.indexNames.contains('savedAt')) { textStore.createIndex('savedAt', 'savedAt'); } const vaeStore = db.objectStoreNames.contains('vae-encoder-latents') ? request.transaction.objectStore('vae-encoder-latents') : db.createObjectStore('vae-encoder-latents', { keyPath: 'key' }); if (!vaeStore.indexNames.contains('savedAt')) { vaeStore.createIndex('savedAt', 'savedAt'); } }; request.onsuccess = () => resolve(request.result); request.onerror = () => { log('text-cache', `IndexedDB open failed: ${request.error?.message || request.error || 'unknown error'}`); resolve(null); }; request.onblocked = () => { log('text-cache', 'IndexedDB open blocked by another page'); resolve(null); }; }); return textContextDbPromise; } async function pruneObjectStoreBySavedAt(db, storeName, maxEntries, stage) { await new Promise((resolve) => { const tx = db.transaction(storeName, 'readwrite'); const store = tx.objectStore(storeName); const entries = []; const request = store.index('savedAt').openCursor(); request.onsuccess = () => { const cursor = request.result; if (!cursor) return; entries.push({ key: cursor.primaryKey, savedAt: Number(cursor.key || 0) }); cursor.continue(); }; tx.oncomplete = () => { if (entries.length <= maxEntries) { resolve(); return; } const deleteTx = db.transaction(storeName, 'readwrite'); const deleteStore = deleteTx.objectStore(storeName); const remove = entries.sort((a, b) => a.savedAt - b.savedAt).slice(0, entries.length - maxEntries); for (const entry of remove) deleteStore.delete(entry.key); deleteTx.oncomplete = () => resolve(); deleteTx.onerror = () => resolve(); }; tx.onerror = () => { log(stage, `IndexedDB prune scan failed for ${storeName}`); resolve(); }; }); } function makePersistentTextContextKey(modelConfig, cacheKey, seqLen, contextDim) { return [ 'ctx-v2', modelConfig.file || 'text-encoder', `seq=${seqLen}`, `dim=${contextDim}`, cacheKey, ].join('|'); } async function loadPersistentTextContext(key, expectedValues) { if (!key) return null; const db = await openTextContextDb(); if (!db) return null; return await new Promise((resolve) => { const tx = db.transaction('text-contexts', 'readonly'); const store = tx.objectStore('text-contexts'); const request = store.get(key); request.onsuccess = () => { const entry = request.result; if (!entry || !entry.data) { resolve(null); return; } const values = entry.data instanceof Uint16Array ? entry.data : (entry.data instanceof ArrayBuffer ? new Uint16Array(entry.data) : null); if (!values || values.length !== expectedValues) { resolve(null); return; } resolve(new Uint16Array(values)); }; request.onerror = () => { log('text-cache', `IndexedDB read failed: ${request.error?.message || request.error || 'unknown error'}`); resolve(null); }; }); } async function savePersistentTextContext(key, values) { if (!key || !(values instanceof Uint16Array)) return false; const db = await openTextContextDb(); if (!db) return false; const saved = await new Promise((resolve) => { const tx = db.transaction('text-contexts', 'readwrite'); const store = tx.objectStore('text-contexts'); const buffer = values.buffer.slice(values.byteOffset, values.byteOffset + values.byteLength); const request = store.put({ key, data: buffer, bytes: values.byteLength, savedAt: Date.now(), }); request.onsuccess = () => resolve(true); request.onerror = () => { log('text-cache', `IndexedDB write failed: ${request.error?.message || request.error || 'unknown error'}`); resolve(false); }; }); if (saved) prunePersistentTextContexts(db).catch((err) => { log('text-cache', `IndexedDB prune failed: ${errorDetail(err)}`); }); return saved; } async function prunePersistentTextContexts(db) { const maxEntries = Math.max(1, Math.trunc(Number(runtimeOptions.persistentTextContextCacheSize || 8))); await pruneObjectStoreBySavedAt(db, 'text-contexts', maxEntries, 'text-cache'); } function rememberTextContext(cacheKey, ctx) { if (!cacheKey) return; textContextCache.set(cacheKey, ctx); while (textContextCache.size > Number(runtimeOptions.textContextCacheSize || 4)) { const oldestKey = textContextCache.keys().next().value; textContextCache.delete(oldestKey); } log('text-encoder', `cached text context ${cacheKey.slice(0, 12)} entries=${textContextCache.size}`); } function makeVaeEncoderCacheKey(params, modelConfig, width, height) { if (!params.cacheVaeEncoder || !params.initImageCacheKey) return ''; return [ modelConfig.file || 'vae-encoder', `${width}x${height}`, String(params.initImageCacheKey), ].join('|'); } function rememberVaeEncoderResult(cacheKey, encoded) { if (!cacheKey) return; vaeEncoderCache.set(cacheKey, { latent: new Uint16Array(encoded.latent), latentWidth: encoded.latentWidth, latentHeight: encoded.latentHeight, }); const maxEntries = Math.max(1, Math.trunc(Number(runtimeOptions.vaeEncoderCacheSize || 3))); while (vaeEncoderCache.size > maxEntries) { const oldestKey = vaeEncoderCache.keys().next().value; vaeEncoderCache.delete(oldestKey); } log('vae-encoder', `cached source latent ${cacheKey.slice(0, 24)} entries=${vaeEncoderCache.size}`); } async function loadPersistentVaeEncoderResult(cacheKey, expectedValues, width, height) { if (!cacheKey || !runtimeOptions.persistentVaeEncoderCache) return null; const db = await openTextContextDb(); if (!db) return null; return await new Promise((resolve) => { const tx = db.transaction('vae-encoder-latents', 'readonly'); const store = tx.objectStore('vae-encoder-latents'); const request = store.get(cacheKey); request.onsuccess = () => { const entry = request.result; if (!entry || !entry.data || Number(entry.latentWidth) !== width || Number(entry.latentHeight) !== height) { resolve(null); return; } const values = entry.data instanceof Uint16Array ? entry.data : (entry.data instanceof ArrayBuffer ? new Uint16Array(entry.data) : null); if (!values || values.length !== expectedValues) { resolve(null); return; } resolve({ latent: new Uint16Array(values), latentWidth: width, latentHeight: height, }); }; request.onerror = () => { log('vae-cache', `IndexedDB read failed: ${request.error?.message || request.error || 'unknown error'}`); resolve(null); }; }); } async function savePersistentVaeEncoderResult(cacheKey, encoded) { if (!cacheKey || !runtimeOptions.persistentVaeEncoderCache || !(encoded.latent instanceof Uint16Array)) return false; const db = await openTextContextDb(); if (!db) return false; const saved = await new Promise((resolve) => { const tx = db.transaction('vae-encoder-latents', 'readwrite'); const store = tx.objectStore('vae-encoder-latents'); const values = encoded.latent; const buffer = values.buffer.slice(values.byteOffset, values.byteOffset + values.byteLength); const request = store.put({ key: cacheKey, data: buffer, latentWidth: encoded.latentWidth, latentHeight: encoded.latentHeight, bytes: values.byteLength, savedAt: Date.now(), }); request.onsuccess = () => resolve(true); request.onerror = () => { log('vae-cache', `IndexedDB write failed: ${request.error?.message || request.error || 'unknown error'}`); resolve(false); }; }); if (saved) { const maxEntries = Math.max(1, Math.trunc(Number(runtimeOptions.persistentVaeEncoderCacheSize || 8))); pruneObjectStoreBySavedAt(db, 'vae-encoder-latents', maxEntries, 'vae-cache').catch((err) => { log('vae-cache', `IndexedDB prune failed: ${errorDetail(err)}`); }); } return saved; } function formatBytes(bytes) { if (!Number.isFinite(bytes)) return 'unknown'; const units = ['B', 'KiB', 'MiB', 'GiB']; let value = bytes; let unit = 0; while (Math.abs(value) >= 1024 && unit < units.length - 1) { value /= 1024; unit += 1; } return `${value.toFixed(unit === 0 ? 0 : 1)} ${units[unit]}`; } function jsMemorySummary() { const memory = performance?.memory; if (!memory) return 'jsHeap=unavailable'; const used = formatBytes(memory.usedJSHeapSize); const total = formatBytes(memory.totalJSHeapSize); const limit = formatBytes(memory.jsHeapSizeLimit); return `jsHeap=${used}/${total} limit=${limit}`; } function logMemory(stage, label) { const deviceMemory = navigator.deviceMemory ? ` deviceMemory=${navigator.deviceMemory}GiB` : ''; log(stage, `memory ${label}: ${jsMemorySummary()}${deviceMemory}`); } function errorDetail(err) { if (!err) return String(err); const parts = []; if (err.name) parts.push(err.name); if (err.message) parts.push(err.message); if (!parts.length) parts.push(String(err)); if (err.stack) parts.push(err.stack); return parts.join(': '); } async function fetchJsonChecked(url, label) { const response = await fetch(url); const body = await response.text(); if (!response.ok) throw new Error(`Failed to fetch ${label}: HTTP ${response.status} ${response.statusText} from ${url}`); return JSON.parse(body); } function float32ToFloat16Bits(value) { if (Number.isNaN(value)) return 0x7e00; if (value === Infinity) return 0x7c00; if (value === -Infinity) return 0xfc00; const sign = value < 0 || Object.is(value, -0) ? 0x8000 : 0; const abs = Math.abs(value); if (abs === 0) return sign; if (abs >= 65504) return sign | 0x7bff; if (abs < 5.960464477539063e-8) return sign; let exponent = Math.floor(Math.log2(abs)); let mantissa = abs / Math.pow(2, exponent) - 1; let halfExp = exponent + 15; if (halfExp <= 0) return sign | Math.round(abs / 5.960464477539063e-8); let halfMant = Math.round(mantissa * 1024); if (halfMant === 1024) { halfMant = 0; halfExp += 1; } if (halfExp >= 31) return sign | 0x7bff; return sign | (halfExp << 10) | (halfMant & 0x03ff); } function float16BitsToFloat32(bits) { const sign = bits & 0x8000 ? -1 : 1; const exponent = (bits >> 10) & 0x1f; const mantissa = bits & 0x03ff; if (exponent === 0) return sign * (mantissa / 1024) * Math.pow(2, -14); if (exponent === 31) return mantissa ? NaN : sign * Infinity; return sign * (1 + mantissa / 1024) * Math.pow(2, exponent - 15); } let halfToImageByteLut = null; function imageByteFromFloat(value) { if (!Number.isFinite(value)) return 0; return Math.max(0, Math.min(255, Math.round((value + 1) * 127.5))); } function getHalfToImageByteLut() { if (halfToImageByteLut) return halfToImageByteLut; const table = new Uint8Array(65536); for (let bits = 0; bits < table.length; bits++) { table[bits] = imageByteFromFloat(float16BitsToFloat32(bits)); } halfToImageByteLut = table; return table; } function f32ToF16Array(values) { const out = new Uint16Array(values.length); for (let i = 0; i < values.length; i++) out[i] = float32ToFloat16Bits(values[i]); return out; } function f16ToF32Array(values) { const out = new Float32Array(values.length); for (let i = 0; i < values.length; i++) out[i] = float16BitsToFloat32(values[i]); return out; } function tensor(type, data, dims) { return new ort.Tensor(type, data, dims); } function alignTo(value, alignment) { return Math.ceil(value / alignment) * alignment; } function getWebGpuDevice() { const device = ort.env?.webgpu?.device; if (!device) throw new Error('ORT WebGPU device is not available'); return device; } function createGpuTensorFromTypedArray(data, dataType, dims, extraUsage = 0) { if (typeof ort.Tensor.fromGpuBuffer !== 'function') { throw new Error('ort.Tensor.fromGpuBuffer is not available in this ONNX Runtime Web build'); } const device = getWebGpuDevice(); const byteLength = data.byteLength; const buffer = device.createBuffer({ size: alignTo(Math.max(byteLength, 4), 16), usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC | extraUsage, }); device.queue.writeBuffer(buffer, 0, data.buffer, data.byteOffset, byteLength); return { tensor: ort.Tensor.fromGpuBuffer(buffer, { dataType, dims }), buffer, byteLength, }; } function getCachedGpuTensorResource(key, factory, stage) { if (staticGpuTensorCache.has(key)) { const cached = staticGpuTensorCache.get(key); log(stage, `reusing cached GPU tensor ${key}`); return cached; } const resource = factory(); resource.__flux2StaticGpuCacheKey = key; staticGpuTensorCache.set(key, resource); log(stage, `cached GPU tensor ${key}: ${formatBytes(resource.byteLength || 0)}`); return resource; } function clearStaticGpuTensorCache() { for (const resource of staticGpuTensorCache.values()) { try { resource.buffer?.destroy?.(); } catch { } } staticGpuTensorCache.clear(); } async function readGpuBuffer(buffer, byteLength) { const device = getWebGpuDevice(); const readback = device.createBuffer({ size: alignTo(Math.max(byteLength, 4), 4), usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ, }); const encoder = device.createCommandEncoder(); encoder.copyBufferToBuffer(buffer, 0, readback, 0, byteLength); device.queue.submit([encoder.finish()]); await readback.mapAsync(GPUMapMode.READ, 0, alignTo(Math.max(byteLength, 4), 4)); const mapped = readback.getMappedRange(0, byteLength); const copy = mapped.slice(0); readback.unmap(); readback.destroy(); return copy; } function getLatentUpdatePipeline() { if (latentUpdatePipeline) return { pipeline: latentUpdatePipeline, layout: latentUpdateBindGroupLayout }; const device = getWebGpuDevice(); if (!device.features?.has?.('shader-f16')) { throw new Error('GPU denoise update requires WebGPU shader-f16 support'); } latentUpdateBindGroupLayout = device.createBindGroupLayout({ entries: [ { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' } }, { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'read-only-storage' } }, { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'uniform' } }, ], }); const pipelineLayout = device.createPipelineLayout({ bindGroupLayouts: [latentUpdateBindGroupLayout] }); const module = device.createShaderModule({ code: ` enable f16; struct Params { dt: f32, length: u32, pad0: u32, pad1: u32, }; @group(0) @binding(0) var x: array; @group(0) @binding(1) var pred: array; @group(0) @binding(2) var params: Params; @compute @workgroup_size(256) fn main(@builtin(global_invocation_id) gid: vec3) { let i = gid.x; if (i >= params.length) { return; } let updated = f32(x[i]) + params.dt * f32(pred[i]); x[i] = f16(clamp(updated, -65504.0, 65504.0)); } `, }); latentUpdatePipeline = device.createComputePipeline({ layout: pipelineLayout, compute: { module, entryPoint: 'main' }, }); return { pipeline: latentUpdatePipeline, layout: latentUpdateBindGroupLayout }; } function dispatchLatentUpdate(xBuffer, predBuffer, length, dt) { const device = getWebGpuDevice(); const { pipeline, layout } = getLatentUpdatePipeline(); const params = new ArrayBuffer(16); const view = new DataView(params); view.setFloat32(0, dt, true); view.setUint32(4, length, true); const paramsBuffer = device.createBuffer({ size: 16, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST, }); device.queue.writeBuffer(paramsBuffer, 0, params); const bindGroup = device.createBindGroup({ layout, entries: [ { binding: 0, resource: { buffer: xBuffer } }, { binding: 1, resource: { buffer: predBuffer } }, { binding: 2, resource: { buffer: paramsBuffer } }, ], }); const encoder = device.createCommandEncoder(); const pass = encoder.beginComputePass(); pass.setPipeline(pipeline); pass.setBindGroup(0, bindGroup); pass.dispatchWorkgroups(Math.ceil(length / 256)); pass.end(); device.queue.submit([encoder.finish()]); return paramsBuffer; } function tensorElementBytes(type) { switch (String(type || '').toLowerCase()) { case 'float32': case 'int32': case 'uint32': return 4; case 'float16': case 'int16': case 'uint16': return 2; case 'int64': case 'uint64': return 8; case 'bool': case 'int8': case 'uint8': return 1; default: return 0; } } function tensorByteEstimate(tensorValue) { if (!tensorValue || typeof tensorValue !== 'object') return 0; if (tensorValue.data?.byteLength) return tensorValue.data.byteLength; if (tensorValue.cpuData?.byteLength) return tensorValue.cpuData.byteLength; const dims = Array.isArray(tensorValue.dims) ? tensorValue.dims : []; if (!dims.length) return 0; const elementBytes = tensorElementBytes(tensorValue.type || tensorValue.dataType); if (!elementBytes) return 0; return dims.reduce((a, b) => a * Number(b || 0), 1) * elementBytes; } function tensorDebugSummary(tensorValue) { if (!tensorValue || typeof tensorValue !== 'object') return 'null'; const ctor = tensorValue.constructor?.name || typeof tensorValue; const type = tensorValue.type || tensorValue.dataType || 'unknown'; const dims = Array.isArray(tensorValue.dims) ? tensorValue.dims.join('x') : 'unknown'; const location = tensorValue.location || tensorValue.dataLocation || 'unknown'; const methods = ['getData', 'dispose', 'release'].filter((name) => typeof tensorValue[name] === 'function').join('|') || 'none'; return `${ctor} type=${type} dims=${dims} location=${location} est=${formatBytes(tensorByteEstimate(tensorValue))} methods=${methods}`; } function yieldToRuntime() { return new Promise((resolve) => setTimeout(resolve, 0)); } async function tensorData(tensorValue) { if (typeof tensorValue?.getData === 'function') return await tensorValue.getData(); if (tensorValue?.data) return tensorValue.data; throw new Error('tensor has no getData() or data field'); } function firstOutput(outputs, preferredName) { if (outputs[preferredName]) return outputs[preferredName]; const keys = Object.keys(outputs); if (!keys.length) throw new Error('session returned no outputs'); return outputs[keys[0]]; } function asFloat16Array(data, label) { if (data instanceof Uint16Array) return data; if (ArrayBuffer.isView(data) && data.constructor?.name === 'Float16Array') { return new Uint16Array(data.buffer, data.byteOffset, data.length); } if (data instanceof Float32Array) return f32ToF16Array(data); throw new Error(`${label} has unsupported data type ${Object.prototype.toString.call(data)}`); } function asFloat32Array(data, label) { if (data instanceof Float32Array) return data; if (data instanceof Uint16Array) return f16ToF32Array(data); if (ArrayBuffer.isView(data) && data.constructor?.name === 'Float16Array') { return f16ToF32Array(new Uint16Array(data.buffer, data.byteOffset, data.length)); } throw new Error(`${label} has unsupported data type ${Object.prototype.toString.call(data)}`); } function transformerUsesFloat32(modelConfig) { const dtype = String(modelConfig.activation_dtype || modelConfig.dtype || '').toLowerCase(); return dtype === 'float32' || dtype === 'fp32'; } function mulberry32(seed) { let state = seed | 0; return function () { state = (state + 0x6D2B79F5) | 0; let value = Math.imul(state ^ (state >>> 15), 1 | state); value = (value + Math.imul(value ^ (value >>> 7), 61 | value)) ^ value; return ((value ^ (value >>> 14)) >>> 0) / 4294967296; }; } function normalSource(seed) { const rng = seed == null ? Math.random : mulberry32(seed); let spare = null; return function () { if (spare !== null) { const value = spare; spare = null; return value; } const u = Math.max(rng(), 1e-7); const v = rng(); const mag = Math.sqrt(-2.0 * Math.log(u)); spare = mag * Math.sin(2.0 * Math.PI * v); return mag * Math.cos(2.0 * Math.PI * v); }; } function generalizedTimeSnrShift(t, mu, sigma) { return Math.exp(mu) / (Math.exp(mu) + Math.pow(1 / t - 1, sigma)); } function computeEmpiricalMu(imageSeqLen, numSteps) { const a1 = 8.73809524e-05; const b1 = 1.89833333; const a2 = 0.00016927; const b2 = 0.45666666; if (imageSeqLen > 4300) return a2 * imageSeqLen + b2; const m200 = a2 * imageSeqLen + b2; const m10 = a1 * imageSeqLen + b1; const a = (m200 - m10) / 190.0; const b = m200 - 200.0 * a; return a * numSteps + b; } function getSchedule(numSteps, imageSeqLen) { const mu = computeEmpiricalMu(imageSeqLen, numSteps); const steps = []; for (let i = 0; i <= numSteps; i++) { const t = 1 - i / numSteps; steps.push(generalizedTimeSnrShift(t, mu, 1.0)); } return steps; } function makeImageIds(latentHeight, latentWidth) { const values = new Float32Array(latentHeight * latentWidth * 4); let offset = 0; for (let h = 0; h < latentHeight; h++) { for (let w = 0; w < latentWidth; w++) { values[offset++] = 0; values[offset++] = h; values[offset++] = w; values[offset++] = 0; } } return values; } function makeTextIds(seqLen) { const values = new Float32Array(seqLen * 4); let offset = 0; for (let i = 0; i < seqLen; i++) { values[offset++] = 0; values[offset++] = 0; values[offset++] = 0; values[offset++] = i; } return values; } function makeNoiseF32(imageSeqLen, channels, seed) { const nextNormal = normalSource(seed); const values = new Float32Array(imageSeqLen * channels); for (let i = 0; i < values.length; i++) values[i] = nextNormal(); return values; } function roundUpToMultiple(value, multiple) { const safeMultiple = Math.max(1, Math.trunc(Number(multiple || 1))); return Math.ceil(Number(value || 0) / safeMultiple) * safeMultiple; } function attentionMaskActiveTokenCount(data, seqLen) { if (data === undefined || data === null || typeof data === 'string') return seqLen; let values = null; if (data instanceof BigInt64Array || data instanceof BigUint64Array) { values = data; } else if (data instanceof ArrayBuffer) { if (data.byteLength !== seqLen * 8) return seqLen; values = new BigInt64Array(data); } else if (ArrayBuffer.isView(data)) { values = data; } else if (Array.isArray(data)) { values = data; } if (!values || values.length < seqLen) return seqLen; for (let index = seqLen - 1; index >= 0; index--) { if (Number(values[index]) !== 0) return index + 1; } return 1; } function coerceLatentF32(data, expectedValues, label) { let values = null; if (data instanceof Float32Array) { values = data; } else if (data instanceof Uint16Array) { values = f16ToF32Array(data); } else if (ArrayBuffer.isView(data)) { values = new Float32Array(data.length); for (let i = 0; i < data.length; i++) values[i] = Number(data[i]); } else if (data instanceof ArrayBuffer) { if (data.byteLength === expectedValues * 4) values = new Float32Array(data); else if (data.byteLength === expectedValues * 2) values = f16ToF32Array(new Uint16Array(data)); } else if (Array.isArray(data)) { values = Float32Array.from(data, Number); } if (!(values instanceof Float32Array)) { throw new Error(`${label} must be a Float32Array, Uint16Array, ArrayBuffer, typed array, or number array`); } if (values.length !== expectedValues) { throw new Error(`${label} has ${values.length} values, expected ${expectedValues}`); } return values; } function logArrayStats(stage, label, values) { const stats = arrayStats(values); const finiteLabel = stats.nonFinite ? ` nonfinite=${stats.nonFinite}` : ''; log(stage, `${label}: min=${stats.min.toFixed(4)} max=${stats.max.toFixed(4)} mean=${stats.mean.toFixed(4)} std=${stats.std.toFixed(4)}${finiteLabel}`); return stats; } function arrayStats(values) { let min = Infinity; let max = -Infinity; let sum = 0; let sumSquares = 0; let finite = 0; let nonFinite = 0; for (let i = 0; i < values.length; i++) { const value = values[i]; if (!Number.isFinite(value)) { nonFinite += 1; continue; } if (value < min) min = value; if (value > max) max = value; sum += value; sumSquares += value * value; finite += 1; } if (!finite) return { min: NaN, max: NaN, mean: NaN, std: NaN, count: values.length, finite, nonFinite }; const mean = sum / finite; const variance = Math.max(0, sumSquares / finite - mean * mean); return { min, max, mean, std: Math.sqrt(variance), count: values.length, finite, nonFinite }; } function assertFiniteStats(stats, label) { if (stats.nonFinite) throw new Error(`${label} contains ${stats.nonFinite} non-finite values`); } function clampFloat16ArrayInPlace(values, limit) { if (!limit || limit <= 0) return 0; let clipped = 0; for (let i = 0; i < values.length; i++) { const value = float16BitsToFloat32(values[i]); if (value > limit) { values[i] = float32ToFloat16Bits(limit); clipped += 1; } else if (value < -limit) { values[i] = float32ToFloat16Bits(-limit); clipped += 1; } } return clipped; } function zeroMaskedTextContextInPlace(values, attentionMask, seqLen, contextDim) { let zeroedTokens = 0; for (let token = 0; token < seqLen; token++) { if (attentionMask[token] !== 0n) continue; values.fill(0, token * contextDim, (token + 1) * contextDim); zeroedTokens += 1; } return zeroedTokens; } async function loadInt64Binary(url, expectedValues, label) { const response = await fetch(url); if (!response.ok) throw new Error(`Failed to fetch ${label}: HTTP ${response.status}`); const buffer = await response.arrayBuffer(); if (buffer.byteLength !== expectedValues * 8) throw new Error(`${label} has ${buffer.byteLength} bytes, expected ${expectedValues * 8}`); return new BigInt64Array(buffer); } function coerceInt64Array(data, expectedValues, label) { let values = null; if (data instanceof BigInt64Array) { values = data; } else if (data instanceof BigUint64Array) { values = new BigInt64Array(data.length); for (let i = 0; i < data.length; i++) values[i] = BigInt.asIntN(64, data[i]); } else if (data instanceof ArrayBuffer) { if (data.byteLength !== expectedValues * 8) throw new Error(`${label} has ${data.byteLength} bytes, expected ${expectedValues * 8}`); values = new BigInt64Array(data); } else if (ArrayBuffer.isView(data)) { values = new BigInt64Array(data.length); for (let i = 0; i < data.length; i++) values[i] = BigInt(Math.trunc(Number(data[i]))); } else if (Array.isArray(data)) { values = new BigInt64Array(data.length); for (let i = 0; i < data.length; i++) values[i] = BigInt(Math.trunc(Number(data[i]))); } if (!(values instanceof BigInt64Array)) { throw new Error(`${label} must be a BigInt64Array, ArrayBuffer, typed array, or number array`); } if (values.length !== expectedValues) { throw new Error(`${label} has ${values.length} values, expected ${expectedValues}`); } return values; } async function loadFloat16Binary(url, expectedValues, label) { const response = await fetch(url); if (!response.ok) throw new Error(`Failed to fetch ${label}: HTTP ${response.status}`); const buffer = await response.arrayBuffer(); if (buffer.byteLength !== expectedValues * 2) throw new Error(`${label} has ${buffer.byteLength} bytes, expected ${expectedValues * 2}`); return new Uint16Array(buffer); } async function loadOptionalFloat16Binary(url, expectedValues, label) { if (!url) return null; try { return await loadFloat16Binary(url, expectedValues, label); } catch (err) { log(label, `optional fp16 binary unavailable: ${errorDetail(err)}`); return null; } } async function loadFloat32Binary(url, expectedValues, label) { const response = await fetch(url); if (!response.ok) throw new Error(`Failed to fetch ${label}: HTTP ${response.status}`); const buffer = await response.arrayBuffer(); if (buffer.byteLength !== expectedValues * 4) throw new Error(`${label} has ${buffer.byteLength} bytes, expected ${expectedValues * 4}`); return new Float32Array(buffer); } async function resolveExternalData(modelConfig, stage) { if (!modelConfig.manifest) { return []; } const manifestUrl = `${modelBaseUrl}/${modelConfig.manifest}`; let manifest; try { manifest = await fetchJsonChecked(manifestUrl, modelConfig.manifest); } catch (err) { log(stage, `no manifest at ${modelConfig.manifest}; treating as inline (no external data)`); return []; } const externalData = []; for (let index = 0; index < manifest.length; index++) { if (abortRequested) throw new Error('ABORTED_BY_CLIENT'); const name = manifest[index]; log(stage, `registering shard ${index + 1}/${manifest.length}: ${name}`); externalData.push({ path: name, data: `${modelBaseUrl}/${name}` }); } return externalData; } async function createSession(modelConfig, stage, freeDimensionOverrides = {}) { const cacheSessions = Boolean(freeDimensionOverrides.__cacheSessions); delete freeDimensionOverrides.__cacheSessions; const preferredOutputLocation = freeDimensionOverrides.__preferredOutputLocation; delete freeDimensionOverrides.__preferredOutputLocation; const requestedProvider = runtimeOptions.executionProvider || modelConfig.execution_provider || 'webgpu'; const graphCaptureRequested = requestedProvider === 'webgpu' && shouldUseGraphCapture(stage); const cacheKey = JSON.stringify({ file: modelConfig.file, executionProvider: requestedProvider, webnnDeviceType: runtimeOptions.webnnDeviceType || 'gpu', enableGraphCapture: graphCaptureRequested, preferredOutputLocation, freeDimensionOverrides, }); if (cacheSessions && sessionCache.has(cacheKey)) { const cached = sessionCache.get(cacheKey); log(stage, `reusing cached WEBGPU session for ${modelConfig.file}; inputs=${(cached.inputNames || []).join(',') || 'unknown'} outputs=${(cached.outputNames || []).join(',') || 'unknown'}`); return cached; } logMemory(stage, 'before resolve external data'); let externalData = await resolveExternalData(modelConfig, stage); let executionProviders; if (requestedProvider === 'webnn') { executionProviders = [{ name: 'webnn', deviceType: runtimeOptions.webnnDeviceType || 'gpu', powerPreference: 'high-performance', }]; } else if (requestedProvider === 'webgpu') { executionProviders = ['webgpu']; } else { throw new Error(`${stage} requested unsupported browser execution provider ${requestedProvider}`); } const baseFreeDimensions = { batch: 1, text_seq: config.text_seq_len || 512, }; const options = { externalData, graphOptimizationLevel: modelConfig.graph_optimization_level || 'disabled', freeDimensionOverrides: { ...baseFreeDimensions, ...freeDimensionOverrides }, enableGraphCapture: graphCaptureRequested, ...(preferredOutputLocation ? { preferredOutputLocation } : {}), enableMemPattern: false, enableCpuMemArena: false, executionProviders, extra: { session: { 'optimization.disable_specified_optimizers': 'MemcpyTransformer,TransformerMemcpy', }, }, }; const quant = modelConfig.quantization ? ` quant=${modelConfig.quantization.algorithm || 'unknown'} accuracy=${modelConfig.quantization.accuracy_level ?? 'n/a'}` : ''; log(stage, `creating ${requestedProvider.toUpperCase()} session for ${modelConfig.file}; freeDims=${JSON.stringify(options.freeDimensionOverrides)} graphCapture=${options.enableGraphCapture}${quant}`); logMemory(stage, 'before session create'); const t0 = performance.now(); let session; try { session = await ort.InferenceSession.create(`${modelBaseUrl}/${modelConfig.file}`, options); } catch (err) { if (!options.enableGraphCapture || !String(errorDetail(err)).includes('graph capture')) { throw err; } log(stage, `graph capture unavailable for ${modelConfig.file}; retrying without graph capture: ${errorDetail(err)}`); options.enableGraphCapture = false; session = await ort.InferenceSession.create(`${modelBaseUrl}/${modelConfig.file}`, options); session.__flux2GraphCaptureFallback = true; } log(stage, `session created in ${((performance.now() - t0) / 1000).toFixed(3)}s; inputs=${(session.inputNames || []).join(',') || 'unknown'} outputs=${(session.outputNames || []).join(',') || 'unknown'}`); externalData = null; options.externalData = null; if (globalThis.gc) globalThis.gc(); logMemory(stage, 'after session create + gc'); if (cacheSessions) { session.__flux2CacheKey = cacheKey; session.__flux2Cached = true; sessionCache.set(cacheKey, session); log(stage, `cached WEBGPU session for ${modelConfig.file}`); } return session; } async function releaseSession(session, stage) { if (!session) return; if (session.__flux2Cached) { log(stage, `releaseSession skipped; cached session retained for ${session.__flux2CacheKey}`); return; } log(stage, `releaseSession start: hasRelease=${typeof session.release === 'function'}`); logMemory(stage, 'before session release'); const t0 = performance.now(); if (typeof session.release === 'function') await session.release(); log(stage, `releaseSession release() returned in ${((performance.now() - t0) / 1000).toFixed(3)}s`); if (globalThis.gc) globalThis.gc(); logMemory(stage, 'after session release + gc'); for (let i = 0; i < 5; i++) { await yieldToRuntime(); if (globalThis.gc) globalThis.gc(); logMemory(stage, `after release yield ${i + 1}/5`); } log(stage, 'session released and runtime yields complete'); } async function releaseTensorValue(tensorValue) { if (!tensorValue || typeof tensorValue !== 'object') return 'none'; try { if (typeof tensorValue.dispose === 'function') { await tensorValue.dispose(); return 'dispose'; } else if (typeof tensorValue.release === 'function') { await tensorValue.release(); return 'release'; } } catch (err) { console.warn(`[runtime] tensor release ignored: ${errorDetail(err)}`); return 'error'; } return 'no-method'; } async function releaseTensorMap(tensors, stage = 'runtime', label = 'tensors') { if (!tensors || typeof tensors !== 'object') { log(stage, `releaseTensorMap ${label}: empty`); return; } const entries = Object.entries(tensors); const totalBytes = entries.reduce((sum, [, value]) => sum + tensorByteEstimate(value), 0); log(stage, `releaseTensorMap ${label}: count=${entries.length} est=${formatBytes(totalBytes)} details=${entries.map(([name, value]) => `${name}:{${tensorDebugSummary(value)}}`).join(' ; ')}`); const counts = {}; const values = Object.values(tensors); for (let i = 0; i < values.length; i++) { const method = await releaseTensorValue(values[i]); counts[method] = (counts[method] || 0) + 1; } log(stage, `releaseTensorMap ${label}: releaseMethods=${JSON.stringify(counts)}`); } async function customKernelStatus(options = {}) { if (!options.force && customKernelState.checked && customKernelState.lastStatus) { return customKernelState.lastStatus; } const base = runtimeOptions.customKernelBaseUrl || '/runtime/custom_lowbit'; const wgslFeatures = navigator.gpu && navigator.gpu.wgslLanguageFeatures ? Array.from(navigator.gpu.wgslLanguageFeatures).sort() : []; const status = { enabled: runtimeOptions.customTransformerMode !== 'off', mode: runtimeOptions.customTransformerMode || 'off', baseUrl: base, gpu: Boolean(navigator.gpu), shaderF16: false, dp4a: wgslFeatures.includes('packed_4x8_integer_dot_product'), scriptsLoaded: typeof globalThis.runCustomSingleStreamMlpBench === 'function', checkpointsLoaded: { singleBlock: typeof globalThis.runCustomSingleStreamMlpBench === 'function', singleBlocksLoop: typeof globalThis.runCustomSingleStreamBlocksLoopBench === 'function', doubleBlock: typeof globalThis.runCustomDoubleStreamBlockBench === 'function', doubleBlocksLoop: typeof globalThis.runCustomDoubleStreamBlocksLoopBench === 'function', fullDenoise: typeof globalThis.runCustomFluxTransformerDenoise === 'function', }, transformerRuntimeLoaded: typeof globalThis.createCustomFluxTransformerRuntime === 'function', bundles: {}, fullTransformer: null, fullImageBackend: typeof globalThis.runCustomFluxTransformerDenoise === 'function' ? 'custom-lowbit-webgpu' : 'onnx-webgpu', notes: [], }; try { if (navigator.gpu) { const adapter = await navigator.gpu.requestAdapter({ powerPreference: 'high-performance' }); status.shaderF16 = Boolean(adapter && adapter.features && adapter.features.has('shader-f16')); } } catch (err) { status.notes.push(`adapter check failed: ${errorDetail(err)}`); } const bundlePaths = { qkv0: 'transformer_img_qkv0/manifest.json', singleLinear1: 'transformer_single_linear1_0/manifest.json', singleLinear2: 'transformer_single_linear2_0/manifest.json', singleBlock0: 'transformer_single_block0/manifest.json', }; const checkLegacyBundles = Boolean(options.checkLegacyBundles ?? runtimeOptions.checkLegacyCustomBundles); if (checkLegacyBundles) { for (const [key, path] of Object.entries(bundlePaths)) { try { const response = await fetch(`${base}/${path}`); if (!response.ok) { status.bundles[key] = { ok: false, status: response.status }; continue; } const manifest = await response.json(); status.bundles[key] = { ok: true, name: manifest.name || '', sourceNode: manifest.source_node || '', shape: manifest.shape || null, }; } catch (err) { status.bundles[key] = { ok: false, error: errorDetail(err) }; } } } else { status.notes.push('legacy checkpoint manifest probe skipped; set legacyStatus=1 for checkpoint diagnostics'); } const diagnosticBundlesReady = !checkLegacyBundles || Object.values(status.bundles).every((item) => item && item.ok); try { const response = await fetch(`${base}/full_transformer/manifest.json`); if (response.ok) { const manifest = await response.json(); status.fullTransformer = { ok: true, name: manifest.name || '', counts: manifest.counts || null, linears: Array.isArray(manifest.linears) ? manifest.linears.length : 0, qkNormGroups: Array.isArray(manifest.qk_norm_groups) ? manifest.qk_norm_groups.length : ( manifest.counts && Number.isFinite(manifest.counts.qk_norm_groups) ? manifest.counts.qk_norm_groups : 0 ), }; } else { status.fullTransformer = { ok: false, status: response.status }; } } catch (err) { status.fullTransformer = { ok: false, error: errorDetail(err) }; } status.available = Boolean( status.gpu && status.shaderF16 && status.dp4a && status.scriptsLoaded && status.transformerRuntimeLoaded && status.checkpointsLoaded.fullDenoise && status.fullTransformer && status.fullTransformer.ok, ); if (!status.available) { status.notes.push('custom transformer is not complete; generation uses ONNX WebGPU correctness backend'); status.fullImageBackend = 'onnx-webgpu'; } else { if (!diagnosticBundlesReady) { status.notes.push('legacy checkpoint manifests are missing; full-image custom denoise remains available'); } status.notes.push(status.checkpointsLoaded.fullDenoise ? 'custom low-bit WebGPU transformer denoise path is available; VAE decode remains on the split ONNX WebGPU path' : 'custom block checkpoints are available, but full denoise is not loaded; generation uses ONNX WebGPU correctness backend'); if (!status.checkpointsLoaded.fullDenoise) status.fullImageBackend = 'onnx-webgpu'; } customKernelState.available = status.available; customKernelState.checked = true; customKernelState.lastStatus = status; log('custom-kernels', JSON.stringify(status)); return status; } async function prepareCustomTransformerAssets(params = {}) { if (typeof globalThis.createCustomFluxTransformerRuntime !== 'function') { throw new Error('CUSTOM_TRANSFORMER_RUNTIME_NOT_LOADED'); } const baseUrl = `${runtimeOptions.customKernelBaseUrl || '/runtime/custom_lowbit'}/full_transformer/`; if (!customTransformerRuntimeCache.has(baseUrl)) { customTransformerRuntimeCache.set(baseUrl, globalThis.createCustomFluxTransformerRuntime({ baseUrl, preloadBinaries: false, })); } const runtime = await customTransformerRuntimeCache.get(baseUrl); if (typeof runtime.prepareAssets !== 'function') { throw new Error('CUSTOM_TRANSFORMER_ASSET_PREPARE_NOT_AVAILABLE'); } const start = performance.now(); const result = await runtime.prepareAssets(params); log('transformer-custom-webgpu', `custom asset prepare returned in ${((performance.now() - start) / 1000).toFixed(3)}s; summary=${JSON.stringify(result.summary || {})}`); return result; } async function getCustomTransformerRuntime() { if (typeof globalThis.createCustomFluxTransformerRuntime !== 'function') { throw new Error('CUSTOM_TRANSFORMER_RUNTIME_NOT_LOADED'); } const baseUrl = `${runtimeOptions.customKernelBaseUrl || '/runtime/custom_lowbit'}/full_transformer/`; if (!customTransformerRuntimeCache.has(baseUrl)) { customTransformerRuntimeCache.set(baseUrl, globalThis.createCustomFluxTransformerRuntime({ baseUrl, preloadBinaries: false, })); } return { baseUrl, runtime: await customTransformerRuntimeCache.get(baseUrl), }; } async function prepareCustomTransformerStageSetup(params = {}) { if (!config || !modelBaseUrl) throw new Error('Engine is not initialized'); const { runtime } = await getCustomTransformerRuntime(); const dispatchWarmup = params.dispatchWarmup === true; const prepareMethod = dispatchWarmup ? runtime.warmDispatch : runtime.prepareStage; if (typeof prepareMethod !== 'function') { throw new Error(dispatchWarmup ? 'CUSTOM_TRANSFORMER_DISPATCH_WARMUP_NOT_AVAILABLE' : 'CUSTOM_TRANSFORMER_STAGE_PREPARE_NOT_AVAILABLE'); } const outputWidth = Number(params.width || config.default_width); const outputHeight = Number(params.height || config.default_height); if (outputWidth % config.latent_downsample || outputHeight % config.latent_downsample) { throw new Error(`width and height must be divisible by ${config.latent_downsample}`); } const renderPlan = makePreviewRenderPlan( outputWidth, outputHeight, params.previewRenderMaxSize, previewRenderEnabledForParams(params), ); const width = renderPlan.renderWidth; const height = renderPlan.renderHeight; const latentWidth = width / config.latent_downsample; const latentHeight = height / config.latent_downsample; const imageSeqLen = latentWidth * latentHeight; const channels = config.latent_channels; const fullTextTokens = Number(config.text_seq_len || 512); const requestedTextTokens = Math.trunc(Number(params.textTokens || params.customTextTokenLimit || fullTextTokens)); const minTextTokens = Math.max(1, Math.trunc(Number(params.minTextTokens || params.customMinTextTokens || 32))); const textTokens = Math.min( fullTextTokens, Math.max(minTextTokens, roundUpToMultiple(Math.max(1, requestedTextTokens), 16)), ); if (imageSeqLen + textTokens > 4608) { throw new Error(`CUSTOM_TRANSFORMER_TOKEN_LIMIT: custom WebGPU path currently supports <=4608 joint tokens, got ${imageSeqLen + textTokens}`); } const contextF16 = new Uint16Array(textTokens * config.context_dim); const initialLatentF32 = new Float32Array(imageSeqLen * channels); const timesteps = Array.isArray(params.customTimesteps) && params.customTimesteps.length > 1 ? params.customTimesteps.map(Number) : getSchedule(params.numSteps || config.num_steps, imageSeqLen); const start = performance.now(); const result = await prepareMethod.call(runtime, { contextF16, initialLatentF32, imageTokens: imageSeqLen, textTokens, latentChannels: channels, imageWidth: latentWidth, timesteps: Array.from(timesteps), linearTileCols: params.customLinearTileCols, projTileCols: params.customProjTileCols, finalTileCols: params.customFinalTileCols, wChunkCols: params.customWChunkCols, linearRowBlock: params.customLinearRowBlock || params.linearRowBlock, singleQ4TileCols: params.customSingleQ4TileCols, singleQ4KChunk: params.customSingleQ4KChunk, singleQ4Dp4aTileCols: params.customSingleQ4Dp4aTileCols, attentionTileKeys: params.customAttentionTileKeys, attentionQueryRows: params.customAttentionQueryRows, singleAttentionKernel: params.customSingleAttentionKernel || params.singleAttentionKernel || params.attentionKernel, textContextCacheKey: "", persistentTextProjectionCache: false, textProjectionUrl: params.textProjectionUrl || '', singleLinear1Output: params.customSingleLinear1Output || params.singleLinear1Output, singleLinear1Q4Kernel: params.customSingleLinear1Q4Kernel || params.singleLinear1Q4Kernel, singleLinear1Q4ActivationScale: params.customSingleLinear1Q4ActivationScale || params.singleLinear1Q4ActivationScale, singleLinear1QkvBackend: params.customSingleLinear1QkvBackend || params.singleLinear1QkvBackend, singleLinear1MlpBackend: params.customSingleLinear1MlpBackend || params.singleLinear1MlpBackend, singleLinear2Backend: params.customSingleLinear2Backend || params.singleLinear2Backend, singleQkNormStorage: params.customSingleQkNormStorage || params.singleQkNormStorage, approxReusePredEvery: params.customApproxReusePredEvery || params.approxReusePredEvery, approxPredictionMode: params.customApproxPredictionMode || params.approxPredictionMode, approxAbScale: params.customApproxAbScale || params.approxAbScale, skipSingleBlocksFromStep: params.customSkipSingleBlocksFromStep || params.skipSingleBlocksFromStep, skipSingleBlocksFromBlock: params.customSkipSingleBlocksFromBlock || params.skipSingleBlocksFromBlock, reuseSingleTailFromStep: params.customReuseSingleTailFromStep || params.reuseSingleTailFromStep, reuseSingleTailFromBlock: params.customReuseSingleTailFromBlock || params.reuseSingleTailFromBlock, profilePhases: false, profileSingleParts: false, stepSubmitFusion: params.customStepSubmitFusion ?? params.stepSubmitFusion, singleComputePass: params.customSingleComputePass || params.singleComputePass, maxDoubleBlocks: params.customMaxDoubleBlocks || params.maxDoubleBlocks, maxSingleBlocks: params.customMaxSingleBlocks || params.maxSingleBlocks, }); log('transformer-custom-webgpu', `custom ${dispatchWarmup ? 'dispatch warmup' : 'stage prepare'} returned in ${((performance.now() - start) / 1000).toFixed(3)}s; load=${JSON.stringify(result.load || {})}`); return { ...result, shape: { width, height, latentWidth, latentHeight, imageSeqLen, textTokens, renderPlan, }, }; } async function prepareCustomTransformerTextProjection(params = {}) { if (!config || !modelBaseUrl) throw new Error('Engine is not initialized'); const { runtime } = await getCustomTransformerRuntime(); if (typeof runtime.runDenoise !== 'function') { throw new Error('CUSTOM_TRANSFORMER_TEXT_PROJECTION_PREPARE_NOT_AVAILABLE'); } const outputWidth = Number(params.width || config.default_width); const outputHeight = Number(params.height || config.default_height); if (outputWidth % config.latent_downsample || outputHeight % config.latent_downsample) { throw new Error(`width and height must be divisible by ${config.latent_downsample}`); } const renderPlan = makePreviewRenderPlan( outputWidth, outputHeight, params.previewRenderMaxSize, previewRenderEnabledForParams(params), ); const width = renderPlan.renderWidth; const height = renderPlan.renderHeight; const latentWidth = width / config.latent_downsample; const latentHeight = height / config.latent_downsample; const imageSeqLen = latentWidth * latentHeight; const channels = config.latent_channels; const fullTextTokens = Number(config.text_seq_len || 512); const inputIds = params.inputIds !== undefined ? coerceInt64Array(params.inputIds, fullTextTokens, 'input ids') : await loadInt64Binary(params.inputIdsUrl, fullTextTokens, 'input ids'); const attentionMask = params.attentionMask !== undefined ? coerceInt64Array(params.attentionMask, fullTextTokens, 'attention mask') : await loadInt64Binary(params.attentionMaskUrl, fullTextTokens, 'attention mask'); const activeTextTokens = attentionMaskActiveTokenCount(attentionMask, fullTextTokens); const requestedTextLimit = Math.max(0, Math.trunc(Number(params.customTextTokenLimit || params.textTokens || 0))); const minTextTokens = Math.max(1, Math.trunc(Number(params.customMinTextTokens || params.minTextTokens || 64))); let textTokens = fullTextTokens; if (requestedTextLimit > 0) { textTokens = Math.min( fullTextTokens, Math.max(minTextTokens, roundUpToMultiple(Math.min(activeTextTokens, requestedTextLimit), 16)), ); } if (imageSeqLen + textTokens > 4608) { throw new Error(`CUSTOM_TRANSFORMER_TOKEN_LIMIT: custom WebGPU path currently supports <=4608 joint tokens, got ${imageSeqLen + textTokens}`); } const ctx = await runTextEncoder({ ...params, inputIds, attentionMask, }); const contextF16 = textTokens < fullTextTokens ? ctx.slice(0, textTokens * config.context_dim) : ctx; const timesteps = Array.isArray(params.customTimesteps) && params.customTimesteps.length > 1 ? params.customTimesteps.map(Number) : getSchedule(params.numSteps || config.num_steps, imageSeqLen); const start = performance.now(); const result = await runtime.runDenoise({ contextF16, initialLatentF32: new Float32Array(imageSeqLen * channels), imageTokens: imageSeqLen, textTokens, latentChannels: channels, imageWidth: latentWidth, timesteps: Array.from(timesteps), linearTileCols: params.customLinearTileCols, projTileCols: params.customProjTileCols, finalTileCols: params.customFinalTileCols, wChunkCols: params.customWChunkCols, linearRowBlock: params.customLinearRowBlock || params.linearRowBlock, singleQ4TileCols: params.customSingleQ4TileCols, singleQ4KChunk: params.customSingleQ4KChunk, singleQ4Dp4aTileCols: params.customSingleQ4Dp4aTileCols, attentionTileKeys: params.customAttentionTileKeys, attentionQueryRows: params.customAttentionQueryRows, singleAttentionKernel: params.customSingleAttentionKernel || params.singleAttentionKernel || params.attentionKernel, textContextCacheKey: params.customTextProjectionCache === false || !params.promptCacheKey ? "" : `${params.promptCacheKey}|txt=${textTokens}`, persistentTextProjectionCache: params.persistentTextContextCache !== false, textProjectionUrl: params.textProjectionUrl || '', singleLinear1Output: params.customSingleLinear1Output || params.singleLinear1Output, singleLinear1Q4Kernel: params.customSingleLinear1Q4Kernel || params.singleLinear1Q4Kernel, singleLinear1Q4ActivationScale: params.customSingleLinear1Q4ActivationScale || params.singleLinear1Q4ActivationScale, singleLinear1QkvBackend: params.customSingleLinear1QkvBackend || params.singleLinear1QkvBackend, singleLinear1MlpBackend: params.customSingleLinear1MlpBackend || params.singleLinear1MlpBackend, singleLinear2Backend: params.customSingleLinear2Backend || params.singleLinear2Backend, singleQkNormStorage: params.customSingleQkNormStorage || params.singleQkNormStorage, approxReusePredEvery: params.customApproxReusePredEvery || params.approxReusePredEvery, approxPredictionMode: params.customApproxPredictionMode || params.approxPredictionMode, approxAbScale: params.customApproxAbScale || params.approxAbScale, skipSingleBlocksFromStep: params.customSkipSingleBlocksFromStep || params.skipSingleBlocksFromStep, skipSingleBlocksFromBlock: params.customSkipSingleBlocksFromBlock || params.skipSingleBlocksFromBlock, reuseSingleTailFromStep: params.customReuseSingleTailFromStep || params.reuseSingleTailFromStep, reuseSingleTailFromBlock: params.customReuseSingleTailFromBlock || params.reuseSingleTailFromBlock, profilePhases: false, profileSingleParts: false, stepSubmitFusion: false, singleComputePass: false, maxDoubleBlocks: params.customMaxDoubleBlocks || params.maxDoubleBlocks, maxSingleBlocks: params.customMaxSingleBlocks || params.maxSingleBlocks, textProjectionOnly: true, }); log('transformer-custom-webgpu', `custom text projection prepare returned in ${((performance.now() - start) / 1000).toFixed(3)}s; load=${JSON.stringify(result.load || {})}`); return { ...result, shape: { width, height, latentWidth, latentHeight, imageSeqLen, textTokens, activeTextTokens, renderPlan, }, textEncoder: lastTextEncoderDetails ? {...lastTextEncoderDetails} : null, }; } async function init(baseUrl, options = {}) { try { clearStaticGpuTensorCache(); modelBaseUrl = baseUrl; runtimeOptions = { enableGraphCapture: Boolean(options.enableGraphCapture), graphCaptureStages: optionList(options.graphCaptureStages), profileWebGpu: Boolean(options.profileWebGpu), profileTopK: Number(options.profileTopK || 12), textContextCacheSize: Number(options.textContextCacheSize || 4), vaeEncoderCacheSize: Number(options.vaeEncoderCacheSize || 3), persistentTextContextCache: options.persistentTextContextCache !== false, persistentTextContextCacheSize: Number(options.persistentTextContextCacheSize || 8), persistentVaeEncoderCache: options.persistentVaeEncoderCache !== false, persistentVaeEncoderCacheSize: Number(options.persistentVaeEncoderCacheSize || 8), checkLegacyCustomBundles: Boolean(options.checkLegacyCustomBundles), enableFusedTransformer: Boolean(options.enableFusedTransformer), executionProvider: options.executionProvider || 'webgpu', webnnDeviceType: options.webnnDeviceType || 'gpu', customKernelBaseUrl: options.customKernelBaseUrl || '/runtime/custom_lowbit', customTransformerMode: options.customTransformerMode || 'off', }; if (runtimeOptions.profileWebGpu) { webGpuProfileEvents = []; ort.env.webgpu.profiling = { mode: 'default', ondata: (data) => { webGpuProfileEvents.push({ ...data, label: activeProfileLabel || 'unscoped' }); }, }; } else { ort.env.webgpu.profiling = { mode: 'off' }; } config = await fetchJsonChecked(`${modelBaseUrl}/flux2-config.json`, 'flux2-config.json'); await customKernelStatus({ force: true }); log('ready', `FLUX.2 staged browser engine initialized; ort=${ortVersion} provider=${runtimeOptions.executionProvider} webnnDevice=${runtimeOptions.webnnDeviceType} gcAvailable=${typeof globalThis.gc === 'function'} gpuAvailable=${!!navigator.gpu} mlAvailable=${!!navigator.ml} graphCapture=${runtimeOptions.enableGraphCapture} graphCaptureStages=${runtimeOptions.graphCaptureStages.join(',') || 'all'} fusedTransformer=${runtimeOptions.enableFusedTransformer} profileWebGpu=${runtimeOptions.profileWebGpu} customTransformerMode=${runtimeOptions.customTransformerMode} customKernels=${customKernelState.available}`); logMemory('ready', 'after config load'); return { backend: 'webgpu', config, customKernels: customKernelState.lastStatus }; } catch (err) { log('error', `Init failed: ${errorDetail(err)}`); throw err; } } async function runTextEncoder(params) { const stage = 'text-encoder'; const modelConfig = config.models.text_encoder; const seqLen = config.text_seq_len; const cacheKey = params.cacheTextContext && params.promptCacheKey ? String(params.promptCacheKey) : ''; const expectedCtxValues = seqLen * config.context_dim; const persistentCacheKey = cacheKey && params.persistentTextContextCache !== false ? makePersistentTextContextKey(modelConfig, cacheKey, seqLen, config.context_dim) : ''; lastTextEncoderDetails = { cacheKey: cacheKey ? cacheKey.slice(0, 12) : '', source: 'miss', persistent: Boolean(persistentCacheKey), persistentMs: 0, }; if (cacheKey && textContextCache.has(cacheKey)) { const cached = textContextCache.get(cacheKey); lastTextEncoderDetails = { ...lastTextEncoderDetails, source: 'memory', }; log(stage, `reusing cached text context ${cacheKey.slice(0, 12)} values=${cached.length}`); return cached; } if (persistentCacheKey) { const persistentStart = performance.now(); const cached = await loadPersistentTextContext(persistentCacheKey, expectedCtxValues); lastTextEncoderDetails.persistentMs = performance.now() - persistentStart; if (cached) { rememberTextContext(cacheKey, cached); lastTextEncoderDetails = { ...lastTextEncoderDetails, source: 'indexeddb', }; log(stage, `reusing IndexedDB text context ${cacheKey.slice(0, 12)} values=${cached.length} read=${lastTextEncoderDetails.persistentMs.toFixed(1)}ms`); return cached; } } if (params.textContextUrl) { const staticStart = performance.now(); const staticCtx = await loadOptionalFloat16Binary(params.textContextUrl, expectedCtxValues, stage); lastTextEncoderDetails.persistentMs = performance.now() - staticStart; if (staticCtx) { if (cacheKey) rememberTextContext(cacheKey, staticCtx); if (persistentCacheKey) { savePersistentTextContext(persistentCacheKey, staticCtx).catch((err) => { log('text-cache', `could not persist static text context: ${errorDetail(err)}`); }); } lastTextEncoderDetails = { ...lastTextEncoderDetails, source: 'static-file', }; log(stage, `loaded static text context ${cacheKey ? cacheKey.slice(0, 12) : ''} values=${staticCtx.length} read=${lastTextEncoderDetails.persistentMs.toFixed(1)}ms`); return staticCtx; } } const inputIds = params.inputIds !== undefined ? coerceInt64Array(params.inputIds, seqLen, 'input ids') : await loadInt64Binary(params.inputIdsUrl, seqLen, 'input ids'); const attentionMask = params.attentionMask !== undefined ? coerceInt64Array(params.attentionMask, seqLen, 'attention mask') : await loadInt64Binary(params.attentionMaskUrl, seqLen, 'attention mask'); let session = null; let feeds = null; let outputs = null; try { session = await createSession(modelConfig, stage, { __cacheSessions: params.cacheSessions }); feeds = { input_ids: tensor('int64', inputIds, [1, seqLen]), attention_mask: tensor('int64', attentionMask, [1, seqLen]), }; logMemory(stage, 'before session.run'); outputs = await profiledRun(session, feeds, null, stage, 'text-encoder'); logMemory(stage, 'after session.run'); const ctxRaw = asFloat16Array(await tensorData(firstOutput(outputs, 'ctx')), 'text encoder output'); const ctx = new Uint16Array(ctxRaw); if (ctx.length !== expectedCtxValues) throw new Error(`text encoder output has ${ctx.length} values, expected ${expectedCtxValues}`); log(stage, `output tensor: ${tensorDebugSummary(firstOutput(outputs, 'ctx'))}`); if (config.zero_padded_context) { const zeroedTokens = zeroMaskedTextContextInPlace(ctx, attentionMask, seqLen, config.context_dim); if (zeroedTokens) log(stage, `zeroed ${zeroedTokens} padded context tokens`); } const ctxClip = config.context_clip ?? 0; const clipped = clampFloat16ArrayInPlace(ctx, ctxClip); if (clipped) log(stage, `clipped ${clipped} context values to +/-${ctxClip}`); logArrayStats(stage, 'context', f16ToF32Array(ctx)); log(stage, `produced context tensor ${ctx.length} fp16 values`); if (cacheKey) { rememberTextContext(cacheKey, ctx); if (persistentCacheKey) { const persistentStart = performance.now(); const saved = await savePersistentTextContext(persistentCacheKey, ctx); lastTextEncoderDetails.persistentMs = performance.now() - persistentStart; lastTextEncoderDetails = { ...lastTextEncoderDetails, source: saved ? 'encoded+indexeddb-save' : 'encoded', }; log(stage, `IndexedDB text context save ${saved ? 'complete' : 'skipped'} in ${lastTextEncoderDetails.persistentMs.toFixed(1)}ms`); } else { lastTextEncoderDetails = { ...lastTextEncoderDetails, source: 'encoded', }; } } return ctx; } finally { await releaseTensorMap(outputs, stage, 'outputs'); await releaseTensorMap(feeds, stage, 'feeds'); outputs = null; feeds = null; await releaseSession(session, stage); } } async function runTransformerGpuResident(ctx, params, width, height) { const stage = 'transformer-gpu'; const modelConfig = config.models.transformer; if (transformerUsesFloat32(modelConfig)) { throw new Error('GPU-resident denoise path currently supports fp16 transformer activations only'); } const latentWidth = width / config.latent_downsample; const latentHeight = height / config.latent_downsample; const imageSeqLen = latentWidth * latentHeight; const channels = config.latent_channels; const latentElements = imageSeqLen * channels; const ctxIds = makeTextIds(config.text_seq_len); const xIds = makeImageIds(latentHeight, latentWidth); const xInitialF32 = params.initialLatentF32 !== undefined ? coerceLatentF32(params.initialLatentF32, latentElements, 'initial latent') : makeNoiseF32(imageSeqLen, channels, params.seed ?? 42); const xInitialF16 = f32ToF16Array(xInitialF32); const timesteps = Array.isArray(params.customTimesteps) && params.customTimesteps.length > 1 ? params.customTimesteps.map(Number) : getSchedule(params.numSteps || config.num_steps, imageSeqLen); let session = null; const gpuResources = []; const paramsBuffers = []; try { session = await createSession(modelConfig, stage, { __cacheSessions: params.cacheSessions, image_seq: imageSeqLen, latent_height: latentHeight, latent_width: latentWidth, height, width, }); const xResource = createGpuTensorFromTypedArray(xInitialF16, 'float16', [1, imageSeqLen, channels]); const predResource = createGpuTensorFromTypedArray(new Uint16Array(latentElements), 'float16', [1, imageSeqLen, channels]); const xIdsResource = getCachedGpuTensorResource( `x_ids:${latentHeight}x${latentWidth}`, () => createGpuTensorFromTypedArray(xIds, 'float32', [1, imageSeqLen, 4]), stage, ); const ctxCacheKey = params.cacheTextContext && params.promptCacheKey ? String(params.promptCacheKey) : ''; const ctxResource = ctxCacheKey ? getCachedGpuTensorResource( `ctx:${ctxCacheKey}`, () => createGpuTensorFromTypedArray(ctx, 'float16', [1, config.text_seq_len, config.context_dim]), stage, ) : createGpuTensorFromTypedArray(ctx, 'float16', [1, config.text_seq_len, config.context_dim]); const ctxIdsResource = getCachedGpuTensorResource( `ctx_ids:${config.text_seq_len}`, () => createGpuTensorFromTypedArray(ctxIds, 'float32', [1, config.text_seq_len, 4]), stage, ); const timestepResource = createGpuTensorFromTypedArray(new Uint16Array(2), 'float16', [1]); gpuResources.push(xResource, predResource, timestepResource); if (!ctxCacheKey) gpuResources.push(ctxResource); const feeds = { x: xResource.tensor, x_ids: xIdsResource.tensor, timesteps: timestepResource.tensor, ctx: ctxResource.tensor, ctx_ids: ctxIdsResource.tensor, }; const fetches = { pred: predResource.tensor }; log(stage, `denoising ${width}x${height}, tokens=${imageSeqLen}, steps=${timesteps.length - 1}, gpu-resident=true`); for (let step = 0; step < timesteps.length - 1; step++) { if (abortRequested) throw new Error('ABORTED_BY_CLIENT'); const tCurr = timesteps[step]; const tPrev = timesteps[step + 1]; const timestepF16 = f32ToF16Array(new Float32Array([tCurr, 0])); getWebGpuDevice().queue.writeBuffer(timestepResource.buffer, 0, timestepF16.buffer, timestepF16.byteOffset, timestepF16.byteLength); const runStart = performance.now(); await profiledRun(session, feeds, fetches, stage, `step-${step + 1}`); const runMs = performance.now() - runStart; const updateStart = performance.now(); paramsBuffers.push(dispatchLatentUpdate(xResource.buffer, predResource.buffer, latentElements, tPrev - tCurr)); const updateMs = performance.now() - updateStart; log(stage, `session.run step ${step + 1} returned in ${(runMs / 1000).toFixed(3)}s; gpu latent update queued in ${(updateMs / 1000).toFixed(3)}s`); } const readStart = performance.now(); const latentBuffer = await readGpuBuffer(xResource.buffer, latentElements * 2); log(stage, `final latent readback returned in ${((performance.now() - readStart) / 1000).toFixed(3)}s`); return { latent: new Uint16Array(latentBuffer), latentWidth, latentHeight }; } finally { for (const buffer of paramsBuffers) { try { buffer.destroy(); } catch { } } for (const resource of gpuResources) { try { resource.buffer.destroy(); } catch { } } await releaseSession(session, stage); } } async function runTransformerFusedGpuResident(ctx, params, width, height) { const stage = 'transformer-fused-gpu'; const modelConfig = config.models.transformer_fused_4step; if (!modelConfig) throw new Error('fused transformer model is not configured'); if (Number(params.numSteps || config.num_steps) !== 4) throw new Error('fused transformer requires exactly 4 denoise steps'); if (transformerUsesFloat32(config.models.transformer)) { throw new Error('fused GPU denoise path currently supports fp16 transformer activations only'); } const latentWidth = width / config.latent_downsample; const latentHeight = height / config.latent_downsample; const imageSeqLen = latentWidth * latentHeight; const channels = config.latent_channels; const latentElements = imageSeqLen * channels; const ctxIds = makeTextIds(config.text_seq_len); const xIds = makeImageIds(latentHeight, latentWidth); const xInitialF32 = params.initialLatentF32 !== undefined ? coerceLatentF32(params.initialLatentF32, latentElements, 'initial latent') : makeNoiseF32(imageSeqLen, channels, params.seed ?? 42); const xInitialF16 = f32ToF16Array(xInitialF32); const timesteps = getSchedule(4, imageSeqLen); let session = null; const gpuResources = []; try { session = await createSession(modelConfig, stage, { __cacheSessions: params.cacheSessions, image_seq: imageSeqLen, latent_height: latentHeight, latent_width: latentWidth, height, width, }); const xResource = createGpuTensorFromTypedArray(xInitialF16, 'float16', [1, imageSeqLen, channels]); const latentResource = createGpuTensorFromTypedArray(new Uint16Array(latentElements), 'float16', [1, imageSeqLen, channels]); const xIdsResource = getCachedGpuTensorResource( `x_ids:${latentHeight}x${latentWidth}`, () => createGpuTensorFromTypedArray(xIds, 'float32', [1, imageSeqLen, 4]), stage, ); const ctxCacheKey = params.cacheTextContext && params.promptCacheKey ? String(params.promptCacheKey) : ''; const ctxResource = ctxCacheKey ? getCachedGpuTensorResource( `ctx:${ctxCacheKey}`, () => createGpuTensorFromTypedArray(ctx, 'float16', [1, config.text_seq_len, config.context_dim]), stage, ) : createGpuTensorFromTypedArray(ctx, 'float16', [1, config.text_seq_len, config.context_dim]); const ctxIdsResource = getCachedGpuTensorResource( `ctx_ids:${config.text_seq_len}`, () => createGpuTensorFromTypedArray(ctxIds, 'float32', [1, config.text_seq_len, 4]), stage, ); gpuResources.push(xResource, latentResource); if (!ctxCacheKey) gpuResources.push(ctxResource); const feeds = { x: xResource.tensor, x_ids: xIdsResource.tensor, ctx: ctxResource.tensor, ctx_ids: ctxIdsResource.tensor, }; for (let step = 0; step < 4; step++) { feeds[`timestep_${step}`] = tensor('float16', f32ToF16Array(new Float32Array([timesteps[step]])), [1]); feeds[`dt_${step}`] = tensor('float32', new Float32Array([timesteps[step + 1] - timesteps[step]]), [1]); } log(stage, `denoising ${width}x${height}, tokens=${imageSeqLen}, steps=4, fused=true gpu-resident=true`); const runStart = performance.now(); await profiledRun(session, feeds, { latent: latentResource.tensor }, stage, 'fused-4step'); log(stage, `fused session.run returned in ${((performance.now() - runStart) / 1000).toFixed(3)}s`); const readStart = performance.now(); const latentBuffer = await readGpuBuffer(latentResource.buffer, latentElements * 2); log(stage, `final latent readback returned in ${((performance.now() - readStart) / 1000).toFixed(3)}s`); return { latent: new Uint16Array(latentBuffer), latentWidth, latentHeight }; } finally { for (const resource of gpuResources) { try { resource.buffer.destroy(); } catch { } } await releaseSession(session, stage); } } async function runTransformerCustomLowbitWebGpu(ctx, params, width, height) { const stage = 'transformer-custom-webgpu'; const latentWidth = width / config.latent_downsample; const latentHeight = height / config.latent_downsample; const imageSeqLen = latentWidth * latentHeight; const channels = config.latent_channels; const fullTextTokens = config.text_seq_len; const activeTextTokens = attentionMaskActiveTokenCount(params.attentionMask, fullTextTokens); const requestedTextLimit = Math.max(0, Math.trunc(Number(params.customTextTokenLimit || 0))); const minTextTokens = Math.max(1, Math.trunc(Number(params.customMinTextTokens || 64))); let textTokens = fullTextTokens; if (requestedTextLimit > 0) { textTokens = Math.min( fullTextTokens, Math.max(minTextTokens, roundUpToMultiple(Math.min(activeTextTokens, requestedTextLimit), 16)), ); } const contextF16 = textTokens < fullTextTokens ? ctx.slice(0, textTokens * config.context_dim) : ctx; const timesteps = Array.isArray(params.customTimesteps) && params.customTimesteps.length > 1 ? params.customTimesteps.map(Number) : getSchedule(params.numSteps || config.num_steps, imageSeqLen); const initialLatentF32 = params.initialLatentF32 !== undefined ? coerceLatentF32(params.initialLatentF32, imageSeqLen * channels, 'initial latent') : makeNoiseF32(imageSeqLen, channels, params.seed ?? 42); const baseUrl = `${runtimeOptions.customKernelBaseUrl || '/runtime/custom_lowbit'}/full_transformer/`; if (imageSeqLen + textTokens > 4608) { throw new Error(`CUSTOM_TRANSFORMER_TOKEN_LIMIT: custom WebGPU path currently supports <=4608 joint tokens, got ${imageSeqLen + textTokens}`); } if (typeof globalThis.createCustomFluxTransformerRuntime !== 'function') { throw new Error('CUSTOM_TRANSFORMER_RUNTIME_NOT_LOADED'); } if (!customTransformerRuntimeCache.has(baseUrl)) { customTransformerRuntimeCache.set(baseUrl, globalThis.createCustomFluxTransformerRuntime({ baseUrl, preloadBinaries: false, })); } const runtime = await customTransformerRuntimeCache.get(baseUrl); const status = runtime.status(); if (!status.ready) { throw new Error(`CUSTOM_TRANSFORMER_RUNTIME_NOT_READY: ${JSON.stringify(status.notes || [])}`); } log(stage, `custom low-bit WebGPU denoise starting ${width}x${height}, imageTokens=${imageSeqLen}, textTokens=${textTokens}/${fullTextTokens}, activeTextTokens=${activeTextTokens}, steps=${timesteps.length - 1}`); const start = performance.now(); const result = await runtime.runDenoise({ contextF16, initialLatentF32, imageTokens: imageSeqLen, textTokens, latentChannels: channels, imageWidth: latentWidth, timesteps: Array.from(timesteps), linearTileCols: params.customLinearTileCols, projTileCols: params.customProjTileCols, finalTileCols: params.customFinalTileCols, wChunkCols: params.customWChunkCols, linearRowBlock: params.customLinearRowBlock || params.linearRowBlock, singleQ4TileCols: params.customSingleQ4TileCols, singleQ4KChunk: params.customSingleQ4KChunk, singleQ4Dp4aTileCols: params.customSingleQ4Dp4aTileCols, attentionTileKeys: params.customAttentionTileKeys, attentionQueryRows: params.customAttentionQueryRows, singleAttentionKernel: params.customSingleAttentionKernel || params.singleAttentionKernel || params.attentionKernel, textContextCacheKey: params.customTextProjectionCache === false || !params.promptCacheKey ? "" : `${params.promptCacheKey}|txt=${textTokens}`, persistentTextProjectionCache: params.persistentTextContextCache !== false, textProjectionUrl: params.textProjectionUrl || '', singleLinear1Output: params.customSingleLinear1Output || params.singleLinear1Output, singleLinear1Q4Kernel: params.customSingleLinear1Q4Kernel || params.singleLinear1Q4Kernel, singleLinear1Q4ActivationScale: params.customSingleLinear1Q4ActivationScale || params.singleLinear1Q4ActivationScale, singleLinear1QkvBackend: params.customSingleLinear1QkvBackend || params.singleLinear1QkvBackend, singleLinear1MlpBackend: params.customSingleLinear1MlpBackend || params.singleLinear1MlpBackend, singleLinear2Backend: params.customSingleLinear2Backend || params.singleLinear2Backend, singleQkNormStorage: params.customSingleQkNormStorage || params.singleQkNormStorage, approxReusePredEvery: params.customApproxReusePredEvery || params.approxReusePredEvery, approxPredictionMode: params.customApproxPredictionMode || params.approxPredictionMode, approxAbScale: params.customApproxAbScale || params.approxAbScale, skipSingleBlocksFromStep: params.customSkipSingleBlocksFromStep || params.skipSingleBlocksFromStep, skipSingleBlocksFromBlock: params.customSkipSingleBlocksFromBlock || params.skipSingleBlocksFromBlock, reuseSingleTailFromStep: params.customReuseSingleTailFromStep || params.reuseSingleTailFromStep, reuseSingleTailFromBlock: params.customReuseSingleTailFromBlock || params.reuseSingleTailFromBlock, profilePhases: params.customProfilePhases ?? params.profilePhases, profileSingleParts: params.customProfileSingleParts ?? params.profileSingleParts, stepSubmitFusion: params.customStepSubmitFusion ?? params.stepSubmitFusion, debugStopAfter: params.debugStopAfter, debugReadbackCount: params.debugReadbackCount, debugImgReadbackCount: params.debugImgReadbackCount, debugTxtReadbackCount: params.debugTxtReadbackCount, singleComputePass: params.customSingleComputePass || params.singleComputePass, maxDoubleBlocks: params.customMaxDoubleBlocks || params.maxDoubleBlocks, maxSingleBlocks: params.customMaxSingleBlocks || params.maxSingleBlocks, }); const latent = result.latentF16 instanceof Uint16Array ? result.latentF16 : new Uint16Array(result.latentF16 || []); if (params.debugStopAfter) { log(stage, `custom low-bit WebGPU debug returned in ${((performance.now() - start) / 1000).toFixed(3)}s; summary=${JSON.stringify(result.summary || {})}`); return { latent, latentWidth, latentHeight, debug: result.debug || null, rawResult: result }; } if (latent.length !== imageSeqLen * channels) { throw new Error(`CUSTOM_TRANSFORMER_BAD_LATENT_LENGTH: got ${latent.length}, expected ${imageSeqLen * channels}`); } log(stage, `custom low-bit WebGPU denoise returned in ${((performance.now() - start) / 1000).toFixed(3)}s; summary=${JSON.stringify(result.summary || {})}`); return { latent, latentWidth, latentHeight, transformerDetails: { verdict: result.verdict || null, config: result.config || null, load: result.load || null, summary: result.summary || null, step_ms: result.step_ms || null, sample: result.sample || null, }, }; } async function runTransformer(ctx, params, width, height) { const stage = 'transformer'; const modelConfig = config.models.transformer; const requestedBackend = params.transformerBackend || runtimeOptions.customTransformerMode || 'onnx-webgpu'; if (requestedBackend === 'custom-lowbit-webgpu') { try { return await runTransformerCustomLowbitWebGpu(ctx, params, width, height); } catch (err) { log(stage, `custom low-bit WebGPU transformer failed: ${errorDetail(err)}`); if (params.requireCustomTransformer) throw err; log(stage, 'falling back to ONNX WebGPU transformer for complete image generation'); } } else if (requestedBackend !== 'onnx-webgpu' && requestedBackend !== 'off') { const status = customKernelState.checked ? customKernelState.lastStatus : await customKernelStatus(); log(stage, `custom transformer backend requested (${requestedBackend}); full custom image backend status=${JSON.stringify(status)}`); if (params.requireCustomTransformer) { throw new Error('CUSTOM_TRANSFORMER_FULL_IMAGE_BACKEND_NOT_READY'); } log(stage, 'falling back to ONNX WebGPU transformer for complete start-to-finish image generation'); } if (params.gpuDenoise && !transformerUsesFloat32(modelConfig)) { if (runtimeOptions.enableFusedTransformer !== false && !params.customTimesteps && Number(params.numSteps || config.num_steps) === 4) { try { return await runTransformerFusedGpuResident(ctx, params, width, height); } catch (err) { log(stage, `fused GPU denoise failed, falling back to per-step GPU loop: ${errorDetail(err)}`); } } try { return await runTransformerGpuResident(ctx, params, width, height); } catch (err) { log(stage, `GPU-resident denoise failed, falling back to CPU tensor loop: ${errorDetail(err)}`); } } const latentWidth = width / config.latent_downsample; const latentHeight = height / config.latent_downsample; const imageSeqLen = latentWidth * latentHeight; const channels = config.latent_channels; const ctxIds = makeTextIds(config.text_seq_len); const xIds = makeImageIds(latentHeight, latentWidth); let x = params.initialLatentF32 !== undefined ? new Float32Array(coerceLatentF32(params.initialLatentF32, imageSeqLen * channels, 'initial latent')) : makeNoiseF32(imageSeqLen, channels, params.seed ?? 42); logArrayStats(stage, 'initial latent', x); const timesteps = Array.isArray(params.customTimesteps) && params.customTimesteps.length > 1 ? params.customTimesteps.map(Number) : getSchedule(params.numSteps || config.num_steps, imageSeqLen); let session = null; let finalLatent = null; try { session = await createSession(modelConfig, stage, { __cacheSessions: params.cacheSessions, image_seq: imageSeqLen, latent_height: latentHeight, latent_width: latentWidth, height, width, }); log(stage, `denoising ${width}x${height}, tokens=${imageSeqLen}, steps=${timesteps.length - 1}`); for (let step = 0; step < timesteps.length - 1; step++) { if (abortRequested) throw new Error('ABORTED_BY_CLIENT'); const tCurr = timesteps[step]; const tPrev = timesteps[step + 1]; const useFloat32 = transformerUsesFloat32(modelConfig); let feeds = null; let outputs = null; try { feeds = { x: useFloat32 ? tensor('float32', x, [1, imageSeqLen, channels]) : tensor('float16', f32ToF16Array(x), [1, imageSeqLen, channels]), x_ids: tensor('float32', xIds, [1, imageSeqLen, 4]), timesteps: useFloat32 ? tensor('float32', new Float32Array([tCurr]), [1]) : tensor('float16', f32ToF16Array(new Float32Array([tCurr])), [1]), ctx: useFloat32 ? tensor('float32', f16ToF32Array(ctx), [1, config.text_seq_len, config.context_dim]) : tensor('float16', ctx, [1, config.text_seq_len, config.context_dim]), ctx_ids: tensor('float32', ctxIds, [1, config.text_seq_len, 4]), }; logMemory(stage, `before session.run step ${step + 1}`); const runStart = performance.now(); outputs = await profiledRun(session, feeds, null, stage, `step-${step + 1}`); log(stage, `session.run step ${step + 1} returned in ${((performance.now() - runStart) / 1000).toFixed(3)}s`); logMemory(stage, `after session.run step ${step + 1}`); const predTensor = firstOutput(outputs, 'pred'); log(stage, `pred tensor step ${step + 1}: ${tensorDebugSummary(predTensor)}`); const pred = asFloat32Array(await tensorData(predTensor), 'transformer output'); const dt = tPrev - tCurr; const predStats = logArrayStats(stage, `pred step ${step + 1}`, pred); assertFiniteStats(predStats, `pred step ${step + 1}`); for (let i = 0; i < x.length; i++) x[i] += dt * pred[i]; const latentStats = logArrayStats(stage, `latent step ${step + 1}`, x); assertFiniteStats(latentStats, `latent step ${step + 1}`); log(stage, `step ${step + 1}/${timesteps.length - 1} complete`); } finally { await releaseTensorMap(outputs, stage, `outputs step ${step + 1}`); await releaseTensorMap(feeds, stage, `feeds step ${step + 1}`); outputs = null; feeds = null; } } assertFiniteStats(logArrayStats(stage, 'final latent', x), 'final latent'); finalLatent = f32ToF16Array(x); x = null; logMemory(stage, 'after final latent conversion'); } finally { await releaseSession(session, stage); session = null; } log(stage, 'transformer resources yielded before VAE creation'); return { latent: finalLatent, latentWidth, latentHeight }; } function latentSequenceToNchw(latent, latentWidth, latentHeight, channels) { const out = new Uint16Array(channels * latentHeight * latentWidth); for (let h = 0; h < latentHeight; h++) { for (let w = 0; w < latentWidth; w++) { const srcBase = (h * latentWidth + w) * channels; for (let c = 0; c < channels; c++) out[(c * latentHeight + h) * latentWidth + w] = latent[srcBase + c]; } } return out; } function latentNchwToSequence(latent, latentWidth, latentHeight, channels) { const out = new Uint16Array(latentWidth * latentHeight * channels); for (let h = 0; h < latentHeight; h++) { for (let w = 0; w < latentWidth; w++) { const dstBase = (h * latentWidth + w) * channels; for (let c = 0; c < channels; c++) out[dstBase + c] = latent[(c * latentHeight + h) * latentWidth + w]; } } return out; } function downsampleLatentSequenceAverage(latent, srcWidth, srcHeight, dstWidth, dstHeight, channels) { if (srcWidth === dstWidth && srcHeight === dstHeight) return latent; const out = new Float32Array(dstWidth * dstHeight * channels); const counts = new Uint16Array(dstWidth * dstHeight); for (let sy = 0; sy < srcHeight; sy++) { const dy = Math.min(dstHeight - 1, Math.floor((sy * dstHeight) / srcHeight)); for (let sx = 0; sx < srcWidth; sx++) { const dx = Math.min(dstWidth - 1, Math.floor((sx * dstWidth) / srcWidth)); const srcBase = (sy * srcWidth + sx) * channels; const dstToken = dy * dstWidth + dx; const dstBase = dstToken * channels; counts[dstToken] += 1; for (let c = 0; c < channels; c++) { out[dstBase + c] += float16BitsToFloat32(latent[srcBase + c]); } } } for (let token = 0; token < counts.length; token++) { const count = Math.max(1, counts[token]); const base = token * channels; for (let c = 0; c < channels; c++) out[base + c] /= count; } return f32ToF16Array(out); } function makePreviewDecodePlan(latentResult, width, height, maxSize) { const requestedMax = Math.max(0, Math.trunc(Number(maxSize || 0))); if (!requestedMax || Math.max(width, height) <= requestedMax) { return { latentResult, decodeWidth: width, decodeHeight: height, previewDecode: null, }; } const scale = requestedMax / Math.max(width, height); const decodeWidth = Math.max(config.latent_downsample, Math.floor((width * scale) / config.latent_downsample) * config.latent_downsample); const decodeHeight = Math.max(config.latent_downsample, Math.floor((height * scale) / config.latent_downsample) * config.latent_downsample); const decodeLatentWidth = decodeWidth / config.latent_downsample; const decodeLatentHeight = decodeHeight / config.latent_downsample; const latent = downsampleLatentSequenceAverage( latentResult.latent, latentResult.latentWidth, latentResult.latentHeight, decodeLatentWidth, decodeLatentHeight, config.latent_channels, ); return { latentResult: { latent, latentWidth: decodeLatentWidth, latentHeight: decodeLatentHeight }, decodeWidth, decodeHeight, previewDecode: { sourceWidth: width, sourceHeight: height, decodeWidth, decodeHeight, maxSize: requestedMax, }, }; } function makePreviewRenderPlan(width, height, maxSize, enabled = true) { const requestedMax = Math.max(0, Math.trunc(Number(maxSize || 0))); if (!enabled || !requestedMax || Math.max(width, height) <= requestedMax) { return { renderWidth: width, renderHeight: height, previewRender: null, }; } const scale = requestedMax / Math.max(width, height); const renderWidth = Math.max(config.latent_downsample, Math.floor((width * scale) / config.latent_downsample) * config.latent_downsample); const renderHeight = Math.max(config.latent_downsample, Math.floor((height * scale) / config.latent_downsample) * config.latent_downsample); return { renderWidth, renderHeight, previewRender: { sourceWidth: width, sourceHeight: height, renderWidth, renderHeight, maxSize: requestedMax, }, }; } function previewRenderEnabledForParams(params = {}) { const hasInitImage = params.initImage !== undefined || params.hasInitImage === true; return !hasInitImage || params.previewRenderInitImage === true; } function coerceImageNchwF16(data, width, height) { const expectedValues = width * height * 3; if (data instanceof Uint16Array) { if (data.length !== expectedValues) throw new Error(`init image has ${data.length} values, expected ${expectedValues}`); return data; } let values; if (data instanceof Float32Array) { values = data; } else if (ArrayBuffer.isView(data)) { values = new Float32Array(data.length); for (let i = 0; i < data.length; i++) values[i] = Number(data[i]); } else if (data instanceof ArrayBuffer) { if (data.byteLength !== expectedValues * 4) throw new Error(`init image has ${data.byteLength} bytes, expected ${expectedValues * 4}`); values = new Float32Array(data); } else if (Array.isArray(data)) { values = Float32Array.from(data, Number); } else { throw new Error('init image must be a normalized NCHW Float32Array/Uint16Array, ArrayBuffer, typed array, or number array'); } if (values.length !== expectedValues) throw new Error(`init image has ${values.length} values, expected ${expectedValues}`); return f32ToF16Array(values); } function img2imgStrengthCurveExponent(curve) { if (curve === 'linear') return 1; if (curve === 'balanced') return 1.25; if (curve === 'strong') return 3; if (curve === 'rewrite') return 6; if (curve === 'edit') return 2.5; return 2.5; } function img2imgEffectiveStrength(strength, curve) { const clamped = Math.max(0, Math.min(1, Number(strength ?? 0.65))); return 1 - Math.pow(1 - clamped, img2imgStrengthCurveExponent(curve)); } function normalizeImg2ImgCurve(curve) { if (curve === 'linear' || curve === 'balanced' || curve === 'edit' || curve === 'strong' || curve === 'rewrite') return curve; return 'edit'; } function normalizeImg2ImgStepPolicy(policy) { return policy === 'adaptive' ? 'adaptive' : 'full'; } function img2imgSchedule(numSteps, imageSeqLen, strength, curve = 'edit', stepPolicy = 'full') { const full = getSchedule(numSteps, imageSeqLen); const normalizedCurve = normalizeImg2ImgCurve(curve); const normalizedStepPolicy = normalizeImg2ImgStepPolicy(stepPolicy); const clamped = Math.max(0, Math.min(1, Number(strength ?? 0.65))); if (clamped <= 0) { return { timesteps: [0], startT: 0, stepCount: 0, requestedSteps: numSteps, strength: clamped, effectiveStrength: 0, curve: normalizedCurve, stepPolicy: normalizedStepPolicy, }; } const effectiveStrength = img2imgEffectiveStrength(clamped, normalizedCurve); const targetSteps = normalizedStepPolicy === 'adaptive' ? Math.max(1, Math.min(numSteps, Math.ceil(numSteps * clamped))) : numSteps; const timesteps = targetSteps === numSteps ? full.map((value) => value * effectiveStrength) : getSchedule(targetSteps, imageSeqLen).map((value) => value * effectiveStrength); return { timesteps, startT: effectiveStrength, stepCount: targetSteps, requestedSteps: numSteps, strength: clamped, effectiveStrength, curve: normalizedCurve, stepPolicy: normalizedStepPolicy, }; } function blendLatentWithNoise(encodedLatentF16, imageSeqLen, channels, seed, startT) { const encoded = f16ToF32Array(encodedLatentF16); const noise = makeNoiseF32(imageSeqLen, channels, seed); const out = new Float32Array(encoded.length); const latentScale = 1 - startT; for (let i = 0; i < out.length; i++) out[i] = encoded[i] * latentScale + noise[i] * startT; return out; } const isLittleEndian = (() => { const buffer = new ArrayBuffer(4); new Uint32Array(buffer)[0] = 0x0a0b0c0d; return new Uint8Array(buffer)[0] === 0x0d; })(); function imageTensorToRgba(imageData, width, height) { const rgba = new Uint8ClampedArray(width * height * 4); const plane = width * height; if (isLittleEndian) { const packed = new Uint32Array(rgba.buffer); if (imageData instanceof Uint16Array) { const lut = getHalfToImageByteLut(); for (let pixel = 0; pixel < plane; pixel++) { const r = lut[imageData[pixel]]; const g = lut[imageData[plane + pixel]]; const b = lut[imageData[plane * 2 + pixel]]; packed[pixel] = 0xff000000 | (b << 16) | (g << 8) | r; } } else { for (let pixel = 0; pixel < plane; pixel++) { const r = imageByteFromFloat(imageData[pixel]); const g = imageByteFromFloat(imageData[plane + pixel]); const b = imageByteFromFloat(imageData[plane * 2 + pixel]); packed[pixel] = 0xff000000 | (b << 16) | (g << 8) | r; } } } else { if (imageData instanceof Uint16Array) { const lut = getHalfToImageByteLut(); for (let pixel = 0; pixel < plane; pixel++) { const out = pixel * 4; rgba[out] = lut[imageData[pixel]]; rgba[out + 1] = lut[imageData[plane + pixel]]; rgba[out + 2] = lut[imageData[plane * 2 + pixel]]; rgba[out + 3] = 255; } } else { for (let pixel = 0; pixel < plane; pixel++) { const out = pixel * 4; rgba[out] = imageByteFromFloat(imageData[pixel]); rgba[out + 1] = imageByteFromFloat(imageData[plane + pixel]); rgba[out + 2] = imageByteFromFloat(imageData[plane * 2 + pixel]); rgba[out + 3] = 255; } } } return rgba; } function canvasToPngBlob(canvas) { return new Promise((resolve, reject) => { canvas.toBlob((blob) => { if (blob) resolve(blob); else reject(new Error('Canvas did not produce a PNG blob')); }, 'image/png'); }); } async function runVaeDecoder(latentResult, width, height, cacheSessions = false, outputOptions = {}) { const stage = 'vae-decoder'; const modelConfig = config.models.vae_decoder; const { latent, latentWidth, latentHeight } = latentResult; log(stage, `latent input sequence: len=${latent.length} bytes=${formatBytes(latent.byteLength)} shape=[1,${latentWidth * latentHeight},${config.latent_channels}]`); logMemory(stage, 'before latentSequenceToNchw'); const z = latentSequenceToNchw(latent, latentWidth, latentHeight, config.latent_channels); log(stage, `VAE z tensor prepared: len=${z.length} bytes=${formatBytes(z.byteLength)} shape=[1,${config.latent_channels},${latentHeight},${latentWidth}]`); logMemory(stage, 'after latentSequenceToNchw'); return runVaeDecoderNchw(z, latentWidth, latentHeight, width, height, stage, modelConfig, cacheSessions, outputOptions); } async function runVaeEncoder(params, width, height, cacheSessions = false) { const stage = 'vae-encoder'; const modelConfig = config.models.vae_encoder; if (!modelConfig) { throw new Error('VAE encoder is not configured in the model bundle; img2img requires flux2-klein-4b-vae-encoder-fp16.onnx'); } const latentWidth = width / config.latent_downsample; const latentHeight = height / config.latent_downsample; const cacheKey = makeVaeEncoderCacheKey(params, modelConfig, width, height); lastVaeEncoderDetails = { source: cacheKey ? 'miss' : 'disabled', cacheKey: cacheKey ? cacheKey.slice(0, 24) : '' }; if (cacheKey && vaeEncoderCache.has(cacheKey)) { const cached = vaeEncoderCache.get(cacheKey); lastVaeEncoderDetails = { ...lastVaeEncoderDetails, source: 'memory' }; log(stage, `reusing cached source latent ${cacheKey.slice(0, 24)} latent=${cached.latentWidth}x${cached.latentHeight}`); return { latent: new Uint16Array(cached.latent), latentWidth: cached.latentWidth, latentHeight: cached.latentHeight, }; } const expectedLatentValues = config.latent_channels * latentHeight * latentWidth; const persistent = await loadPersistentVaeEncoderResult(cacheKey, expectedLatentValues, latentWidth, latentHeight); if (persistent) { rememberVaeEncoderResult(cacheKey, persistent); lastVaeEncoderDetails = { ...lastVaeEncoderDetails, source: 'indexeddb' }; log(stage, `loaded cached source latent from IndexedDB ${cacheKey.slice(0, 24)} latent=${persistent.latentWidth}x${persistent.latentHeight}`); return persistent; } const image = coerceImageNchwF16(params.initImage, width, height); let session = null; let feeds = null; let outputs = null; try { logMemory(stage, 'before VAE encoder session create'); session = await createSession(modelConfig, stage, { __cacheSessions: cacheSessions, height, width, latent_height: latentHeight, latent_width: latentWidth, }); feeds = { image: tensor('float16', image, [1, 3, height, width]), }; log(stage, `VAE encode starting: image=${width}x${height} latent=${latentWidth}x${latentHeight} feed=${tensorDebugSummary(feeds.image)}`); logMemory(stage, 'before session.run'); const runStart = performance.now(); outputs = await profiledRun(session, feeds, null, stage, 'vae-encoder'); log(stage, `VAE encoder session.run returned in ${((performance.now() - runStart) / 1000).toFixed(3)}s`); logMemory(stage, 'after session.run'); const zTensor = firstOutput(outputs, 'z'); log(stage, `encoded latent tensor: ${tensorDebugSummary(zTensor)}`); const zRaw = asFloat16Array(await tensorData(zTensor), 'VAE encoder output'); const zNchw = new Uint16Array(zRaw); if (zNchw.length !== expectedLatentValues) throw new Error(`VAE encoder output has ${zNchw.length} values, expected ${expectedLatentValues}`); const latent = latentNchwToSequence(zNchw, latentWidth, latentHeight, config.latent_channels); const encoded = { latent, latentWidth, latentHeight }; rememberVaeEncoderResult(cacheKey, encoded); savePersistentVaeEncoderResult(cacheKey, encoded).then((saved) => { if (saved) log('vae-cache', `IndexedDB source latent save complete ${cacheKey.slice(0, 24)}`); }).catch((err) => { log('vae-cache', `IndexedDB save failed: ${errorDetail(err)}`); }); if (cacheKey) lastVaeEncoderDetails = { ...lastVaeEncoderDetails, source: 'encoded' }; return encoded; } finally { await releaseTensorMap(outputs, stage, 'outputs'); await releaseTensorMap(feeds, stage, 'feeds'); outputs = null; feeds = null; await releaseSession(session, stage); } } async function runVaeDecoderNchw(z, latentWidth, latentHeight, width, height, stage = 'vae-decoder', modelConfig = config.models.vae_decoder, cacheSessions = false, outputOptions = {}) { const postStageKeys = Array.isArray(config.vae_decoder_post_stage_keys) ? config.vae_decoder_post_stage_keys : []; const hasPostStages = postStageKeys.some((key) => config.models[key]); if (config.models.vae_decoder_pre && config.models.vae_decoder_attn_chunk && (config.models.vae_decoder_post || hasPostStages)) { return runVaeDecoderSplitNchw(z, latentWidth, latentHeight, width, height, cacheSessions, outputOptions); } if (!modelConfig) { throw new Error('VAE decoder is not configured: missing both split decoder stages and unified vae_decoder model'); } let session = null; let feeds = null; let outputs = null; try { logMemory(stage, 'before VAE session create'); session = await createSession(modelConfig, stage, { __cacheSessions: cacheSessions, latent_height: latentHeight, latent_width: latentWidth, height, width, }); feeds = { z: tensor('float16', z, [1, config.latent_channels, latentHeight, latentWidth]), }; log(stage, `VAE run starting: latent=${latentWidth}x${latentHeight} output=${width}x${height} feed=${tensorDebugSummary(feeds.z)}`); logMemory(stage, 'before session.run'); const runStart = performance.now(); outputs = await profiledRun(session, feeds, null, stage, 'vae'); log(stage, `VAE session.run returned in ${((performance.now() - runStart) / 1000).toFixed(3)}s`); logMemory(stage, 'after session.run'); const imageTensor = firstOutput(outputs, 'image'); log(stage, `image tensor: ${tensorDebugSummary(imageTensor)}`); const image = await tensorData(imageTensor); log(stage, `image tensor data materialized: ${image.constructor?.name || typeof image} bytes=${formatBytes(image.byteLength || 0)}`); logMemory(stage, 'after image getData'); return await imageTensorDataToOutput(image, width, height, stage, outputOptions); } finally { await releaseTensorMap(outputs, stage, 'outputs'); await releaseTensorMap(feeds, stage, 'feeds'); outputs = null; feeds = null; await releaseSession(session, stage); } } async function imageTensorDataToOutput(image, width, height, stage, outputOptions = {}) { if (outputOptions.warmDispatchOnly === true) { const bytes = image?.byteLength || 0; log(stage, `warm dispatch decoded image tensor; skipping RGBA/PNG output: bytes=${formatBytes(bytes)}`); return { imageBytes: bytes, pngBytes: 0, warmDispatchOnly: true, timings: { imageTensorToRgbaMs: 0, pngEncodeMs: 0 } }; } const convertStart = performance.now(); const rgba = imageTensorToRgba(image, width, height); const imageTensorToRgbaMs = performance.now() - convertStart; log(stage, `RGBA buffer prepared: bytes=${formatBytes(rgba.byteLength)}`); logMemory(stage, 'after imageTensorToRgba'); const imageDataStart = performance.now(); const imageData = new ImageData(rgba, width, height); const imageDataMs = performance.now() - imageDataStart; if (outputOptions.returnImageData && !outputOptions.returnImageBlob) { log(stage, `returning ImageData without PNG encode: ${formatBytes(rgba.byteLength)}`); return { imageData, imageBytes: rgba.byteLength, pngBytes: 0, timings: { imageTensorToRgbaMs, imageDataMs, pngEncodeMs: 0 } }; } const canvas = document.createElement('canvas'); canvas.width = width; canvas.height = height; const context = canvas.getContext('2d'); context.putImageData(imageData, 0, 0); const pngStart = performance.now(); const png = await canvasToPngBlob(canvas); const pngEncodeMs = performance.now() - pngStart; log(stage, `encoded PNG ${png.size} bytes`); return { imageBlob: png, imageData: outputOptions.returnImageData ? imageData : undefined, imageBytes: rgba.byteLength, pngBytes: png.size, timings: { imageTensorToRgbaMs, imageDataMs, pngEncodeMs }, }; } function vaeSplitTensorDims(name, width, height, decoderWidth, decoderHeight, seqLen, headDim, length = null) { const tensorName = String(name || ''); if (tensorName === 'z') return [1, config.latent_channels, decoderHeight / 2, decoderWidth / 2]; if (tensorName.includes('/decoder/block_1/Add_output_0')) return [1, headDim, decoderHeight, decoderWidth]; if (tensorName.includes('/decoder/attn_1/MatMul_1_output_0')) return [1, 1, seqLen, headDim]; if (tensorName.includes('/decoder/upsample/conv/Conv_output_0')) return [1, headDim, decoderHeight * 2, decoderWidth * 2]; if (tensorName.includes('/decoder/upsample/conv_1/Conv_output_0')) return [1, headDim, decoderHeight * 4, decoderWidth * 4]; if (tensorName.includes('/decoder/block.0_2/Add_output_0')) return [1, Math.max(1, Math.round(Number(length || 0) / ((width / 2) * (height / 2)))), height / 2, width / 2]; if (tensorName.includes('/decoder/block.1_2/Add_output_0')) return [1, Math.max(1, Math.round(Number(length || 0) / ((width / 2) * (height / 2)))), height / 2, width / 2]; if (tensorName.includes('/decoder/block.2_2/Add_output_0')) return [1, Math.max(1, Math.round(Number(length || 0) / ((width / 2) * (height / 2)))), height / 2, width / 2]; if (tensorName.includes('/decoder/upsample/conv_2/Conv_output_0')) return [1, Math.max(1, Math.round(Number(length || 0) / (width * height))), height, width]; if (tensorName.includes('/decoder/block.0/nin_shortcut_1/Conv_output_0')) return [1, Math.max(1, Math.round(Number(length || 0) / (width * height))), height, width]; if (tensorName.includes('/decoder/block.0/norm1_3/Add_output_0')) return [1, Math.max(1, Math.round(Number(length || 0) / (width * height))), height, width]; if (tensorName.includes('/decoder/block.0/conv1_3/Conv_output_0')) return [1, Math.max(1, Math.round(Number(length || 0) / (width * height))), height, width]; if (tensorName.includes('/decoder/block.0/conv2_3/Conv_output_0')) return [1, Math.max(1, Math.round(Number(length || 0) / (width * height))), height, width]; if (tensorName.includes('/decoder/block.0_3/Add_output_0')) return [1, Math.max(1, Math.round(Number(length || 0) / (width * height))), height, width]; if (tensorName.includes('/decoder/block.1/norm1_3/Add_output_0')) return [1, Math.max(1, Math.round(Number(length || 0) / (width * height))), height, width]; if (tensorName.includes('/decoder/block.1/conv1_3/Conv_output_0')) return [1, Math.max(1, Math.round(Number(length || 0) / (width * height))), height, width]; if (tensorName.includes('/decoder/block.1/norm2_3/Add_output_0')) return [1, Math.max(1, Math.round(Number(length || 0) / (width * height))), height, width]; if (tensorName.includes('/decoder/block.1/conv2_3/Conv_output_0')) return [1, Math.max(1, Math.round(Number(length || 0) / (width * height))), height, width]; if (tensorName.includes('/decoder/block.1_3/Add_output_0')) return [1, Math.max(1, Math.round(Number(length || 0) / (width * height))), height, width]; if (tensorName.includes('/decoder/block.2/norm1_3/Add_output_0')) return [1, Math.max(1, Math.round(Number(length || 0) / (width * height))), height, width]; if (tensorName.includes('/decoder/block.2/conv1_3/Conv_output_0')) return [1, Math.max(1, Math.round(Number(length || 0) / (width * height))), height, width]; if (tensorName.includes('/decoder/block.2/norm2_3/Add_output_0')) return [1, Math.max(1, Math.round(Number(length || 0) / (width * height))), height, width]; if (tensorName.includes('/decoder/block.2/conv2_3/Conv_output_0')) return [1, Math.max(1, Math.round(Number(length || 0) / (width * height))), height, width]; if (tensorName.includes('/decoder/block.2_3/Add_output_0')) return [1, Math.max(1, Math.round(Number(length || 0) / (width * height))), height, width]; if (tensorName === 'image') return [1, 3, height, width]; if (length && width > 0 && height > 0 && length % (width * height) === 0) return [1, length / (width * height), height, width]; throw new Error(`No split VAE tensor shape rule for ${tensorName}`); } async function runVaeDecoderPostStages(stage, postStageKeys, residualName, residual, attnName, attn, dimsContext) { const tensors = new Map([ [residualName, residual], [attnName, attn], ]); const futureInputsByStage = postStageKeys.map((_, index) => { const names = new Set(); for (let future = index + 1; future < postStageKeys.length; future++) { const futureConfig = config.models[postStageKeys[future]]; for (const inputName of futureConfig?.inputs || []) names.add(inputName); } return names; }); let finalImage = null; for (let index = 0; index < postStageKeys.length; index++) { if (abortRequested) throw new Error('ABORTED_BY_CLIENT'); const key = postStageKeys[index]; const stageConfig = config.models[key]; if (!stageConfig) throw new Error(`Missing VAE post stage config ${key}`); const inputNames = stageConfig.inputs || []; const outputName = stageConfig.outputs?.[0]; if (!inputNames.length || !outputName) throw new Error(`Invalid VAE post stage config ${key}`); let postSession = null; let postFeeds = null; let postOutputs = null; try { postSession = await createSession(stageConfig, `${stage}:post:${index + 1}/${postStageKeys.length}`, { __cacheSessions: dimsContext.cacheSessions, latent_height: dimsContext.latentHeight, latent_width: dimsContext.latentWidth, height: dimsContext.height, width: dimsContext.width, kv_seq: dimsContext.seqLen, chunk_seq: dimsContext.chunkSize, }); postFeeds = {}; for (const inputName of inputNames) { const data = tensors.get(inputName); if (!data) throw new Error(`VAE post stage ${key} missing input ${inputName}`); postFeeds[inputName] = tensor( 'float16', data, vaeSplitTensorDims( inputName, dimsContext.width, dimsContext.height, dimsContext.decoderWidth, dimsContext.decoderHeight, dimsContext.seqLen, dimsContext.headDim, data.length, ), ); } log(stage, `post stage ${index + 1}/${postStageKeys.length} starting: ${stageConfig.file} inputs=${inputNames.map((name) => `${name}:${formatBytes(tensors.get(name)?.byteLength || 0)}`).join(', ')}`); logMemory(stage, `before post stage ${index + 1}`); const t0 = performance.now(); postOutputs = await profiledRun(postSession, postFeeds, null, `${stage}:post:${index + 1}`, `stage-${index + 1}`); const outputTensor = postOutputs[outputName] || firstOutput(postOutputs, outputName); log(stage, `post stage ${index + 1}/${postStageKeys.length} returned in ${((performance.now() - t0) / 1000).toFixed(3)}s; output=${tensorDebugSummary(outputTensor)}`); const futureInputs = futureInputsByStage[index]; await releaseTensorMap(postFeeds, `${stage}:post:${index + 1}`, 'feeds-after-run'); postFeeds = null; for (const inputName of inputNames) { if (inputName !== outputName && !futureInputs.has(inputName)) tensors.delete(inputName); } log(stage, `post stage ${index + 1}/${postStageKeys.length} released unneeded inputs before materializing output; retained=${Array.from(tensors.keys()).map((name) => `${name}:${formatBytes(tensors.get(name)?.byteLength || 0)}`).join(', ') || 'none'}`); await yieldToRuntime(); const outputData = asFloat16Array(await tensorData(outputTensor), outputName); log(stage, `post stage ${index + 1}/${postStageKeys.length} materialized ${outputName}: ${formatBytes(outputData.byteLength)}`); logMemory(stage, `after post stage ${index + 1} getData`); tensors.set(outputName, outputData); if (outputName === 'image') finalImage = outputData; } finally { await releaseTensorMap(postOutputs, `${stage}:post:${index + 1}`, 'outputs'); await releaseTensorMap(postFeeds, `${stage}:post:${index + 1}`, 'feeds'); await releaseSession(postSession, `${stage}:post:${index + 1}`); await yieldToRuntime(); } } if (!finalImage) { finalImage = tensors.get('image'); } if (!finalImage) throw new Error('VAE post stages did not produce image'); return finalImage; } async function releaseTensorMapValues(tensors, stage, label) { if (!tensors || !tensors.size) { log(stage, `releaseTensorMapValues ${label}: empty`); return; } const counts = {}; for (const tensorValue of tensors.values()) { const method = await releaseTensorValue(tensorValue); counts[method] = (counts[method] || 0) + 1; } tensors.clear(); log(stage, `releaseTensorMapValues ${label}: releaseMethods=${JSON.stringify(counts)}`); } async function runVaeDecoderPostStagesGpuHandoff(stage, postStageKeys, tensors, dimsContext) { const futureInputsByStage = postStageKeys.map((_, index) => { const names = new Set(); for (let future = index + 1; future < postStageKeys.length; future++) { const futureConfig = config.models[postStageKeys[future]]; for (const inputName of futureConfig?.inputs || []) names.add(inputName); } return names; }); let finalTensor = null; for (let index = 0; index < postStageKeys.length; index++) { if (abortRequested) throw new Error('ABORTED_BY_CLIENT'); const key = postStageKeys[index]; const stageConfig = config.models[key]; if (!stageConfig) throw new Error(`Missing VAE post stage config ${key}`); const inputNames = stageConfig.inputs || []; const outputName = stageConfig.outputs?.[0]; if (!inputNames.length || !outputName) throw new Error(`Invalid VAE post stage config ${key}`); let postSession = null; let postOutputs = null; try { postSession = await createSession(stageConfig, `${stage}:gpu-post:${index + 1}/${postStageKeys.length}`, { __cacheSessions: dimsContext.cacheSessions, __preferredOutputLocation: { [outputName]: 'gpu-buffer' }, latent_height: dimsContext.latentHeight, latent_width: dimsContext.latentWidth, height: dimsContext.height, width: dimsContext.width, kv_seq: dimsContext.seqLen, chunk_seq: dimsContext.chunkSize, }); const postFeeds = {}; for (const inputName of inputNames) { const inputTensor = tensors.get(inputName); if (!inputTensor) throw new Error(`VAE GPU post stage ${key} missing input ${inputName}`); postFeeds[inputName] = inputTensor; } log(stage, `GPU post stage ${index + 1}/${postStageKeys.length} starting: ${stageConfig.file} inputs=${inputNames.map((name) => `${name}:{${tensorDebugSummary(tensors.get(name))}}`).join(', ')}`); const t0 = performance.now(); postOutputs = await profiledRun(postSession, postFeeds, null, `${stage}:gpu-post:${index + 1}`, `stage-${index + 1}`); const outputTensor = postOutputs[outputName] || firstOutput(postOutputs, outputName); log(stage, `GPU post stage ${index + 1}/${postStageKeys.length} returned in ${((performance.now() - t0) / 1000).toFixed(3)}s; output=${tensorDebugSummary(outputTensor)}`); const futureInputs = futureInputsByStage[index]; for (const inputName of inputNames) { if (inputName !== outputName && !futureInputs.has(inputName)) { await releaseTensorValue(tensors.get(inputName)); tensors.delete(inputName); } } tensors.set(outputName, outputTensor); if (outputName === 'image') finalTensor = outputTensor; postOutputs = null; await yieldToRuntime(); } finally { await releaseTensorMap(postOutputs, `${stage}:gpu-post:${index + 1}`, 'outputs-after-error'); await releaseSession(postSession, `${stage}:gpu-post:${index + 1}`); } } if (!finalTensor) { finalTensor = tensors.get('image'); } if (!finalTensor) throw new Error('VAE GPU post stages did not produce image'); return finalTensor; } async function runVaeDecoderSplitNchwGpuHandoff(z, latentWidth, latentHeight, width, height, cacheSessions = false, outputOptions = {}) { const stage = 'vae-decoder-split-gpu'; const preConfig = config.models.vae_decoder_pre; const chunkConfig = config.models.vae_decoder_attn_chunk; const postStageKeys = (Array.isArray(config.vae_decoder_post_stage_keys) ? config.vae_decoder_post_stage_keys : []).filter((key) => config.models[key]); if (!postStageKeys.length) throw new Error('GPU VAE handoff currently expects staged post decoder graphs'); const postConfig = config.models.vae_decoder_post || config.models[postStageKeys[0]]; const decoderHeight = latentHeight * 2; const decoderWidth = latentWidth * 2; const seqLen = decoderHeight * decoderWidth; const headDim = Number(chunkConfig.head_dim || 512); const residualName = postConfig?.residual_input || postConfig?.inputs?.[0] || '/decoder/block_1/Add_output_0'; const attnName = postConfig?.attn_input || postConfig?.inputs?.[1] || '/decoder/attn_1/MatMul_1_output_0'; const preOutputs = preConfig.outputs || []; const qName = preOutputs.find((name) => String(name).includes('/Mul_6_output_0')) || '/decoder/attn_1/Mul_6_output_0'; const ktName = preOutputs.find((name) => String(name).includes('/Mul_7_output_0')) || '/decoder/attn_1/Mul_7_output_0'; const vName = preOutputs.find((name) => String(name).includes('/Reshape_2_output_0')) || '/decoder/attn_1/Reshape_2_output_0'; const chunkOutputName = chunkConfig.outputs?.[0] || 'attn_chunk'; const tensors = new Map(); let preSession = null; let preFeeds = null; let preRunOutputs = null; let chunkSession = null; let chunkOutputs = null; try { log(stage, `GPU handoff pre-attention run: latent=${latentWidth}x${latentHeight} decoder_mid=${decoderWidth}x${decoderHeight} seq=${seqLen}`); preSession = await createSession(preConfig, `${stage}:pre`, { __cacheSessions: cacheSessions, __preferredOutputLocation: { [residualName]: 'gpu-buffer', [qName]: 'gpu-buffer', [ktName]: 'gpu-buffer', [vName]: 'gpu-buffer', }, latent_height: latentHeight, latent_width: latentWidth, height, width, }); preFeeds = { z: tensor('float16', z, [1, config.latent_channels, latentHeight, latentWidth]) }; preRunOutputs = await profiledRun(preSession, preFeeds, null, `${stage}:pre`, 'pre'); for (const name of [residualName, qName, ktName, vName]) { const outputTensor = preRunOutputs[name]; if (!outputTensor) throw new Error(`GPU VAE pre missing output ${name}`); tensors.set(name, outputTensor); } log(stage, `GPU pre outputs ready: ${Array.from(tensors.entries()).map(([name, value]) => `${name}:{${tensorDebugSummary(value)}}`).join(' ; ')}`); preRunOutputs = null; await releaseTensorMap(preFeeds, `${stage}:pre`, 'feeds'); preFeeds = null; await releaseSession(preSession, `${stage}:pre`); preSession = null; chunkSession = await createSession(chunkConfig, `${stage}:attn`, { __cacheSessions: cacheSessions, __preferredOutputLocation: { [chunkOutputName]: 'gpu-buffer' }, chunk_seq: seqLen, kv_seq: seqLen, batch: 1, heads: 1, }); const chunkFeeds = { q_chunk: tensors.get(qName), kt: tensors.get(ktName), v: tensors.get(vName), }; log(stage, `GPU attention full chunk starting: seq=${seqLen} q=${tensorDebugSummary(chunkFeeds.q_chunk)} kt=${tensorDebugSummary(chunkFeeds.kt)} v=${tensorDebugSummary(chunkFeeds.v)}`); const attnStart = performance.now(); chunkOutputs = await profiledRun(chunkSession, chunkFeeds, null, `${stage}:attn`, 'full-chunk'); const attnTensor = chunkOutputs[chunkOutputName] || firstOutput(chunkOutputs, chunkOutputName); log(stage, `GPU attention full chunk returned in ${((performance.now() - attnStart) / 1000).toFixed(3)}s; output=${tensorDebugSummary(attnTensor)}`); for (const name of [qName, ktName, vName]) { await releaseTensorValue(tensors.get(name)); tensors.delete(name); } tensors.set(attnName, attnTensor); chunkOutputs = null; await releaseSession(chunkSession, `${stage}:attn`); chunkSession = null; const imageTensor = await runVaeDecoderPostStagesGpuHandoff(stage, postStageKeys, tensors, { latentWidth, latentHeight, width, height, decoderWidth, decoderHeight, seqLen, headDim, chunkSize: seqLen, cacheSessions, }); log(stage, `GPU final image tensor: ${tensorDebugSummary(imageTensor)}`); const image = await tensorData(imageTensor); log(stage, `GPU final image data materialized: ${image.constructor?.name || typeof image} bytes=${formatBytes(image.byteLength || 0)}`); return await imageTensorDataToOutput(image, width, height, stage, outputOptions); } finally { await releaseTensorMap(preRunOutputs, `${stage}:pre`, 'outputs-after-error'); await releaseTensorMap(preFeeds, `${stage}:pre`, 'feeds'); await releaseSession(preSession, `${stage}:pre`); await releaseTensorMap(chunkOutputs, `${stage}:attn`, 'outputs-after-error'); await releaseSession(chunkSession, `${stage}:attn`); await releaseTensorMapValues(tensors, stage, 'gpu-handoff-tensors'); } } async function runVaeDecoderSplitNchw(z, latentWidth, latentHeight, width, height, cacheSessions = false, outputOptions = {}) { const stage = 'vae-decoder-split'; const preConfig = config.models.vae_decoder_pre; const chunkConfig = config.models.vae_decoder_attn_chunk; const postStageKeys = (Array.isArray(config.vae_decoder_post_stage_keys) ? config.vae_decoder_post_stage_keys : []).filter((key) => config.models[key]); const postConfig = config.models.vae_decoder_post || (postStageKeys.length ? config.models[postStageKeys[0]] : null); const decoderHeight = latentHeight * 2; const decoderWidth = latentWidth * 2; const seqLen = decoderHeight * decoderWidth; const headDim = Number(chunkConfig.head_dim || 512); const baseChunkSize = Number(chunkConfig.chunk_size || 1024); const requestedChunkSize = Math.trunc(Number(outputOptions.vaeAttentionChunkSize || 0)); const maxAutoChunkSeq = Math.trunc(Number(outputOptions.maxAutoVaeAttentionChunkSeq || 0)); const chunkSize = Math.min( seqLen, requestedChunkSize > 0 ? requestedChunkSize : (seqLen <= maxAutoChunkSeq ? seqLen : baseChunkSize), ); const chunkCount = Math.ceil(seqLen / chunkSize); if (chunkSize !== baseChunkSize) { log(stage, `attention chunk size override: ${baseChunkSize} -> ${chunkSize} for seq=${seqLen}`); } const residualName = postConfig?.residual_input || postConfig?.inputs?.[0] || '/decoder/block_1/Add_output_0'; const attnName = postConfig?.attn_input || postConfig?.inputs?.[1] || '/decoder/attn_1/MatMul_1_output_0'; const preOutputs = preConfig.outputs || []; const qName = preOutputs.find((name) => String(name).includes('/Mul_6_output_0')) || '/decoder/attn_1/Mul_6_output_0'; const ktName = preOutputs.find((name) => String(name).includes('/Mul_7_output_0')) || '/decoder/attn_1/Mul_7_output_0'; const vName = preOutputs.find((name) => String(name).includes('/Reshape_2_output_0')) || '/decoder/attn_1/Reshape_2_output_0'; const chunkOutputName = chunkConfig.outputs?.[0] || 'attn_chunk'; const maxGpuHandoffSeq = Number(outputOptions.maxGpuHandoffSeq || 4096); const canTryGpuHandoff = outputOptions.gpuHandoff === true && runtimeOptions.executionProvider === 'webgpu' && postStageKeys.length > 0 && seqLen <= maxGpuHandoffSeq; if (canTryGpuHandoff) { try { return await runVaeDecoderSplitNchwGpuHandoff(z, latentWidth, latentHeight, width, height, cacheSessions, outputOptions); } catch (err) { log(stage, `GPU VAE handoff failed; falling back to CPU-staged split VAE: ${errorDetail(err)}`); } } let preSession = null; let preFeeds = null; let preRunOutputs = null; let residual = null; let q = null; let kt = null; let v = null; try { log(stage, `pre-attention run: latent=${latentWidth}x${latentHeight} decoder_mid=${decoderWidth}x${decoderHeight} seq=${seqLen}`); preSession = await createSession(preConfig, `${stage}:pre`, { __cacheSessions: cacheSessions, latent_height: latentHeight, latent_width: latentWidth, height, width, }); preFeeds = { z: tensor('float16', z, [1, config.latent_channels, latentHeight, latentWidth]) }; preRunOutputs = await profiledRun(preSession, preFeeds, null, `${stage}:pre`, 'pre'); residual = asFloat16Array(await tensorData(preRunOutputs[residualName]), residualName).slice(); q = asFloat16Array(await tensorData(preRunOutputs[qName]), qName).slice(); kt = asFloat16Array(await tensorData(preRunOutputs[ktName]), ktName).slice(); v = asFloat16Array(await tensorData(preRunOutputs[vName]), vName).slice(); log(stage, `pre outputs materialized: residual=${formatBytes(residual.byteLength)} q=${formatBytes(q.byteLength)} kt=${formatBytes(kt.byteLength)} v=${formatBytes(v.byteLength)}`); } finally { await releaseTensorMap(preRunOutputs, `${stage}:pre`, 'outputs'); await releaseTensorMap(preFeeds, `${stage}:pre`, 'feeds'); await releaseSession(preSession, `${stage}:pre`); } let attn = new Uint16Array(seqLen * headDim); let chunkSession = null; try { chunkSession = await createSession(chunkConfig, `${stage}:attn`, { __cacheSessions: cacheSessions, chunk_seq: chunkSize, kv_seq: seqLen, batch: 1, heads: 1, }); for (let start = 0, chunkIndex = 0; start < seqLen; start += chunkSize, chunkIndex++) { if (abortRequested) throw new Error('ABORTED_BY_CLIENT'); let chunkFeeds = null; let chunkOutputs = null; const validRows = Math.min(chunkSize, seqLen - start); const end = start + validRows; try { let qChunk = q.slice(start * headDim, end * headDim); if (validRows < chunkSize) { const padded = new Uint16Array(chunkSize * headDim); padded.set(qChunk); qChunk = padded; } chunkFeeds = { q_chunk: tensor('float16', qChunk, [1, 1, chunkSize, headDim]), kt: tensor('float16', kt, [1, 1, headDim, seqLen]), v: tensor('float16', v, [1, 1, seqLen, headDim]), }; log(stage, `attention chunk ${chunkIndex + 1}/${chunkCount}: q=${start}:${end}${validRows < chunkSize ? ` padded=${chunkSize - validRows}` : ''}`); const t0 = performance.now(); chunkOutputs = await profiledRun(chunkSession, chunkFeeds, null, `${stage}:attn`, `chunk-${chunkIndex + 1}`); const chunk = asFloat16Array(await tensorData(chunkOutputs[chunkOutputName]), chunkOutputName); attn.set(validRows < chunkSize ? chunk.slice(0, validRows * headDim) : chunk, start * headDim); log(stage, `attention chunk ${chunkIndex + 1} returned in ${((performance.now() - t0) / 1000).toFixed(3)}s; output=${formatBytes(chunk.byteLength)}`); } finally { await releaseTensorMap(chunkOutputs, `${stage}:attn`, `chunk-${chunkIndex + 1}-outputs`); await releaseTensorMap(chunkFeeds, `${stage}:attn`, `chunk-${chunkIndex + 1}-feeds`); await yieldToRuntime(); } } } finally { await releaseSession(chunkSession, `${stage}:attn`); q = null; kt = null; v = null; } if (postStageKeys.length) { const image = await runVaeDecoderPostStages(stage, postStageKeys, residualName, residual, attnName, attn, { latentWidth, latentHeight, width, height, decoderWidth, decoderHeight, seqLen, headDim, chunkSize, cacheSessions, }); residual = null; attn = null; return await imageTensorDataToOutput(image, width, height, stage, outputOptions); } else { let postSession = null; let postFeeds = null; let postOutputs = null; try { postSession = await createSession(postConfig, `${stage}:post`, { __cacheSessions: cacheSessions, latent_height: latentHeight, latent_width: latentWidth, height, width, kv_seq: seqLen, chunk_seq: chunkSize, }); postFeeds = { [residualName]: tensor('float16', residual, [1, headDim, decoderHeight, decoderWidth]), [attnName]: tensor('float16', attn, [1, 1, seqLen, headDim]), }; residual = null; log(stage, `post-attention run starting: attn=${formatBytes(attn.byteLength)}`); postOutputs = await profiledRun(postSession, postFeeds, null, `${stage}:post`, 'post'); const imageTensor = firstOutput(postOutputs, 'image'); log(stage, `image tensor: ${tensorDebugSummary(imageTensor)}`); const image = await tensorData(imageTensor); log(stage, `image tensor data materialized: ${image.constructor?.name || typeof image} bytes=${formatBytes(image.byteLength || 0)}`); return await imageTensorDataToOutput(image, width, height, stage, outputOptions); } finally { await releaseTensorMap(postOutputs, `${stage}:post`, 'outputs'); await releaseTensorMap(postFeeds, `${stage}:post`, 'feeds'); await releaseSession(postSession, `${stage}:post`); } } } function makeVaePrepareShape(params = {}) { const outputWidth = Number(params.width || config.default_width); const outputHeight = Number(params.height || config.default_height); if (outputWidth % config.latent_downsample || outputHeight % config.latent_downsample) { throw new Error(`width and height must be divisible by ${config.latent_downsample}`); } const renderPlan = makePreviewRenderPlan( outputWidth, outputHeight, params.previewRenderMaxSize, previewRenderEnabledForParams(params), ); const width = renderPlan.renderWidth; const height = renderPlan.renderHeight; const requestedMax = Math.max(0, Math.trunc(Number(params.previewDecodeMaxSize || 0))); let decodeWidth = width; let decodeHeight = height; let previewDecode = null; if (requestedMax && Math.max(width, height) > requestedMax) { const scale = requestedMax / Math.max(width, height); decodeWidth = Math.max(config.latent_downsample, Math.floor((width * scale) / config.latent_downsample) * config.latent_downsample); decodeHeight = Math.max(config.latent_downsample, Math.floor((height * scale) / config.latent_downsample) * config.latent_downsample); previewDecode = { sourceWidth: width, sourceHeight: height, decodeWidth, decodeHeight, maxSize: requestedMax, }; } return { outputWidth, outputHeight, renderWidth: width, renderHeight: height, decodeWidth, decodeHeight, latentWidth: decodeWidth / config.latent_downsample, latentHeight: decodeHeight / config.latent_downsample, previewRender: renderPlan.previewRender, previewDecode, }; } function vaeDecoderChunkSizeForShape(latentWidth, latentHeight, outputOptions = {}) { const chunkConfig = config.models.vae_decoder_attn_chunk; if (!chunkConfig) return 0; const decoderHeight = latentHeight * 2; const decoderWidth = latentWidth * 2; const seqLen = decoderHeight * decoderWidth; const baseChunkSize = Number(chunkConfig.chunk_size || 1024); const requestedChunkSize = Math.trunc(Number(outputOptions.vaeAttentionChunkSize || 0)); const maxAutoChunkSeq = Math.trunc(Number(outputOptions.maxAutoVaeAttentionChunkSeq || 0)); return Math.min( seqLen, requestedChunkSize > 0 ? requestedChunkSize : (seqLen <= maxAutoChunkSeq ? seqLen : baseChunkSize), ); } async function prepareVaeDecoderSessionsForShape(latentWidth, latentHeight, width, height, outputOptions = {}) { const stage = 'vae-decoder-prepare'; const postStageKeys = (Array.isArray(config.vae_decoder_post_stage_keys) ? config.vae_decoder_post_stage_keys : []).filter((key) => config.models[key]); if (config.models.vae_decoder_pre && config.models.vae_decoder_attn_chunk && (config.models.vae_decoder_post || postStageKeys.length)) { const preConfig = config.models.vae_decoder_pre; const chunkConfig = config.models.vae_decoder_attn_chunk; const postConfig = config.models.vae_decoder_post || (postStageKeys.length ? config.models[postStageKeys[0]] : null); const decoderHeight = latentHeight * 2; const decoderWidth = latentWidth * 2; const seqLen = decoderHeight * decoderWidth; const headDim = Number(chunkConfig.head_dim || 512); const chunkSize = vaeDecoderChunkSizeForShape(latentWidth, latentHeight, outputOptions); const created = []; const createCached = async (modelConfig, stageName, freeDims) => { const session = await createSession(modelConfig, stageName, { __cacheSessions: true, ...freeDims, }); created.push(stageName); await releaseSession(session, stageName); }; await createCached(preConfig, `${stage}:pre`, { latent_height: latentHeight, latent_width: latentWidth, height, width, }); await createCached(chunkConfig, `${stage}:attn`, { chunk_seq: chunkSize, kv_seq: seqLen, batch: 1, heads: 1, }); if (postStageKeys.length) { for (let index = 0; index < postStageKeys.length; index++) { const key = postStageKeys[index]; await createCached(config.models[key], `${stage}:post:${index + 1}/${postStageKeys.length}`, { latent_height: latentHeight, latent_width: latentWidth, height, width, kv_seq: seqLen, chunk_seq: chunkSize, }); } } else if (postConfig) { await createCached(postConfig, `${stage}:post`, { latent_height: latentHeight, latent_width: latentWidth, height, width, kv_seq: seqLen, chunk_seq: chunkSize, }); } return { split: true, sessions: created.length, stages: created, seqLen, chunkSize, chunkCount: Math.ceil(seqLen / chunkSize), headDim, }; } if (!config.models.vae_decoder) { throw new Error('VAE decoder is not configured: missing both split decoder stages and unified vae_decoder model'); } const session = await createSession(config.models.vae_decoder, stage, { __cacheSessions: true, latent_height: latentHeight, latent_width: latentWidth, height, width, }); await releaseSession(session, stage); return { split: false, sessions: 1, stages: [stage], seqLen: latentWidth * latentHeight * 4, chunkSize: 0, chunkCount: 0, }; } async function prepareVaeSessions(params = {}) { if (!config || !modelBaseUrl) throw new Error('Engine is not initialized'); const stage = 'vae-prepare'; const start = performance.now(); const shape = makeVaePrepareShape(params); const outputOptions = { vaeAttentionChunkSize: params.vaeAttentionChunkSize, maxAutoVaeAttentionChunkSeq: params.maxAutoVaeAttentionChunkSeq, }; const lutStart = performance.now(); getHalfToImageByteLut(); const imageByteLutMs = performance.now() - lutStart; const decoder = await prepareVaeDecoderSessionsForShape( shape.latentWidth, shape.latentHeight, shape.decodeWidth, shape.decodeHeight, outputOptions, ); let encoder = null; if (params.includeEncoder === true && config.models.vae_encoder) { const encoderSession = await createSession(config.models.vae_encoder, `${stage}:encoder`, { __cacheSessions: true, height: shape.renderHeight, width: shape.renderWidth, latent_height: shape.renderHeight / config.latent_downsample, latent_width: shape.renderWidth / config.latent_downsample, }); await releaseSession(encoderSession, `${stage}:encoder`); encoder = { sessions: 1, width: shape.renderWidth, height: shape.renderHeight, latentWidth: shape.renderWidth / config.latent_downsample, latentHeight: shape.renderHeight / config.latent_downsample, }; } let dispatch = null; if (params.dispatchWarmup === true) { const dispatchStart = performance.now(); const latent = new Uint16Array(shape.latentWidth * shape.latentHeight * config.latent_channels); const output = await runVaeDecoder( { latent, latentWidth: shape.latentWidth, latentHeight: shape.latentHeight }, shape.decodeWidth, shape.decodeHeight, true, { ...outputOptions, warmDispatchOnly: true, returnImageBlob: false, returnImageData: false, }, ); dispatch = { elapsed_ms: performance.now() - dispatchStart, imageBytes: output.imageBytes || 0, pngBytes: output.pngBytes || 0, }; } const elapsedMs = performance.now() - start; log(stage, `prepared VAE sessions in ${(elapsedMs / 1000).toFixed(3)}s; decoder=${JSON.stringify(decoder)} encoder=${JSON.stringify(encoder)} dispatch=${JSON.stringify(dispatch)}`); return { verdict: 'vae-sessions-prepared', elapsed_ms: elapsedMs, shape, decoder, encoder, dispatch, imageByteLutMs, }; } async function prepareVaeEncoderLatent(params = {}) { if (!config || !modelBaseUrl) throw new Error('Engine is not initialized'); const width = Number(params.width || config.default_width); const height = Number(params.height || config.default_height); if (width % config.latent_downsample || height % config.latent_downsample) { throw new Error(`width and height must be divisible by ${config.latent_downsample}`); } if (params.initImage === undefined) { throw new Error('prepareVaeEncoderLatent requires initImage'); } const start = performance.now(); const encoded = await runVaeEncoder( { ...params, cacheVaeEncoder: params.cacheVaeEncoder !== false, }, width, height, true, ); const elapsedMs = performance.now() - start; return { verdict: 'vae-encoder-latent-prepared', elapsed_ms: elapsedMs, source: lastVaeEncoderDetails?.source || 'unknown', cacheKey: lastVaeEncoderDetails?.cacheKey || '', latentWidth: encoded.latentWidth, latentHeight: encoded.latentHeight, values: encoded.latent.length, }; } async function generateImage(params) { if (!config || !modelBaseUrl) throw new Error('Engine is not initialized'); abortRequested = false; const outputWidth = params.width || config.default_width; const outputHeight = params.height || config.default_height; if (outputWidth % config.latent_downsample || outputHeight % config.latent_downsample) { throw new Error(`width and height must be divisible by ${config.latent_downsample}`); } const renderPlan = makePreviewRenderPlan(outputWidth, outputHeight, params.previewRenderMaxSize, previewRenderEnabledForParams(params)); const width = renderPlan.renderWidth; const height = renderPlan.renderHeight; try { const totalStart = performance.now(); const timings = {}; let stageStart = performance.now(); const customStatus = await customKernelStatus(); timings.customKernelStatusSeconds = Number(((performance.now() - stageStart) / 1000).toFixed(3)); logMemory('generate', 'start'); const img2imgImageSeqLen = params.initImage !== undefined ? (width / config.latent_downsample) * (height / config.latent_downsample) : 0; const precomputedImg2imgSchedule = params.initImage !== undefined ? img2imgSchedule(params.numSteps || config.num_steps, img2imgImageSeqLen, params.strength, params.strengthCurve, params.img2imgStepPolicy) : null; stageStart = performance.now(); let ctx = null; if (!precomputedImg2imgSchedule || precomputedImg2imgSchedule.stepCount > 0) { ctx = await runTextEncoder(params); timings.textEncoderSeconds = Number(((performance.now() - stageStart) / 1000).toFixed(3)); timings.textEncoderCacheSource = lastTextEncoderDetails?.source || 'unknown'; timings.textEncoderCacheSeconds = Number(((lastTextEncoderDetails?.persistentMs || 0) / 1000).toFixed(3)); logMemory('generate', 'after text encoder'); } else { timings.textEncoderSeconds = 0; timings.textEncoderCacheSource = 'skipped-source-only'; timings.textEncoderCacheSeconds = 0; log('generate', 'skipping text encoder for img2img strength=0 source-only decode'); } let transformerParams = params; let img2imgInfo = null; let encodedLatentResult = null; if (params.initImage !== undefined) { stageStart = performance.now(); const encoded = await runVaeEncoder(params, width, height, params.cacheSessions); timings.vaeEncoderSeconds = Number(((performance.now() - stageStart) / 1000).toFixed(3)); timings.vaeEncoderCacheSource = lastVaeEncoderDetails?.source || 'unknown'; const imageSeqLen = encoded.latentWidth * encoded.latentHeight; const schedule = precomputedImg2imgSchedule || img2imgSchedule(params.numSteps || config.num_steps, imageSeqLen, params.strength, params.strengthCurve, params.img2imgStepPolicy); img2imgInfo = { strength: schedule.strength, effectiveStrength: Number(schedule.effectiveStrength.toFixed(6)), curve: schedule.curve, startT: Number(schedule.startT.toFixed(6)), denoiseSteps: schedule.stepCount, requestedSteps: schedule.requestedSteps, stepPolicy: schedule.stepPolicy, }; if (schedule.stepCount <= 0) { transformerParams = params; encodedLatentResult = encoded; } else { transformerParams = { ...params, customTimesteps: schedule.timesteps, initialLatentF32: blendLatentWithNoise(encoded.latent, imageSeqLen, config.latent_channels, params.seed ?? 42, schedule.startT), }; } log('img2img', `encoded init image; strength=${img2imgInfo.strength} effective=${img2imgInfo.effectiveStrength} curve=${img2imgInfo.curve} startT=${img2imgInfo.startT} denoiseSteps=${img2imgInfo.denoiseSteps}`); } stageStart = performance.now(); let latentResult; if (encodedLatentResult) { latentResult = encodedLatentResult; timings.transformerSeconds = 0; } else { latentResult = await runTransformer(ctx, transformerParams, width, height); timings.transformerSeconds = Number(((performance.now() - stageStart) / 1000).toFixed(3)); } const transformerDetails = latentResult.transformerDetails || null; ctx = null; if (globalThis.gc) globalThis.gc(); logMemory('generate', 'after transformer + ctx drop'); stageStart = performance.now(); const decodePlan = makePreviewDecodePlan(latentResult, width, height, params.previewDecodeMaxSize); const outputOptions = { returnImageBlob: Boolean(params.returnImageBlob || params.returnImageObjectUrl || params.resultUrl), returnImageData: Boolean(params.returnImageData), gpuHandoff: params.gpuVaeHandoff === true, maxGpuHandoffSeq: params.maxGpuHandoffSeq, vaeAttentionChunkSize: params.vaeAttentionChunkSize, maxAutoVaeAttentionChunkSeq: params.maxAutoVaeAttentionChunkSeq, }; const imageOutput = await runVaeDecoder(decodePlan.latentResult, decodePlan.decodeWidth, decodePlan.decodeHeight, params.cacheSessions, outputOptions); timings.vaeSeconds = Number(((performance.now() - stageStart) / 1000).toFixed(3)); if (imageOutput.timings) { timings.imageTensorToRgbaSeconds = Number(((imageOutput.timings.imageTensorToRgbaMs || 0) / 1000).toFixed(3)); timings.imageDataSeconds = Number(((imageOutput.timings.imageDataMs || 0) / 1000).toFixed(3)); timings.pngEncodeSeconds = Number(((imageOutput.timings.pngEncodeMs || 0) / 1000).toFixed(3)); } latentResult = null; if (globalThis.gc) globalThis.gc(); logMemory('generate', 'after VAE + latent drop'); if (params.resultUrl) await fetch(params.resultUrl, { method: 'POST', body: imageOutput.imageBlob }); timings.totalBrowserSeconds = Number(((performance.now() - totalStart) / 1000).toFixed(3)); const response = { width: outputWidth, height: outputHeight, renderWidth: width, renderHeight: height, seed: params.seed ?? 42, numSteps: params.numSteps || config.num_steps, mode: params.initImage !== undefined ? 'image-to-image' : 'text-to-image', img2img: img2imgInfo, pngBytes: imageOutput.pngBytes || 0, imageBytes: imageOutput.imageBytes || imageOutput.pngBytes || 0, imageDataWidth: decodePlan.decodeWidth, imageDataHeight: decodePlan.decodeHeight, previewRender: renderPlan.previewRender, previewDecode: decodePlan.previewDecode, staged: params.initImage !== undefined ? (img2imgInfo?.denoiseSteps > 0 ? ['text_encoder', 'vae_encoder', 'transformer', 'vae_decoder'] : ['vae_encoder', 'vae_decoder']) : ['text_encoder', 'transformer', 'vae_decoder'], transformerBackend: params.transformerBackend || runtimeOptions.customTransformerMode || 'onnx-webgpu', customKernels: customStatus, timings, transformerDetails, }; if (params.returnImageData) response.imageData = imageOutput.imageData; if (params.returnImageBlob) response.imageBlob = imageOutput.imageBlob; if (params.returnImageObjectUrl) response.imageObjectUrl = URL.createObjectURL(imageOutput.imageBlob); return response; } catch (err) { log('error', errorDetail(err)); throw err; } } async function encodeText(params) { if (!config || !modelBaseUrl) throw new Error('Engine is not initialized'); const ctx = await runTextEncoder(params); const ctxF32 = f16ToF32Array(ctx); const stats = logArrayStats('text-encoder-contract', 'context', ctxF32); if (params.resultUrl) { const bytes = ctxF32.byteOffset === 0 && ctxF32.byteLength === ctxF32.buffer.byteLength ? ctxF32.buffer : ctxF32.slice().buffer; await fetch(params.resultUrl, { method: 'POST', body: bytes }); } return { shape: [1, config.text_seq_len, config.context_dim], dtype: 'float32', stats, }; } async function decodeLatent(params) { if (!config || !modelBaseUrl) throw new Error('Engine is not initialized'); const width = params.width || config.default_width; const height = params.height || config.default_height; const latentWidth = width / config.latent_downsample; const latentHeight = height / config.latent_downsample; if (!Number.isInteger(latentWidth) || !Number.isInteger(latentHeight)) { throw new Error(`width and height must be divisible by ${config.latent_downsample}`); } const z = await loadFloat16Binary(params.latentUrl, config.latent_channels * latentHeight * latentWidth, 'latent'); const imageOutput = await runVaeDecoderNchw(z, latentWidth, latentHeight, width, height, 'vae-decoder', config.models.vae_decoder, false, { returnImageBlob: true, returnImageData: false, }); if (params.resultUrl) await fetch(params.resultUrl, { method: 'POST', body: imageOutput.imageBlob }); return { width, height, pngBytes: imageOutput.pngBytes, imageBytes: imageOutput.imageBytes, staged: ['vae_decoder'] }; } async function transformerContract(params) { if (!config || !modelBaseUrl) throw new Error('Engine is not initialized'); const stage = 'transformer-contract'; const modelConfig = config.models.transformer; const latentHeight = Number(params.latentHeight || params.latent_height || 2); const latentWidth = Number(params.latentWidth || params.latent_width || 2); const imageSeqLen = latentHeight * latentWidth; const channels = config.latent_channels; const x = await loadFloat32Binary(params.xUrl, imageSeqLen * channels, 'contract latent'); const ctx = await loadFloat32Binary(params.ctxUrl, config.text_seq_len * config.context_dim, 'contract context'); const ctxIds = makeTextIds(config.text_seq_len); const xIds = makeImageIds(latentHeight, latentWidth); const timestep = Number(params.timestep ?? getSchedule(params.numSteps || config.num_steps, imageSeqLen)[0]); let session = null; try { session = await createSession(modelConfig, stage, { image_seq: imageSeqLen, latent_height: latentHeight, latent_width: latentWidth, height: latentHeight * config.latent_downsample, width: latentWidth * config.latent_downsample, }); const useFloat32 = transformerUsesFloat32(modelConfig); const outputs = await profiledRun(session, { x: useFloat32 ? tensor('float32', x, [1, imageSeqLen, channels]) : tensor('float16', f32ToF16Array(x), [1, imageSeqLen, channels]), x_ids: tensor('float32', xIds, [1, imageSeqLen, 4]), timesteps: useFloat32 ? tensor('float32', new Float32Array([timestep]), [1]) : tensor('float16', f32ToF16Array(new Float32Array([timestep])), [1]), ctx: useFloat32 ? tensor('float32', ctx, [1, config.text_seq_len, config.context_dim]) : tensor('float16', f32ToF16Array(ctx), [1, config.text_seq_len, config.context_dim]), ctx_ids: tensor('float32', ctxIds, [1, config.text_seq_len, 4]), }, null, stage, 'contract'); const pred = asFloat32Array(await tensorData(firstOutput(outputs, 'pred')), 'transformer contract output'); const stats = logArrayStats(stage, 'pred', pred); if (params.resultUrl) { const bytes = pred.byteOffset === 0 && pred.byteLength === pred.buffer.byteLength ? pred.buffer : pred.slice().buffer; await fetch(params.resultUrl, { method: 'POST', body: bytes }); } return { shape: [1, imageSeqLen, channels], dtype: 'float32', timestep, stats, }; } finally { await releaseSession(session, stage); } } function deterministicContractContextF32(seqLen, dim) { const values = new Float32Array(seqLen * dim); for (let row = 0; row < seqLen; row++) { for (let col = 0; col < dim; col++) { const code = ((row * 37 + col * 17 + 11) % 257) - 128; values[row * dim + col] = code / 256; } } return values; } function compareFloatArrays(actual, expected, maxItems = 16) { let finite = 0; let maxAbs = 0; let sumAbs = 0; let sumSq = 0; let maxRel = 0; const count = Math.min(actual.length, expected.length); for (let i = 0; i < count; i++) { const a = actual[i]; const e = expected[i]; if (Number.isFinite(a) && Number.isFinite(e)) finite += 1; const diff = Math.abs(a - e); maxAbs = Math.max(maxAbs, diff); sumAbs += diff; sumSq += diff * diff; maxRel = Math.max(maxRel, diff / Math.max(1e-6, Math.abs(e))); } return { count, finite, max_abs: maxAbs, mean_abs: count ? sumAbs / count : 0, rms: count ? Math.sqrt(sumSq / count) : 0, max_rel: maxRel, actual_sample: Array.from(actual.slice(0, Math.min(maxItems, actual.length))), expected_sample: Array.from(expected.slice(0, Math.min(maxItems, expected.length))), }; } async function customTransformerContract(params = {}) { if (!config || !modelBaseUrl) throw new Error('Engine is not initialized'); const stage = 'custom-transformer-contract'; const latentHeight = Number(params.latentHeight || params.latent_height || 2); const latentWidth = Number(params.latentWidth || params.latent_width || 2); const imageSeqLen = latentHeight * latentWidth; const channels = config.latent_channels; const x = makeNoiseF32(imageSeqLen, channels, Number(params.seed ?? 123)); const ctxF32 = deterministicContractContextF32(config.text_seq_len, config.context_dim); const ctxF16 = f32ToF16Array(ctxF32); const timestep = Number(params.timestep ?? getSchedule(params.numSteps || config.num_steps, imageSeqLen)[0]); const ctxIds = makeTextIds(config.text_seq_len); const xIds = makeImageIds(latentHeight, latentWidth); const modelConfig = config.models.transformer; if (params.debugStopAfter) { const custom = await runTransformerCustomLowbitWebGpu(ctxF16, { seed: Number(params.seed ?? 123), numSteps: params.numSteps || config.num_steps, customTimesteps: [timestep, 0], customAttentionTileKeys: params.customAttentionTileKeys, customAttentionQueryRows: params.customAttentionQueryRows, customSingleQ4TileCols: params.customSingleQ4TileCols, customSingleQ4Dp4aTileCols: params.customSingleQ4Dp4aTileCols, customSingleLinear1Output: params.customSingleLinear1Output || params.singleLinear1Output, customSingleLinear1Q4Kernel: params.customSingleLinear1Q4Kernel || params.singleLinear1Q4Kernel, customSingleLinear1Q4ActivationScale: params.customSingleLinear1Q4ActivationScale || params.singleLinear1Q4ActivationScale, customSingleLinear1QkvBackend: params.customSingleLinear1QkvBackend || params.singleLinear1QkvBackend, customSingleLinear1MlpBackend: params.customSingleLinear1MlpBackend || params.singleLinear1MlpBackend, customSingleLinear2Backend: params.customSingleLinear2Backend || params.singleLinear2Backend, debugStopAfter: params.debugStopAfter, debugReadbackCount: params.debugReadbackCount, debugImgReadbackCount: params.debugImgReadbackCount, debugTxtReadbackCount: params.debugTxtReadbackCount, }, latentWidth * config.latent_downsample, latentHeight * config.latent_downsample); return { shape: [1, imageSeqLen, channels], timestep, latentHeight, latentWidth, debug: custom.debug || null, }; } let session = null; try { session = await createSession(modelConfig, stage, { image_seq: imageSeqLen, latent_height: latentHeight, latent_width: latentWidth, height: latentHeight * config.latent_downsample, width: latentWidth * config.latent_downsample, }); const useFloat32 = transformerUsesFloat32(modelConfig); const onnxOutputs = await profiledRun(session, { x: useFloat32 ? tensor('float32', x, [1, imageSeqLen, channels]) : tensor('float16', f32ToF16Array(x), [1, imageSeqLen, channels]), x_ids: tensor('float32', xIds, [1, imageSeqLen, 4]), timesteps: useFloat32 ? tensor('float32', new Float32Array([timestep]), [1]) : tensor('float16', f32ToF16Array(new Float32Array([timestep])), [1]), ctx: useFloat32 ? tensor('float32', ctxF32, [1, config.text_seq_len, config.context_dim]) : tensor('float16', ctxF16, [1, config.text_seq_len, config.context_dim]), ctx_ids: tensor('float32', ctxIds, [1, config.text_seq_len, 4]), }, null, stage, 'onnx-reference'); const onnxPred = asFloat32Array(await tensorData(firstOutput(onnxOutputs, 'pred')), 'onnx contract pred'); const custom = await runTransformerCustomLowbitWebGpu(ctxF16, { seed: Number(params.seed ?? 123), numSteps: params.numSteps || config.num_steps, customTimesteps: [timestep, 0], customAttentionTileKeys: params.customAttentionTileKeys, customAttentionQueryRows: params.customAttentionQueryRows, customSingleQ4TileCols: params.customSingleQ4TileCols, customSingleQ4Dp4aTileCols: params.customSingleQ4Dp4aTileCols, customSingleLinear1Output: params.customSingleLinear1Output || params.singleLinear1Output, customSingleLinear1Q4Kernel: params.customSingleLinear1Q4Kernel || params.singleLinear1Q4Kernel, customSingleLinear1Q4ActivationScale: params.customSingleLinear1Q4ActivationScale || params.singleLinear1Q4ActivationScale, customSingleLinear1QkvBackend: params.customSingleLinear1QkvBackend || params.singleLinear1QkvBackend, customSingleLinear1MlpBackend: params.customSingleLinear1MlpBackend || params.singleLinear1MlpBackend, customSingleLinear2Backend: params.customSingleLinear2Backend || params.singleLinear2Backend, debugStopAfter: params.debugStopAfter, debugReadbackCount: params.debugReadbackCount, debugImgReadbackCount: params.debugImgReadbackCount, debugTxtReadbackCount: params.debugTxtReadbackCount, }, latentWidth * config.latent_downsample, latentHeight * config.latent_downsample); if (params.debugStopAfter) { return { shape: [1, imageSeqLen, channels], timestep, latentHeight, latentWidth, debug: custom.debug || null, }; } const customLatent = f16ToF32Array(custom.latent); const customPred = new Float32Array(customLatent.length); for (let i = 0; i < customPred.length; i++) { customPred[i] = (x[i] - customLatent[i]) / timestep; } const comparison = compareFloatArrays(customPred, onnxPred, Number(params.sample ?? 16)); log(stage, `comparison=${JSON.stringify(comparison)}`); return { shape: [1, imageSeqLen, channels], timestep, latentHeight, latentWidth, comparison, }; } finally { await releaseSession(session, stage); } } function abortGeneration() { abortRequested = true; } window.flux2Engine = { init, generateImage, encodeText, decodeLatent, transformerContract, customTransformerContract, customKernelStatus, prepareCustomTransformerAssets, prepareCustomTransformerStageSetup, prepareCustomTransformerTextProjection, prepareVaeSessions, prepareVaeEncoderLatent, abortGeneration };