mirror of
https://github.com/snapotter-hq/SnapOtter.git
synced 2026-08-03 07:46:42 +02:00
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:
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
});
|
||||
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -1558,6 +1558,9 @@ export const ar: TranslationKeys = {
|
||||
progressLabelBatch: "جاري تكبير {count} صورة",
|
||||
},
|
||||
"erase-object": {
|
||||
qualityFast: "سريع",
|
||||
qualityHq: "جودة عالية",
|
||||
qualityHint: "الجودة العالية تستخدم نموذج انتشار للأجسام الكبيرة والقوام التفصيلي.",
|
||||
brushMode: "فرشاة",
|
||||
lassoMode: "لاسو",
|
||||
brushSize: "حجم الفرشاة",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -1388,6 +1388,9 @@ export const hi: TranslationKeys = {
|
||||
progressLabelBatch: "{count} इमेज अपस्केल हो रही हैं",
|
||||
},
|
||||
"erase-object": {
|
||||
qualityFast: "तेज़",
|
||||
qualityHq: "उच्च गुणवत्ता",
|
||||
qualityHint: "उच्च गुणवत्ता बड़ी वस्तुओं और विस्तृत बनावट के लिए डिफ्यूज़न मॉडल का उपयोग करती है।",
|
||||
brushMode: "ब्रश",
|
||||
lassoMode: "लासो",
|
||||
brushSize: "ब्रश साइज़",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -1530,6 +1530,9 @@ export const ja: TranslationKeys = {
|
||||
progressLabelBatch: "{count}枚の画像をアップスケール中",
|
||||
},
|
||||
"erase-object": {
|
||||
qualityFast: "高速",
|
||||
qualityHq: "高品質",
|
||||
qualityHint: "高品質は大きなオブジェクトや細かいテクスチャに拡散モデルを使用します。",
|
||||
brushMode: "ブラシ",
|
||||
lassoMode: "投げ縄",
|
||||
brushSize: "ブラシサイズ",
|
||||
|
||||
@@ -1512,6 +1512,9 @@ export const ko: TranslationKeys = {
|
||||
progressLabelBatch: "{count}개 이미지 업스케일 중",
|
||||
},
|
||||
"erase-object": {
|
||||
qualityFast: "빠름",
|
||||
qualityHq: "고품질",
|
||||
qualityHint: "고품질은 큰 객체와 세밀한 텍스처에 디퓨전 모델을 사용합니다.",
|
||||
brushMode: "브러시",
|
||||
lassoMode: "올가미",
|
||||
brushSize: "브러시 크기",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -1571,6 +1571,10 @@ export const ru: TranslationKeys = {
|
||||
progressLabelBatch: "Увеличение {count} изображений",
|
||||
},
|
||||
"erase-object": {
|
||||
qualityFast: "Быстро",
|
||||
qualityHq: "Высокое качество",
|
||||
qualityHint:
|
||||
"Высокое качество использует диффузионную модель для больших объектов и детальных текстур.",
|
||||
brushMode: "Кисть",
|
||||
lassoMode: "Лассо",
|
||||
brushSize: "Размер кисти",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -1548,6 +1548,9 @@ export const th: TranslationKeys = {
|
||||
progressLabelBatch: "กำลังขยาย {count} ภาพ",
|
||||
},
|
||||
"erase-object": {
|
||||
qualityFast: "เร็ว",
|
||||
qualityHq: "คุณภาพสูง",
|
||||
qualityHint: "คุณภาพสูงใช้โมเดล diffusion สำหรับวัตถุขนาดใหญ่และพื้นผิวที่มีรายละเอียด",
|
||||
brushMode: "แปรง",
|
||||
lassoMode: "ลาสโซ",
|
||||
brushSize: "ขนาดแปรง",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -1571,6 +1571,10 @@ export const uk: TranslationKeys = {
|
||||
progressLabelBatch: "Збільшення {count} зображень",
|
||||
},
|
||||
"erase-object": {
|
||||
qualityFast: "Швидко",
|
||||
qualityHq: "Висока якість",
|
||||
qualityHint:
|
||||
"Висока якість використовує дифузійну модель для великих об'єктів і детальних текстур.",
|
||||
brushMode: "Пензель",
|
||||
lassoMode: "Ласо",
|
||||
brushSize: "Розмір пензля",
|
||||
|
||||
@@ -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ọ",
|
||||
|
||||
@@ -1336,6 +1336,9 @@ export const zhCN: TranslationKeys = {
|
||||
progressLabelBatch: "正在放大 {count} 张图片",
|
||||
},
|
||||
"erase-object": {
|
||||
qualityFast: "快速",
|
||||
qualityHq: "高质量",
|
||||
qualityHint: "高质量使用扩散模型处理大型物体和精细纹理。",
|
||||
brushMode: "画笔",
|
||||
lassoMode: "套索",
|
||||
brushSize: "画笔大小",
|
||||
|
||||
@@ -1336,6 +1336,9 @@ export const zhTW: TranslationKeys = {
|
||||
progressLabelBatch: "正在放大{count}張影像",
|
||||
},
|
||||
"erase-object": {
|
||||
qualityFast: "快速",
|
||||
qualityHq: "高品質",
|
||||
qualityHint: "高品質使用擴散模型處理大型物件與精細紋理。",
|
||||
brushMode: "筆刷",
|
||||
lassoMode: "套索",
|
||||
brushSize: "筆刷大小",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
@@ -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");
|
||||
});
|
||||
});
|
||||
@@ -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();
|
||||
|
||||
@@ -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", () => {
|
||||
|
||||
@@ -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");
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user