86 lines
2.6 KiB
TypeScript
86 lines
2.6 KiB
TypeScript
import { spawn } from "node:child_process";
|
||
import path from "node:path";
|
||
import { downloadCanvasMedia } from "../../lib/video-canvas/download-media";
|
||
import {
|
||
videoUpscaleSchema,
|
||
type VideoUpscaleInput,
|
||
} from "../../lib/video-canvas/video-upscale-contract";
|
||
export async function upscaleJob(
|
||
value: VideoUpscaleInput,
|
||
_tenant: string,
|
||
dir: string,
|
||
progress: (stage: string) => Promise<void>
|
||
) {
|
||
const input = videoUpscaleSchema.parse(value);
|
||
const file = path.join(dir, "input.bin");
|
||
await downloadCanvasMedia(input.url, file);
|
||
await progress("GPU 正在进行 AI 超分");
|
||
return new Promise<any>((resolve, reject) => {
|
||
const child = spawn(
|
||
process.env.UPSCALE_PYTHON || "/opt/video-upscale/venv/bin/python",
|
||
[
|
||
process.env.UPSCALE_SCRIPT || path.join(__dirname, "upscale.py"),
|
||
"--input",
|
||
file,
|
||
"--output",
|
||
path.join(dir, "result.mp4"),
|
||
"--resolution",
|
||
input.resolution,
|
||
"--model",
|
||
process.env.UPSCALE_MODEL ||
|
||
"/opt/video-upscale/models/realesr-general-x4v3.pth",
|
||
],
|
||
{ stdio: ["ignore", "pipe", "pipe"], detached: true }
|
||
);
|
||
let buffer = "",
|
||
result: any,
|
||
error = "",
|
||
stages = Promise.resolve();
|
||
const timer = setTimeout(() => {
|
||
try {
|
||
process.kill(-child.pid!, "SIGKILL");
|
||
} catch {}
|
||
}, 2 * 3600000);
|
||
child.stdout.on("data", (chunk) => {
|
||
buffer += chunk;
|
||
const lines = buffer.split("\n");
|
||
buffer = lines.pop()!;
|
||
for (const line of lines)
|
||
try {
|
||
const value = JSON.parse(line);
|
||
if (value.result) result = value.result;
|
||
if (value.stage)
|
||
stages = stages
|
||
.then(() => progress(value.stage))
|
||
.catch(() => {
|
||
try {
|
||
process.kill(-child.pid!, "SIGKILL");
|
||
} catch {}
|
||
});
|
||
} catch {}
|
||
});
|
||
child.stderr.on("data", (chunk) => {
|
||
error = (error + chunk).slice(-4000);
|
||
});
|
||
child.on("error", (e) => {
|
||
clearTimeout(timer);
|
||
reject(e);
|
||
});
|
||
child.on("close", async (code) => {
|
||
clearTimeout(timer);
|
||
await stages.catch(() => {});
|
||
if (code !== 0 || !result) {
|
||
console.error(
|
||
"Upscale inference failed",
|
||
error.replace(/https?:\/\/\S+/g, "[url]")
|
||
);
|
||
reject(
|
||
Error(
|
||
"AI 超分失败，请检查视频是否为恒定帧率、120 秒以内，且放大倍率为 1–4 倍"
|
||
)
|
||
);
|
||
} else resolve({ ...result, duration: result.durationSeconds });
|
||
});
|
||
});
|
||
}
|