Merge origin/master: bf16 precision, chunked attention, seg-reuse, mimalloc purge

Adopts master as the standard for all overlapping work. Master's landed
optimizations supersede server-edition's own fp16 experiment:

- Unified --precision fp32|bf16 knob (bgfilter/precision.py) driving all
  three models: ViTMatte weight cast + both BiRefNets via autocast, with a
  hardware gate (falls back to fp32 off native-bf16 CPUs) and an
  AutocastCPU fp32 shim for torchvision deform_conv2d.
- Query-chunked ViTMatte global attention (bgfilter/attn_chunk.py), exact
  and bitwise-identical, caps the N^2 spike (~19 -> ~4 GB at 2048).
- Cross-check HR-matting forward reused as the seg mask (birefnet backend
  only), skipping the primary seg model; MIMALLOC_PURGE_DELAY=0.
- inference_mode and the detect_background_color removal converge with
  server-edition's earlier equivalents.

Conflict resolution (favoring master, preserving server-only features):
- vitmatte_infer/segmentation: dropped server's device-derived fp16 for
  master's precision path, kept resolve_model_source (local weights).
- service.py: cross-check SegmentationSettings now passes precision so the
  HTTP service honors bf16 like the CLI's _second_opinion does.

Verified on CPU: default fp32 pipeline loads 2 models (seg-reuse active)
and bf16 path runs (autocast + deform_conv2d shim) — both exit 0.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-07-06 16:02:44 +08:00
14 changed files with 316 additions and 24 deletions
+18 -9
View File
@@ -5,6 +5,7 @@ import os
import numpy as np
from PIL import Image
from .precision import ensure_deform_conv2d_autocast_cpu, resolve_dtype
from .settings import SegmentationSettings
from .weights import resolve_model_source
@@ -37,17 +38,20 @@ class BiRefNetSegmenter:
) from exc
self.torch = torch
self.device = self._resolve_device(settings.device)
# fp16 only pays off on CUDA; CPU lacks native fp16 kernels (~100x slower),
# so the dtype follows the device. The fp16 checkpoint upcasts on CPU.
self.dtype = torch.float16 if self.device.type == "cuda" else torch.float32
self.model = AutoModelForImageSegmentation.from_pretrained(
resolve_model_source(settings.model_name),
trust_remote_code=True,
dtype=self.dtype,
resolve_model_source(settings.model_name), trust_remote_code=True
)
self.model.eval()
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(
@@ -68,8 +72,13 @@ class BiRefNetSegmenter:
def mask(self, rgb: np.ndarray) -> np.ndarray:
"""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, dtype=self.dtype)
with self.torch.inference_mode():
tensor = self.transform(image).unsqueeze(0).to(self.device)
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]: