diff --git a/README.md b/README.md index 55e71e3..54b3bd4 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,13 @@ download (~0.9 GB) and one inference pass per image; disable with `--no-cross-check` (see the `cross_check` config section, and `docs/hair_gap_artifacts.md` for the analysis behind it). +When the cross-check is on, that same HR-matting forward is **reused as the +segmentation mask by default**, skipping the primary seg model entirely (one +less model to load, ~20 s faster per image on CPU). Pilot-validated +(TestImage3 / FixImage1): trimap 99.8% identical, no structural change to +fingers, hair wisps or thin lines. Disable with `--no-cross-check-as-seg` to +run the dedicated seg model instead. + The segmentation trimap defaults to `directional` mode (chroma + seg + a hue-direction split: it keeps a background-coloured garment such as a white shirt while dropping a background-hued residual such as blue trapped between hair strands). Switch with @@ -90,8 +97,9 @@ 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), `--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` +`fp32`; `bf16` speeds up all three models and 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`. diff --git a/bgfilter/__init__.py b/bgfilter/__init__.py index 8d43e99..e002b31 100644 --- a/bgfilter/__init__.py +++ b/bgfilter/__init__.py @@ -1,3 +1,12 @@ """Flat-background character matting pipeline.""" +import os + +# Windows PyTorch allocates CPU tensors through a bundled mimalloc whose lazy +# purge keeps ~10 GB of already-freed forward activations committed after a +# BiRefNet@2048 pass, stacking under ViTMatte's attention peak (measured +# 2048x2048 bf16: peak 25.2 -> 21.4 GB with purge 0, and slightly *faster*). +# Must be set before the torch DLL loads; harmless where mimalloc is not used. +os.environ.setdefault("MIMALLOC_PURGE_DELAY", "0") + __version__ = "0.1.0" diff --git a/bgfilter/cli.py b/bgfilter/cli.py index 4285232..bcf6ee9 100644 --- a/bgfilter/cli.py +++ b/bgfilter/cli.py @@ -21,7 +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)"), + precision: str | None = typer.Option(None, "--precision", help="Compute precision for all models (ViTMatte cast + BiRefNet autocast): fp32 (default) | bf16 (faster + 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"), @@ -31,6 +31,7 @@ def main( fg_safe_radius_ratio: float | None = typer.Option(None, "--fg-safe-radius-ratio", min=0.0), despill: bool | None = typer.Option(None, "--despill/--no-despill"), cross_check: bool | None = typer.Option(None, "--cross-check/--no-cross-check", help="Second-opinion veto of background-hued residue between hair strands (default: on; costs one extra model inference)"), + cross_check_as_seg: bool | None = typer.Option(None, "--cross-check-as-seg/--no-cross-check-as-seg", help="Reuse the cross-check forward as the segmentation mask, skipping the primary seg model (default: on; saves ~20s on CPU; needs cross-check on)"), trimap_mode: str | None = typer.Option(None, "--trimap-mode", help="Trimap mode (segmentation pipeline): directional | seg | directional-hard-bg"), seg_backend: str | None = typer.Option(None, "--seg-backend", help="Segmentation backend: birefnet (default) | anime-seg"), ) -> None: @@ -49,6 +50,7 @@ def main( fg_safe_radius_ratio=fg_safe_radius_ratio, despill=despill, cross_check=cross_check, + cross_check_as_seg=cross_check_as_seg, trimap_mode=trimap_mode, seg_backend=seg_backend, ) diff --git a/bgfilter/config.py b/bgfilter/config.py index c1bac61..688aa1d 100644 --- a/bgfilter/config.py +++ b/bgfilter/config.py @@ -105,9 +105,16 @@ def override_settings(settings: PipelineSettings, **overrides: Any) -> PipelineS cross_check_updates: dict[str, Any] = {} if overrides.get("cross_check") is not None: cross_check_updates["enabled"] = overrides["cross_check"] + if overrides.get("cross_check_as_seg") is not None: + cross_check_updates["reuse_as_seg"] = overrides["cross_check_as_seg"] seg_updates: dict[str, Any] = {} if overrides.get("device") is not None: seg_updates["device"] = overrides["device"] + # --precision is a single knob for all three models (ViTMatte handled via + # model_updates above; per-model values remain settable in the config file). + if overrides.get("precision") is not None: + seg_updates["precision"] = overrides["precision"] + cross_check_updates["precision"] = overrides["precision"] if overrides.get("seg_backend") is not None: backend = overrides["seg_backend"] if backend not in _SEG_MODELS: diff --git a/bgfilter/pipeline.py b/bgfilter/pipeline.py index 8a7ee5f..9599482 100644 --- a/bgfilter/pipeline.py +++ b/bgfilter/pipeline.py @@ -56,6 +56,7 @@ class MattingPipeline: model_name=check.model_name, device=self.settings.segmentation.device, input_size=check.input_size, + precision=check.precision, ) ) return self._cross_checker.mask(rgb) @@ -108,8 +109,15 @@ def _run_image( bg_confidence, model = compute_bg_confidence( rgb, settings=settings.chroma, screen_color=screen_color ) + second_alpha = None if settings.segmentation.enabled: - seg_mask = pipeline._segment(rgb) + if settings.cross_check.enabled and settings.cross_check.reuse_as_seg: + # One HR-matting forward serves both the trimap topology and the + # cross-check second opinion; the primary seg model is never loaded. + second_alpha = pipeline._second_opinion(rgb) + seg_mask = second_alpha + else: + seg_mask = pipeline._segment(rgb) mode = settings.trimap.mode if mode == "seg": trimap, trimap_stats = fuse_trimap(seg_mask, bg_confidence, settings.trimap) @@ -136,10 +144,10 @@ def _run_image( raw_alpha=raw_alpha if alpha_source == "vitmatte" else None, ) alpha = clean_alpha(alpha, trimap, settings.alpha_post) - second_alpha = None if settings.cross_check.enabled: lab = convert_color_spaces(rgb)[1] - second_alpha = pipeline._second_opinion(rgb) + if second_alpha is None: + second_alpha = pipeline._second_opinion(rgb) alpha = cross_check_alpha( alpha, second_alpha, diff --git a/bgfilter/precision.py b/bgfilter/precision.py new file mode 100644 index 0000000..942d83a --- /dev/null +++ b/bgfilter/precision.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +import warnings + + +def resolve_dtype(torch, device, precision: str): + """Map a precision string to a torch dtype, gated on hardware support. + + bf16 is only worth it on hardware with native kernels (any modern GPU, or a + CPU with AVX512-BF16/AMX). Without them PyTorch silently routes bf16 matmuls + onto a fallback that is orders of magnitude slower than fp32, so this gate + falls back to fp32 with a warning instead. The CPU check is the same one + PyTorch uses to route bf16 matmuls onto oneDNN's native kernels. + """ + if precision == "fp32": + return torch.float32 + if precision != "bf16": + raise RuntimeError(f"Unsupported precision '{precision}'. Use 'fp32' or 'bf16'.") + if device.type == "cuda": + if torch.cuda.is_bf16_supported(): + return torch.bfloat16 + reason = "this GPU has no bf16 support" + else: + 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 + + +# Keep the Library object alive: dropping it would un-register the impl. +_deform_conv2d_shim: object | None = None + + +def ensure_deform_conv2d_autocast_cpu(torch) -> None: + """Give torchvision's deform_conv2d an AutocastCPU wrapper (fp32 fallback). + + torchvision (<= 0.27 at least) ships no AutocastCPU registration and no + BFloat16 CPU kernel for ``torchvision::deform_conv2d``, so running a model + that uses it (BiRefNet) under ``torch.autocast("cpu", bf16)`` crashes with + "deformable_im2col not implemented for 'BFloat16'". This registers a + dispatcher-level wrapper that casts the op's inputs to fp32 and runs it + outside autocast — upstream/downstream convs keep their bf16 speed, only + this one op pays a cast. No-op if called twice or if a future torchvision + registers its own wrapper. + """ + global _deform_conv2d_shim + if _deform_conv2d_shim is not None: + return + import torchvision # noqa: F401 (registers torchvision::deform_conv2d) + + lib = torch.library.Library("torchvision", "IMPL", "AutocastCPU") + + def _deform_conv2d_fp32(input, weight, offset, mask, bias, + stride_h, stride_w, pad_h, pad_w, + dil_h, dil_w, n_weight_grps, n_offset_grps, use_mask): + with torch.autocast("cpu", enabled=False): + return torch.ops.torchvision.deform_conv2d( + input.float(), weight.float(), offset.float(), mask.float(), + bias.float(), stride_h, stride_w, pad_h, pad_w, + dil_h, dil_w, n_weight_grps, n_offset_grps, use_mask, + ) + + try: + lib.impl("deform_conv2d", _deform_conv2d_fp32) + except RuntimeError: + # torchvision grew its own AutocastCPU wrapper; ours is unneeded. + return + _deform_conv2d_shim = lib diff --git a/bgfilter/segmentation.py b/bgfilter/segmentation.py index 6c6f8b0..2de1564 100644 --- a/bgfilter/segmentation.py +++ b/bgfilter/segmentation.py @@ -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]: diff --git a/bgfilter/settings.py b/bgfilter/settings.py index 9d3e501..72063b6 100644 --- a/bgfilter/settings.py +++ b/bgfilter/settings.py @@ -74,6 +74,15 @@ class CrossCheckSettings: enabled: bool = True model_name: str = "ZhengPeng7/BiRefNet_HR-matting" input_size: int = 2048 + # Reuse the cross-check forward as the segmentation mask: one HR-matting + # inference serves both trimap topology and the second opinion, skipping the + # primary seg model entirely (saves its load + forward). Pilot-validated on + # TestImage3/FixImage1: trimap 99.8% identical, no structural change. Only + # takes effect when both cross_check and segmentation are enabled. + reuse_as_seg: bool = True + # Compute precision of the cross-check forward: "fp32" or "bf16" (autocast; + # needs bf16-capable hardware, else falls back to fp32 with a warning). + precision: str = "fp32" proj_min: float = 3.0 # bg-hue projection above which a pixel is suspect l_min: float = 45.0 # Lab lightness below which a pixel is exempt (dark hair) feather_sigma: float = 2.0 # Gaussian feather of the zone boundary, in px @@ -134,6 +143,10 @@ class SegmentationSettings: model_name: str = "ZhengPeng7/BiRefNet" device: str = "cpu" input_size: int = 1024 + # Compute precision of the BiRefNet forward: "fp32" or "bf16" (autocast; needs + # bf16-capable hardware, else falls back to fp32 with a warning). The ONNX + # anime-seg backend ignores this. + precision: str = "fp32" @dataclass(frozen=True) diff --git a/bgfilter/vitmatte_infer.py b/bgfilter/vitmatte_infer.py index f282701..70e2a20 100644 --- a/bgfilter/vitmatte_infer.py +++ b/bgfilter/vitmatte_infer.py @@ -1,10 +1,9 @@ from __future__ import annotations -import warnings - import numpy as np from PIL import Image +from .precision import resolve_dtype from .settings import ModelSettings @@ -23,7 +22,7 @@ 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.dtype = self._resolve_dtype(settings.precision) + self.dtype = resolve_dtype(self.torch, self.device, settings.precision) self.model.to(device=self.device, dtype=self.dtype) self.model.eval() @@ -34,30 +33,6 @@ 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") diff --git a/configs/default.yaml b/configs/default.yaml index 8893748..8f3b3eb 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -8,8 +8,10 @@ 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. + # "bf16" halves ViTMatte activation memory and is ~3x faster on native-bf16 + # CPUs (visually identical alpha); needs bf16-capable hardware (modern GPU / + # AVX512-BF16 CPU), else falls back to fp32. --precision bf16 sets bf16 here + # AND on the BiRefNet models below in one go. precision: fp32 segmentation: @@ -20,6 +22,9 @@ segmentation: model_name: ZhengPeng7/BiRefNet # matches backend; --seg-backend anime-seg auto-selects skytnt/anime-seg device: cpu input_size: 1024 + # BiRefNet forward precision: "bf16" runs via autocast (~1.4x on native-bf16 + # CPUs, visually identical), else falls back to fp32. anime-seg ignores this. + precision: fp32 trimap: sure_bg_threshold: 0.92 @@ -59,6 +64,15 @@ cross_check: enabled: true model_name: ZhengPeng7/BiRefNet_HR-matting input_size: 2048 + # Reuse the cross-check forward as the seg mask: one HR-matting inference + # drives both trimap topology and the veto, skipping the primary seg model + # (saves its ~9s forward + ~0.9 GB load; CPU pipeline ~66s -> ~45s). Pilot: + # trimap 99.8% identical, no structural change (fingers/wisps/bowstring OK). + # Disable with --no-cross-check-as-seg. No effect when cross_check is off. + reuse_as_seg: true + # Cross-check forward precision, same semantics as segmentation.precision + # (bf16 via autocast: 2048 pass ~35s -> ~26s on a native-bf16 CPU). + precision: fp32 proj_min: 3.0 # bg-hue projection above which a pixel is suspect l_min: 45.0 # Lab lightness below which a pixel is exempt (dark hair) feather_sigma: 2.0 # zone-boundary feather, px