Add local queued video upscale with Real-ESRGAN and RIFE

scripts/comfy-host-agent.mjs

@@ -1,3 +1,4 @@
+import { createUpscaleHost } from './upscale-host.mjs'
 import { createGpuReservation } from './gpu-reservation.mjs'
 import { createGpuProxy } from './gpu-proxy.mjs'
 import { createYueGpHost } from './yuegp-host.mjs'
@@ -144,7 +145,7 @@ let proxyTarget = 0
 
 function ensureProxyListening() {
   if (proxyServer) return
-  proxyServer = createGpuProxy({ target: () => proxyTarget, reservation: gpuReservation, authorized, markWork, externalBusy: () => yueGp.busy() })
+  proxyServer = createGpuProxy({ target: () => proxyTarget, reservation: gpuReservation, authorized, markWork, externalBusy: () => yueGp.busy() || upscale.busy() })
   proxyServer.on('error', (error) => {
     console.log(JSON.stringify({ src: 'comfy-host-agent', event: 'proxy-error', error: String(error.message || error) }))
   })
@@ -813,7 +814,7 @@ function purgeDesktopFiles(body) {
 }
 
 const gpuReservation = createGpuReservation({ idle: async () => {
-  if (yueGp.busy()) return false
+  if (yueGp.busy() || upscale.busy()) return false
   if ((await trainingLock()).busy) return false
   const healthy = await syncProxy()
   if (healthy) {
@@ -839,9 +840,20 @@ const yueGp = createYueGpHost({
   }
 })
 
+const upscale = createUpscaleHost({ leaseValid: token => gpuReservation.isOwner(token) })
+
 async function handleControl(req, res) {
   if (!authorized(req)) return json(res, 401, { ok: false, error: 'unauthorized' })
   const url = new URL(req.url || '/', 'http://localhost')
+  if (url.pathname.startsWith('/upscale/')) {
+    const match = url.pathname.match(/^\/upscale\/jobs\/([a-zA-Z0-9-]{12,80})(\/input|\/video|\/cancel)?$/)
+    if (req.method === 'POST' && url.pathname === '/upscale/jobs') return json(res, 200, upscale.start(await readJson(req), String(req.headers['x-aigen-gpu-lease'] || '')))
+    if (match && req.method === 'PUT' && match[2] === '/input') { await upscale.upload(match[1], req); return json(res, 200, { uploaded: true }) }
+    if (match && req.method === 'POST' && match[2] === '/cancel') return json(res, 200, await upscale.cancel(match[1]))
+    if (match && req.method === 'GET' && match[2] === '/video') { const path = upscale.output(match[1]); return path ? streamFile(res, path) : json(res, 404, { error: 'Video not ready' }) }
+    if (match && req.method === 'GET' && !match[2]) { const job = upscale.read(match[1]); return json(res, job ? 200 : 404, job || { error: 'Upscale job not found' }) }
+    return json(res, 404, { error: 'Unknown upscale endpoint' })
+  }
   if (url.pathname.startsWith('/yuegp/')) {
     const match = url.pathname.match(/^\/yuegp\/jobs\/([a-zA-Z0-9-]{12,80})(\/audio|\/cancel)?$/)
     if (req.method === 'GET' && url.pathname === '/yuegp/status') return json(res, 200, { configured: yueGp.configured(), busy: yueGp.busy(), backend: 'yuegp' })
@@ -888,10 +900,12 @@ async function handleControl(req, res) {
       proxyPort,
       gpu: gpuReservation.availability(),
       training: { busy: lastTraining.busy },
-      yuegp: { busy: yueGp.busy(), configured: yueGp.configured() }
+      yuegp: { busy: yueGp.busy(), configured: yueGp.configured() },
+      upscale: { busy: upscale.busy(), engine: 'realesrgan-rife', local: true }
     })
   }
   if (req.method === 'POST' && url.pathname === '/start') {
+    if (upscale.busy()) return json(res, 409, { message: 'Local upscale is using the GPU.' })
     if (yueGp.busy()) return json(res, 409, { message: 'YuEGP is using the GPU.' })
     const training = await trainingLock()
     if (training.busy) {
@@ -983,12 +997,12 @@ const server = http.createServer(async (req, res) => {
   if (!authorized(req)) return json(res, 401, { ok: false, error: 'unauthorized' })
   try {
     const path = new URL(req.url || '/', 'http://localhost').pathname
-    if (req.method === 'POST' && !path.startsWith('/gpu/')) {
+    if (['POST', 'PUT'].includes(req.method) && !path.startsWith('/gpu/')) {
       await gpuReservation.permit(String(req.headers['x-aigen-gpu-lease'] || ''), () => handleControl(req, res))
     } else await handleControl(req, res)
   } catch (error) {
     req.resume()
-    if (String(req.url || '').startsWith('/yuegp/') && !res.headersSent) return json(res, error.statusCode || 400, { error: error.message || 'YuEGP request failed' })
+    if ((String(req.url || '').startsWith('/yuegp/') || String(req.url || '').startsWith('/upscale/')) && !res.headersSent) return json(res, error.statusCode || 400, { error: error.message || 'YuEGP request failed' })
     if (!res.headersSent) json(res, error.statusCode || 400, { ok: false, message: error.statusCode === 409 ? 'GPU is in use. Waiting for availability.' : 'GPU coordination request failed.' })
   }
 })

scripts/setup-upscale.ps1

@@ -0,0 +1,30 @@
+param([string]$Root = $(if ($env:UPSCALE_ROOT) { $env:UPSCALE_ROOT } else { Join-Path (Split-Path $PSScriptRoot) '../AIGen-Upscale' }))
+$ErrorActionPreference = 'Stop'
+$ProgressPreference = 'SilentlyContinue'
+$Root = [IO.Path]::GetFullPath($Root)
+New-Item -ItemType Directory -Force -Path $Root | Out-Null
+function Install-Package($Name, $Url, $Executable) {
+  $Destination = Join-Path $Root $Name
+  $Ready = Test-Path (Join-Path $Destination $Executable)
+  if ($Name -eq 'esrgan') { $Ready = $Ready -and (Test-Path (Join-Path $Destination 'models/realesrgan-x4plus.bin')) -and (Test-Path (Join-Path $Destination 'models/realesrgan-x4plus.param')) }
+  if ($Name -eq 'rife') { $Ready = $Ready -and (Test-Path (Join-Path $Destination 'rife-v4*/flownet.bin')) -and (Test-Path (Join-Path $Destination 'rife-v4*/flownet.param')) }
+  if ($Name -eq 'ffmpeg') { $Ready = $Ready -and (Test-Path (Join-Path $Destination 'ffprobe.exe')) }
+  if ($Ready) { return }
+  $Zip = Join-Path $Root "$Name.zip"
+  Write-Output "Downloading $Name tools and weights..."
+  Invoke-WebRequest -UseBasicParsing -Uri $Url -OutFile $Zip -TimeoutSec 1800
+  $Unpack = Join-Path $Root "$Name-unpack"
+  New-Item -ItemType Directory -Force -Path $Unpack | Out-Null
+  Expand-Archive -LiteralPath $Zip -DestinationPath $Unpack -Force
+  $Exe = Get-ChildItem -LiteralPath $Unpack -Recurse -Filter $Executable | Select-Object -First 1
+  if (!$Exe) { throw "Package is missing $Executable" }
+  New-Item -ItemType Directory -Force -Path $Destination | Out-Null
+  Copy-Item -Path (Join-Path $Exe.DirectoryName '*') -Destination $Destination -Recurse -Force
+}
+Install-Package 'esrgan' 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesrgan-ncnn-vulkan-20220424-windows.zip' 'realesrgan-ncnn-vulkan.exe'
+Install-Package 'rife' 'https://github.com/nihui/rife-ncnn-vulkan/releases/download/20221029/rife-ncnn-vulkan-20221029-windows.zip' 'rife-ncnn-vulkan.exe'
+Install-Package 'ffmpeg' 'https://www.gyan.dev/ffmpeg/builds/ffmpeg-release-essentials.zip' 'ffmpeg.exe'
+if (!(Test-Path (Join-Path $Root 'ffmpeg/ffprobe.exe'))) { throw 'FFprobe is missing.' }
+if (!(Test-Path (Join-Path $Root 'esrgan/models/realesrgan-x4plus.bin'))) { throw 'Real-ESRGAN weights are missing.' }
+if (!(Get-ChildItem -LiteralPath (Join-Path $Root 'rife') -Directory -Filter 'rife-v4*')) { throw 'RIFE v4 weights are missing.' }
+Write-Output "Local upscale tools and models ready: $Root"

scripts/upscale-host.mjs

@@ -0,0 +1,100 @@
+import { spawn } from 'node:child_process'
+import { existsSync, mkdirSync, readFileSync, writeFileSync, renameSync, createWriteStream, appendFileSync, readdirSync } from 'node:fs'
+import { pipeline } from 'node:stream/promises'
+import { Transform } from 'node:stream'
+import { join, resolve } from 'node:path'
+import { fileURLToPath } from 'node:url'
+import { upscaleOptions } from '../shared/video-upscale.mjs'
+
+export function createUpscaleHost({ leaseValid, spawnProcess = spawn, root = process.env.UPSCALE_ROOT || fileURLToPath(new URL('../../AIGen-Upscale', import.meta.url)) }) {
+  root = resolve(root)
+  const data = join(root, 'jobs')
+  mkdirSync(data, { recursive: true })
+  let active = null
+  const directory = id => {
+    if (!/^[a-zA-Z0-9-]{12,80}$/.test(id || '')) throw new Error('Invalid upscale ID.')
+    return join(data, id)
+  }
+  const persist = state => { const path = join(directory(state.id), 'status.json'); writeFileSync(path + '.tmp', JSON.stringify(state)); renameSync(path + '.tmp', path) }
+  // A previous worker's parent watchdog terminates its process tree before another GPU job starts.
+  let holdUntil = 0
+  for (const id of readdirSync(data)) {
+    try {
+      const state = JSON.parse(readFileSync(join(directory(id), 'status.json'), 'utf8'))
+      if (state.status === 'running') { holdUntil = Date.now() + 15000; state.status = 'error'; state.error = 'Upscale host restarted. Original video is safe.'; persist(state) }
+    } catch { /* Not a valid job record. */ }
+  }
+  const read = id => {
+    const path = join(directory(id), 'status.json')
+    return existsSync(path) ? JSON.parse(readFileSync(path, 'utf8')) : null
+  }
+  const stop = () => new Promise((resolve, reject) => {
+    if (!active?.child?.pid) return resolve()
+    const child = active.child
+    child.once('close', resolve)
+    const killer = spawn('taskkill', ['/PID', String(child.pid), '/T', '/F'], { windowsHide: true, stdio: 'ignore' })
+    killer.once('error', reject)
+    killer.once('close', code => { if (code && child.exitCode === null) reject(new Error('Could not stop upscale worker. GPU remains reserved.')) })
+  })
+  return {
+    busy: () => Boolean(active) || Date.now() < holdUntil,
+    read,
+    output: id => read(id)?.status === 'complete' ? join(directory(id), 'output.mp4') : null,
+    async upload(id, stream) {
+      if (read(id)) throw new Error('This upscale job has already been submitted.')
+      const dir = directory(id); mkdirSync(dir, { recursive: true })
+      let bytes = 0
+      await pipeline(stream, new Transform({ transform(chunk, _, callback) {
+        bytes += chunk.length
+        callback(bytes > 8 * 1024 ** 3 ? new Error('Video exceeds the 8 GB upload limit.') : null, chunk)
+      } }), createWriteStream(join(dir, 'input.upload'), { flags: 'w' }))
+      if (bytes < 32) throw new Error('Video file is empty.')
+      renameSync(join(dir, 'input.upload'), join(dir, 'input.mp4'))
+    },
+    start(body, lease) {
+      const options = upscaleOptions(body), dir = directory(body.id)
+      const previous = read(body.id)
+      if (previous) return previous
+      if (active || Date.now() < holdUntil) throw Object.assign(new Error('Local upscaler is busy.'), { statusCode: 409 })
+      if (!leaseValid(lease)) throw new Error('GPU reservation expired.')
+      if (!existsSync(join(dir, 'input.mp4'))) throw new Error('Upload a completed video first.')
+      const state = { id: body.id, status: 'running', startedAt: Date.now(), message: 'Preparing local upscaler', progress: 0 }
+      writeFileSync(join(dir, 'request.json'), JSON.stringify({ ...options, root, parentPid: process.pid }))
+      persist(state)
+      const child = spawnProcess(process.execPath, [fileURLToPath(new URL('./upscale-worker.mjs', import.meta.url)), join(dir, 'request.json')], { windowsHide: true, shell: false, stdio: ['ignore', 'pipe', 'pipe'] })
+      active = { child, state }
+      let lines = '', tail = '', cancelled = false
+      const watchdog = setInterval(() => { if (!leaseValid(lease)) { state.error = 'GPU reservation expired.'; void stop().catch(() => {}) } }, 5000)
+      child.stderr.on('data', chunk => { appendFileSync(join(dir, 'worker.log'), chunk); tail = (tail + chunk).slice(-2000) })
+      child.stdout.on('data', chunk => {
+        appendFileSync(join(dir, 'worker.log'), chunk)
+        lines += chunk
+        const parts = lines.split('\n'); lines = parts.pop().slice(-64000)
+        for (const line of parts) if (line.startsWith('AIGEN_EVENT ')) {
+          try { const event = JSON.parse(line.slice(12)); for (const key of ['message', 'progress', 'error', 'width', 'height', 'sourceWidth', 'sourceHeight', 'fps', 'duration', 'engine', 'elapsedMs']) if (event[key] !== undefined) state[key] = event[key]; persist(state) } catch {}
+        }
+      })
+      active.cancel = async () => { cancelled = true; await stop() }
+      let finished = false
+      const finish = error => {
+        if (finished) return
+        finished = true; clearInterval(watchdog)
+        state.status = cancelled ? 'cancelled' : !error && !state.error && existsSync(join(dir, 'output.mp4')) ? 'complete' : 'error'
+        state.elapsedMs = Date.now() - state.startedAt
+        if (state.status === 'error') state.error ||= error?.message || tail || 'Upscale failed. Original video is safe.'
+        persist(state); active = null
+      }
+      child.once('error', finish)
+      child.once('close', code => finish(code ? new Error(`Upscaler exited (${code}). ${tail}`) : null))
+      return state
+    },
+    async cancel(id) {
+      if (active?.state.id === id) await active.cancel()
+      if (!read(id)) {
+        mkdirSync(directory(id), { recursive: true })
+        persist({ id, status: 'cancelled', message: 'Cancelled before processing', elapsedMs: 0 })
+      }
+      return read(id)
+    }
+  }
+}

scripts/upscale-worker.mjs

@@ -0,0 +1,84 @@
+import { spawn } from 'node:child_process'
+import { readFileSync, mkdirSync, readdirSync, rmSync, existsSync, writeFileSync } from 'node:fs'
+import { join, dirname, resolve } from 'node:path'
+import { fileURLToPath } from 'node:url'
+import { upscaleOptions, upscaleDimensions } from '../shared/video-upscale.mjs'
+
+export async function upscaleVideo(requestPath, { runCommand, report } = {}) {
+const request = JSON.parse(readFileSync(requestPath, 'utf8'))
+const root = request.root, dir = dirname(requestPath), source = join(dir, 'input.mp4')
+const options = upscaleOptions(request)
+const started = Date.now()
+let child
+function event(message, extra = {}) { if (report) report({ message, ...extra }); else process.stdout.write(`AIGEN_EVENT ${JSON.stringify({ message, ...extra })}\n`) }
+function execute(exe, args, cwd = dir) {
+  return new Promise((resolve, reject) => {
+    child = spawn(exe, args, { cwd, windowsHide: true, shell: false, stdio: ['ignore', 'pipe', 'pipe'] })
+    let out = '', tail = ''
+    child.stdout.on('data', b => { out = (out + b).slice(-2_000_000) })
+    child.stderr.on('data', b => { tail = (tail + b).slice(-6000) })
+    child.once('error', reject)
+    child.once('close', code => { child = null; code === 0 ? resolve(out) : reject(new Error(/out of memory|allocate|vkAllocateMemory/i.test(tail) ? 'GPU out of memory. Try 2x with Keep original FPS.' : `${exe.split(/[\\/]/).pop()} failed: ${tail.slice(-1500)}`)) })
+  })
+}
+const run = runCommand || execute
+// Host cancellation kills this entire tree. Also stop descendants if the host disappears.
+const watchdog = setInterval(() => {
+  try { process.kill(Number(request.parentPid), 0) } catch {
+    spawn('taskkill', ['/PID', String(process.pid), '/T', '/F'], { windowsHide: true })
+  }
+}, 2000)
+try {
+  event('Checking local tools and weights', { progress: 1 })
+  await run('powershell.exe', ['-NoProfile', '-NonInteractive', '-ExecutionPolicy', 'Bypass', '-File', fileURLToPath(new URL('./setup-upscale.ps1', import.meta.url)), '-Root', root])
+  const ffmpeg = join(root, 'ffmpeg', 'ffmpeg.exe'), ffprobe = join(root, 'ffmpeg', 'ffprobe.exe')
+  const probe = JSON.parse(await run(ffprobe, ['-v', 'error', '-show_streams', '-show_format', '-of', 'json', source]))
+  const video = probe.streams.find(s => s.codec_type === 'video')
+  if (!video) throw new Error('Input has no video stream.')
+  const rate = String(video.avg_frame_rate).split('/').map(Number)
+  const fps = rate[0] / (rate[1] || 1), duration = Number(video.duration || probe.format.duration)
+  if (!Number.isFinite(fps) || fps <= 0 || fps > 240 || !Number.isFinite(duration) || duration <= 0) throw new Error('Cannot read source frame rate or duration.')
+  const target = upscaleDimensions(video.width, video.height, options)
+  const outputFps = options.fps === 'keep' ? fps : options.fps
+  const interpolate = options.fps !== 'keep'
+  const multiplier = interpolate ? Math.max(2, Math.ceil(outputFps / fps)) : 1
+  const totalFrames = Math.max(1, Math.round(duration * fps))
+  const chunkFrames = Math.max(2, Math.round(fps * 2))
+  const rifeDir = join(root, 'rife')
+  const rifeModel = readdirSync(rifeDir).filter(n => /^rife-v4/.test(n)).sort().at(-1)
+  const info = { sourceWidth: video.width, sourceHeight: video.height, ...target, duration, fps: outputFps, engine: 'Real-ESRGAN' + (interpolate ? ' + RIFE' : '') }
+  event('Upscaling video', { ...info, progress: 3 })
+  const segments = []
+  for (let start = 0, index = 0; start < totalFrames; start += chunkFrames, index++) {
+    const count = Math.min(chunkFrames, totalFrames - start)
+    const work = join(dir, `chunk-${index}`)
+    for (const folder of ['in', 'esrgan', 'sized', 'rife']) mkdirSync(join(work, folder), { recursive: true })
+    const pattern = folder => join(work, folder, '%08d.png')
+    await run(ffmpeg, ['-v', 'error', '-ss', String(start / fps), '-i', source, '-an', '-vf', `fps=${fps},tpad=stop_mode=clone:stop_duration=1`, '-frames:v', String(count + 1), '-start_number', '1', pattern('in')])
+    event(`Upscaling ${Math.round(start / fps)}–${Math.min(Math.round((start + count) / fps), Math.ceil(duration))}s`, { progress: 3 + 90 * start / totalFrames })
+    await run(join(root, 'esrgan', 'realesrgan-ncnn-vulkan.exe'), ['-i', join(work, 'in'), '-o', join(work, 'esrgan'), '-m', join(root, 'esrgan', 'models'), '-n', 'realesrgan-x4plus', '-s', '4', '-t', String(process.env.UPSCALE_TILE || 128), '-g', process.env.UPSCALE_GPU_ID || '0', '-j', '1:1:1', '-f', 'png'])
+    await run(ffmpeg, ['-v', 'error', '-i', pattern('esrgan'), '-vf', `scale=${target.width}:${target.height}:flags=lanczos${options.enhance === 'sharpen' ? ',unsharp=5:5:0.25:5:5:0' : ''}`, '-start_number', '1', pattern('sized')])
+    if (interpolate) {
+      event(`Smoothing motion at ${Math.round(start / fps)}s`, { progress: 3 + 90 * (start + count / 2) / totalFrames })
+      await run(join(rifeDir, 'rife-ncnn-vulkan.exe'), ['-i', join(work, 'sized'), '-o', join(work, 'rife'), '-m', join(rifeDir, rifeModel), '-n', String((count + 1) * multiplier), '-g', process.env.UPSCALE_GPU_ID || '0', '-j', '1:1:1', '-u', '-f', '%08d.png'], rifeDir)
+    }
+    const outputCount = Math.round((start + count) / fps * outputFps) - Math.round(start / fps * outputFps)
+    const segment = `part-${index}.mp4`
+    await run(ffmpeg, ['-v', 'error', '-framerate', String(fps * multiplier), '-i', pattern(interpolate ? 'rife' : 'sized'), '-vf', `fps=${outputFps},setsar=1`, '-frames:v', String(outputCount), '-an', '-c:v', 'libx264', '-preset', 'medium', '-crf', '17', '-pix_fmt', 'yuv420p', '-video_track_timescale', '90000', join(dir, segment)])
+    segments.push(`file '${segment}'`)
+    rmSync(work, { recursive: true, force: true }) // Fixed job-local directory, never an input path.
+  }
+  writeFileSync(join(dir, 'concat.txt'), segments.join('\n'))
+  event('Joining video and preserving original audio', { progress: 96 })
+  await run(ffmpeg, ['-v', 'error', '-f', 'concat', '-safe', '1', '-i', join(dir, 'concat.txt'), '-i', source, '-map', '0:v:0', '-map', '1:a?', '-c:v', 'copy', '-c:a', 'copy', '-map_metadata', '1', '-t', String(duration), '-movflags', '+faststart', join(dir, 'output.mp4')])
+  if (!existsSync(join(dir, 'output.mp4'))) throw new Error('Upscaler did not produce a video.')
+  event('Upscale complete', { ...info, progress: 100, elapsedMs: Date.now() - started })
+} catch (error) {
+  event(error.message, { error: error.message, elapsedMs: Date.now() - started })
+  throw error
+} finally { clearInterval(watchdog) }
+
+}
+if (process.argv[1] && resolve(process.argv[1]) === fileURLToPath(import.meta.url)) {
+  upscaleVideo(process.argv[2]).catch(() => { process.exitCode = 1 })
+}