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

packages/ai/python/dispatcher.py

@@ -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",

packages/ai/python/inpaint_hq.py

@@ -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()

packages/ai/python/tests/test_inpaint_hq.py

@@ -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

packages/ai/src/feature-gate.ts

@@ -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",

packages/ai/src/index.ts

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

packages/ai/src/inpainting.ts

@@ -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,
   });