Add local queued video upscale with Real-ESRGAN and RIFE

2026-09-08 17:55:02 -05:00
parent 4c97980646
commit 1e0cff9dea
24 changed files with 747 additions and 16 deletions
@@ -27,6 +27,7 @@ function fixture({ rows = [], lives = [], queue = { running: 0, pending: 0 }, hi
  const timers = []
  const started = []
  const modules = {
+    './videoUpscale': { startUpscaleJob: async item => { started.push({ upscale: true, sourceId: item.payload.upscale.sourceId }) } },
    '~/server/utils/sharedGpu': { withSharedGpuStart: async fn => fn(), maintainSharedGpu: async () => {}, acquireSharedGpu: async () => false, sharedGpuWaitReason: () => 'GPU is in use. Waiting for availability.' },
    '~/server/utils/jobs': {
      listJobs: () => lives,
@@ -93,6 +94,20 @@ test('standalone YuEGP retains its queue slot while Comfy is empty for longer th
  assert.equal(f.started.length, 0)
  assert.equal(f.rows()[1].status, 'waiting')
 })
+test('standalone upscale retains its queue slot while Comfy is empty for longer than the zombie timeout', async () => {
+  const live = liveJob('upscale-active', 'video', { upscale: true, promptId: undefined, startedAt: Date.now() - 3600000 })
+  const f = fixture({ rows: [row('upscale-active'), row('next', 'waiting', undefined)], lives: [live] })
+  await f.api.kickStudioQueue()
+  assert.equal(live.status, 'running')
+  assert.equal(f.started.length, 0)
+  assert.equal(f.rows()[1].status, 'waiting')
+})
+test('upscale queue row dispatches to the standalone worker before any video generation path', async () => {
+  const item = { ...row('upscale-waiting', 'waiting', undefined), kind: 'video', payload: { prompt: 'existing video', workflow: 'ltx-cherry', extensions: [], upscale: { sourceId: 'finished-master', scale: 2, fps: 'keep', target: 'preserve', enhance: 'faithful' } } }
+  const f = fixture({ rows: [item], queue: null })
+  await f.api.startStudioJob(item)
+  assert.deepEqual(f.started, [{ upscale: true, sourceId: 'finished-master' }])
+})
 async function finishes(promise) {
  let timer
  try {
@@ -0,0 +1,108 @@
+import test from 'node:test'
+import assert from 'node:assert/strict'
+import { EventEmitter } from 'node:events'
+import { PassThrough, Readable } from 'node:stream'
+import { mkdtempSync, mkdirSync, writeFileSync, readFileSync, existsSync } from 'node:fs'
+import { tmpdir } from 'node:os'
+import { join } from 'node:path'
+import { upscaleOptions, upscaleDimensions, upscaleName } from '../shared/video-upscale.mjs'
+import { createUpscaleHost } from '../scripts/upscale-host.mjs'
+import { upscaleVideo } from '../scripts/upscale-worker.mjs'
+import { execFile } from 'node:child_process'
+import { promisify } from 'node:util'
+const exec = promisify(execFile)
+
+test('upscale defaults, delivery sizes, portrait and square aspect, and collision names', () => {
+  assert.deepEqual(upscaleOptions(), { scale: 2, target: 'preserve', fps: 'keep', enhance: 'faithful' })
+  assert.deepEqual(upscaleDimensions(1344, 768, upscaleOptions()), { width: 2688, height: 1536 })
+  assert.deepEqual(upscaleDimensions(768, 1344, upscaleOptions({ target: '1080p' })), { width: 1098, height: 1920 })
+  assert.deepEqual(upscaleDimensions(960, 960, upscaleOptions({ target: '4k' })), { width: 3840, height: 3840 })
+  assert.equal(upscaleName('master', 4, 3), 'master_up4x-3.mp4')
+  for (const value of [{ scale: 3 }, { fps: 30 }, { target: 'crop' }, { enhance: 'creative' }]) assert.throws(() => upscaleOptions(value))
+})
+function fixture() {
+  const root = mkdtempSync(join(tmpdir(), 'aigen-upscale-'))
+  const calls = [], child = new EventEmitter()
+  child.stdout = new PassThrough(); child.stderr = new PassThrough()
+  const host = createUpscaleHost({ root, leaseValid: token => token === 'lease', spawnProcess: (...args) => { calls.push(args); return child } })
+  return { root, calls, child, host }
+}
+test('standalone upscale requires a lease, uploads locally, deduplicates, and reports OOM after exit', async () => {
+  const f = fixture(), id = 'test-upscale-12345'
+  await f.host.upload(id, Readable.from(Buffer.alloc(64)))
+  assert.throws(() => f.host.start({ id }, 'invalid'), /reservation/)
+  f.host.start({ id }, 'lease'); f.host.start({ id }, 'lease')
+  assert.equal(f.calls.length, 1)
+  assert.equal(f.calls[0][2].windowsHide, true)
+  assert.match(f.calls[0][1][0], /upscale-worker/)
+  assert.equal(f.host.busy(), true)
+  f.child.stdout.write('AIGEN_EVENT {"error":"GPU out of ')
+  f.child.stdout.write('memory"}\n')
+  assert.equal(f.host.output(id), null)
+  f.child.emit('close', 1)
+  assert.equal(f.host.busy(), false)
+  assert.equal(f.host.read(id).status, 'error')
+  assert.match(f.host.read(id).error, /out of memory/)
+  assert.equal(readFileSync(join(f.root, 'jobs', id, 'input.mp4')).length, 64)
+})
+test('cancel before submission is a durable tombstone and cannot later start a GPU worker', async () => {
+  const f = fixture(), id = 'cancel-upscale-12345'
+  await f.host.cancel(id)
+  assert.equal(f.host.start({ id }, 'lease').status, 'cancelled')
+  assert.equal(f.calls.length, 0)
+  await assert.rejects(f.host.upload('../escape', Readable.from('x')), /Invalid/)
+})
+test('two-minute fractional-FPS master stays one job; ESRGAN precedes RIFE and audio is stream-copied', async () => {
+  const root = mkdtempSync(join(tmpdir(), 'aigen-upscale-plan-'))
+  mkdirSync(join(root, 'rife', 'rife-v4.6'), { recursive: true })
+  const request = join(root, 'request.json'), events = [], calls = []
+  writeFileSync(request, JSON.stringify({ root, parentPid: process.pid, fps: 50 }))
+  let outputFrames = 0
+  await upscaleVideo(request, { report: e => events.push(e), runCommand: async (exe, args) => {
+    calls.push({ exe, args })
+    if (exe.endsWith('ffprobe.exe')) return JSON.stringify({ streams: [{ codec_type: 'video', width: 1344, height: 768, avg_frame_rate: '30000/1001', duration: '120.02' }] })
+    if (args.includes('libx264')) outputFrames += Number(args[args.indexOf('-frames:v') + 1])
+    if (args.at(-1) === join(root, 'output.mp4')) writeFileSync(args.at(-1), 'mock result')
+    return ''
+  } })
+  assert.equal(outputFrames, Math.round(Math.round(120.02 * 30000 / 1001) / (30000 / 1001) * 50))
+  const engines = calls.filter(c => /ncnn-vulkan/.test(c.exe))
+  assert.ok(engines.length > 100)
+  for (let i = 0; i < engines.length; i += 2) { assert.match(engines[i].exe, /realesrgan/); assert.match(engines[i + 1].exe, /rife/) }
+  const mux = calls.at(-1).args
+  assert.equal(mux[mux.indexOf('-c:a') + 1], 'copy')
+  assert.ok(mux.includes('1:a?'))
+  assert.ok(!mux.includes('-shortest'))
+  assert.equal(events.at(-1).progress, 100)
+  assert.ok(existsSync(join(root, 'output.mp4')))
+})
+
+const cpuFfmpeg = process.env.UPSCALE_TEST_FFMPEG || join(process.cwd(), '.data', 'video-test-tools', 'ffmpeg.exe')
+for (const hasAudio of [true, false]) test(`CPU-only fixture: chunk boundary, frame count and ${hasAudio ? 'bit-identical audio' : 'video-only output'}`, { skip: !existsSync(cpuFfmpeg) }, async () => {
+  const root = mkdtempSync(join(tmpdir(), 'aigen-upscale-cpu-'))
+  mkdirSync(join(root, 'rife', 'rife-v4.6'), { recursive: true })
+  const source = join(root, 'input.mp4'), result = join(root, 'output.mp4')
+  const ffmpeg = async args => (await exec(cpuFfmpeg, args, { windowsHide: true, maxBuffer: 4_000_000 })).stdout
+  await ffmpeg(['-v', 'error', '-f', 'lavfi', '-i', 'testsrc2=size=48x32:rate=24:duration=2.5', ...(hasAudio ? ['-f', 'lavfi', '-i', 'sine=frequency=440:duration=2.5'] : []), '-c:v', 'libx264', ...(hasAudio ? ['-c:a', 'aac'] : []), source])
+  const original = readFileSync(source)
+  writeFileSync(join(root, 'request.json'), JSON.stringify({ root, parentPid: process.pid, fps: hasAudio ? 50 : 'keep' }))
+  await upscaleVideo(join(root, 'request.json'), { report: () => {}, runCommand: async (exe, args) => {
+    if (exe === 'powershell.exe') return ''
+    if (exe.endsWith('ffprobe.exe')) return JSON.stringify({ streams: [{ codec_type: 'video', width: 48, height: 32, avg_frame_rate: '24/1', duration: '2.5' }] })
+    if (exe.endsWith('realesrgan-ncnn-vulkan.exe')) {
+      // Substitute a CPU scaler only in this test. Never execute the neural engine.
+      return ffmpeg(['-v', 'error', '-i', join(args[args.indexOf('-i') + 1], '%08d.png'), '-vf', 'scale=192:128', '-threads', '1', join(args[args.indexOf('-o') + 1], '%08d.png')])
+    }
+    if (exe.endsWith('rife-ncnn-vulkan.exe')) {
+      return ffmpeg(['-v', 'error', '-framerate', '24', '-i', join(args[args.indexOf('-i') + 1], '%08d.png'), '-vf', 'fps=72', '-frames:v', args[args.indexOf('-n') + 1], '-threads', '1', join(args[args.indexOf('-o') + 1], '%08d.png')])
+    }
+    return ffmpeg(args)
+  } })
+  assert.deepEqual(readFileSync(source), original)
+  const hashAudio = file => ffmpeg(['-v', 'error', '-i', file, '-map', '0:a:0', '-c:a', 'copy', '-f', 'hash', '-'])
+  if (hasAudio) assert.equal(await hashAudio(result), await hashAudio(source))
+  else await assert.rejects(hashAudio(result), /matches no streams/)
+  const frames = await ffmpeg(['-v', 'error', '-i', result, '-map', '0:v:0', '-f', 'framemd5', '-'])
+  assert.equal(frames.split('\n').filter(line => /^0,/.test(line)).length, hasAudio ? 125 : 60)
+  assert.match(frames, /dimensions 0: 96x64/)
+})