85 lines
6.5 KiB
JavaScript
85 lines
6.5 KiB
JavaScript
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 })
|
||
}
|