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.
This commit is contained in:
SnapOtter
2026-04-28 02:20:48 +08:00
parent 4f6fd707a1
commit c6c78dc76f
3 changed files with 74 additions and 14 deletions
+4
View File
@@ -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>,