class Upscaler { static DEFAULT_TIMEOUTS = { IMAGE: 300000, VIDEO: 3600000, METADATA: 10000 }; static DEFAULT_DELAYS = { INIT: 500, NETWORK: 300 }; static isSupported() { try { const hasWorker = typeof Worker !== "undefined"; const hasOffscreen = typeof OffscreenCanvas !== "undefined" || (typeof document !== "undefined" && !!document.createElement("canvas").transferControlToOffscreen); const hasBlob = typeof Blob !== "undefined"; return hasWorker && hasOffscreen && hasBlob; } catch { return false; } } static isVideoSupported() { try { return ( typeof VideoEncoder !== "undefined" && typeof VideoDecoder !== "undefined" ); } catch { return false; } } constructor(options = {}) { const weightsBaseUrl = options.weightsBaseUrl || "/weights"; this.networks = options.networks || { small: { name: "anime4k/cnn-2x-l", weightsUrl: `${weightsBaseUrl}/cnn-8.json`, }, medium: { name: "anime4k/cnn-2x-16", weightsUrl: `${weightsBaseUrl}/cnn-16.json`, }, large: { name: "anime4k/cnn-2x-28", weightsUrl: `${weightsBaseUrl}/cnn-28.json`, }, }; this.weightsBaseUrl = weightsBaseUrl; this.networkSize = options.networkSize || "medium"; this.workerUrl = options.workerUrl || "/worker/main.js"; this.timeouts = { ...Upscaler.DEFAULT_TIMEOUTS, ...(options.timeouts || {}), }; this.delays = { ...Upscaler.DEFAULT_DELAYS, ...(options.delays || {}) }; this.imageScale = typeof options.imageScale === "number" && options.imageScale > 0 ? options.imageScale : 2; this.videoScale = typeof options.videoScale === "number" && options.videoScale > 0 ? options.videoScale : 2; this.weightsCache = new Map(); this.workerInstance = null; this.messageHandlers = {}; this.progressCallback = null; this.processingType = null; this.currentReject = null; this.currentTimeoutId = null; this.realtimeLoopId = null; this.realtimeState = null; this.init(); } init({ prewarm = true } = {}) { if (!Upscaler.isSupported()) { throw new Error("Upscaler is not supported in this environment"); } if (prewarm) { this.getWorker().postMessage({ cmd: "isSupported" }); } return this; } getWorker() { if (!this.workerInstance) { this.workerInstance = new Worker(this.workerUrl); this.workerInstance.onmessage = (event) => this.handleWorkerMessage(event); this.workerInstance.onerror = (error) => { console.error("[Upscaler] Worker error:", error); this.rejectCurrent( new Error( error.message || "Worker crashed unexpectedly", ), ); }; this.workerInstance.onmessageerror = (error) => { console.error("[Upscaler] Worker message error:", error); this.rejectCurrent( new Error("Worker message serialization error"), ); }; } return this.workerInstance; } rejectCurrent(error) { if (this.currentReject) { const reject = this.currentReject; this.currentReject = null; if (this.currentTimeoutId) { clearTimeout(this.currentTimeoutId); this.currentTimeoutId = null; } // Clean up whichever blob handler is currently registered delete this.messageHandlers.videoBlob; delete this.messageHandlers.imageBlob; reject(error); } } static extractProgressValue(data) { const value = data.data ?? data.progress ?? data.value ?? data.percentage; return typeof value === "number" ? Math.min(100, Math.max(0, Math.round(value))) : null; } handleBlobResponse(data) { const blobType = this.processingType === "video" ? "videoBlob" : "imageBlob"; const handler = this.messageHandlers[blobType]; if (data.data instanceof Blob && handler) { handler({ [blobType]: data.data }); } } requestBlob() { const cmd = this.processingType === "video" ? "getVideoBlob" : "getImageBlob"; this.getWorker().postMessage({ cmd }); } handleWorkerMessage(event) { const { data } = event; console.debug("Worker message:", data); if (!data.cmd) return; const { cmd } = data; if (cmd === "error") { const errorMsg = data.data || data.message || "Unknown worker error"; console.error("[Upscaler] Worker reported error:", errorMsg); this.rejectCurrent(new Error(errorMsg)); } else if (cmd === "progress") { const progress = Upscaler.extractProgressValue(data); if (progress !== null && this.progressCallback) { this.progressCallback(progress); } if (this.messageHandlers.progress) this.messageHandlers.progress(data); } else if (cmd === "finished") { if (data.data instanceof Blob) { this.handleBlobResponse(data); } else { this.requestBlob(); } } else if (this.messageHandlers[cmd]) { this.messageHandlers[cmd](data); } } async loadWeights(networkSize) { this.validateNetworkSize(networkSize); if (this.weightsCache.has(networkSize)) { return this.weightsCache.get(networkSize); } const network = this.networks[networkSize]; const response = await fetch(network.weightsUrl); if (!response.ok) throw new Error(`Failed to fetch weights: ${response.statusText}`); const weights = await response.json(); this.weightsCache.set(networkSize, weights); return weights; } static loadImageMetadata(arrayBuffer, mimeType) { return new Promise((resolve, reject) => { const blob = new Blob([arrayBuffer], { type: mimeType }); const url = URL.createObjectURL(blob); const img = new Image(); img.onload = () => { URL.revokeObjectURL(url); resolve({ width: img.width, height: img.height }); }; img.onerror = () => { URL.revokeObjectURL(url); reject(new Error("Failed to load image")); }; img.src = url; }); } static loadVideoMetadata(file, timeoutMs) { return new Promise((resolve, reject) => { const url = URL.createObjectURL(file); const video = document.createElement("video"); video.src = url; const timeout = setTimeout(() => reject(new Error("Timeout")), timeoutMs); video.onloadedmetadata = () => { clearTimeout(timeout); URL.revokeObjectURL(url); resolve({ width: video.videoWidth, height: video.videoHeight, duration: video.duration, }); }; video.onerror = () => { clearTimeout(timeout); URL.revokeObjectURL(url); reject(new Error("Failed to load video")); }; }); } createOffscreenCanvas(width, height, scale = 2) { const factor = typeof scale === "number" && scale > 0 ? scale : 1; const canvas = document.createElement("canvas"); canvas.width = width * factor; canvas.height = height * factor; if (typeof canvas.transferControlToOffscreen !== "function") { throw new Error("OffscreenCanvas is not supported in this environment"); } return canvas.transferControlToOffscreen(); } createBlobPromise(blobType, timeout) { return new Promise((resolve, reject) => { const timeoutId = setTimeout(() => { this.currentReject = null; this.currentTimeoutId = null; if (this.messageHandlers[blobType]) { delete this.messageHandlers[blobType]; reject(new Error(`${blobType} processing timeout`)); } }, timeout); this.currentReject = reject; this.currentTimeoutId = timeoutId; this.messageHandlers[blobType] = (msg) => { delete this.messageHandlers[blobType]; clearTimeout(timeoutId); this.currentReject = null; this.currentTimeoutId = null; const blob = msg[blobType] || msg.data; if (blob instanceof Blob) { resolve(blob); } else { reject(new Error(`Invalid ${blobType} format`)); } }; }); } validateNetworkSize(size) { if (!this.networks[size]) { throw new Error( "Invalid networkSize: use one of " + Object.keys(this.networks).join(", "), ); } } async upscaleImage(imageFile, networkSize, scaleOverride) { if (!imageFile) throw new Error("Image file is required"); const size = networkSize || this.networkSize; this.validateNetworkSize(size); // Reject any stale promise from a previous aborted operation this.rejectCurrent(new Error("Aborted by new operation")); this.processingType = "image"; const arrayBuffer = await imageFile.arrayBuffer(); const mimeType = imageFile.type || "image/jpeg"; const { width, height } = await Upscaler.loadImageMetadata( arrayBuffer, mimeType, ); const weights = await this.loadWeights(size); const network = this.networks[size]; const scale = typeof scaleOverride === "number" && scaleOverride > 0 ? scaleOverride : this.imageScale; const [canvasUp, canvasOrig] = [ this.createOffscreenCanvas(width, height, scale), this.createOffscreenCanvas(width, height, 1), ]; const blobPromise = this.createBlobPromise( "imageBlob", this.timeouts.IMAGE, ); this.getWorker().postMessage( { cmd: "init", data: { imageArrayBuffer: arrayBuffer, imageMimeType: mimeType, upscaled: canvasUp, original: canvasOrig, resolution: { width, height }, network_name: network.name, weights, }, }, [canvasUp, canvasOrig], ); setTimeout(() => { this.getWorker().postMessage({ cmd: "network", data: { name: network.name, imageArrayBuffer: arrayBuffer, imageMimeType: mimeType, weights, }, }); setTimeout(() => { this.getWorker().postMessage({ cmd: "getImageBlob", data: { imageArrayBuffer: arrayBuffer, imageMimeType: mimeType }, }); }, this.delays.NETWORK); }, this.delays.INIT); return blobPromise; } async upscaleVideo(videoFile, networkSize, scaleOverride) { if (!videoFile) throw new Error("Video file is required"); const size = networkSize || this.networkSize; this.validateNetworkSize(size); if (!Upscaler.isVideoSupported()) { throw new Error("WebCodecs API not supported. Requires Chrome 94+"); } // Reject any stale promise from a previous aborted operation this.rejectCurrent(new Error("Aborted by new operation")); this.processingType = "video"; const { width, height, duration } = await Upscaler.loadVideoMetadata( videoFile, this.timeouts.METADATA, ); const weights = await this.loadWeights(size); const network = this.networks[size]; const scale = typeof scaleOverride === "number" && scaleOverride > 0 ? scaleOverride : this.videoScale; const [canvasOut, canvasIn] = [ this.createOffscreenCanvas(width, height, scale), this.createOffscreenCanvas(width, height, 1), ]; const blobPromise = this.createBlobPromise( "videoBlob", this.timeouts.VIDEO, ); this.getWorker().postMessage( { cmd: "process", file: videoFile, fileSize: videoFile.size, duration, adjustedResolution: { adjustedInputWidth: width, adjustedInputHeight: height, adjustedOutputWidth: width * scale, adjustedOutputHeight: height * scale, }, upscaled: canvasOut, original: canvasIn, weights, network_name: network.name, skipDemuxProgress: true, }, [canvasOut, canvasIn], ); return blobPromise; } async startRealtimeUpscale( videoElement, canvasElement, networkSize, scaleOverride, ) { if (!videoElement || !canvasElement) { throw new Error("videoElement and canvasElement are required"); } const size = networkSize || this.networkSize; this.validateNetworkSize(size); const network = this.networks[size]; const weights = await this.loadWeights(size); const scale = typeof scaleOverride === "number" && scaleOverride > 0 ? scaleOverride : this.videoScale; const ensureMetadata = () => { return new Promise((resolve, reject) => { if (videoElement.readyState >= 1) { resolve(); return; } const onLoaded = () => { cleanup(); resolve(); }; const onError = () => { cleanup(); reject(new Error("Failed to load video metadata")); }; const cleanup = () => { videoElement.removeEventListener("loadedmetadata", onLoaded); videoElement.removeEventListener("error", onError); }; videoElement.addEventListener("loadedmetadata", onLoaded); videoElement.addEventListener("error", onError); }); }; await ensureMetadata(); const width = videoElement.videoWidth; const height = videoElement.videoHeight; if (!width || !height) { throw new Error("Invalid video dimensions"); } const firstInit = !canvasElement._upscalerTransferred; let offscreen = null; if (firstInit) { canvasElement.width = width * scale; canvasElement.height = height * scale; if (typeof canvasElement.transferControlToOffscreen !== "function") { throw new Error("OffscreenCanvas is not supported in this environment"); } offscreen = canvasElement.transferControlToOffscreen(); canvasElement._upscalerTransferred = true; } const captureCanvas = document.createElement("canvas"); captureCanvas.width = width; captureCanvas.height = height; const captureCtx = captureCanvas.getContext("2d"); if (this.realtimeLoopId) { cancelAnimationFrame(this.realtimeLoopId); this.realtimeLoopId = null; } const initPayload = { cmd: "realtimeInit", data: { upscaled: firstInit ? offscreen : null, resolution: { width, height, scale, outputWidth: width * scale, outputHeight: height * scale, }, network_name: network.name, weights, }, }; if (firstInit && offscreen) { this.getWorker().postMessage(initPayload, [offscreen]); } else { this.getWorker().postMessage(initPayload); } this.realtimeState = { running: true, video: videoElement, canvas: canvasElement, scale, busy: false, frameIndex: 0, captureCanvas, captureCtx, }; const loop = async () => { if (!this.realtimeState || !this.realtimeState.running) return; const state = this.realtimeState; const v = state.video; const haveCurrentData = v.HAVE_CURRENT_DATA || 2; if ( !v.paused && !v.ended && !state.busy && v.readyState >= haveCurrentData && v.videoWidth > 0 && v.videoHeight > 0 ) { state.busy = true; try { const captureCanvas = state.captureCanvas; const captureCtx = state.captureCtx; if (!captureCanvas || !captureCtx) { throw new Error("Missing capture canvas"); } captureCtx.drawImage(v, 0, 0, captureCanvas.width, captureCanvas.height); const frame = await createImageBitmap(captureCanvas); state.frameIndex += 1; this.getWorker().postMessage( { cmd: "realtimeFrame", frame, }, [frame], ); } catch (e) { console.warn("realtimeFrame error", e); } finally { if (this.realtimeState) { this.realtimeState.busy = false; } } } this.realtimeLoopId = requestAnimationFrame(loop); }; this.realtimeLoopId = requestAnimationFrame(loop); } stopRealtimeUpscale() { if (this.realtimeLoopId) { cancelAnimationFrame(this.realtimeLoopId); this.realtimeLoopId = null; } if (this.realtimeState) { this.realtimeState.running = false; this.realtimeState = null; } } async downloadUpscaled(file, networkSize, filename = null) { const isVideo = file.type.startsWith("video"); const size = networkSize || this.networkSize; const blob = isVideo ? await this.upscaleVideo(file, size) : await this.upscaleImage(file, size); const url = URL.createObjectURL(blob); const link = document.createElement("a"); link.href = url; link.download = filename || `upscaled_${Date.now()}.${isVideo ? "mp4" : "png"}`; link.click(); URL.revokeObjectURL(url); return blob; } getSupportedNetworks() { return Object.keys(this.networks); } static async getGPUCapability() { try { if (navigator.gpu) { const adapter = await navigator.gpu.requestAdapter(); if (adapter) { const buffer = adapter.limits.maxStorageBufferBindingSize / 512; const texture = 8192 * 8192; return Math.min(buffer, texture); } } const canvas = document.createElement("canvas"); const gl = canvas.getContext("webgl2"); if (gl) return gl.getParameter(gl.MAX_TEXTURE_SIZE) ** 2; return 8388608; } catch { return 8388608; } } static getBackendType() { try { if (navigator.gpu) return "webgpu"; if (document.createElement("canvas").getContext("webgl2")) return "webgl"; return "unknown"; } catch { return "unknown"; } } getAppState() { return { backend: Upscaler.getBackendType(), isProcessing: false, progress: 0, width: 0, height: 0, }; } onProgress(callback) { if (typeof callback === "function") { this.progressCallback = callback; } } async upscaleVideoBatch(files, options = {}) { const { networkSize, directoryHandle, onFileStart, onFileProgress, onFileComplete, onFileError, } = options; const results = []; const size = networkSize || this.networkSize; for (let i = 0; i < files.length; i++) { const file = files[i]; if (!file.type.startsWith("video/")) { const err = new Error(`Skipped (not a video): ${file.name}`); if (onFileError) onFileError(file, i, err); results.push({ file, error: err, success: false }); continue; } try { if (onFileStart) onFileStart(file, i, files.length); this.onProgress((progress) => { if (onFileProgress) onFileProgress(file, i, progress); }); const blob = await this.upscaleVideo(file, size); if (directoryHandle) { const outName = this._buildOutputName(file.name); const fh = await directoryHandle.getFileHandle(outName, { create: true, }); const writable = await fh.createWritable(); await writable.write(blob); await writable.close(); } if (onFileComplete) onFileComplete(file, i, blob); results.push({ file, blob, success: true }); } catch (error) { if (onFileError) onFileError(file, i, error); results.push({ file, error, success: false }); } } this.progressCallback = null; return results; } _buildOutputName(originalName) { const dot = originalName.lastIndexOf("."); const base = dot > 0 ? originalName.slice(0, dot) : originalName; return `${base}_upscaled.mp4`; } dispose() { this.stopRealtimeUpscale(); if (this.workerInstance) { this.workerInstance.terminate(); this.workerInstance = null; } this.messageHandlers = {}; this.progressCallback = null; this.processingType = null; this.weightsCache.clear(); } } window.Upscaler = Upscaler;