mirror of
https://github.com/heygen-com/hyperframes.git
synced 2026-09-03 12:54:29 +00:00
feat(cli): add remove-background command for transparent video
Adds `hyperframes remove-background` — a local-AI subcommand that mattes a video or image with the u2net_human_seg ONNX model and emits a transparent WebM (VP9-alpha), ProRes 4444 .mov, or RGBA PNG. Drops directly into any composition's <video> tag — no green screen, no API keys, no upload. Auto-picks the fastest available execution provider via onnxruntime-node: CoreML on Apple Silicon, CUDA when HYPERFRAMES_CUDA=1, CPU otherwise. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
James
@@ -0,0 +1,172 @@
|
||||
/**
|
||||
* u2net_human_seg inference: RGB frame → RGBA frame (alpha = human mask).
|
||||
*
|
||||
* Pre/postprocessing matches rembg's u2net session
|
||||
* (https://github.com/danielgatis/rembg/blob/main/rembg/sessions/u2net.py)
|
||||
* so output should be pixel-equivalent to `rembg new_session("u2net_human_seg")`.
|
||||
*/
|
||||
import type { InferenceSession, Tensor } from "onnxruntime-node";
|
||||
import type sharpType from "sharp";
|
||||
import { ensureModel, selectProviders, type Device, type ModelId } from "./manager.js";
|
||||
|
||||
const INPUT_SIZE = 320;
|
||||
const INPUT_PLANE = INPUT_SIZE * INPUT_SIZE;
|
||||
const MEAN = [0.485, 0.456, 0.406] as const;
|
||||
const STD = [1.0, 1.0, 1.0] as const;
|
||||
|
||||
type Sharp = typeof sharpType;
|
||||
interface OrtModule {
|
||||
InferenceSession: typeof InferenceSession;
|
||||
Tensor: typeof Tensor;
|
||||
}
|
||||
|
||||
export interface Session {
|
||||
/** Run inference on one RGB frame, return RGBA bytes (H*W*4). */
|
||||
process(rgb: Buffer, width: number, height: number): Promise<Buffer>;
|
||||
/** ORT EP that was actually selected. */
|
||||
provider: string;
|
||||
close(): Promise<void>;
|
||||
}
|
||||
|
||||
export interface CreateSessionOptions {
|
||||
model?: ModelId;
|
||||
device?: Device;
|
||||
onProgress?: (message: string) => void;
|
||||
}
|
||||
|
||||
export async function createSession(options: CreateSessionOptions = {}): Promise<Session> {
|
||||
const ort = (await import("onnxruntime-node")) as unknown as OrtModule;
|
||||
const sharpMod = await import("sharp");
|
||||
const sharp = sharpMod.default as Sharp;
|
||||
|
||||
const choice = selectProviders(options.device ?? "auto");
|
||||
const path = await ensureModel(options.model, { onProgress: options.onProgress });
|
||||
|
||||
options.onProgress?.(`Loading model on ${choice.label}...`);
|
||||
|
||||
const tryCreate = (providers: string[]) =>
|
||||
ort.InferenceSession.create(path, {
|
||||
executionProviders: providers,
|
||||
graphOptimizationLevel: "all",
|
||||
});
|
||||
|
||||
let session: InferenceSession;
|
||||
let providerUsed = choice.label;
|
||||
try {
|
||||
session = await tryCreate(choice.providers);
|
||||
} catch (err) {
|
||||
if (choice.providers[0] === "cpu") throw err;
|
||||
options.onProgress?.(
|
||||
`${choice.label} provider failed (${(err as Error).message}); falling back to CPU.`,
|
||||
);
|
||||
session = await tryCreate(["cpu"]);
|
||||
providerUsed = "CPU";
|
||||
}
|
||||
|
||||
const inputName = session.inputNames[0];
|
||||
const outputName = session.outputNames[0];
|
||||
if (!inputName || !outputName) {
|
||||
throw new Error("ONNX session is missing input or output bindings");
|
||||
}
|
||||
|
||||
// Pre-allocated per-frame buffers reused across every process() call.
|
||||
// At 1080p this saves ~9 MB of allocations per frame. rgbaBuf is sized
|
||||
// lazily on the first call (we don't know W/H until then).
|
||||
const inputData = new Float32Array(3 * INPUT_PLANE);
|
||||
const maskBuf = Buffer.allocUnsafe(INPUT_PLANE);
|
||||
let rgbaBuf: Buffer | null = null;
|
||||
|
||||
return {
|
||||
provider: providerUsed,
|
||||
async process(rgb, width, height) {
|
||||
const tensor = await preprocess(sharp, ort, rgb, width, height, inputData);
|
||||
const outputs = await session.run({ [inputName]: tensor });
|
||||
const output = outputs[outputName];
|
||||
if (!output) throw new Error(`Model did not return output '${outputName}'`);
|
||||
const expectedBytes = width * height * 4;
|
||||
if (!rgbaBuf || rgbaBuf.length !== expectedBytes) {
|
||||
rgbaBuf = Buffer.allocUnsafe(expectedBytes);
|
||||
}
|
||||
return await postprocess(sharp, output, rgb, width, height, maskBuf, rgbaBuf);
|
||||
},
|
||||
async close() {
|
||||
await session.release();
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
async function preprocess(
|
||||
sharp: Sharp,
|
||||
ort: OrtModule,
|
||||
rgb: Buffer,
|
||||
width: number,
|
||||
height: number,
|
||||
inputData: Float32Array,
|
||||
): Promise<Tensor> {
|
||||
const resized = await sharp(rgb, { raw: { width, height, channels: 3 } })
|
||||
.resize(INPUT_SIZE, INPUT_SIZE, { kernel: "lanczos3", fit: "fill" })
|
||||
.raw()
|
||||
.toBuffer();
|
||||
|
||||
// rembg's normalize divides by `np.max(im_ary)` (NOT 255). Match exactly so
|
||||
// we hit the same operating point as the model's training distribution.
|
||||
let maxPixel = 0;
|
||||
for (let i = 0; i < resized.length; i++) {
|
||||
if (resized[i]! > maxPixel) maxPixel = resized[i]!;
|
||||
}
|
||||
if (maxPixel === 0) maxPixel = 1;
|
||||
|
||||
for (let y = 0; y < INPUT_SIZE; y++) {
|
||||
for (let x = 0; x < INPUT_SIZE; x++) {
|
||||
const src = (y * INPUT_SIZE + x) * 3;
|
||||
const dst = y * INPUT_SIZE + x;
|
||||
inputData[dst] = (resized[src]! / maxPixel - MEAN[0]) / STD[0];
|
||||
inputData[INPUT_PLANE + dst] = (resized[src + 1]! / maxPixel - MEAN[1]) / STD[1];
|
||||
inputData[2 * INPUT_PLANE + dst] = (resized[src + 2]! / maxPixel - MEAN[2]) / STD[2];
|
||||
}
|
||||
}
|
||||
|
||||
return new ort.Tensor("float32", inputData, [1, 3, INPUT_SIZE, INPUT_SIZE]);
|
||||
}
|
||||
|
||||
async function postprocess(
|
||||
sharp: Sharp,
|
||||
output: Tensor,
|
||||
rgb: Buffer,
|
||||
width: number,
|
||||
height: number,
|
||||
maskBuf: Buffer,
|
||||
rgbaBuf: Buffer,
|
||||
): Promise<Buffer> {
|
||||
const raw = output.data as Float32Array;
|
||||
|
||||
let lo = Infinity;
|
||||
let hi = -Infinity;
|
||||
for (let i = 0; i < INPUT_PLANE; i++) {
|
||||
const v = raw[i]!;
|
||||
if (v < lo) lo = v;
|
||||
if (v > hi) hi = v;
|
||||
}
|
||||
const range = hi - lo || 1;
|
||||
|
||||
for (let i = 0; i < INPUT_PLANE; i++) {
|
||||
const norm = (raw[i]! - lo) / range;
|
||||
maskBuf[i] = Math.max(0, Math.min(255, Math.round(norm * 255)));
|
||||
}
|
||||
|
||||
// lanczos3 keeps soft edges; nearest leaves visible jaggies on hair.
|
||||
const fullMask = await sharp(maskBuf, {
|
||||
raw: { width: INPUT_SIZE, height: INPUT_SIZE, channels: 1 },
|
||||
})
|
||||
.resize(width, height, { kernel: "lanczos3", fit: "fill" })
|
||||
.raw()
|
||||
.toBuffer();
|
||||
|
||||
for (let i = 0; i < width * height; i++) {
|
||||
rgbaBuf[i * 4] = rgb[i * 3]!;
|
||||
rgbaBuf[i * 4 + 1] = rgb[i * 3 + 1]!;
|
||||
rgbaBuf[i * 4 + 2] = rgb[i * 3 + 2]!;
|
||||
rgbaBuf[i * 4 + 3] = fullMask[i]!;
|
||||
}
|
||||
return rgbaBuf;
|
||||
}
|
||||
@@ -0,0 +1,81 @@
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||
|
||||
describe("background-removal/manager — selectProviders", () => {
|
||||
beforeEach(() => {
|
||||
vi.resetModules();
|
||||
delete process.env["HYPERFRAMES_CUDA"];
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
it("returns CPU explicitly when --device cpu", async () => {
|
||||
vi.doMock("node:os", () => ({
|
||||
platform: () => "darwin",
|
||||
arch: () => "arm64",
|
||||
homedir: () => "/tmp",
|
||||
}));
|
||||
const { selectProviders } = await import("./manager.js");
|
||||
const choice = selectProviders("cpu");
|
||||
expect(choice.providers).toEqual(["cpu"]);
|
||||
expect(choice.label).toBe("CPU");
|
||||
});
|
||||
|
||||
it("auto picks CoreML on darwin-arm64", async () => {
|
||||
vi.doMock("node:os", () => ({
|
||||
platform: () => "darwin",
|
||||
arch: () => "arm64",
|
||||
homedir: () => "/tmp",
|
||||
}));
|
||||
const { selectProviders } = await import("./manager.js");
|
||||
const choice = selectProviders("auto");
|
||||
expect(choice.providers).toEqual(["coreml", "cpu"]);
|
||||
expect(choice.label).toBe("CoreML");
|
||||
});
|
||||
|
||||
it("auto falls back to CPU on linux without HYPERFRAMES_CUDA", async () => {
|
||||
vi.doMock("node:os", () => ({
|
||||
platform: () => "linux",
|
||||
arch: () => "x64",
|
||||
homedir: () => "/tmp",
|
||||
}));
|
||||
const { selectProviders } = await import("./manager.js");
|
||||
const choice = selectProviders("auto");
|
||||
expect(choice.providers).toEqual(["cpu"]);
|
||||
expect(choice.label).toBe("CPU");
|
||||
});
|
||||
|
||||
it("auto picks CUDA on linux when HYPERFRAMES_CUDA=1", async () => {
|
||||
process.env["HYPERFRAMES_CUDA"] = "1";
|
||||
vi.doMock("node:os", () => ({
|
||||
platform: () => "linux",
|
||||
arch: () => "x64",
|
||||
homedir: () => "/tmp",
|
||||
}));
|
||||
const { selectProviders } = await import("./manager.js");
|
||||
const choice = selectProviders("auto");
|
||||
expect(choice.providers).toEqual(["cuda", "cpu"]);
|
||||
expect(choice.label).toBe("CUDA");
|
||||
});
|
||||
|
||||
it("--device coreml on linux throws", async () => {
|
||||
vi.doMock("node:os", () => ({
|
||||
platform: () => "linux",
|
||||
arch: () => "x64",
|
||||
homedir: () => "/tmp",
|
||||
}));
|
||||
const { selectProviders } = await import("./manager.js");
|
||||
expect(() => selectProviders("coreml")).toThrow(/CoreML execution provider not available/);
|
||||
});
|
||||
|
||||
it("--device cuda without env var throws", async () => {
|
||||
vi.doMock("node:os", () => ({
|
||||
platform: () => "linux",
|
||||
arch: () => "x64",
|
||||
homedir: () => "/tmp",
|
||||
}));
|
||||
const { selectProviders } = await import("./manager.js");
|
||||
expect(() => selectProviders("cuda")).toThrow(/CUDA execution provider not available/);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,96 @@
|
||||
import { existsSync, mkdirSync } from "node:fs";
|
||||
import { homedir, platform, arch } from "node:os";
|
||||
import { join } from "node:path";
|
||||
import { downloadFile } from "../utils/download.js";
|
||||
|
||||
export const MODELS_DIR = join(homedir(), ".cache", "hyperframes", "background-removal", "models");
|
||||
|
||||
export const DEFAULT_MODEL = "u2net_human_seg" as const;
|
||||
export type ModelId = typeof DEFAULT_MODEL;
|
||||
|
||||
const MODEL_URLS: Record<ModelId, string> = {
|
||||
u2net_human_seg:
|
||||
"https://github.com/danielgatis/rembg/releases/download/v0.0.0/u2net_human_seg.onnx",
|
||||
};
|
||||
|
||||
export const MODEL_MEMORY_MB: Record<ModelId, number> = {
|
||||
u2net_human_seg: 1500,
|
||||
};
|
||||
|
||||
export const DEVICES = ["auto", "cpu", "coreml", "cuda"] as const;
|
||||
export type Device = (typeof DEVICES)[number];
|
||||
|
||||
export function isDevice(value: unknown): value is Device {
|
||||
return typeof value === "string" && (DEVICES as readonly string[]).includes(value);
|
||||
}
|
||||
|
||||
export interface ProviderChoice {
|
||||
providers: string[];
|
||||
label: "CoreML" | "CUDA" | "CPU";
|
||||
}
|
||||
|
||||
export function selectProviders(device: Device = "auto"): ProviderChoice {
|
||||
if (device === "cpu") return { providers: ["cpu"], label: "CPU" };
|
||||
|
||||
const available = listAvailableProviders();
|
||||
const hasCoreML = available.includes("coreml");
|
||||
const hasCUDA = available.includes("cuda");
|
||||
|
||||
if (device === "coreml") {
|
||||
if (!hasCoreML) {
|
||||
throw new Error(
|
||||
"CoreML execution provider not available. Install onnxruntime-node on Apple Silicon, or use --device cpu.",
|
||||
);
|
||||
}
|
||||
return { providers: ["coreml", "cpu"], label: "CoreML" };
|
||||
}
|
||||
if (device === "cuda") {
|
||||
if (!hasCUDA) {
|
||||
throw new Error(
|
||||
"CUDA execution provider not available. Use --device cpu or install an onnxruntime-node build with CUDA support.",
|
||||
);
|
||||
}
|
||||
return { providers: ["cuda", "cpu"], label: "CUDA" };
|
||||
}
|
||||
|
||||
if (hasCoreML && platform() === "darwin" && arch() === "arm64") {
|
||||
return { providers: ["coreml", "cpu"], label: "CoreML" };
|
||||
}
|
||||
if (hasCUDA) return { providers: ["cuda", "cpu"], label: "CUDA" };
|
||||
return { providers: ["cpu"], label: "CPU" };
|
||||
}
|
||||
|
||||
let _cachedProviders: string[] | undefined;
|
||||
export function listAvailableProviders(): string[] {
|
||||
if (_cachedProviders) return _cachedProviders;
|
||||
|
||||
// The npm onnxruntime-node ships with CPU on every platform and bundles the
|
||||
// CoreML EP only on darwin-arm64. CUDA is opt-in via a separate gpu build —
|
||||
// gate behind an env var so we don't try to bind to a missing EP.
|
||||
const out: string[] = ["cpu"];
|
||||
if (platform() === "darwin" && arch() === "arm64") out.push("coreml");
|
||||
if (process.env["HYPERFRAMES_CUDA"] === "1") out.push("cuda");
|
||||
_cachedProviders = out;
|
||||
return out;
|
||||
}
|
||||
|
||||
export function modelPath(model: ModelId = DEFAULT_MODEL): string {
|
||||
return join(MODELS_DIR, `${model}.onnx`);
|
||||
}
|
||||
|
||||
export async function ensureModel(
|
||||
model: ModelId = DEFAULT_MODEL,
|
||||
options?: { onProgress?: (message: string) => void },
|
||||
): Promise<string> {
|
||||
const dest = modelPath(model);
|
||||
if (existsSync(dest)) return dest;
|
||||
|
||||
mkdirSync(MODELS_DIR, { recursive: true });
|
||||
options?.onProgress?.(`Downloading ${model} weights (~168 MB)...`);
|
||||
await downloadFile(MODEL_URLS[model], dest);
|
||||
|
||||
if (!existsSync(dest)) {
|
||||
throw new Error(`Model download failed: ${model}`);
|
||||
}
|
||||
return dest;
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { inferOutputFormat, inferInputKind, buildEncoderArgs } from "./pipeline.js";
|
||||
|
||||
describe("background-removal/pipeline — inferOutputFormat", () => {
|
||||
it("maps .webm → webm", () => {
|
||||
expect(inferOutputFormat("/tmp/out.webm")).toBe("webm");
|
||||
});
|
||||
it("maps .mov → mov", () => {
|
||||
expect(inferOutputFormat("/tmp/out.mov")).toBe("mov");
|
||||
});
|
||||
it("maps .png → png", () => {
|
||||
expect(inferOutputFormat("/tmp/out.png")).toBe("png");
|
||||
});
|
||||
it("rejects unknown extensions", () => {
|
||||
expect(() => inferOutputFormat("/tmp/out.mp4")).toThrow(/Unsupported output extension/);
|
||||
});
|
||||
});
|
||||
|
||||
describe("background-removal/pipeline — inferInputKind", () => {
|
||||
it("recognizes mp4/mov/webm/mkv/avi as video", () => {
|
||||
for (const ext of [".mp4", ".mov", ".webm", ".mkv", ".avi"]) {
|
||||
expect(inferInputKind(`/tmp/clip${ext}`)).toBe("video");
|
||||
}
|
||||
});
|
||||
it("recognizes jpg/png/webp as image", () => {
|
||||
for (const ext of [".jpg", ".jpeg", ".png", ".webp"]) {
|
||||
expect(inferInputKind(`/tmp/img${ext}`)).toBe("image");
|
||||
}
|
||||
});
|
||||
it("rejects unknown extensions", () => {
|
||||
expect(() => inferInputKind("/tmp/file.gif")).toThrow(/Unsupported input/);
|
||||
});
|
||||
});
|
||||
|
||||
describe("background-removal/pipeline — buildEncoderArgs", () => {
|
||||
it("webm preset emits VP9 + alpha_mode metadata", () => {
|
||||
const args = buildEncoderArgs("webm", 1920, 1080, 30, "/tmp/out.webm");
|
||||
expect(args).toContain("libvpx-vp9");
|
||||
expect(args).toContain("yuva420p");
|
||||
// The alpha_mode metadata must be present; without it Chrome ignores the alpha plane.
|
||||
const idx = args.indexOf("-metadata:s:v:0");
|
||||
expect(idx).toBeGreaterThan(-1);
|
||||
expect(args[idx + 1]).toBe("alpha_mode=1");
|
||||
expect(args[args.length - 1]).toBe("/tmp/out.webm");
|
||||
});
|
||||
|
||||
it("mov preset emits ProRes 4444 + yuva444p10le", () => {
|
||||
const args = buildEncoderArgs("mov", 1920, 1080, 30, "/tmp/out.mov");
|
||||
expect(args).toContain("prores_ks");
|
||||
expect(args).toContain("4444");
|
||||
expect(args).toContain("yuva444p10le");
|
||||
});
|
||||
|
||||
it("png preset emits a single RGBA frame", () => {
|
||||
const args = buildEncoderArgs("png", 1920, 1080, 30, "/tmp/out.png");
|
||||
expect(args).toContain("-frames:v");
|
||||
expect(args).toContain("rgba");
|
||||
});
|
||||
|
||||
it("threads input dimensions and fps into raw video header", () => {
|
||||
const args = buildEncoderArgs("webm", 640, 480, 24, "/tmp/o.webm");
|
||||
const sIdx = args.indexOf("-s");
|
||||
expect(args[sIdx + 1]).toBe("640x480");
|
||||
const rIdx = args.indexOf("-r");
|
||||
expect(args[rIdx + 1]).toBe("24");
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,324 @@
|
||||
/**
|
||||
* Background-removal rendering pipeline.
|
||||
*
|
||||
* Decode source frames via ffmpeg → run inference per frame → encode the RGBA
|
||||
* stream via a second ffmpeg process. Output formats:
|
||||
* .webm → VP9 with alpha (HTML5-native, ~1 MB / 4s @ 1080p)
|
||||
* .mov → ProRes 4444 with alpha (editing round-trip)
|
||||
* .png → single RGBA still (only when input is also a single image)
|
||||
*
|
||||
* The encode flags for VP9-with-alpha mirror the `chunkEncoder.ts` pattern in
|
||||
* @hyperframes/engine — `-pix_fmt yuva420p` plus the
|
||||
* `-metadata:s:v:0 alpha_mode=1` tag are what make Chrome's `<video>` element
|
||||
* decode the alpha plane.
|
||||
*/
|
||||
import { spawn } from "node:child_process";
|
||||
import { extname } from "node:path";
|
||||
import { hasFFmpeg, hasFFprobe } from "../whisper/manager.js";
|
||||
import { createSession, type Session } from "./inference.js";
|
||||
import { type Device, type ModelId } from "./manager.js";
|
||||
|
||||
export type OutputFormat = "webm" | "mov" | "png";
|
||||
|
||||
export interface RenderOptions {
|
||||
inputPath: string;
|
||||
outputPath: string;
|
||||
device?: Device;
|
||||
model?: ModelId;
|
||||
onProgress?: (event: ProgressEvent) => void;
|
||||
}
|
||||
|
||||
export type ProgressEvent =
|
||||
| { kind: "info"; message: string }
|
||||
| { kind: "metadata"; width: number; height: number; fps: number; frameCount: number }
|
||||
| { kind: "frame"; index: number; total: number; avgMsPerFrame: number };
|
||||
|
||||
export interface RenderResult {
|
||||
outputPath: string;
|
||||
framesProcessed: number;
|
||||
durationSeconds: number;
|
||||
avgMsPerFrame: number;
|
||||
provider: string;
|
||||
format: OutputFormat;
|
||||
}
|
||||
|
||||
const VIDEO_EXTENSIONS = new Set([".mp4", ".mov", ".webm", ".mkv", ".avi"]);
|
||||
const IMAGE_EXTENSIONS = new Set([".jpg", ".jpeg", ".png", ".webp"]);
|
||||
|
||||
interface MediaInfo {
|
||||
width: number;
|
||||
height: number;
|
||||
fps: number;
|
||||
frameCount: number;
|
||||
}
|
||||
|
||||
export function inferOutputFormat(outputPath: string): OutputFormat {
|
||||
const ext = extname(outputPath).toLowerCase();
|
||||
if (ext === ".webm") return "webm";
|
||||
if (ext === ".mov") return "mov";
|
||||
if (ext === ".png") return "png";
|
||||
throw new Error(
|
||||
`Unsupported output extension: ${ext}. Use .webm (VP9 alpha), .mov (ProRes 4444), or .png.`,
|
||||
);
|
||||
}
|
||||
|
||||
export function inferInputKind(inputPath: string): "video" | "image" {
|
||||
const ext = extname(inputPath).toLowerCase();
|
||||
if (VIDEO_EXTENSIONS.has(ext)) return "video";
|
||||
if (IMAGE_EXTENSIONS.has(ext)) return "image";
|
||||
throw new Error(
|
||||
`Unsupported input: ${ext}. Use a video (mp4/mov/webm/mkv/avi) or image (jpg/png/webp).`,
|
||||
);
|
||||
}
|
||||
|
||||
interface EngineMetadata {
|
||||
width: number;
|
||||
height: number;
|
||||
fps: number;
|
||||
durationSeconds: number;
|
||||
}
|
||||
|
||||
async function probeMedia(inputPath: string): Promise<MediaInfo> {
|
||||
const isImage = inferInputKind(inputPath) === "image";
|
||||
const engine = (await import("@hyperframes/engine")) as {
|
||||
extractMediaMetadata: (path: string) => Promise<EngineMetadata>;
|
||||
};
|
||||
const meta = await engine.extractMediaMetadata(inputPath);
|
||||
|
||||
if (isImage) {
|
||||
return { width: meta.width, height: meta.height, fps: 0, frameCount: 1 };
|
||||
}
|
||||
|
||||
const fps = meta.fps || 30;
|
||||
const frameCount = meta.durationSeconds ? Math.round(meta.durationSeconds * fps) : 0;
|
||||
return { width: meta.width, height: meta.height, fps, frameCount };
|
||||
}
|
||||
|
||||
export function buildEncoderArgs(
|
||||
format: OutputFormat,
|
||||
width: number,
|
||||
height: number,
|
||||
fps: number,
|
||||
outputPath: string,
|
||||
): string[] {
|
||||
const base = [
|
||||
"-y",
|
||||
"-f",
|
||||
"rawvideo",
|
||||
"-pix_fmt",
|
||||
"rgba",
|
||||
"-s",
|
||||
`${width}x${height}`,
|
||||
"-r",
|
||||
String(fps || 30),
|
||||
"-i",
|
||||
"-",
|
||||
];
|
||||
|
||||
if (format === "webm") {
|
||||
return [
|
||||
...base,
|
||||
"-c:v",
|
||||
"libvpx-vp9",
|
||||
"-b:v",
|
||||
"0",
|
||||
"-crf",
|
||||
"30",
|
||||
"-deadline",
|
||||
"good",
|
||||
"-row-mt",
|
||||
"1",
|
||||
"-auto-alt-ref",
|
||||
"0",
|
||||
"-pix_fmt",
|
||||
"yuva420p",
|
||||
"-metadata:s:v:0",
|
||||
"alpha_mode=1",
|
||||
"-an",
|
||||
outputPath,
|
||||
];
|
||||
}
|
||||
if (format === "mov") {
|
||||
return [
|
||||
...base,
|
||||
"-c:v",
|
||||
"prores_ks",
|
||||
"-profile:v",
|
||||
"4444",
|
||||
"-vendor",
|
||||
"apl0",
|
||||
"-pix_fmt",
|
||||
"yuva444p10le",
|
||||
"-an",
|
||||
outputPath,
|
||||
];
|
||||
}
|
||||
return [...base, "-frames:v", "1", "-pix_fmt", "rgba", "-update", "1", outputPath];
|
||||
}
|
||||
|
||||
async function* readFrames(
|
||||
stream: NodeJS.ReadableStream,
|
||||
frameBytes: number,
|
||||
): AsyncGenerator<Buffer> {
|
||||
let buffered: Buffer = Buffer.alloc(0);
|
||||
for await (const chunk of stream) {
|
||||
buffered =
|
||||
buffered.length === 0 ? (chunk as Buffer) : Buffer.concat([buffered, chunk as Buffer]);
|
||||
while (buffered.length >= frameBytes) {
|
||||
// Copy because the next concat would clobber the underlying memory.
|
||||
yield Buffer.from(buffered.subarray(0, frameBytes));
|
||||
buffered = buffered.subarray(frameBytes);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export async function render(options: RenderOptions): Promise<RenderResult> {
|
||||
if (!hasFFmpeg() || !hasFFprobe()) {
|
||||
throw new Error("ffmpeg and ffprobe are required. Install: brew install ffmpeg");
|
||||
}
|
||||
|
||||
const format = inferOutputFormat(options.outputPath);
|
||||
const inputKind = inferInputKind(options.inputPath);
|
||||
|
||||
if (inputKind === "image" && format !== "png") {
|
||||
throw new Error(
|
||||
`Image input requires a .png output (got ${extname(options.outputPath)}). Use a video input for .webm/.mov.`,
|
||||
);
|
||||
}
|
||||
if (inputKind === "video" && format === "png") {
|
||||
throw new Error(
|
||||
`Video input requires a .webm or .mov output (got .png). Use an image input for .png.`,
|
||||
);
|
||||
}
|
||||
|
||||
const media = await probeMedia(options.inputPath);
|
||||
|
||||
options.onProgress?.({
|
||||
kind: "metadata",
|
||||
width: media.width,
|
||||
height: media.height,
|
||||
fps: media.fps,
|
||||
frameCount: media.frameCount,
|
||||
});
|
||||
|
||||
const session = await createSession({
|
||||
model: options.model,
|
||||
device: options.device,
|
||||
onProgress: (msg) => options.onProgress?.({ kind: "info", message: msg }),
|
||||
});
|
||||
|
||||
try {
|
||||
const start = Date.now();
|
||||
const framesProcessed = await runPipeline(options, session, media, format);
|
||||
const durationSeconds = (Date.now() - start) / 1000;
|
||||
const avgMsPerFrame = framesProcessed ? (durationSeconds * 1000) / framesProcessed : 0;
|
||||
|
||||
return {
|
||||
outputPath: options.outputPath,
|
||||
framesProcessed,
|
||||
durationSeconds,
|
||||
avgMsPerFrame,
|
||||
provider: session.provider,
|
||||
format,
|
||||
};
|
||||
} finally {
|
||||
await session.close();
|
||||
}
|
||||
}
|
||||
|
||||
const RECENT_WINDOW = 30;
|
||||
|
||||
async function runPipeline(
|
||||
options: RenderOptions,
|
||||
session: Session,
|
||||
media: MediaInfo,
|
||||
format: OutputFormat,
|
||||
): Promise<number> {
|
||||
const { inputPath, outputPath } = options;
|
||||
const { width, height, fps, frameCount } = media;
|
||||
const frameBytes = width * height * 3;
|
||||
|
||||
const decoder = spawn(
|
||||
"ffmpeg",
|
||||
["-loglevel", "error", "-i", inputPath, "-f", "rawvideo", "-pix_fmt", "rgb24", "-an", "-"],
|
||||
{ stdio: ["ignore", "pipe", "pipe"] },
|
||||
);
|
||||
|
||||
let decoderStderr = "";
|
||||
decoder.stderr?.on("data", (d: Buffer) => {
|
||||
decoderStderr += d.toString();
|
||||
});
|
||||
const decoderExit = waitForExit(decoder, "ffmpeg decoder", () => decoderStderr);
|
||||
|
||||
const encoder = spawn("ffmpeg", buildEncoderArgs(format, width, height, fps || 30, outputPath), {
|
||||
stdio: ["pipe", "ignore", "pipe"],
|
||||
});
|
||||
let encoderStderr = "";
|
||||
encoder.stderr?.on("data", (d: Buffer) => {
|
||||
encoderStderr += d.toString();
|
||||
});
|
||||
const encoderExit = waitForExit(encoder, "ffmpeg encoder", () => encoderStderr);
|
||||
|
||||
let processed = 0;
|
||||
const total = frameCount;
|
||||
|
||||
// Running average over the last RECENT_WINDOW frames.
|
||||
const recentMs = new Array<number>(RECENT_WINDOW).fill(0);
|
||||
let recentSum = 0;
|
||||
let recentSlot = 0;
|
||||
let recentCount = 0;
|
||||
|
||||
try {
|
||||
for await (const rgb of readFrames(decoder.stdout!, frameBytes)) {
|
||||
const t0 = Date.now();
|
||||
const rgba = await session.process(rgb, width, height);
|
||||
const elapsed = Date.now() - t0;
|
||||
|
||||
recentSum += elapsed - recentMs[recentSlot]!;
|
||||
recentMs[recentSlot] = elapsed;
|
||||
recentSlot = (recentSlot + 1) % RECENT_WINDOW;
|
||||
if (recentCount < RECENT_WINDOW) recentCount++;
|
||||
|
||||
if (!encoder.stdin!.write(rgba)) {
|
||||
await new Promise<void>((resolve) => encoder.stdin!.once("drain", () => resolve()));
|
||||
}
|
||||
|
||||
processed++;
|
||||
options.onProgress?.({
|
||||
kind: "frame",
|
||||
index: processed,
|
||||
total,
|
||||
avgMsPerFrame: recentSum / recentCount,
|
||||
});
|
||||
}
|
||||
} catch (err) {
|
||||
decoder.kill("SIGKILL");
|
||||
encoder.kill("SIGKILL");
|
||||
throw err;
|
||||
}
|
||||
|
||||
encoder.stdin!.end();
|
||||
await Promise.all([decoderExit, encoderExit]);
|
||||
|
||||
if (processed === 0) {
|
||||
throw new Error(
|
||||
`No frames produced from ${inputPath}. Decoder stderr:\n${decoderStderr.slice(-400)}`,
|
||||
);
|
||||
}
|
||||
|
||||
return processed;
|
||||
}
|
||||
|
||||
function waitForExit(
|
||||
proc: ReturnType<typeof spawn>,
|
||||
label: string,
|
||||
getStderr: () => string,
|
||||
): Promise<void> {
|
||||
return new Promise<void>((resolve, reject) => {
|
||||
proc.on("error", reject);
|
||||
proc.on("exit", (code) => {
|
||||
if (code === 0 || code === null) resolve();
|
||||
else reject(new Error(`${label} exited with code ${code}: ${getStderr().slice(-400)}`));
|
||||
});
|
||||
});
|
||||
}
|
||||
Reference in New Issue
Block a user