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
This commit is contained in:
SnapOtter
2026-07-19 20:47:35 +08:00
committed by GitHub
parent 84c18eb82c
commit 1bac663a2e
41 changed files with 954 additions and 35 deletions
+33 -7
View File
@@ -1,6 +1,6 @@
import { randomUUID } from "node:crypto";
import { inpaint } from "@snapotter/ai";
import { getBundleForTool, TOOL_BUNDLE_MAP } from "@snapotter/shared";
import { FEATURE_BUNDLES, getBundleForTool, TOOL_BUNDLE_MAP } from "@snapotter/shared";
import type { FastifyInstance, FastifyReply, FastifyRequest } from "fastify";
import sharp from "sharp";
import { z } from "zod";
@@ -8,7 +8,7 @@ import { enqueueToolJob } from "../../jobs/enqueue.js";
import { INVALID_SAVE_MODE_ERROR, parseSaveModeField } from "../../jobs/types.js";
import { autoOrient } from "../../lib/auto-orient.js";
import { stripInternalPaths } from "../../lib/errors.js";
import { isToolInstalled } from "../../lib/feature-status.js";
import { isFeatureInstalled, isToolInstalled } from "../../lib/feature-status.js";
import { validateImageBuffer } from "../../lib/file-validation.js";
import { decodeToSharpCompat, needsCliDecode } from "../../lib/format-decoders.js";
import { encodeJxl } from "../../lib/format-encoders.js";
@@ -24,8 +24,12 @@ const settingsSchema = z.object({
.enum(["auto", "png", "jpg", "jpeg", "webp", "tiff", "gif", "avif", "heic", "heif", "jxl"])
.default("auto"),
quality: z.number().int().min(1).max(100).default(95),
// "fast" = LaMa (always available); "hq" = diffusion, gated behind inpaint-hq.
qualityMode: z.enum(["fast", "hq"]).default("fast"),
});
const HQ_BUNDLE_ID = "inpaint-hq";
/**
* Object eraser / inpainting route.
* Accepts an image and a mask image, erases masked areas using LaMa.
@@ -60,6 +64,7 @@ export function registerEraseObject(app: FastifyInstance) {
let saveModeRaw: string | null = null;
let format = "png";
let quality = 95;
let qualityMode = "fast";
let imageKey: string | null = null;
let maskKey: string | null = null;
@@ -88,6 +93,8 @@ export function registerEraseObject(app: FastifyInstance) {
format = (part.value as string) || "png";
} else if (part.fieldname === "quality") {
quality = Number(part.value) || 95;
} else if (part.fieldname === "qualityMode") {
qualityMode = (part.value as string) || "fast";
}
}
} catch (err) {
@@ -123,8 +130,8 @@ export function registerEraseObject(app: FastifyInstance) {
return reply.status(400).send({ error: `Invalid mask: ${maskValidation.reason}` });
}
// Validate format and quality via Zod
const settingsResult = settingsSchema.safeParse({ format, quality });
// Validate format, quality, and quality mode via Zod
const settingsResult = settingsSchema.safeParse({ format, quality, qualityMode });
if (!settingsResult.success) {
return reply.status(400).send({
error: "Invalid settings",
@@ -135,6 +142,21 @@ export function registerEraseObject(app: FastifyInstance) {
}
format = settingsResult.data.format;
quality = settingsResult.data.quality;
qualityMode = settingsResult.data.qualityMode;
// High-Quality mode needs the optional diffusion bundle on top of the base
// (LaMa) bundle already checked above. Fail loud with the standard install
// contract; never silently downgrade HQ to the fast path.
if (qualityMode === "hq" && !isFeatureInstalled(HQ_BUNDLE_ID)) {
const hqBundle = FEATURE_BUNDLES[HQ_BUNDLE_ID];
return reply.status(501).send({
error: "Feature not installed",
code: "FEATURE_NOT_INSTALLED",
feature: HQ_BUNDLE_ID,
featureName: hqBundle?.name ?? "High-Quality Inpainting",
estimatedSize: hqBundle?.estimatedSize ?? "unknown",
});
}
if (format === "auto") {
const detected = await resolveOutputFormat(imageBuffer, filename);
@@ -176,7 +198,7 @@ export function registerEraseObject(app: FastifyInstance) {
pool: "ai",
inputRefs: [imageKey, maskKey],
filename,
settings: { format, quality },
settings: { format, quality, qualityMode },
clientJobId: clientJobId ?? undefined,
fileId: fileId ?? undefined,
saveMode,
@@ -198,8 +220,12 @@ registerAiJobHandler("erase-object", async (input, data, ctx) => {
const format = settings.format;
const quality = settings.quality;
const resultBuffer = await inpaint(input, maskBuffer, ctx.scratchDir, (percent, stage) =>
ctx.report(percent, stage),
const resultBuffer = await inpaint(
input,
maskBuffer,
ctx.scratchDir,
(percent, stage) => ctx.report(percent, stage),
settings.qualityMode,
);
// Convert to requested output format
@@ -1,13 +1,18 @@
import { Download, Lasso, Paintbrush, Redo, Trash2 } from "lucide-react";
import { Download, Lasso, Loader2, Paintbrush, Redo, Sparkles, Trash2, Zap } from "lucide-react";
import { useEffect, useRef, useState } from "react";
import { ProgressCard } from "@/components/common/progress-card";
import { useTranslation } from "@/contexts/i18n-context";
import { useAuth } from "@/hooks/use-auth";
import { formatHeaders } from "@/lib/api";
import { format } from "@/lib/format";
import { format, formatFileSize } from "@/lib/format";
import { generateId } from "@/lib/utils";
import { useFeaturesStore } from "@/stores/features-store";
import { useFileStore } from "@/stores/file-store";
import type { EraserCanvasRef } from "./eraser-canvas";
type QualityMode = "fast" | "hq";
const HQ_BUNDLE_ID = "inpaint-hq";
const OUTPUT_FORMATS = [
"png",
"jpg",
@@ -160,6 +165,22 @@ export function EraseObjectSettings({
const [outputFormat, setOutputFormat] = useState("png");
const [quality, setQuality] = useState(95);
const [qualityMode, setQualityMode] = useState<QualityMode>("fast");
// High-Quality (diffusion) mode is backed by the optional inpaint-hq bundle.
// Mirrors the OCR quality control: pick the mode, and if the pack is missing
// show the standard install prompt instead of silently running the fast path.
const { hasPermission } = useAuth();
const hqBundle = useFeaturesStore((s) => s.bundles.find((b) => b.id === HQ_BUNDLE_ID));
const hqInstalled = hqBundle?.status === "installed";
const installBundle = useFeaturesStore((s) => s.installBundle);
const hqInstalling = useFeaturesStore((s) => s.installing[HQ_BUNDLE_ID]);
const hqQueued = useFeaturesStore((s) => s.queued.includes(HQ_BUNDLE_ID));
const hqInstallError = useFeaturesStore((s) => s.errors[HQ_BUNDLE_ID]);
const needsHqPack = qualityMode === "hq" && !hqInstalled;
const isAdmin = hasPermission("features:manage");
const hqSizeBytes = hqBundle?.missingDownloadBytes ?? hqBundle?.downloadBytes;
const hqSize = hqSizeBytes ? formatFileSize(hqSizeBytes) : (hqBundle?.estimatedSize ?? "5-7 GB");
const processOneFile = (
entryIndex: number,
@@ -199,6 +220,7 @@ export function EraseObjectSettings({
formData.append("clientJobId", clientJobId);
formData.append("format", outputFormat);
formData.append("quality", String(quality));
formData.append("qualityMode", qualityMode);
const xhr = new XMLHttpRequest();
xhr.timeout = 600_000;
@@ -323,6 +345,7 @@ export function EraseObjectSettings({
formData.append("clientJobId", clientJobId);
formData.append("format", outputFormat);
formData.append("quality", String(quality));
formData.append("qualityMode", qualityMode);
const xhr = new XMLHttpRequest();
xhr.timeout = 600_000;
@@ -482,6 +505,81 @@ export function EraseObjectSettings({
</button>
</div>
{/* Quality: Fast (LaMa, always available) vs High quality (diffusion, inpaint-hq) */}
<div>
<div className="flex gap-1 rounded-lg bg-muted p-1">
<button
type="button"
data-testid="eraser-quality-fast"
aria-pressed={qualityMode === "fast"}
disabled={processing}
onClick={() => setQualityMode("fast")}
className={`flex-1 flex items-center justify-center gap-1.5 py-1.5 rounded-md text-xs font-medium transition-colors disabled:opacity-50 ${
qualityMode === "fast"
? "bg-background text-foreground shadow-sm"
: "text-muted-foreground hover:text-foreground"
}`}
>
<Zap className="h-3.5 w-3.5" />
{t.toolSettings["erase-object"].qualityFast}
</button>
<button
type="button"
data-testid="eraser-quality-hq"
aria-pressed={qualityMode === "hq"}
disabled={processing}
onClick={() => setQualityMode("hq")}
className={`flex-1 flex items-center justify-center gap-1.5 py-1.5 rounded-md text-xs font-medium transition-colors disabled:opacity-50 ${
qualityMode === "hq"
? "bg-background text-foreground shadow-sm"
: "text-muted-foreground hover:text-foreground"
}`}
>
<Sparkles className="h-3.5 w-3.5" />
{t.toolSettings["erase-object"].qualityHq}
</button>
</div>
{qualityMode === "hq" && (
<p className="mt-1 text-[10px] text-muted-foreground">
{t.toolSettings["erase-object"].qualityHint}
</p>
)}
{needsHqPack && (
<div className="mt-2 rounded-lg border border-border bg-muted/40 p-3 text-start">
<p className="text-xs text-muted-foreground">
{format(t.features.requiresDownload, { size: hqSize })}
</p>
{isAdmin ? (
<button
type="button"
data-testid="eraser-install-hq"
onClick={() => installBundle(HQ_BUNDLE_ID)}
disabled={!!hqInstalling || hqQueued}
className="mt-2 inline-flex items-center gap-1.5 rounded-md bg-primary px-3 py-1.5 text-xs font-medium text-primary-foreground disabled:opacity-50"
>
{hqInstalling || hqQueued ? (
<Loader2 className="h-3.5 w-3.5 animate-spin" />
) : (
<Download className="h-3.5 w-3.5" />
)}
{hqInstalling || hqQueued
? t.settings.aiFeatures.installing
: format(t.features.enableButton, {
name: hqBundle?.name ?? "High-Quality Inpainting",
})}
</button>
) : (
<p className="mt-1 text-xs text-muted-foreground">
{t.features.notEnabledDescription}
</p>
)}
{hqInstallError && <p className="mt-1 text-xs text-destructive">{hqInstallError}</p>}
</div>
)}
</div>
{/* Brush size (brush mode only) */}
{mode === "brush" && (
<div>
@@ -606,7 +704,7 @@ export function EraseObjectSettings({
type="button"
data-testid="erase-object-submit"
onClick={maskedFileCount > 1 ? handleProcessAll : handleProcess}
disabled={!hasFile || (!hasStrokes && maskedFileCount === 0) || processing}
disabled={!hasFile || (!hasStrokes && maskedFileCount === 0) || processing || needsHqPack}
className="w-full py-2.5 rounded-lg bg-primary text-primary-foreground font-medium disabled:opacity-50 disabled:cursor-not-allowed flex items-center justify-center gap-2"
>
{maskedFileCount > 1
+6 -1
View File
@@ -333,9 +333,14 @@ for model in models:
kwargs = {"repo_id": repo_id, "local_dir": local_dir}
# Only download specific file if specified
# Restrict the snapshot when specified. "file" pins one file (single-file
# models like an ONNX weight); "allowPatterns" narrows a multi-file model
# (e.g. a diffusers pipeline) to the fp16 weight variant + configs so the
# bundle does not ship unused fp32/.bin weights.
if "file" in model:
kwargs["allow_patterns"] = [model["file"]]
elif "allowPatterns" in model:
kwargs["allow_patterns"] = model["allowPatterns"]
# Handle non-default repo types (e.g. "space")
if "repoType" in model:
+48
View File
@@ -203,6 +203,54 @@
"smokeImports": ["onnxruntime"],
"enablesTools": ["erase-object", "colorize", "ai-canvas-expand"]
},
"inpaint-hq": {
"name": "High-Quality Inpainting",
"description": "Diffusion-based object removal for large objects and detailed textures",
"estimatedSize": "5-7 GB",
"archives": {
"amd64-gpu": {
"file": "v2.0.0/inpaint-hq-amd64-gpu.tar.gz",
"sha256": "da65c2b7e678b50815c7c2d8ae4594ad7bbf4c1e84682bfc91b46948cdff71de",
"compressedSize": 7283076233,
"extractedSize": 10745029298
},
"arm64-cpu": {
"file": "v2.0.0/inpaint-hq-arm64-cpu.tar.gz",
"sha256": "b854b101a6190dde94d9d0039129941122c62df7c9a6429d425231cf4fa4630f",
"compressedSize": 5130985609,
"extractedSize": 7159491153
}
},
"packages": {
"common": [
"diffusers==0.31.0",
"transformers==4.46.3",
"accelerate==1.2.1",
"safetensors==0.4.5",
"huggingface-hub[hf_xet]==0.36.2"
],
"amd64": [
"torch==2.7.0+cu126 torchvision==0.22.0+cu126 --index-url https://download.pytorch.org/whl/cu126"
],
"arm64": [
"torch==2.7.0 torchvision==0.22.0 --index-url https://download.pytorch.org/whl/cpu"
]
},
"pipFlags": {},
"postInstall": [],
"models": [
{
"id": "sd15-inpainting",
"downloadFn": "hf_snapshot",
"args": ["stable-diffusion-v1-5/stable-diffusion-inpainting", "sd15-inpainting"],
"allowPatterns": ["*.json", "*.txt", "tokenizer/*", "*.fp16.safetensors"],
"path": "sd15-inpainting/unet/diffusion_pytorch_model.fp16.safetensors",
"minSize": 1000000000
}
],
"smokeImports": ["diffusers", "torch"],
"enablesTools": ["erase-object"]
},
"upscale-enhance": {
"name": "Upscale & Enhance",
"description": "AI upscaling, face enhancement, and noise removal",
+2
View File
@@ -64,6 +64,7 @@ ALLOWED_SCRIPTS = {
"face_landmarks",
"gif_remove_bg",
"inpaint",
"inpaint_hq",
"install_feature",
"noise_removal",
"ocr_preprocess",
@@ -114,6 +115,7 @@ TOOL_BUNDLE_MAP = {
"face_landmarks": "face-detection",
"red_eye_removal": "face-detection",
"inpaint": "object-eraser-colorize",
"inpaint_hq": "inpaint-hq",
"outpaint": "object-eraser-colorize",
"colorize": "object-eraser-colorize",
"upscale": "upscale-enhance",
+226
View File
@@ -0,0 +1,226 @@
"""High-quality object erasing via Stable Diffusion 1.5 inpainting (diffusers).
This is the optional "High Quality" backend for the Object Eraser, gated behind
the `inpaint-hq` feature bundle. The default fast path stays `inpaint.py` (LaMa).
Design: reuse the crop-and-composite geometry from `inpaint.py`
(`inpaint_array` dilates the mask, crops a padded HD window, runs a model on the
crop, and blends only the masked region back into the untouched original). The
only difference here is the model step: a diffusion pipeline replaces the LaMa
ONNX session. Diffusion synthesizes plausible texture over large/structured
regions where a non-diffusion model smears, which is exactly #141's open case.
Heavy imports (torch/diffusers) are lazy so the base AI dispatcher stays lean and
the geometry stays unit-testable with an injected fake pipeline.
"""
import json
import os
import sys
import inpaint # reuse dilate/crop/composite geometry (inpaint_array)
def emit_progress(percent, stage):
"""Emit structured progress to stderr for bridge.ts to capture."""
print(json.dumps({"progress": percent, "stage": stage}), file=sys.stderr, flush=True)
# Model directory: the inpaint-hq bundle downloads the SD1.5 inpainting model
# (diffusers layout) here via hf_snapshot. MODELS_PATH is set by the bridge to
# DATA_DIR/ai/models; /opt/models is the baked fallback for other model kinds.
_MODELS_BASE = os.environ.get("MODELS_PATH", "/opt/models")
SD_MODEL_DIR = os.environ.get("SD15_INPAINT_DIR", os.path.join(_MODELS_BASE, "sd15-inpainting"))
# Diffusion runs at SD1.5's native 512. Crops are resized to this and the result
# resized back, so the crop-HD property (small objects keep resolution) still holds.
MODEL_SIZE = 512
# Inference defaults. Overridable by env for tuning without a rebuild. An empty
# prompt with a "keep it background" negative prompt biases toward clean removal
# (continue the surroundings) rather than hallucinating a new object.
STEPS = int(os.environ.get("SD15_INPAINT_STEPS", "28"))
GUIDANCE = float(os.environ.get("SD15_INPAINT_GUIDANCE", "7.0"))
PROMPT = os.environ.get("SD15_INPAINT_PROMPT", "")
NEGATIVE_PROMPT = os.environ.get(
"SD15_INPAINT_NEGATIVE",
"object, person, text, watermark, artifact, blurry, distorted, extra limbs",
)
# Fixed seed so a given input erases deterministically (stable, reproducible,
# testable) instead of changing on every run.
SEED = int(os.environ.get("SD15_INPAINT_SEED", "0"))
def make_run_model(pipe, device, steps=STEPS, guidance=GUIDANCE, prompt=PROMPT,
negative_prompt=NEGATIVE_PROMPT, seed=SEED, progress=None):
"""Build a run_model(crop_img, crop_mask) backed by a diffusion pipeline.
Matches inpaint.py's run_model contract: crop in (HxWx3 uint8 RGB), inpainted
crop out (same HxWx3). Resizes the crop to the model's 512, runs the pipe with
the (dilated) mask, and resizes the result back to the native crop size.
`pipe` is any callable with the diffusers inpaint signature; injecting a fake
keeps this unit-testable without torch/diffusers or a real model.
"""
import inspect
import cv2
import numpy as np
from PIL import Image
supports_step_cb = "callback_on_step_end" in inspect.signature(pipe.__call__).parameters
def _make_generator():
try:
import torch
return torch.Generator(device=device).manual_seed(seed)
except Exception:
return None
def run_model(crop_img, crop_mask):
h, w = crop_img.shape[:2]
interp = cv2.INTER_AREA if (w > MODEL_SIZE or h > MODEL_SIZE) else cv2.INTER_LINEAR
img_resized = cv2.resize(crop_img, (MODEL_SIZE, MODEL_SIZE), interpolation=interp)
mask_resized = cv2.resize(
crop_mask, (MODEL_SIZE, MODEL_SIZE), interpolation=cv2.INTER_NEAREST
)
pil_img = Image.fromarray(img_resized)
pil_mask = Image.fromarray((mask_resized > 127).astype(np.uint8) * 255)
kwargs = dict(
prompt=prompt,
negative_prompt=negative_prompt,
image=pil_img,
mask_image=pil_mask,
num_inference_steps=steps,
guidance_scale=guidance,
height=MODEL_SIZE,
width=MODEL_SIZE,
)
gen = _make_generator()
if gen is not None:
kwargs["generator"] = gen
if supports_step_cb and progress is not None:
def _cb(_pipe, step, _timestep, cbk):
progress(int(45 + 30 * (step + 1) / max(1, steps)), "Erasing objects")
return cbk
kwargs["callback_on_step_end"] = _cb
out = pipe(**kwargs).images[0]
out_arr = np.array(out.convert("RGB"))
# Diffusion emits exactly MODEL_SIZE; resize back to the native crop.
if out_arr.shape[:2] != (h, w):
out_arr = cv2.resize(out_arr, (w, h), interpolation=cv2.INTER_LANCZOS4)
return out_arr
return run_model
def _resolve_device():
"""cuda when torch can actually use it, else cpu (mirrors the other torch tools)."""
try:
from gpu import torch_gpu_available
return "cuda" if torch_gpu_available() else "cpu"
except Exception:
return "cpu"
def _load_pipeline(model_dir, device):
"""Load the SD1.5 inpainting pipeline from the local bundle dir (never downloads)."""
# Check the model exists before importing the heavy stack, so a missing
# bundle fails fast with an actionable message instead of an ImportError.
if not os.path.isdir(model_dir):
raise FileNotFoundError(
f"High-quality inpainting model not found at {model_dir}. "
"Install the 'High-Quality Inpainting' feature bundle first."
)
import torch
from diffusers import StableDiffusionInpaintPipeline
dtype = torch.float16 if device == "cuda" else torch.float32
# Prefer the fp16 weight variant when the bundle ships it: it halves the
# download and loads on GPU (fp16) or CPU (cast up to fp32) alike. Fall back
# to non-variant (fp32) weights when only those are present.
fp16_unet = os.path.join(model_dir, "unet", "diffusion_pytorch_model.fp16.safetensors")
variant = "fp16" if os.path.exists(fp16_unet) else None
pipe = StableDiffusionInpaintPipeline.from_pretrained(
model_dir,
torch_dtype=dtype,
variant=variant,
safety_checker=None,
requires_safety_checker=False,
local_files_only=True,
)
pipe = pipe.to(device)
pipe.set_progress_bar_config(disable=True)
# Keep peak memory modest so mid-range GPUs and CPU hosts do not OOM at 512.
try:
pipe.enable_attention_slicing()
except Exception:
pass
return pipe
def main():
input_path = sys.argv[1]
mask_path = sys.argv[2]
output_path = sys.argv[3]
try:
emit_progress(5, "Preparing")
from PIL import Image
import numpy as np
try:
import cv2 # noqa: F401
import torch # noqa: F401
import diffusers # noqa: F401
except ImportError as e:
print(json.dumps({
"success": False,
"error": (
f"Missing dependency: {e}. The High-Quality Inpainting bundle "
"provides diffusers/torch; install it and retry."
),
}))
sys.exit(1)
emit_progress(15, "Loading model")
device = _resolve_device()
pipe = _load_pipeline(SD_MODEL_DIR, device)
emit_progress(35, "Loading images")
img = Image.open(input_path).convert("RGB")
mask = Image.open(mask_path).convert("L")
img_array = np.array(img)
mask_array = np.array(mask)
if mask_array.shape[:2] != img_array.shape[:2]:
import cv2
mask_array = cv2.resize(
mask_array,
(img_array.shape[1], img_array.shape[0]),
interpolation=cv2.INTER_NEAREST,
)
run_model = make_run_model(pipe, device, progress=emit_progress)
result = inpaint.inpaint_array(
img_array, mask_array, run_model, progress=emit_progress
)
emit_progress(90, "Saving")
Image.fromarray(result).save(output_path)
print(json.dumps({"success": True, "method": "sd15-inpainting"}))
except Exception as e:
print(json.dumps({"success": False, "error": str(e)}))
sys.exit(1)
if __name__ == "__main__":
main()
+124
View File
@@ -0,0 +1,124 @@
"""Unit tests for the SD1.5 diffusion inpainting sidecar (inpaint_hq.py).
The diffusion pipeline is injected as a fake, so no torch/diffusers or model is
needed: only the resize/call/resize-back contract and its reuse of inpaint.py's
crop-and-composite geometry are exercised. Skips where numpy/cv2 are absent, as
on CI integration shards (matches test_inpaint_geometry.py).
"""
import os
import sys
import pytest
np = pytest.importorskip("numpy")
cv2 = pytest.importorskip("cv2")
pytest.importorskip("PIL")
from PIL import Image # noqa: E402
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
import inpaint # noqa: E402
import inpaint_hq # noqa: E402
class _FakePipeResult:
def __init__(self, image):
self.images = [image]
class FakePipe:
"""Stand-in for a diffusers inpaint pipeline. Fills the whole 512 canvas with
a constant colour so callers can assert what landed in the masked region.
Records the last call's image/mask sizes and whether a step callback ran."""
def __init__(self, fill=(255, 0, 255)):
self.fill = fill
self.calls = []
self.callback_ran = False
def __call__(
self,
prompt=None,
negative_prompt=None,
image=None,
mask_image=None,
num_inference_steps=1,
guidance_scale=7.0,
height=512,
width=512,
generator=None,
callback_on_step_end=None,
):
self.calls.append(
{
"image_size": image.size if image is not None else None,
"mask_size": mask_image.size if mask_image is not None else None,
"steps": num_inference_steps,
}
)
if callback_on_step_end is not None:
callback_on_step_end(self, 0, 0, {})
self.callback_ran = True
out = Image.new("RGB", (width, height), self.fill)
return _FakePipeResult(out)
def test_make_run_model_resizes_to_model_and_back():
pipe = FakePipe(fill=(10, 20, 30))
run_model = inpaint_hq.make_run_model(pipe, "cpu", steps=3)
crop = np.zeros((100, 120, 3), np.uint8)
mask = np.zeros((100, 120), np.uint8)
mask[30:70, 40:80] = 255
out = run_model(crop, mask)
# Output is resized back to the native crop size, RGB.
assert out.shape == (100, 120, 3)
# The pipe saw a 512x512 image and mask (PIL size is (w, h)).
assert pipe.calls[-1]["image_size"] == (inpaint_hq.MODEL_SIZE, inpaint_hq.MODEL_SIZE)
assert pipe.calls[-1]["mask_size"] == (inpaint_hq.MODEL_SIZE, inpaint_hq.MODEL_SIZE)
assert pipe.calls[-1]["steps"] == 3
# The constant fill is what came back (resized), so the centre is that colour.
assert tuple(int(v) for v in out[50, 60]) == (10, 20, 30)
def test_make_run_model_invokes_progress_callback():
pipe = FakePipe()
seen = []
run_model = inpaint_hq.make_run_model(
pipe, "cpu", steps=2, progress=lambda pct, stage: seen.append((pct, stage))
)
run_model(np.zeros((60, 60, 3), np.uint8), _center_mask(60, 60))
assert pipe.callback_ran is True
assert seen and all(0 <= p <= 100 for p, _ in seen)
def test_inpaint_array_with_diffusion_leaves_far_pixels_untouched():
# Reuse inpaint.py's crop/composite via a diffusion run_model. The fill only
# lands inside the (feathered) mask; everything far from it stays identical.
rng = np.random.RandomState(0)
img = rng.randint(0, 256, (400, 500, 3), np.uint8)
mask = np.zeros((400, 500), np.uint8)
cv2.circle(mask, (250, 200), 50, 255, -1)
pipe = FakePipe(fill=(255, 0, 255))
run_model = inpaint_hq.make_run_model(pipe, "cpu", steps=1)
out = inpaint.inpaint_array(img, mask, run_model)
assert out.shape == img.shape
# Corner far from the mask is byte-identical to the original.
assert np.array_equal(out[0:60, 0:60], img[0:60, 0:60])
# Mask centre received the magenta fill.
assert out[200, 250, 0] > 200 and out[200, 250, 2] > 200
def test_load_pipeline_missing_dir_raises_actionable_error():
with pytest.raises(FileNotFoundError) as exc:
inpaint_hq._load_pipeline("/nonexistent/sd15-inpainting", "cpu")
assert "feature bundle" in str(exc.value).lower()
def _center_mask(h, w):
m = np.zeros((h, w), np.uint8)
m[h // 4 : 3 * h // 4, w // 4 : 3 * w // 4] = 255
return m
+1
View File
@@ -17,6 +17,7 @@ export const SCRIPT_BUNDLE_MAP: Record<string, string> = {
face_landmarks: "face-detection",
red_eye_removal: "face-detection",
inpaint: "object-eraser-colorize",
inpaint_hq: "inpaint-hq",
outpaint: "object-eraser-colorize",
colorize: "object-eraser-colorize",
upscale: "upscale-enhance",
+1 -1
View File
@@ -27,7 +27,7 @@ export { enhanceFaces } from "./face-enhancement.js";
export type { FaceLandmarkPoint, FaceLandmarks, FaceLandmarksResult } from "./face-landmarks.js";
export { detectFaceLandmarks } from "./face-landmarks.js";
export { missingBundleForScript, SCRIPT_BUNDLE_MAP } from "./feature-gate.js";
export { inpaint } from "./inpainting.js";
export { type InpaintQuality, inpaint } from "./inpainting.js";
export { noiseRemoval } from "./noise-removal.js";
export type {
OcrExecutionMetadata,
+12 -1
View File
@@ -3,11 +3,21 @@ import { join } from "node:path";
import sharp from "sharp";
import { type ProgressCallback, parseStdoutJson, runPythonWithProgress } from "./bridge.js";
/**
* Inpainting backend. "fast" is the always-available LaMa ONNX path
* (`inpaint.py`); "hq" is the optional diffusion path (`inpaint_hq.py`), gated
* behind the `inpaint-hq` feature bundle. The route decides which mode to pass;
* the sidecar/per-request feature gate independently rejects "hq" when the
* bundle is absent.
*/
export type InpaintQuality = "fast" | "hq";
export async function inpaint(
inputBuffer: Buffer,
maskBuffer: Buffer,
outputDir: string,
onProgress?: ProgressCallback,
quality: InpaintQuality = "fast",
): Promise<Buffer> {
const inputPath = join(outputDir, "input_inpaint.png");
const maskPath = join(outputDir, "mask_inpaint.png");
@@ -18,7 +28,8 @@ export async function inpaint(
await writeFile(inputPath, pngInput);
await writeFile(maskPath, pngMask);
const { stdout } = await runPythonWithProgress("inpaint.py", [inputPath, maskPath, outputPath], {
const script = quality === "hq" ? "inpaint_hq.py" : "inpaint.py";
const { stdout } = await runPythonWithProgress(script, [inputPath, maskPath, outputPath], {
onProgress,
});
+19 -3
View File
@@ -69,6 +69,14 @@ export const FEATURE_BUNDLES: Record<string, FeatureBundleInfo> = {
estimatedSize: "1-2 GB",
enablesTools: ["erase-object", "colorize", "ai-canvas-expand"],
},
"inpaint-hq": {
id: "inpaint-hq",
name: "High-Quality Inpainting",
description:
"Diffusion-based object removal for large objects, detailed textures, and structured backgrounds",
estimatedSize: "5-7 GB",
enablesTools: ["erase-object"],
},
"upscale-enhance": {
id: "upscale-enhance",
name: "Upscale & Enhance",
@@ -107,14 +115,22 @@ export const FEATURE_BUNDLES: Record<string, FeatureBundleInfo> = {
export const TOOL_OPTIONAL_BUNDLE_MAP: Readonly<Record<string, string>> = {
ocr: "ocr",
"ocr-pdf": "ocr",
// High-Quality (diffusion) inpainting upgrades Object Eraser without gating
// it: the base LaMa model in `object-eraser-colorize` stays the tool's
// required primary. `getRequiredBundlesForTool("erase-object")` is unchanged;
// HQ availability is a separate, explicit `inpaint-hq` install check.
"erase-object": "inpaint-hq",
};
export const TOOL_BUNDLE_MAP: Record<string, string> = {};
for (const [bundleId, bundle] of Object.entries(FEATURE_BUNDLES)) {
for (const toolId of bundle.enablesTools) {
if (!TOOL_OPTIONAL_BUNDLE_MAP[toolId]) {
TOOL_BUNDLE_MAP[toolId] = bundleId;
}
// An optional pack must never claim a tool's required-primary slot, but a
// different, non-optional bundle still can. Skip only when THIS bundle is
// the tool's optional pack; the first non-optional bundle to list the tool
// wins. This is behavior-identical for every tool that has no optional pack.
if (TOOL_OPTIONAL_BUNDLE_MAP[toolId] === bundleId) continue;
if (!TOOL_BUNDLE_MAP[toolId]) TOOL_BUNDLE_MAP[toolId] = bundleId;
}
}
+3
View File
@@ -1558,6 +1558,9 @@ export const ar: TranslationKeys = {
progressLabelBatch: "جاري تكبير {count} صورة",
},
"erase-object": {
qualityFast: "سريع",
qualityHq: "جودة عالية",
qualityHint: "الجودة العالية تستخدم نموذج انتشار للأجسام الكبيرة والقوام التفصيلي.",
brushMode: "فرشاة",
lassoMode: "لاسو",
brushSize: "حجم الفرشاة",
+3
View File
@@ -1575,6 +1575,9 @@ export const de: TranslationKeys = {
progressLabelBatch: "{count} Bilder werden hochskaliert",
},
"erase-object": {
qualityFast: "Schnell",
qualityHq: "Hohe Qualität",
qualityHint: "Hohe Qualität nutzt ein Diffusionsmodell für große Objekte und feine Texturen.",
brushMode: "Pinsel",
lassoMode: "Lasso",
brushSize: "Pinselgröße",
+3
View File
@@ -1522,6 +1522,9 @@ export const en = {
progressLabelBatch: "Upscaling {count} images",
},
"erase-object": {
qualityFast: "Fast",
qualityHq: "High quality",
qualityHint: "High quality uses a diffusion model for large objects and detailed textures.",
brushMode: "Brush",
lassoMode: "Lasso",
brushSize: "Brush Size",
+4
View File
@@ -1558,6 +1558,10 @@ export const es: TranslationKeys = {
progressLabelBatch: "Escalando {count} imágenes",
},
"erase-object": {
qualityFast: "Rápido",
qualityHq: "Alta calidad",
qualityHint:
"La alta calidad usa un modelo de difusión para objetos grandes y texturas detalladas.",
brushMode: "Pincel",
lassoMode: "Lazo",
brushSize: "Tamaño del pincel",
+4
View File
@@ -1582,6 +1582,10 @@ export const fr: TranslationKeys = {
progressLabelBatch: "Agrandissement de {count} images",
},
"erase-object": {
qualityFast: "Rapide",
qualityHq: "Haute qualité",
qualityHint:
"La haute qualité utilise un modèle de diffusion pour les grands objets et les textures détaillées.",
brushMode: "Pinceau",
lassoMode: "Lasso",
brushSize: "Taille du pinceau",
+3
View File
@@ -1388,6 +1388,9 @@ export const hi: TranslationKeys = {
progressLabelBatch: "{count} इमेज अपस्केल हो रही हैं",
},
"erase-object": {
qualityFast: "तेज़",
qualityHq: "उच्च गुणवत्ता",
qualityHint: "उच्च गुणवत्ता बड़ी वस्तुओं और विस्तृत बनावट के लिए डिफ्यूज़न मॉडल का उपयोग करती है।",
brushMode: "ब्रश",
lassoMode: "लासो",
brushSize: "ब्रश साइज़",
+4
View File
@@ -1568,6 +1568,10 @@ export const id: TranslationKeys = {
progressLabelBatch: "Memperbesar {count} gambar",
},
"erase-object": {
qualityFast: "Cepat",
qualityHq: "Kualitas tinggi",
qualityHint:
"Kualitas tinggi menggunakan model difusi untuk objek besar dan tekstur mendetail.",
brushMode: "Kuas",
lassoMode: "Laso",
brushSize: "Ukuran Kuas",
+4
View File
@@ -1573,6 +1573,10 @@ export const it: TranslationKeys = {
progressLabelBatch: "Ingrandimento di {count} immagini",
},
"erase-object": {
qualityFast: "Veloce",
qualityHq: "Alta qualità",
qualityHint:
"L'alta qualità usa un modello di diffusione per oggetti grandi e texture dettagliate.",
brushMode: "Pennello",
lassoMode: "Lazo",
brushSize: "Dimensione pennello",
+3
View File
@@ -1530,6 +1530,9 @@ export const ja: TranslationKeys = {
progressLabelBatch: "{count}枚の画像をアップスケール中",
},
"erase-object": {
qualityFast: "高速",
qualityHq: "高品質",
qualityHint: "高品質は大きなオブジェクトや細かいテクスチャに拡散モデルを使用します。",
brushMode: "ブラシ",
lassoMode: "投げ縄",
brushSize: "ブラシサイズ",
+3
View File
@@ -1512,6 +1512,9 @@ export const ko: TranslationKeys = {
progressLabelBatch: "{count}개 이미지 업스케일 중",
},
"erase-object": {
qualityFast: "빠름",
qualityHq: "고품질",
qualityHint: "고품질은 큰 객체와 세밀한 텍스처에 디퓨전 모델을 사용합니다.",
brushMode: "브러시",
lassoMode: "올가미",
brushSize: "브러시 크기",
+4
View File
@@ -1572,6 +1572,10 @@ export const nl: TranslationKeys = {
progressLabelBatch: "{count} afbeeldingen opschalen",
},
"erase-object": {
qualityFast: "Snel",
qualityHq: "Hoge kwaliteit",
qualityHint:
"Hoge kwaliteit gebruikt een diffusiemodel voor grote objecten en gedetailleerde texturen.",
brushMode: "Penseel",
lassoMode: "Lasso",
brushSize: "Penseelgrootte",
+4
View File
@@ -1573,6 +1573,10 @@ export const pl: TranslationKeys = {
progressLabelBatch: "Powiększanie {count} obrazów",
},
"erase-object": {
qualityFast: "Szybko",
qualityHq: "Wysoka jakość",
qualityHint:
"Wysoka jakość używa modelu dyfuzyjnego do dużych obiektów i szczegółowych tekstur.",
brushMode: "Pędzel",
lassoMode: "Lasso",
brushSize: "Rozmiar pędzla",
+4
View File
@@ -1571,6 +1571,10 @@ export const ptBR: TranslationKeys = {
progressLabelBatch: "Ampliando {count} imagens",
},
"erase-object": {
qualityFast: "Rápido",
qualityHq: "Alta qualidade",
qualityHint:
"A alta qualidade usa um modelo de difusão para objetos grandes e texturas detalhadas.",
brushMode: "Pincel",
lassoMode: "Laço",
brushSize: "Tamanho do pincel",
+4
View File
@@ -1571,6 +1571,10 @@ export const ru: TranslationKeys = {
progressLabelBatch: "Увеличение {count} изображений",
},
"erase-object": {
qualityFast: "Быстро",
qualityHq: "Высокое качество",
qualityHint:
"Высокое качество использует диффузионную модель для больших объектов и детальных текстур.",
brushMode: "Кисть",
lassoMode: "Лассо",
brushSize: "Размер кисти",
+4
View File
@@ -1568,6 +1568,10 @@ export const sv: TranslationKeys = {
progressLabelBatch: "Uppskalar {count} bilder",
},
"erase-object": {
qualityFast: "Snabb",
qualityHq: "Hög kvalitet",
qualityHint:
"Hög kvalitet använder en diffusionsmodell för stora objekt och detaljerade texturer.",
brushMode: "Pensel",
lassoMode: "Lasso",
brushSize: "Penselstorlek",
+3
View File
@@ -1548,6 +1548,9 @@ export const th: TranslationKeys = {
progressLabelBatch: "กำลังขยาย {count} ภาพ",
},
"erase-object": {
qualityFast: "เร็ว",
qualityHq: "คุณภาพสูง",
qualityHint: "คุณภาพสูงใช้โมเดล diffusion สำหรับวัตถุขนาดใหญ่และพื้นผิวที่มีรายละเอียด",
brushMode: "แปรง",
lassoMode: "ลาสโซ",
brushSize: "ขนาดแปรง",
+4
View File
@@ -1571,6 +1571,10 @@ export const tr: TranslationKeys = {
progressLabelBatch: "{count} görüntü büyütülüyor",
},
"erase-object": {
qualityFast: "Hızlı",
qualityHq: "Yüksek kalite",
qualityHint:
"Yüksek kalite, büyük nesneler ve ayrıntılı dokular için bir difüzyon modeli kullanır.",
brushMode: "Fırça",
lassoMode: "Kement",
brushSize: "Fırça Boyutu",
+4
View File
@@ -1571,6 +1571,10 @@ export const uk: TranslationKeys = {
progressLabelBatch: "Збільшення {count} зображень",
},
"erase-object": {
qualityFast: "Швидко",
qualityHq: "Висока якість",
qualityHint:
"Висока якість використовує дифузійну модель для великих об'єктів і детальних текстур.",
brushMode: "Пензель",
lassoMode: "Ласо",
brushSize: "Розмір пензля",
+3
View File
@@ -1571,6 +1571,9 @@ export const vi: TranslationKeys = {
progressLabelBatch: "Đang phóng to {count} ảnh",
},
"erase-object": {
qualityFast: "Nhanh",
qualityHq: "Chất lượng cao",
qualityHint: "Chất lượng cao dùng mô hình khuếch tán cho vật thể lớn và kết cấu chi tiết.",
brushMode: "Cọ",
lassoMode: "Lasso",
brushSize: "Kích thước cọ",
+3
View File
@@ -1336,6 +1336,9 @@ export const zhCN: TranslationKeys = {
progressLabelBatch: "正在放大 {count} 张图片",
},
"erase-object": {
qualityFast: "快速",
qualityHq: "高质量",
qualityHint: "高质量使用扩散模型处理大型物体和精细纹理。",
brushMode: "画笔",
lassoMode: "套索",
brushSize: "画笔大小",
+3
View File
@@ -1336,6 +1336,9 @@ export const zhTW: TranslationKeys = {
progressLabelBatch: "正在放大{count}張影像",
},
"erase-object": {
qualityFast: "快速",
qualityHq: "高品質",
qualityHint: "高品質使用擴散模型處理大型物件與精細紋理。",
brushMode: "筆刷",
lassoMode: "套索",
brushSize: "筆刷大小",
+2 -1
View File
@@ -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",
+2 -1
View File
@@ -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",
+38
View File
@@ -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);
});
});
@@ -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",
}),
);
@@ -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);
});
});
+102
View File
@@ -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");
});
});
+3 -2
View File
@@ -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();
+1 -1
View File
@@ -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", () => {
+49 -13
View File
@@ -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");
});
});