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>
This commit is contained in:
2026-07-06 15:15:12 +08:00
parent dfc85566c7
commit 9a9084b3d6
10 changed files with 158 additions and 36 deletions
+14 -1
View File
@@ -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]: