feat(erase-object): optional high-quality diffusion inpainting bundle (#566)

Adds an opt-in High Quality mode to the Object Eraser, backed by a new inpaint-hq feature bundle (Stable Diffusion 1.5 inpainting via diffusers). The default fast LaMa path is unchanged. Both arch archives are published to deepsafe/feature-bundles and the manifest carries their real sha256/sizes.

Verified end to end: a fresh container pulls the bundle from HuggingFace, checksum-verifies it, extracts torch/diffusers plus the fp16 model, and the HQ sidecar erases a large object with a plausible fill.

Refs #141

tests/e2e-docker/feature-lifecycle.spec.ts

@@ -129,12 +129,13 @@ test.describe("Feature listing baseline", () => {
     const res = await request.get(`${API}/api/v1/features`, { headers });
     expect(res.ok()).toBeTruthy();
     const data = (await res.json()) as FeatureResponse;
-    expect(data.bundles).toHaveLength(7);
+    expect(data.bundles).toHaveLength(8);
 
     const expectedIds = [
       "background-removal",
       "face-detection",
       "object-eraser-colorize",
+      "inpaint-hq",
       "upscale-enhance",
       "photo-restoration",
       "ocr",

tests/e2e-docker/features.spec.ts

@@ -48,12 +48,13 @@ test.describe("Feature API", () => {
     });
     expect(response.ok()).toBeTruthy();
     const data = (await response.json()) as { bundles: BundleInfo[] };
-    expect(data.bundles).toHaveLength(7);
+    expect(data.bundles).toHaveLength(8);
 
     const expectedBundles = [
       "background-removal",
       "face-detection",
       "object-eraser-colorize",
+      "inpaint-hq",
       "upscale-enhance",
       "photo-restoration",
       "ocr",

tests/e2e/erase-object.spec.ts

@@ -231,4 +231,42 @@ test.describe("Erase Object tool", () => {
     // Both files now have masks -> batch submit button.
     await expect(page.getByTestId("erase-object-submit")).toHaveText("Erase All (2)");
   });
+
+  test("High Quality mode is gated on the inpaint-hq pack and blocks submit until installed", async ({
+    loggedInPage: page,
+  }) => {
+    // gotoEraser mocks only object-eraser-colorize as installed, so the optional
+    // inpaint-hq (diffusion) pack reads as missing.
+    await gotoEraser(page);
+    await uploadFile(page, fixturePath("image/valid/test-200x150.png"));
+
+    // The Fast/High-Quality toggle is present; Fast is the default.
+    await expect(page.getByTestId("eraser-quality-fast")).toHaveAttribute("aria-pressed", "true");
+    await expect(page.getByTestId("eraser-quality-hq")).toBeVisible();
+
+    // Paint a stroke so the ONLY thing gating submit is the quality mode.
+    const canvas = page.locator("canvas");
+    await canvas.waitFor({ state: "visible", timeout: 5_000 });
+    const box = await canvas.boundingBox();
+    if (!box) throw new Error("Canvas not found");
+    await page.mouse.move(box.x + box.width / 2, box.y + box.height / 2);
+    await page.mouse.down();
+    await page.mouse.move(box.x + box.width / 2 + 30, box.y + box.height / 2);
+    await page.mouse.up();
+
+    // Fast mode with a stroke: submit is enabled.
+    await expect(page.getByTestId("erase-object-submit")).toBeEnabled();
+
+    // Switch to High Quality: the pack is missing, so submit is blocked (never a
+    // silent downgrade to the fast path) and the install prompt appears.
+    await page.getByTestId("eraser-quality-hq").click();
+    await expect(page.getByTestId("eraser-quality-hq")).toHaveAttribute("aria-pressed", "true");
+    await expect(page.getByTestId("erase-object-submit")).toBeDisabled();
+    await expect(page.getByTestId("eraser-install-hq")).toBeVisible();
+
+    // Back to Fast: submit re-enables and the prompt is gone.
+    await page.getByTestId("eraser-quality-fast").click();
+    await expect(page.getByTestId("erase-object-submit")).toBeEnabled();
+    await expect(page.getByTestId("eraser-install-hq")).toHaveCount(0);
+  });
 });

tests/integration/tools/image/ai-async-route-coverage.test.ts

@@ -283,7 +283,7 @@ describe("custom async AI image routes", () => {
           expect.stringContaining("subject.png"),
           expect.stringContaining("mask.png"),
         ]),
-        settings: { format: "webp", quality: 72 },
+        settings: { format: "webp", quality: 72, qualityMode: "fast" },
         kind: "ai-tool",
       }),
     );

