import { CUSTOM_MODEL_URL_PREFIX, getModelCache } from 'lib/fetch-progress'; const CUSTOM_MODELS_KEY = 'upscaler_custom_models'; function sanitizeScale(scale) { const parsed = parseInt(scale, 10); if (!Number.isFinite(parsed) || parsed < 1) return 1; if (parsed > 16) return 16; return parsed; } function sanitizeRange(range) { const parsed = parseInt(range, 10); if (parsed === 255) return 255; return 1; } function sanitizeLayout(layout) { const normalized = String(layout || '').toLowerCase(); if (normalized === 'nhwc') return 'nhwc'; return 'nchw'; } function sanitizeMultipleOf(multipleOf) { const parsed = parseInt(multipleOf, 10); if (!Number.isFinite(parsed) || parsed < 1) return 1; if (parsed > 256) return 256; return parsed; } function sanitizeMaxTileSize(maxTileSize) { if (maxTileSize == null || maxTileSize === '') return null; const parsed = parseInt(maxTileSize, 10); if (!Number.isFinite(parsed) || parsed < 1) return null; if (parsed > 4096) return 4096; return parsed; } function sanitizePrecision(precision) { return String(precision || '').toLowerCase() === 'fp16' ? 'fp16' : 'fp32'; } function sanitizeUpscaleBefore(value) { if (value === true || value === 'true' || value === 1 || value === '1') return true; return false; } function sanitizeTileBlend(value) { return value === 'gaussian' ? 'gaussian' : 'overlapCrop'; } function normalizeLabel(label) { const trimmed = (label || '').trim(); return trimmed || 'Custom ONNX'; } function toModelRecord(raw = {}) { const id = String(raw.id || ''); if (!id) return null; const scale = sanitizeScale(raw.scale); const range = sanitizeRange(raw.range); const layout = sanitizeLayout(raw.layout); const multipleOf = sanitizeMultipleOf(raw.multipleOf); const maxTileSize = sanitizeMaxTileSize(raw.maxTileSize); const precision = sanitizePrecision(raw.precision); const upscaleBefore = sanitizeUpscaleBefore(raw.upscaleBefore); const tileBlend = sanitizeTileBlend(raw.tileBlend); const sizeBytes = Number.isFinite(raw.sizeBytes) ? Math.max(0, raw.sizeBytes) : 0; const sizeMB = Number((sizeBytes / (1024 * 1024)).toFixed(1)); return { id, url: `${CUSTOM_MODEL_URL_PREFIX}${id}`, label: normalizeLabel(raw.label), scale, range, layout, multipleOf, maxTileSize, precision, upscaleBefore, tileBlend, sizeBytes, sizeMB, custom: true, }; } function readStoredRecords() { try { const json = localStorage.getItem(CUSTOM_MODELS_KEY); if (!json) return []; const parsed = JSON.parse(json); if (!Array.isArray(parsed)) return []; return parsed .map(toModelRecord) .filter(Boolean); } catch { return []; } } function persistRecords(records) { const payload = records.map(({ id, label, scale, range, layout, multipleOf, maxTileSize, precision, upscaleBefore, tileBlend, sizeBytes }) => ({ id, label, scale, range, layout, multipleOf, maxTileSize, precision, upscaleBefore, tileBlend, sizeBytes, })); localStorage.setItem(CUSTOM_MODELS_KEY, JSON.stringify(payload)); } export function listCustomModels() { return readStoredRecords(); } export function getUploadCustomOptionHTML() { return ''; } export function getCustomModelByUrl(url) { if (!url || typeof url !== 'string') return null; return readStoredRecords().find((model) => model.url === url) || null; } export async function saveCustomModel({ file, label, scale, range, layout, multipleOf, maxTileSize, precision, upscaleBefore, tileBlend }) { if (!(file instanceof File)) { throw new Error('Missing ONNX file for custom model upload.'); } const idSeed = `${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 8)}`; const id = `custom-${idSeed}`; const url = `${CUSTOM_MODEL_URL_PREFIX}${id}`; const bytes = await file.arrayBuffer(); const cache = await getModelCache(); if (!cache) { throw new Error('Browser Cache API is unavailable, cannot store custom model.'); } await cache.put(url, new Response(bytes, { headers: { 'content-type': 'application/octet-stream', 'content-length': String(bytes.byteLength), }, })); const records = readStoredRecords(); const model = toModelRecord({ id, label, scale, range, layout, multipleOf, maxTileSize, precision, upscaleBefore, tileBlend, sizeBytes: bytes.byteLength, }); records.unshift(model); persistRecords(records); return model; } /** * Update an existing custom model's metadata (label / scale / range / layout / * multipleOf). The cached ONNX file and id/url are preserved. Unspecified * fields are left unchanged. Returns the updated model record, or `null` if * no model matches `url`. */ export function updateCustomModelByUrl(url, updates = {}) { if (!url || typeof url !== 'string') return null; const records = readStoredRecords(); const index = records.findIndex((entry) => entry.url === url); if (index === -1) return null; const current = records[index]; const merged = toModelRecord({ id: current.id, label: 'label' in updates ? updates.label : current.label, scale: 'scale' in updates ? updates.scale : current.scale, range: 'range' in updates ? updates.range : current.range, layout: 'layout' in updates ? updates.layout : current.layout, multipleOf: 'multipleOf' in updates ? updates.multipleOf : current.multipleOf, maxTileSize: 'maxTileSize' in updates ? updates.maxTileSize : current.maxTileSize, precision: 'precision' in updates ? updates.precision : current.precision, upscaleBefore: 'upscaleBefore' in updates ? updates.upscaleBefore : current.upscaleBefore, tileBlend: 'tileBlend' in updates ? updates.tileBlend : current.tileBlend, sizeBytes: current.sizeBytes, }); records[index] = merged; persistRecords(records); return merged; } export async function deleteCustomModelByUrl(url) { if (!url || typeof url !== 'string') return false; const records = readStoredRecords(); const model = records.find((entry) => entry.url === url); if (!model) return false; const cache = await getModelCache(); await cache?.delete(url); persistRecords(records.filter((entry) => entry.id !== model.id)); return true; }