109 lines
7.0 KiB
JavaScript
109 lines
7.0 KiB
JavaScript
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/)
|
|
})
|