From c6c78dc76fefc8d230523944347bffdcb3d94f48 Mon Sep 17 00:00:00 2001 From: SnapOtter Date: Tue, 28 Apr 2026 02:20:48 +0800 Subject: [PATCH] fix: prevent OOM kills during background removal on CPU Skip alpha matting on CPU (pymatting's sparse matrices are the main memory hog), auto-downscale images above 2048px before sending to rembg, and retry with the lighter u2net model when OOM is detected. --- packages/ai/python/remove_bg.py | 13 +++-- packages/ai/src/background-removal.ts | 71 ++++++++++++++++++++---- tests/unit/ai/background-removal.test.ts | 4 ++ 3 files changed, 74 insertions(+), 14 deletions(-) diff --git a/packages/ai/python/remove_bg.py b/packages/ai/python/remove_bg.py index 6fc39b16..defe5442 100644 --- a/packages/ai/python/remove_bg.py +++ b/packages/ai/python/remove_bg.py @@ -94,18 +94,23 @@ def main(): input_data = f.read() emit_progress(30, "Analyzing image") + use_alpha_matting = device != "cpu" try: output_data = remove( input_data, session=session, - alpha_matting=True, + alpha_matting=use_alpha_matting, alpha_matting_foreground_threshold=240, alpha_matting_background_threshold=10, ) except Exception as e: - raise RuntimeError( - f"Alpha matting failed: {e}. Try again without alpha matting or with a different model." - ) from e + if use_alpha_matting: + emit_progress(35, "Retrying without alpha matting") + output_data = remove(input_data, session=session, alpha_matting=False) + else: + raise RuntimeError( + f"Background removal failed: {e}" + ) from e emit_progress(80, "Background removed") diff --git a/packages/ai/src/background-removal.ts b/packages/ai/src/background-removal.ts index b3b35c87..f459a026 100644 --- a/packages/ai/src/background-removal.ts +++ b/packages/ai/src/background-removal.ts @@ -10,6 +10,9 @@ export interface RemoveBackgroundOptions { backgroundColor?: string; } +const MAX_REMBG_PX = Number(process.env.MAX_REMBG_PX) || 2048; +const OOM_FALLBACK_MODEL = "u2net"; + export async function removeBackground( inputBuffer: Buffer, outputDir: string, @@ -20,29 +23,77 @@ export async function removeBackground( const inputPath = join(tmpdir(), `rembg_in_${id}.png`); const outputPath = join(outputDir, `rembg_out_${id}.png`); - const pngBuffer = await sharp(inputBuffer).png().toBuffer(); + const meta = await sharp(inputBuffer).metadata(); + const origW = meta.width ?? 0; + const origH = meta.height ?? 0; + const longest = Math.max(origW, origH); + const needsDownscale = longest > MAX_REMBG_PX; + + let pipeline = sharp(inputBuffer); + if (needsDownscale) { + pipeline = pipeline.resize({ + width: origW >= origH ? MAX_REMBG_PX : undefined, + height: origH > origW ? MAX_REMBG_PX : undefined, + fit: "inside", + withoutEnlargement: true, + }); + } + const pngBuffer = await pipeline.png().toBuffer(); await writeFile(inputPath, pngBuffer); + try { - const meta = await sharp(inputBuffer).metadata(); - const megapixels = ((meta.width ?? 0) * (meta.height ?? 0)) / 1_000_000; + const megapixels = (origW * origH) / 1_000_000; const baseTimeout = options.model?.startsWith("birefnet") ? 600000 : 300000; const timeout = Math.max(baseTimeout, megapixels * 30 * 1000); + + const rawMask = await runAndParse(inputPath, outputPath, options, onProgress, timeout); + + if (needsDownscale) { + return sharp(rawMask).resize({ width: origW, height: origH, fit: "fill" }).png().toBuffer(); + } + + return rawMask; + } finally { + await unlink(inputPath).catch(() => {}); + await unlink(outputPath).catch(() => {}); + } +} + +async function runAndParse( + inputPath: string, + outputPath: string, + options: RemoveBackgroundOptions, + onProgress: ProgressCallback | undefined, + timeout: number, +): Promise { + try { const { stdout } = await runPythonWithProgress( "remove_bg.py", [inputPath, outputPath, JSON.stringify(options)], { onProgress, timeout }, ); - const result = parseStdoutJson(stdout); if (!result.success) { throw new Error(result.error || "Background removal failed"); } + return readFile(outputPath); + } catch (err) { + const isOom = err instanceof Error && err.message.includes("out of memory"); + const canFallback = isOom && options.model !== OOM_FALLBACK_MODEL; - const outputBuffer = await readFile(outputPath); - return outputBuffer; - } finally { - // Clean up temp files - await unlink(inputPath).catch(() => {}); - await unlink(outputPath).catch(() => {}); + if (!canFallback) throw err; + + onProgress?.(5, `Retrying with lighter model (${OOM_FALLBACK_MODEL})`); + const fallbackOpts = { ...options, model: OOM_FALLBACK_MODEL }; + const { stdout } = await runPythonWithProgress( + "remove_bg.py", + [inputPath, outputPath, JSON.stringify(fallbackOpts)], + { onProgress, timeout: 300000 }, + ); + const result = parseStdoutJson(stdout); + if (!result.success) { + throw new Error(result.error || "Background removal failed"); + } + return readFile(outputPath); } } diff --git a/tests/unit/ai/background-removal.test.ts b/tests/unit/ai/background-removal.test.ts index 6cd9b68a..3566aef2 100644 --- a/tests/unit/ai/background-removal.test.ts +++ b/tests/unit/ai/background-removal.test.ts @@ -4,6 +4,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; vi.mock("sharp", () => { const mockSharp = vi.fn(() => ({ png: vi.fn().mockReturnThis(), + resize: vi.fn().mockReturnThis(), toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")), metadata: vi.fn().mockResolvedValue({ width: 800, height: 600 }), })); @@ -42,6 +43,7 @@ beforeEach(() => { () => ({ png: vi.fn().mockReturnThis(), + resize: vi.fn().mockReturnThis(), toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")), metadata: vi.fn().mockResolvedValue({ width: 800, height: 600 }), }) as unknown as ReturnType, @@ -189,6 +191,7 @@ describe("removeBackground", () => { () => ({ png: vi.fn().mockReturnThis(), + resize: vi.fn().mockReturnThis(), toBuffer: vi.fn().mockRejectedValue(new Error("Invalid image")), metadata: vi.fn().mockResolvedValue({ width: 800, height: 600 }), }) as unknown as ReturnType, @@ -225,6 +228,7 @@ describe("removeBackground", () => { () => ({ png: vi.fn().mockReturnThis(), + resize: vi.fn().mockReturnThis(), toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")), metadata: vi.fn().mockResolvedValue({ width: 6000, height: 4000 }), }) as unknown as ReturnType,