mirror of
https://github.com/heygen-com/hyperframes.git
synced 2026-09-03 04:38:33 +00:00
Merge pull request #612 from heygen-com/feat/cli-remove-background
feat(cli): add remove-background command for transparent video
This commit is contained in:
James Russo
@@ -33,10 +33,12 @@
|
||||
"giget": "^3.2.0",
|
||||
"hono": "^4.0.0",
|
||||
"mime-types": "^3.0.2",
|
||||
"onnxruntime-node": "^1.20.0",
|
||||
"open": "^10.0.0",
|
||||
"postcss": "^8.5.8",
|
||||
"prettier": "^3.8.1",
|
||||
"puppeteer-core": "^24.39.1"
|
||||
"puppeteer-core": "^24.39.1",
|
||||
"sharp": "^0.34.0"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@clack/prompts": "^1.1.0",
|
||||
|
||||
@@ -0,0 +1,18 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { MEAN, STD } from "./inference.js";
|
||||
|
||||
// Regression: the u2net_human_seg model was trained with ImageNet
|
||||
// normalization. Drifting away from these exact values changes the input
|
||||
// tensor at every pixel and shifts the predicted alpha mask noticeably
|
||||
// (Miguel reproduced 8,317 pixel changes with delta up to 78/255 when std
|
||||
// was set to (1, 1, 1)). Reference:
|
||||
// https://github.com/danielgatis/rembg/blob/main/rembg/sessions/u2net_human_seg.py#L33
|
||||
describe("background-removal/inference — rembg u2net_human_seg parity", () => {
|
||||
it("MEAN matches U2netHumanSegSession reference", () => {
|
||||
expect(MEAN).toEqual([0.485, 0.456, 0.406]);
|
||||
});
|
||||
|
||||
it("STD matches U2netHumanSegSession reference (ImageNet, not the base u2net's (1,1,1))", () => {
|
||||
expect(STD).toEqual([0.229, 0.224, 0.225]);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,176 @@
|
||||
/**
|
||||
* 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;
|
||||
|
||||
// Must match rembg's U2netHumanSegSession.predict — ImageNet mean/std, NOT the
|
||||
// (1.0, 1.0, 1.0) std used by the general-purpose u2net session.
|
||||
// https://github.com/danielgatis/rembg/blob/main/rembg/sessions/u2net_human_seg.py#L33
|
||||
export const MEAN = [0.485, 0.456, 0.406] as const;
|
||||
export const STD = [0.229, 0.224, 0.225] 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,108 @@
|
||||
import { describe, expect, it } from "vitest";
|
||||
import { EventEmitter } from "node:events";
|
||||
import type { spawn } from "node:child_process";
|
||||
import { inferOutputFormat, inferInputKind, buildEncoderArgs, waitForExit } 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");
|
||||
});
|
||||
});
|
||||
|
||||
// Regression: a previous version of waitForExit treated `code === null` as
|
||||
// success. Per Node's child_process docs, that's the signal-killed case —
|
||||
// reporting it as success means a SIGTERM/SIGKILL'd ffmpeg encoder produces
|
||||
// a "successful" render with a missing or truncated output file.
|
||||
describe("background-removal/pipeline — waitForExit signal handling", () => {
|
||||
function fakeProc(): ReturnType<typeof spawn> {
|
||||
return new EventEmitter() as unknown as ReturnType<typeof spawn>;
|
||||
}
|
||||
|
||||
it("resolves on a clean exit (code=0, signal=null)", async () => {
|
||||
const proc = fakeProc();
|
||||
const promise = waitForExit(proc, "ffmpeg encoder", () => "");
|
||||
proc.emit("exit", 0, null);
|
||||
await expect(promise).resolves.toBeUndefined();
|
||||
});
|
||||
|
||||
it("rejects when killed by signal (code=null, signal='SIGTERM')", async () => {
|
||||
const proc = fakeProc();
|
||||
const promise = waitForExit(proc, "ffmpeg encoder", () => "tail of stderr");
|
||||
proc.emit("exit", null, "SIGTERM");
|
||||
await expect(promise).rejects.toThrow(/killed by SIGTERM/);
|
||||
await expect(promise).rejects.toThrow(/tail of stderr/);
|
||||
});
|
||||
|
||||
it("rejects on non-zero exit code", async () => {
|
||||
const proc = fakeProc();
|
||||
const promise = waitForExit(proc, "ffmpeg encoder", () => "");
|
||||
proc.emit("exit", 1, null);
|
||||
await expect(promise).rejects.toThrow(/exited with code 1/);
|
||||
});
|
||||
|
||||
it("rejects on SIGKILL", async () => {
|
||||
const proc = fakeProc();
|
||||
const promise = waitForExit(proc, "ffmpeg encoder", () => "");
|
||||
proc.emit("exit", null, "SIGKILL");
|
||||
await expect(promise).rejects.toThrow(/killed by SIGKILL/);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,333 @@
|
||||
/**
|
||||
* 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;
|
||||
}
|
||||
|
||||
export function waitForExit(
|
||||
proc: ReturnType<typeof spawn>,
|
||||
label: string,
|
||||
getStderr: () => string,
|
||||
): Promise<void> {
|
||||
return new Promise<void>((resolve, reject) => {
|
||||
proc.on("error", reject);
|
||||
// Per Node docs the exit callback is (code, signal): on a normal exit
|
||||
// `code` is the numeric exit status and `signal` is null; on a
|
||||
// signal-killed exit `code` is null and `signal` is the signal name.
|
||||
// Treating null-code as success would silently report SIGTERM/SIGKILL
|
||||
// as a successful render.
|
||||
proc.on("exit", (code, signal) => {
|
||||
if (code === 0 && !signal) {
|
||||
resolve();
|
||||
return;
|
||||
}
|
||||
const cause = signal ? `killed by ${signal}` : `exited with code ${code}`;
|
||||
reject(new Error(`${label} ${cause}: ${getStderr().slice(-400)}`));
|
||||
});
|
||||
});
|
||||
}
|
||||
@@ -38,6 +38,7 @@ const subCommands = {
|
||||
compositions: () => import("./commands/compositions.js").then((m) => m.default),
|
||||
benchmark: () => import("./commands/benchmark.js").then((m) => m.default),
|
||||
browser: () => import("./commands/browser.js").then((m) => m.default),
|
||||
"remove-background": () => import("./commands/remove-background.js").then((m) => m.default),
|
||||
transcribe: () => import("./commands/transcribe.js").then((m) => m.default),
|
||||
tts: () => import("./commands/tts.js").then((m) => m.default),
|
||||
docs: () => import("./commands/docs.js").then((m) => m.default),
|
||||
|
||||
@@ -0,0 +1,181 @@
|
||||
import { defineCommand } from "citty";
|
||||
import { resolve } from "node:path";
|
||||
import { existsSync } from "node:fs";
|
||||
import * as clack from "@clack/prompts";
|
||||
import { c } from "../ui/colors.js";
|
||||
import { isDevice, DEVICES } from "../background-removal/manager.js";
|
||||
import type { Example } from "./_examples.js";
|
||||
|
||||
export const examples: Example[] = [
|
||||
[
|
||||
"Remove background from a video, output transparent VP9 WebM (default)",
|
||||
"hyperframes remove-background avatar.mp4 -o transparent.webm",
|
||||
],
|
||||
[
|
||||
"Output ProRes 4444 .mov for editing round-trip",
|
||||
"hyperframes remove-background avatar.mp4 -o transparent.mov",
|
||||
],
|
||||
[
|
||||
"Remove background from a single image, output transparent PNG",
|
||||
"hyperframes remove-background portrait.jpg -o cutout.png",
|
||||
],
|
||||
[
|
||||
"Force CPU (skip CoreML/CUDA)",
|
||||
"hyperframes remove-background avatar.mp4 -o transparent.webm --device cpu",
|
||||
],
|
||||
["Show detected providers without rendering", "hyperframes remove-background --info"],
|
||||
];
|
||||
|
||||
export default defineCommand({
|
||||
meta: {
|
||||
name: "remove-background",
|
||||
description:
|
||||
"Remove background from a video or image using a local AI model — outputs transparent WebM, ProRes 4444, or PNG",
|
||||
},
|
||||
args: {
|
||||
input: {
|
||||
type: "positional",
|
||||
description: "Source video (.mp4/.mov/.webm/.mkv) or image (.jpg/.png/.webp)",
|
||||
required: false,
|
||||
},
|
||||
output: {
|
||||
type: "string",
|
||||
description: "Output path. Format inferred from extension: .webm (default), .mov, .png",
|
||||
alias: "o",
|
||||
},
|
||||
device: {
|
||||
type: "string",
|
||||
description: `Execution provider: ${DEVICES.join(", ")}`,
|
||||
default: "auto",
|
||||
},
|
||||
info: {
|
||||
type: "boolean",
|
||||
description: "Print detected execution providers and exit (no render)",
|
||||
default: false,
|
||||
},
|
||||
json: {
|
||||
type: "boolean",
|
||||
description: "Output result as JSON",
|
||||
default: false,
|
||||
},
|
||||
},
|
||||
async run({ args }) {
|
||||
if (args.info) {
|
||||
return showInfo(args.json);
|
||||
}
|
||||
if (!args.input) {
|
||||
console.error(
|
||||
c.error(
|
||||
"Input file is required. Run `hyperframes remove-background --info` for providers.",
|
||||
),
|
||||
);
|
||||
process.exit(1);
|
||||
}
|
||||
if (!args.output) {
|
||||
console.error(c.error("--output (-o) is required. Use a .webm, .mov, or .png path."));
|
||||
process.exit(1);
|
||||
}
|
||||
if (!isDevice(args.device)) {
|
||||
console.error(
|
||||
c.error(`Invalid --device '${String(args.device)}'. Use: ${DEVICES.join(", ")}.`),
|
||||
);
|
||||
process.exit(1);
|
||||
}
|
||||
|
||||
const inputPath = resolve(args.input);
|
||||
const outputPath = resolve(args.output);
|
||||
|
||||
const { render } = await import("../background-removal/pipeline.js");
|
||||
|
||||
const spin = args.json ? null : clack.spinner();
|
||||
spin?.start("Preparing background-removal pipeline...");
|
||||
|
||||
try {
|
||||
const result = await render({
|
||||
inputPath,
|
||||
outputPath,
|
||||
device: args.device,
|
||||
onProgress: (event) => {
|
||||
if (event.kind === "info") {
|
||||
spin?.message(event.message);
|
||||
} else if (event.kind === "metadata") {
|
||||
const dims = `${event.width}×${event.height}`;
|
||||
const frames = event.frameCount ? ` · ${event.frameCount} frames` : "";
|
||||
spin?.message(`Source ${dims} @ ${event.fps.toFixed(0)}fps${frames}`);
|
||||
} else if (event.kind === "frame") {
|
||||
const pct = event.total ? ` (${Math.floor((100 * event.index) / event.total)}%)` : "";
|
||||
spin?.message(
|
||||
`Frame ${event.index}${event.total ? `/${event.total}` : ""}${pct} — ${Math.round(event.avgMsPerFrame)}ms/frame avg`,
|
||||
);
|
||||
}
|
||||
},
|
||||
});
|
||||
|
||||
if (args.json) {
|
||||
console.log(
|
||||
JSON.stringify({
|
||||
ok: true,
|
||||
outputPath: result.outputPath,
|
||||
framesProcessed: result.framesProcessed,
|
||||
durationSeconds: Number(result.durationSeconds.toFixed(2)),
|
||||
avgMsPerFrame: Number(result.avgMsPerFrame.toFixed(1)),
|
||||
provider: result.provider,
|
||||
format: result.format,
|
||||
}),
|
||||
);
|
||||
} else {
|
||||
const fpsThroughput = result.durationSeconds
|
||||
? (result.framesProcessed / result.durationSeconds).toFixed(1)
|
||||
: "n/a";
|
||||
spin?.stop(
|
||||
c.success(
|
||||
`Removed background from ${c.accent(String(result.framesProcessed))} frames in ${result.durationSeconds.toFixed(1)}s (${fpsThroughput} fps, ${c.accent(result.provider)}) → ${c.accent(result.outputPath)}`,
|
||||
),
|
||||
);
|
||||
}
|
||||
} catch (err) {
|
||||
const message = err instanceof Error ? err.message : String(err);
|
||||
if (args.json) {
|
||||
console.log(JSON.stringify({ ok: false, error: message }));
|
||||
} else {
|
||||
spin?.stop(c.error(`Background removal failed: ${message}`));
|
||||
}
|
||||
process.exit(1);
|
||||
}
|
||||
},
|
||||
});
|
||||
|
||||
async function showInfo(json: boolean): Promise<void> {
|
||||
const { selectProviders, listAvailableProviders, DEFAULT_MODEL, MODEL_MEMORY_MB, modelPath } =
|
||||
await import("../background-removal/manager.js");
|
||||
|
||||
const providers = listAvailableProviders();
|
||||
const auto = selectProviders("auto");
|
||||
const cached = existsSync(modelPath());
|
||||
|
||||
if (json) {
|
||||
console.log(
|
||||
JSON.stringify({
|
||||
defaultModel: DEFAULT_MODEL,
|
||||
modelCached: cached,
|
||||
modelPath: modelPath(),
|
||||
peakMemoryMb: MODEL_MEMORY_MB[DEFAULT_MODEL],
|
||||
availableProviders: providers,
|
||||
autoProvider: auto.label,
|
||||
}),
|
||||
);
|
||||
return;
|
||||
}
|
||||
|
||||
console.log(c.bold("hyperframes remove-background — system info"));
|
||||
console.log("");
|
||||
console.log(` ${c.dim("Default model:")} ${c.accent(DEFAULT_MODEL)}`);
|
||||
console.log(` ${c.dim("Peak memory:")} ~${MODEL_MEMORY_MB[DEFAULT_MODEL]} MB`);
|
||||
console.log(
|
||||
` ${c.dim("Weights cached:")} ${cached ? c.success("yes") : c.dim("no (will download on first run)")}`,
|
||||
);
|
||||
console.log(` ${c.dim("Cache path:")} ${modelPath()}`);
|
||||
console.log("");
|
||||
console.log(` ${c.dim("Available providers:")} ${providers.join(", ")}`);
|
||||
console.log(` ${c.dim("Auto-selected:")} ${c.accent(auto.label)}`);
|
||||
}
|
||||
@@ -60,6 +60,7 @@ const GROUPS: Group[] = [
|
||||
"Transcribe audio/video to word-level timestamps, or import an existing transcript",
|
||||
],
|
||||
["tts", "Generate speech audio from text using a local AI model (Kokoro-82M)"],
|
||||
["remove-background", "Remove background from a video or image to produce transparent media"],
|
||||
],
|
||||
},
|
||||
{
|
||||
@@ -101,7 +102,7 @@ const STATIC_EXAMPLES: Record<string, Example[]> = {
|
||||
|
||||
// ── Render root help ───────────────────────────────────────────────────────
|
||||
function renderRootHelp(): string {
|
||||
const NAME_COL = 16;
|
||||
const NAME_COL = 19;
|
||||
const CMD_COL = 46;
|
||||
const lines: string[] = [];
|
||||
|
||||
|
||||
@@ -205,8 +205,16 @@ export async function ensureModel(
|
||||
}
|
||||
|
||||
export function hasFFmpeg(): boolean {
|
||||
return hasBinary("ffmpeg");
|
||||
}
|
||||
|
||||
export function hasFFprobe(): boolean {
|
||||
return hasBinary("ffprobe");
|
||||
}
|
||||
|
||||
function hasBinary(name: string): boolean {
|
||||
try {
|
||||
execFileSync("ffmpeg", ["-version"], { stdio: "ignore", timeout: 5000 });
|
||||
execFileSync(name, ["-version"], { stdio: "ignore", timeout: 5000 });
|
||||
return true;
|
||||
} catch {
|
||||
return false;
|
||||
|
||||
Reference in New Issue
Block a user