Cut pipeline cost: seg-reuse, unified bf16 knob, eager mimalloc purge

Three optimizations from profiling the cross-check-dominated pipeline
(9700X CPU, all pilot-validated on TestImage3/FixImage1):

- Reuse the cross-check HR-matting@2048 forward as the segmentation mask
  (cross_check.reuse_as_seg, default ON; --no-cross-check-as-seg to opt
  out). Skips the BiRefNet@1024 load+forward entirely: ~66s -> ~45s,
  one less 0.9GB model. Trimap 99.8% identical, no structural change.

- --precision bf16 now fans out to all three models: ViTMatte keeps its
  weight cast; both BiRefNets run their forward under autocast with a
  dispatcher-level AutocastCPU fp32 shim for torchvision::deform_conv2d
  (no bf16 CPU kernel, no autocast wrapper upstream). Shared hardware
  gate in bgfilter/precision.py falls back to fp32 off native-bf16
  hardware. TestImage3: 51.9s -> 33.1s; alpha diff max 0.15, none >0.25.

- MIMALLOC_PURGE_DELAY=0 (bgfilter/__init__.py, before torch loads):
  Windows torch's bundled mimalloc lazily retains ~10GB of freed
  BiRefNet activations, stacking under ViTMatte's attention peak.
  2048x2048 bf16: peak 25.2 -> 21.4GB and slightly faster (73 -> 63s).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

README.md

@@ -38,6 +38,13 @@ download (~0.9 GB) and one inference pass per image; disable with
 `--no-cross-check` (see the `cross_check` config section, and
 `docs/hair_gap_artifacts.md` for the analysis behind it).
 
+When the cross-check is on, that same HR-matting forward is **reused as the
+segmentation mask by default**, skipping the primary seg model entirely (one
+less model to load, ~20 s faster per image on CPU). Pilot-validated
+(TestImage3 / FixImage1): trimap 99.8% identical, no structural change to
+fingers, hair wisps or thin lines. Disable with `--no-cross-check-as-seg` to
+run the dedicated seg model instead.
+
 The segmentation trimap defaults to `directional` mode (chroma + seg + a hue-direction
 split: it keeps a background-coloured garment such as a white shirt while dropping a
 background-hued residual such as blue trapped between hair strands). Switch with
@@ -90,8 +97,9 @@ No `--config` is needed — the built-in defaults are identical to
 `configs/default.yaml`. Pass `--config configs\default.yaml` only after you edit that
 file to tune the detailed parameters. Runtime choices stay on the command line:
 `--device` (default **CPU**; use `--device cuda` for GPU), `--precision` (default
-`fp32`; `bf16` halves the matting model's activation memory with visually identical
-alpha — needs bf16-capable hardware, falls back to fp32 elsewhere), `--seg-backend`
+`fp32`; `bf16` speeds up all three models and halves the matting model's activation
+memory with visually identical alpha — needs bf16-capable hardware, falls back to
+fp32 elsewhere), `--seg-backend`
 (default `birefnet`; `anime-seg` for anime characters), `--screen-color` (default:
 auto-detect the flat background), and `--trimap-mode`.
 

bgfilter/__init__.py

@@ -1,3 +1,12 @@
 """Flat-background character matting pipeline."""
 
+import os
+
+# Windows PyTorch allocates CPU tensors through a bundled mimalloc whose lazy
+# purge keeps ~10 GB of already-freed forward activations committed after a
+# BiRefNet@2048 pass, stacking under ViTMatte's attention peak (measured
+# 2048x2048 bf16: peak 25.2 -> 21.4 GB with purge 0, and slightly *faster*).
+# Must be set before the torch DLL loads; harmless where mimalloc is not used.
+os.environ.setdefault("MIMALLOC_PURGE_DELAY", "0")
+
 __version__ = "0.1.0"

bgfilter/cli.py

