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:
@@ -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]:
|
||||
|
||||
Reference in New Issue
Block a user