tests/integration/tools/image/erase-object-hq-gate.test.ts

@@ -0,0 +1,109 @@
+/**
+ * Integration tests for the Object Eraser "High Quality" (diffusion) gate.
+ *
+ * erase-object always requires its base bundle (object-eraser-colorize, LaMa).
+ * The optional qualityMode=hq additionally requires the inpaint-hq bundle. The
+ * route must 501 loudly for the missing HQ pack (never silently fall back to
+ * the fast path), while qualityMode=fast keeps working with only the base
+ * bundle installed.
+ *
+ * DATA_DIR is set to an isolated temp dir BEFORE importing feature-status (it
+ * reads DATA_DIR at module load) so we control which bundles read as installed.
+ */
+import { randomUUID } from "node:crypto";
+import { mkdirSync, rmSync, writeFileSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { afterAll, beforeAll, describe, expect, it } from "vitest";
+
+const testRoot = join(tmpdir(), `snapotter-erase-hq-guard-${randomUUID()}`);
+const aiDir = join(testRoot, "ai");
+const installedPath = join(aiDir, "installed.json");
+
+process.env.DATA_DIR = testRoot;
+process.env.FEATURE_MANIFEST_PATH = join(process.cwd(), "docker/feature-manifest.json");
+
+mkdirSync(join(aiDir, "models"), { recursive: true });
+writeFileSync(installedPath, JSON.stringify({ bundles: {} }), "utf-8");
+
+const { invalidateCache } = await import("../../../../apps/api/src/lib/feature-status.js");
+const { fixtures, readFixture } = await import("../../../fixtures/index.js");
+const { buildTestApp, createMultipartPayload, loginAsAdmin } = await import("../../test-server.js");
+
+type TestAppType = Awaited<ReturnType<typeof buildTestApp>>;
+
+const PNG = readFixture(fixtures.image.base.png200);
+
+let testApp: TestAppType;
+let app: TestAppType["app"];
+let adminToken: string;
+
+function setInstalled(bundleIds: string[]): void {
+  const bundles: Record<string, { version: string; installedAt: string; models: string[] }> = {};
+  for (const id of bundleIds) {
+    bundles[id] = { version: "1.0.0-test", installedAt: "2026-01-01T00:00:00.000Z", models: [] };
+  }
+  writeFileSync(installedPath, JSON.stringify({ bundles }), "utf-8");
+  invalidateCache();
+}
+
+async function postErase(qualityMode: "fast" | "hq") {
+  const { body, contentType } = createMultipartPayload([
+    { name: "file", filename: "test.png", contentType: "image/png", content: PNG },
+    { name: "mask", filename: "mask.png", contentType: "image/png", content: PNG },
+    { name: "qualityMode", content: qualityMode },
+  ]);
+  return app.inject({
+    method: "POST",
+    url: "/api/v1/tools/image/erase-object",
+    headers: { authorization: `Bearer ${adminToken}`, "content-type": contentType },
+    body,
+  });
+}
+
+beforeAll(async () => {
+  testApp = await buildTestApp();
+  app = testApp.app;
+  adminToken = await loginAsAdmin(app);
+}, 30_000);
+
+afterAll(async () => {
+  await testApp.cleanup();
+  rmSync(testRoot, { recursive: true, force: true });
+}, 10_000);
+
+describe("Object Eraser HQ (inpaint-hq) feature gate", () => {
+  it("501s naming the base bundle when nothing is installed", async () => {
+    setInstalled([]);
+    const res = await postErase("fast");
+    expect(res.statusCode).toBe(501);
+    const json = JSON.parse(res.body);
+    expect(json.code).toBe("FEATURE_NOT_INSTALLED");
+    expect(json.feature).toBe("object-eraser-colorize");
+  });
+
+  it("501s naming inpaint-hq when HQ is requested but only the base is installed", async () => {
+    setInstalled(["object-eraser-colorize"]);
+    const res = await postErase("hq");
+    expect(res.statusCode).toBe(501);
+    const json = JSON.parse(res.body);
+    expect(json.code).toBe("FEATURE_NOT_INSTALLED");
+    expect(json.feature).toBe("inpaint-hq");
+    expect(json.featureName).toBe("High-Quality Inpainting");
+  });
+
+  it("accepts fast mode with only the base bundle installed (no HQ needed)", async () => {
+    setInstalled(["object-eraser-colorize"]);
+    const res = await postErase("fast");
+    // The route enqueues and returns 202; it never 501s in fast mode.
+    expect(res.statusCode).not.toBe(501);
+    expect(res.statusCode).toBe(202);
+  });
+
+  it("accepts HQ mode once both the base and inpaint-hq bundles are installed", async () => {
+    setInstalled(["object-eraser-colorize", "inpaint-hq"]);
+    const res = await postErase("hq");
+    expect(res.statusCode).not.toBe(501);
+    expect(res.statusCode).toBe(202);
+  });
+});

tests/unit/ai/inpaint.test.ts

@@ -0,0 +1,102 @@
+import { readFile, writeFile } from "node:fs/promises";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+
+vi.mock("sharp", () => {
+  const mockSharp = vi.fn(() => ({
+    png: vi.fn().mockReturnThis(),
+    toBuffer: vi.fn().mockResolvedValue(Buffer.from("mock-png-data")),
+  }));
+  return { default: mockSharp };
+});
+
+vi.mock("node:fs/promises", () => ({
+  readFile: vi.fn().mockResolvedValue(Buffer.from("mock-output-data")),
+  writeFile: vi.fn().mockResolvedValue(undefined),
+}));
+
+vi.mock("../../../packages/ai/src/bridge.js", () => ({
+  runPythonWithProgress: vi.fn(),
+  parseStdoutJson: vi.fn(),
+}));
+
+import { parseStdoutJson, runPythonWithProgress } from "../../../packages/ai/src/bridge.js";
+import { inpaint } from "../../../packages/ai/src/inpainting.js";
+
+const IMG = Buffer.from("fake-image");
+const MASK = Buffer.from("fake-mask");
+const DIR = "/tmp/test-inpaint";
+
+beforeEach(() => {
+  vi.clearAllMocks();
+  vi.mocked(readFile).mockResolvedValue(Buffer.from("mock-output-data"));
+  vi.mocked(writeFile).mockResolvedValue(undefined);
+  vi.mocked(runPythonWithProgress).mockResolvedValue({ stdout: '{"success": true}', stderr: "" });
+  vi.mocked(parseStdoutJson).mockReturnValue({ success: true, method: "lama-onnx" });
+});
+
+afterEach(() => {
+  vi.restoreAllMocks();
+});
+
+describe("inpaint quality-mode script selection", () => {
+  const ARGS = [`${DIR}/input_inpaint.png`, `${DIR}/mask_inpaint.png`, `${DIR}/output_inpaint.png`];
+
+  it("runs the LaMa script (inpaint.py) by default", async () => {
+    await inpaint(IMG, MASK, DIR);
+    expect(runPythonWithProgress).toHaveBeenCalledWith("inpaint.py", ARGS, expect.any(Object));
+  });
+
+  it("runs the LaMa script when quality is explicitly 'fast'", async () => {
+    await inpaint(IMG, MASK, DIR, undefined, "fast");
+    expect(runPythonWithProgress).toHaveBeenCalledWith("inpaint.py", ARGS, expect.any(Object));
+  });
+
+  it("runs the diffusion script (inpaint_hq.py) when quality is 'hq'", async () => {
+    await inpaint(IMG, MASK, DIR, undefined, "hq");
+    // A regression here would silently run LaMa while the UI reported High Quality.
+    expect(runPythonWithProgress).toHaveBeenCalledWith("inpaint_hq.py", ARGS, expect.any(Object));
+    expect(runPythonWithProgress).not.toHaveBeenCalledWith(
+      "inpaint.py",
+      expect.anything(),
+      expect.anything(),
+    );
+  });
+});
+
+describe("inpaint contract", () => {
+  it("writes the input and mask as PNGs and returns the output buffer", async () => {
+    const out = await inpaint(IMG, MASK, DIR, undefined, "hq");
+    expect(writeFile).toHaveBeenCalledWith(
+      `${DIR}/input_inpaint.png`,
+      Buffer.from("mock-png-data"),
+    );
+    expect(writeFile).toHaveBeenCalledWith(`${DIR}/mask_inpaint.png`, Buffer.from("mock-png-data"));
+    expect(readFile).toHaveBeenCalledWith(`${DIR}/output_inpaint.png`);
+    expect(out).toEqual(Buffer.from("mock-output-data"));
+  });
+
+  it("forwards onProgress to the bridge", async () => {
+    const onProgress = vi.fn();
+    await inpaint(IMG, MASK, DIR, onProgress, "hq");
+    expect(runPythonWithProgress).toHaveBeenCalledWith(
+      "inpaint_hq.py",
+      expect.any(Array),
+      expect.objectContaining({ onProgress }),
+    );
+  });
+
+  it("throws the Python error (no silent fallback) when the script fails", async () => {
+    vi.mocked(parseStdoutJson).mockReturnValue({
+      success: false,
+      error: "High-quality inpainting model not found",
+    });
+    await expect(inpaint(IMG, MASK, DIR, undefined, "hq")).rejects.toThrow(
+      "High-quality inpainting model not found",
+    );
+  });
+
+  it("throws a fallback message when the script fails without an error string", async () => {
+    vi.mocked(parseStdoutJson).mockReturnValue({ success: false });
+    await expect(inpaint(IMG, MASK, DIR, undefined, "hq")).rejects.toThrow("Inpainting failed");
+  });
+});

tests/unit/features/feature-manifest.test.ts

@@ -37,11 +37,12 @@ describe("Feature manifest structure", () => {
     expect(manifest.basePackages).toBeInstanceOf(Array);
   });
 
-  it("all 7 bundles are defined", () => {
-    expect(Object.keys(bundles)).toHaveLength(7);
+  it("all 8 bundles are defined", () => {
+    expect(Object.keys(bundles)).toHaveLength(8);
     expect(bundles["background-removal"]).toBeDefined();
     expect(bundles["face-detection"]).toBeDefined();
     expect(bundles["object-eraser-colorize"]).toBeDefined();
+    expect(bundles["inpaint-hq"]).toBeDefined();
     expect(bundles["upscale-enhance"]).toBeDefined();
     expect(bundles["photo-restoration"]).toBeDefined();
     expect(bundles.ocr).toBeDefined();

tests/unit/features/feature-status.test.ts

@@ -1207,7 +1207,7 @@ describe("Composite state - getFeatureStates", () => {
     for (const state of states) {
       expect(state.status).toBe("not_installed");
     }
-    expect(states.length).toBe(7);
+    expect(states.length).toBe(8);
   });
 
   it("installed bundle with valid models returns installed with version", () => {

tests/unit/shared/features.test.ts

@@ -7,6 +7,7 @@ import {
   PYTHON_SIDECAR_TOOLS,
   TOOL_BUNDLE_MAP,
   TOOL_EXTRA_BUNDLES,
+  TOOL_OPTIONAL_BUNDLE_MAP,
 } from "@snapotter/shared";
 import { describe, expect, it } from "vitest";
 
@@ -30,39 +31,58 @@ describe("Feature bundles", () => {
     expect(tools).not.toContain("upscale");
   });
 
-  it("all 7 bundles are defined", () => {
-    expect(Object.keys(FEATURE_BUNDLES)).toHaveLength(7);
+  it("all 8 bundles are defined", () => {
+    expect(Object.keys(FEATURE_BUNDLES)).toHaveLength(8);
     expect(FEATURE_BUNDLES["background-removal"]).toBeDefined();
     expect(FEATURE_BUNDLES["face-detection"]).toBeDefined();
     expect(FEATURE_BUNDLES["object-eraser-colorize"]).toBeDefined();
+    expect(FEATURE_BUNDLES["inpaint-hq"]).toBeDefined();
     expect(FEATURE_BUNDLES["upscale-enhance"]).toBeDefined();
     expect(FEATURE_BUNDLES["photo-restoration"]).toBeDefined();
     expect(FEATURE_BUNDLES.ocr).toBeDefined();
     expect(FEATURE_BUNDLES.transcription).toBeDefined();
   });
 
-  it("TOOL_BUNDLE_MAP covers sidecar tools without an optional capability pack", () => {
+  it("every sidecar tool is reachable; only built-in-fast tools skip the required map", () => {
     const mappedTools = Object.keys(TOOL_BUNDLE_MAP);
     for (const toolId of PYTHON_SIDECAR_TOOLS) {
-      if (getOptionalBundleForTool(toolId)) {
+      // Reachable via a required primary and/or an optional upgrade pack.
+      expect(
+        getBundleForTool(toolId) !== null || getOptionalBundleForTool(toolId) !== null,
+        `${toolId} has no bundle at all`,
+      ).toBe(true);
+      // A tool ABSENT from TOOL_BUNDLE_MAP must be a built-in-fast tool whose only
+      // bundle is an optional pack (e.g. OCR's Fast tier + accurate pack). A tool
+      // with a required base stays mapped even if it also has an optional upgrade
+      // pack (e.g. erase-object's LaMa base + inpaint-hq diffusion pack).
+      if (!mappedTools.includes(toolId)) {
         expect(
-          mappedTools,
-          `${toolId} must remain available without its optional pack`,
-        ).not.toContain(toolId);
-      } else {
-        expect(mappedTools, `${toolId} missing from TOOL_BUNDLE_MAP`).toContain(toolId);
+          getOptionalBundleForTool(toolId),
+          `${toolId} is neither required-mapped nor a built-in-fast optional-pack tool`,
+        ).not.toBeNull();
       }
     }
   });
 });
 
 describe("Feature bundle edge cases", () => {
-  it("no duplicate tools across bundles", () => {
-    const allTools: string[] = [];
+  it("no tool appears in two non-optional bundles (an optional pack may re-list its tool)", () => {
+    const firstBundle = new Map<string, string>();
     for (const bundle of Object.values(FEATURE_BUNDLES)) {
       for (const tool of bundle.enablesTools) {
-        expect(allTools, `Tool ${tool} appears in multiple bundles`).not.toContain(tool);
-        allTools.push(tool);
+        const prior = firstBundle.get(tool);
+        if (prior === undefined) {
+          firstBundle.set(tool, bundle.id);
+          continue;
+        }
+        // The only allowed overlap: a tool's optional upgrade pack re-lists a
+        // tool its primary bundle already enables (e.g. inpaint-hq over
+        // erase-object). Any other pairing is an accidental duplicate.
+        const optional = TOOL_OPTIONAL_BUNDLE_MAP[tool];
+        expect(
+          optional !== undefined && (prior === optional || bundle.id === optional),
+          `Tool ${tool} appears in two non-optional bundles (${prior}, ${bundle.id})`,
+        ).toBe(true);
       }
     }
   });
@@ -158,3 +178,19 @@ describe("TOOL_EXTRA_BUNDLES", () => {
     }
   });
 });
+
+describe("inpaint-hq optional upgrade for erase-object", () => {
+  it("keeps object-eraser-colorize as the required primary; inpaint-hq stays optional", () => {
+    // The HQ diffusion pack upgrades Object Eraser but must not gate it: the base
+    // LaMa bundle remains the tool's required primary, and HQ is a separate,
+    // explicit install check (mirrors OCR's Fast tier + optional accurate pack).
+    expect(FEATURE_BUNDLES["inpaint-hq"]).toBeDefined();
+    expect(TOOL_BUNDLE_MAP["erase-object"]).toBe("object-eraser-colorize");
+    expect(TOOL_OPTIONAL_BUNDLE_MAP["erase-object"]).toBe("inpaint-hq");
+    expect(getBundleForTool("erase-object")?.id).toBe("object-eraser-colorize");
+    expect(getOptionalBundleForTool("erase-object")?.id).toBe("inpaint-hq");
+    // erase-object must NOT require inpaint-hq (fast path works without it).
+    expect(getRequiredBundlesForTool("erase-object")).toEqual(["object-eraser-colorize"]);
+    expect(getRequiredBundlesForTool("erase-object")).not.toContain("inpaint-hq");
+  });
+});