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