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

@@ -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(
+            if use_alpha_matting:
-                f"Alpha matting failed: {e}. Try again without alpha matting or with a different model."
+                emit_progress(35, "Retrying without alpha matting")
-            ) from e
+                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")
 

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 = (origW * origH) / 1_000_000;
-    const megapixels = ((meta.width ?? 0) * (meta.height ?? 0)) / 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<Buffer> {
+  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);
+    if (!canFallback) throw err;
-    return outputBuffer;
+
-  } finally {
+    onProgress?.(5, `Retrying with lighter model (${OOM_FALLBACK_MODEL})`);
-    // Clean up temp files
+    const fallbackOpts = { ...options, model: OOM_FALLBACK_MODEL };
-    await unlink(inputPath).catch(() => {});
+    const { stdout } = await runPythonWithProgress(
-    await unlink(outputPath).catch(() => {});
+      "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);
   }
 }

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<typeof sharp>,
@@ -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<typeof sharp>,
@@ -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<typeof sharp>,