diff --git a/README.md b/README.md index 661d955..55e71e3 100644 --- a/README.md +++ b/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 diff --git a/bgfilter/cli.py b/bgfilter/cli.py index 9d1ac93..4285232 100644 --- a/bgfilter/cli.py +++ b/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, diff --git a/bgfilter/config.py b/bgfilter/config.py index f395c49..c1bac61 100644 --- a/bgfilter/config.py +++ b/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 = { diff --git a/bgfilter/settings.py b/bgfilter/settings.py index 0274f32..9d3e501 100644 --- a/bgfilter/settings.py +++ b/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) diff --git a/bgfilter/vitmatte_infer.py b/bgfilter/vitmatte_infer.py index a12ee95..f282701 100644 --- a/bgfilter/vitmatte_infer.py +++ b/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) diff --git a/configs/default.yaml b/configs/default.yaml index 48fbe9a..8893748 100644 --- a/configs/default.yaml +++ b/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)