mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
Siddharth Kumar Sah
fb1d366d5f
BiRefNet-Lite times out on first load (~60s+ for 973MB model). U2-Net works in 2 seconds. Users can still select BiRefNet for higher quality when they're willing to wait. Added timing hints in model descriptions and increased timeout for BiRefNet models.
34 lines
977 B
TypeScript
34 lines
977 B
TypeScript
import { runPythonScript } from "./bridge.js";
|
|
import { writeFile, readFile } from "node:fs/promises";
|
|
import { join } from "node:path";
|
|
|
|
export interface RemoveBackgroundOptions {
|
|
model?: string;
|
|
backgroundColor?: string;
|
|
}
|
|
|
|
export async function removeBackground(
|
|
inputBuffer: Buffer,
|
|
outputDir: string,
|
|
options: RemoveBackgroundOptions = {},
|
|
): Promise<Buffer> {
|
|
const inputPath = join(outputDir, "input_bg.png");
|
|
const outputPath = join(outputDir, "output_bg.png");
|
|
|
|
await writeFile(inputPath, inputBuffer);
|
|
// BiRefNet models need longer timeout (up to 10 min for first load)
|
|
const timeout = options.model?.startsWith("birefnet") ? 600000 : 300000;
|
|
const { stdout } = await runPythonScript("remove_bg.py", [
|
|
inputPath,
|
|
outputPath,
|
|
JSON.stringify(options),
|
|
], timeout);
|
|
|
|
const result = JSON.parse(stdout);
|
|
if (!result.success) {
|
|
throw new Error(result.error || "Background removal failed");
|
|
}
|
|
|
|
return readFile(outputPath);
|
|
}
|