Add bf16 precision option for ViTMatte (--precision, default fp32)

bf16 halves the matting model's activation memory -- its full-resolution
attention is the pipeline's memory peak -- with visually identical alpha
(measured: 0 px alpha deviation > 0.25 on samples; cross-check veto
behaviour unchanged, region IoU 0.984).

Guarded by a hardware check so it never lands on a slow emulation path:
CUDA requires torch.cuda.is_bf16_supported(); CPU requires the same
oneDNN native-bf16 gate PyTorch uses for matmul routing (AVX512-BF16/
AMX). Without support it warns and falls back to fp32 -- on a Zen2 EPYC
the fallback kernels measured 17-370x slower than fp32, so silent bf16
there would be a performance landmine.

Segmentation stays fp32: torchvision deform_conv2d (used by BiRefNet)
has no bf16 CPU kernel, and the segmenter is not the memory peak.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

README.md

@@ -89,9 +89,11 @@ D:\MiniConda\envs\lightML\python.exe -m bgfilter.cli `
 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), `--seg-backend` (default
-`birefnet`; `anime-seg` for anime characters), `--screen-color` (default: auto-detect
-the flat background), and `--trimap-mode`.
+`--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`
+(default `birefnet`; `anime-seg` for anime characters), `--screen-color` (default:
+auto-detect the flat background), and `--trimap-mode`.
 
 ## Batch
 

bgfilter/cli.py

@@ -21,6 +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)"),
     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"),
@@ -38,6 +39,7 @@ def main(
             load_settings(config),
             model_name=model_name,
             device=device,
+            precision=precision,
             screen_color=screen_color,
             matting_method=matting_method,
             fallback_to_chroma_alpha=fallback_to_chroma_alpha,

bgfilter/config.py

@@ -84,7 +84,7 @@ def override_settings(settings: PipelineSettings, **overrides: Any) -> PipelineS
 
     model_updates = {
         key: overrides[key]
-        for key in ["model_name", "device", "matting_method", "fallback_to_chroma_alpha"]
+        for key in ["model_name", "device", "matting_method", "fallback_to_chroma_alpha", "precision"]
         if overrides.get(key) is not None
     }
     trimap_updates = {

bgfilter/settings.py

@@ -119,6 +119,12 @@ class ModelSettings:
     device: str = "cpu"
     matting_method: str = "vitmatte"
     fallback_to_chroma_alpha: bool = False
+    # ViTMatte compute precision: "fp32" or "bf16". bf16 halves the matting
+    # model's activation memory (its full-resolution attention is the pipeline's
+    # memory peak) with visually identical alpha, but needs bf16-capable
+    # hardware (any modern GPU, or a CPU with AVX512-BF16/AMX) — otherwise it
+    # falls back to fp32 with a warning. Segmentation always stays fp32.
+    precision: str = "fp32"
 
 
 @dataclass(frozen=True)

bgfilter/vitmatte_infer.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import warnings
+
 import numpy as np
 from PIL import Image
 
@@ -21,7 +23,8 @@ 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.model.to(self.device)
+        self.dtype = self._resolve_dtype(settings.precision)
+        self.model.to(device=self.device, dtype=self.dtype)
         self.model.eval()
 
     def _resolve_device(self, requested: str):
@@ -31,11 +34,38 @@ 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")
         inputs = self.processor(images=image, trimaps=trimap_image, return_tensors="pt")
-        inputs = {key: value.to(self.device) for key, value in inputs.items()}
+        inputs = {
+            key: value.to(self.device, dtype=self.dtype) if value.is_floating_point() else value.to(self.device)
+            for key, value in inputs.items()
+        }
         with self.torch.inference_mode():
             outputs = self.model(**inputs)
 

configs/default.yaml

@@ -8,6 +8,9 @@ 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.
+  precision: fp32
 
 segmentation:
   # enabled: true  = single-segmenter pipeline (one seg model drives the trimap)