@@ -21,7 +21,7 @@ def main(
     config: Path | None = typer.Option(None, "--config", exists=True, file_okay=True, dir_okay=False),
     model_name: str | None = typer.Option(None, "--model-name"),
     device: str | None = typer.Option(None, "--device"),
-    precision: str | None = typer.Option(None, "--precision", help="ViTMatte precision: fp32 (default) | bf16 (halves matting activation memory; needs bf16-capable hardware, else falls back to fp32)"),
+    precision: str | None = typer.Option(None, "--precision", help="Compute precision for all models (ViTMatte cast + BiRefNet autocast): fp32 (default) | bf16 (faster + halves matting activation memory; needs bf16-capable hardware, else falls back to fp32)"),
     screen_color: str | None = typer.Option(None, "--screen-color", help="Background colour prior as #RRGGBB (default: auto-detect the flat background colour)"),
     matting_method: str | None = typer.Option(None, "--matting-method"),
     fallback_to_chroma_alpha: bool | None = typer.Option(None, "--fallback-to-chroma-alpha/--no-fallback-to-chroma-alpha"),
@@ -31,6 +31,7 @@ def main(
     fg_safe_radius_ratio: float | None = typer.Option(None, "--fg-safe-radius-ratio", min=0.0),
     despill: bool | None = typer.Option(None, "--despill/--no-despill"),
     cross_check: bool | None = typer.Option(None, "--cross-check/--no-cross-check", help="Second-opinion veto of background-hued residue between hair strands (default: on; costs one extra model inference)"),
+    cross_check_as_seg: bool | None = typer.Option(None, "--cross-check-as-seg/--no-cross-check-as-seg", help="Reuse the cross-check forward as the segmentation mask, skipping the primary seg model (default: on; saves ~20s on CPU; needs cross-check on)"),
     trimap_mode: str | None = typer.Option(None, "--trimap-mode", help="Trimap mode (segmentation pipeline): directional | seg | directional-hard-bg"),
     seg_backend: str | None = typer.Option(None, "--seg-backend", help="Segmentation backend: birefnet (default) | anime-seg"),
 ) -> None:
@@ -49,6 +50,7 @@ def main(
             fg_safe_radius_ratio=fg_safe_radius_ratio,
             despill=despill,
             cross_check=cross_check,
+            cross_check_as_seg=cross_check_as_seg,
             trimap_mode=trimap_mode,
             seg_backend=seg_backend,
         )

bgfilter/config.py

@@ -105,9 +105,16 @@ def override_settings(settings: PipelineSettings, **overrides: Any) -> PipelineS
     cross_check_updates: dict[str, Any] = {}
     if overrides.get("cross_check") is not None:
         cross_check_updates["enabled"] = overrides["cross_check"]
+    if overrides.get("cross_check_as_seg") is not None:
+        cross_check_updates["reuse_as_seg"] = overrides["cross_check_as_seg"]
     seg_updates: dict[str, Any] = {}
     if overrides.get("device") is not None:
         seg_updates["device"] = overrides["device"]
+    # --precision is a single knob for all three models (ViTMatte handled via
+    # model_updates above; per-model values remain settable in the config file).
+    if overrides.get("precision") is not None:
+        seg_updates["precision"] = overrides["precision"]
+        cross_check_updates["precision"] = overrides["precision"]
     if overrides.get("seg_backend") is not None:
         backend = overrides["seg_backend"]
         if backend not in _SEG_MODELS:

bgfilter/pipeline.py

@@ -56,6 +56,7 @@ class MattingPipeline:
                     model_name=check.model_name,
                     device=self.settings.segmentation.device,
                     input_size=check.input_size,
+                    precision=check.precision,
                 )
             )
         return self._cross_checker.mask(rgb)
