ArtPlayer/packages/artplayer-plugin-websr/src/Upscaler.js

export default 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.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)
    }
    return this.workerInstance
  }

  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
    if (!data.cmd)
      return

    const { cmd } = data

    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 TypeError('OffscreenCanvas is not supported in this environment')
    }
    return canvas.transferControlToOffscreen()
  }

  createBlobPromise(blobType, timeout) {
    return new Promise((resolve, reject) => {
      const timeoutId = setTimeout(() => {
        if (this.messageHandlers[blobType]) {
          delete this.messageHandlers[blobType]
          reject(new Error(`${blobType} processing timeout`))
        }
      }, timeout)

      this.messageHandlers[blobType] = (msg) => {
        delete this.messageHandlers[blobType]
        clearTimeout(timeoutId)
        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)

    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+')
    }

    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
        }
        let onLoaded, onError
        const cleanup = () => {
          videoElement.removeEventListener('loadedmetadata', onLoaded)
          videoElement.removeEventListener('error', onError)
        }
        onLoaded = () => {
          cleanup()
          resolve()
        }
        onError = () => {
          cleanup()
          reject(new Error('Failed to load video metadata'))
        }
        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 TypeError('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
    }
  }

  dispose() {
    this.stopRealtimeUpscale()
    if (this.workerInstance) {
      this.workerInstance.terminate()
      this.workerInstance = null
    }
    this.messageHandlers = {}
    this.progressCallback = null
    this.processingType = null
    this.weightsCache.clear()
  }
}