@@ -108,8 +109,15 @@ def _run_image(
     bg_confidence, model = compute_bg_confidence(
         rgb, settings=settings.chroma, screen_color=screen_color
     )
+    second_alpha = None
     if settings.segmentation.enabled:
-        seg_mask = pipeline._segment(rgb)
+        if settings.cross_check.enabled and settings.cross_check.reuse_as_seg:
+            # One HR-matting forward serves both the trimap topology and the
+            # cross-check second opinion; the primary seg model is never loaded.
+            second_alpha = pipeline._second_opinion(rgb)
+            seg_mask = second_alpha
+        else:
+            seg_mask = pipeline._segment(rgb)
         mode = settings.trimap.mode
         if mode == "seg":
             trimap, trimap_stats = fuse_trimap(seg_mask, bg_confidence, settings.trimap)
@@ -136,10 +144,10 @@ def _run_image(
         raw_alpha=raw_alpha if alpha_source == "vitmatte" else None,
     )
     alpha = clean_alpha(alpha, trimap, settings.alpha_post)
-    second_alpha = None
     if settings.cross_check.enabled:
         lab = convert_color_spaces(rgb)[1]
-        second_alpha = pipeline._second_opinion(rgb)
+        if second_alpha is None:
+            second_alpha = pipeline._second_opinion(rgb)
         alpha = cross_check_alpha(
             alpha,
             second_alpha,

bgfilter/precision.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import warnings
+
+
+def resolve_dtype(torch, device, precision: str):
+    """Map a precision string to a torch dtype, gated on hardware support.
+
+    bf16 is only worth it on hardware with native kernels (any modern GPU, or a
+    CPU with AVX512-BF16/AMX). Without them PyTorch silently routes bf16 matmuls
+    onto a fallback that is orders of magnitude slower than fp32, so this gate
+    falls back to fp32 with a warning instead. The CPU check is the same one
+    PyTorch uses to route bf16 matmuls onto oneDNN's native kernels.
+    """
+    if precision == "fp32":
+        return torch.float32
+    if precision != "bf16":
+        raise RuntimeError(f"Unsupported precision '{precision}'. Use 'fp32' or 'bf16'.")
+    if device.type == "cuda":
+        if torch.cuda.is_bf16_supported():
+            return torch.bfloat16
+        reason = "this GPU has no bf16 support"
+    else:
+        try:
+            native = torch.ops.mkldnn._is_mkldnn_bf16_supported()
+        except (AttributeError, RuntimeError):
+            native = False
+        if native:
+            return torch.bfloat16
+        reason = "this CPU has no native bf16 support (needs AVX512-BF16/AMX)"
+    warnings.warn(f"precision 'bf16' requested but {reason}; running fp32 instead.")
+    return torch.float32
+
+
+# Keep the Library object alive: dropping it would un-register the impl.
+_deform_conv2d_shim: object | None = None
+
+
+def ensure_deform_conv2d_autocast_cpu(torch) -> None:
+    """Give torchvision's deform_conv2d an AutocastCPU wrapper (fp32 fallback).
+
+    torchvision (<= 0.27 at least) ships no AutocastCPU registration and no
+    BFloat16 CPU kernel for ``torchvision::deform_conv2d``, so running a model
+    that uses it (BiRefNet) under ``torch.autocast("cpu", bf16)`` crashes with
+    "deformable_im2col not implemented for 'BFloat16'". This registers a
+    dispatcher-level wrapper that casts the op's inputs to fp32 and runs it
+    outside autocast — upstream/downstream convs keep their bf16 speed, only
+    this one op pays a cast. No-op if called twice or if a future torchvision
+    registers its own wrapper.
+    """
+    global _deform_conv2d_shim
+    if _deform_conv2d_shim is not None:
+        return
+    import torchvision  # noqa: F401  (registers torchvision::deform_conv2d)
+
+    lib = torch.library.Library("torchvision", "IMPL", "AutocastCPU")
+
+    def _deform_conv2d_fp32(input, weight, offset, mask, bias,
+                            stride_h, stride_w, pad_h, pad_w,
+                            dil_h, dil_w, n_weight_grps, n_offset_grps, use_mask):
+        with torch.autocast("cpu", enabled=False):
+            return torch.ops.torchvision.deform_conv2d(
+                input.float(), weight.float(), offset.float(), mask.float(),
+                bias.float(), stride_h, stride_w, pad_h, pad_w,
+                dil_h, dil_w, n_weight_grps, n_offset_grps, use_mask,
+            )
+
+    try:
+        lib.impl("deform_conv2d", _deform_conv2d_fp32)
+    except RuntimeError:
+        # torchvision grew its own AutocastCPU wrapper; ours is unneeded.
+        return
+    _deform_conv2d_shim = lib

bgfilter/segmentation.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import numpy as np
 from PIL import Image
 
+from .precision import ensure_deform_conv2d_autocast_cpu, resolve_dtype
 from .settings import SegmentationSettings
 
 
@@ -41,6 +42,13 @@ class BiRefNetSegmenter:
         self.model.float()  # checkpoint ships as fp16; force fp32 to match inputs
         self.device = self._resolve_device(settings.device)
         self.model.to(self.device)
+        # bf16 runs via autocast (weights stay fp32): conv/matmul take the fast
+        # bf16 kernels while ops without them keep fp32. deform_conv2d has no
+        # bf16 CPU kernel and no AutocastCPU wrapper in torchvision, so it needs
+        # a dispatcher-level fp32 shim or the forward crashes.
+        self.dtype = resolve_dtype(torch, self.device, settings.precision)
+        if self.dtype == torch.bfloat16 and self.device.type == "cpu":
+            ensure_deform_conv2d_autocast_cpu(torch)
 
         size = settings.input_size
         self.transform = transforms.Compose(
@@ -62,7 +70,12 @@ class BiRefNetSegmenter:
         """Return a soft foreground mask, float32 0..1, at the input resolution."""
         image = Image.fromarray(rgb.astype(np.uint8), mode="RGB")
         tensor = self.transform(image).unsqueeze(0).to(self.device)
-        with self.torch.inference_mode():
+        autocast = self.torch.autocast(
+            self.device.type,
+            dtype=self.torch.bfloat16,
+            enabled=self.dtype == self.torch.bfloat16,
+        )
+        with self.torch.inference_mode(), autocast:
             pred = self.model(tensor)[-1].sigmoid()
         pred = pred[0, 0].detach().float().cpu().numpy()
         if pred.shape != rgb.shape[:2]:

bgfilter/settings.py

@@ -74,6 +74,15 @@ class CrossCheckSettings:
     enabled: bool = True
     model_name: str = "ZhengPeng7/BiRefNet_HR-matting"
     input_size: int = 2048
+    # Reuse the cross-check forward as the segmentation mask: one HR-matting
+    # inference serves both trimap topology and the second opinion, skipping the
+    # primary seg model entirely (saves its load + forward). Pilot-validated on
+    # TestImage3/FixImage1: trimap 99.8% identical, no structural change. Only
+    # takes effect when both cross_check and segmentation are enabled.
+    reuse_as_seg: bool = True
+    # Compute precision of the cross-check forward: "fp32" or "bf16" (autocast;
+    # needs bf16-capable hardware, else falls back to fp32 with a warning).
+    precision: str = "fp32"
     proj_min: float = 3.0       # bg-hue projection above which a pixel is suspect
     l_min: float = 45.0         # Lab lightness below which a pixel is exempt (dark hair)
     feather_sigma: float = 2.0  # Gaussian feather of the zone boundary, in px
@@ -134,6 +143,10 @@ class SegmentationSettings:
     model_name: str = "ZhengPeng7/BiRefNet"
     device: str = "cpu"
     input_size: int = 1024
+    # Compute precision of the BiRefNet forward: "fp32" or "bf16" (autocast; needs
+    # bf16-capable hardware, else falls back to fp32 with a warning). The ONNX
+    # anime-seg backend ignores this.
+    precision: str = "fp32"
 
 
 @dataclass(frozen=True)

bgfilter/vitmatte_infer.py

@@ -1,10 +1,9 @@
 from __future__ import annotations
 
-import warnings
-
 import numpy as np
 from PIL import Image
 
+from .precision import resolve_dtype
 from .settings import ModelSettings
 
 
@@ -23,7 +22,7 @@ class ViTMatteRunner:
         self.processor = VitMatteImageProcessor.from_pretrained(settings.model_name)
         self.model = VitMatteForImageMatting.from_pretrained(settings.model_name)
         self.device = self._resolve_device(settings.device)
-        self.dtype = self._resolve_dtype(settings.precision)
+        self.dtype = resolve_dtype(self.torch, self.device, settings.precision)
         self.model.to(device=self.device, dtype=self.dtype)
         self.model.eval()
 
@@ -34,30 +33,6 @@ class ViTMatteRunner:
             )
         return self.torch.device(requested)
 
-    def _resolve_dtype(self, precision: str):
-        torch = self.torch
-        if precision == "fp32":
-            return torch.float32
-        if precision != "bf16":
-            raise RuntimeError(f"Unsupported precision '{precision}'. Use 'fp32' or 'bf16'.")
-        if self.device.type == "cuda":
-            if torch.cuda.is_bf16_supported():
-                return torch.bfloat16
-            reason = "this GPU has no bf16 support"
-        else:
-            # Same gate PyTorch uses to route bf16 matmuls onto oneDNN's native
-            # kernels; without the hardware (AVX512-BF16/AMX) they hit a fallback
-            # that is orders of magnitude slower than fp32.
-            try:
-                native = torch.ops.mkldnn._is_mkldnn_bf16_supported()
-            except (AttributeError, RuntimeError):
-                native = False
-            if native:
-                return torch.bfloat16
-            reason = "this CPU has no native bf16 support (needs AVX512-BF16/AMX)"
-        warnings.warn(f"precision 'bf16' requested but {reason}; running fp32 instead.")
-        return torch.float32
-
     def predict(self, rgb: np.ndarray, trimap: np.ndarray) -> np.ndarray:
         image = Image.fromarray(rgb.astype(np.uint8), mode="RGB")
         trimap_image = Image.fromarray(trimap.astype(np.uint8), mode="L")

configs/default.yaml

@@ -8,8 +8,10 @@ model:
   device: cpu
   matting_method: vitmatte
   fallback_to_chroma_alpha: false
-  # "bf16" halves ViTMatte activation memory (visually identical alpha); needs
-  # bf16-capable hardware (modern GPU / AVX512-BF16 CPU), else falls back to fp32.
+  # "bf16" halves ViTMatte activation memory and is ~3x faster on native-bf16
+  # CPUs (visually identical alpha); needs bf16-capable hardware (modern GPU /
+  # AVX512-BF16 CPU), else falls back to fp32. --precision bf16 sets bf16 here
+  # AND on the BiRefNet models below in one go.
   precision: fp32
 
 segmentation:
@@ -20,6 +22,9 @@ segmentation:
   model_name: ZhengPeng7/BiRefNet  # matches backend; --seg-backend anime-seg auto-selects skytnt/anime-seg
   device: cpu
   input_size: 1024
+  # BiRefNet forward precision: "bf16" runs via autocast (~1.4x on native-bf16
+  # CPUs, visually identical), else falls back to fp32. anime-seg ignores this.
+  precision: fp32
 
 trimap:
   sure_bg_threshold: 0.92
@@ -59,6 +64,15 @@ cross_check:
   enabled: true
   model_name: ZhengPeng7/BiRefNet_HR-matting
   input_size: 2048
+  # Reuse the cross-check forward as the seg mask: one HR-matting inference
+  # drives both trimap topology and the veto, skipping the primary seg model
+  # (saves its ~9s forward + ~0.9 GB load; CPU pipeline ~66s -> ~45s). Pilot:
+  # trimap 99.8% identical, no structural change (fingers/wisps/bowstring OK).
+  # Disable with --no-cross-check-as-seg. No effect when cross_check is off.
+  reuse_as_seg: true
+  # Cross-check forward precision, same semantics as segmentation.precision
+  # (bf16 via autocast: 2048 pass ~35s -> ~26s on a native-bf16 CPU).
+  precision: fp32
   proj_min: 3.0        # bg-hue projection above which a pixel is suspect
   l_min: 45.0          # Lab lightness below which a pixel is exempt (dark hair)
   feather_sigma: 2.0   # zone-boundary feather, px