diff --git a/README.md b/README.md index 842badb..4d01f9f 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,16 @@ 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 reuse is a same-family swap, so it +applies to the `birefnet` backend only: with `--seg-backend anime-seg` the +anime segmenter keeps its own forward and the cross-check veto still runs +independently on top. + 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 @@ -89,9 +99,14 @@ 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` 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; large inputs additionally get ViTMatte's global attention computed +in query chunks by default — exact, bitwise-identical, caps the memory spike at +~4 GB instead of ~19 GB at 2048x2048, see `model.attn_query_chunk`), `--seg-backend` +(default `birefnet`; `anime-seg` for anime characters), `--screen-color` (default: +auto-detect the flat background), and `--trimap-mode`. ## Batch @@ -102,6 +117,26 @@ D:\MiniConda\envs\lightML\python.exe -m bgfilter.cli ` --debug-dir Outputs\debug ``` +## Performance (CPU reference numbers) + +Measured on a Ryzen 9700X (Zen 5, native bf16), 32 GB RAM, default settings +(cross-check on, reused as seg; chunked global attention on): + +| input | precision | warm / image | peak memory | +|---|---|---|---| +| 1024x1536 | fp32 | ~45 s | — | +| 1024x1536 | bf16 | ~33 s | — | +| 2048x2048 | bf16 | ~53 s | 11.4 GB (batch), 8.1 GB (single image) | + +The memory ceiling is the cross-check BiRefNet_HR forward (fixed +`input_size` 2048 regardless of the image size) — genuine live activations +of a full-resolution dense prediction net. Before chunked attention and the +mimalloc fix, the same 2048x2048 bf16 run peaked at 25.2 GB; a fully +un-optimized fp32 run would need an estimated 55-60 GB (ViTMatte's un-chunked +N^2 attention alone ~39 GB). CPUs without native bf16 (e.g. Zen 2) auto-fall +back to fp32 — on such machines prefer `--no-cross-check` if memory or +latency is tight. + ## Chroma-alpha debug mode `--matting-method chroma` skips ViTMatte and uses chroma confidence directly as the 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/attn_chunk.py b/bgfilter/attn_chunk.py new file mode 100644 index 0000000..1f796df --- /dev/null +++ b/bgfilter/attn_chunk.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +"""Query-chunked attention for ViTMatte's VitDet backbone. + +ViTMatte's backbone runs 4 *global* attention blocks over every patch token. +At 2048x2048 input that is 16384 tokens, and the stock transformers forward +materializes the full [heads x N x N] attention map (plus a second copy when +the decomposed relative-position bias is added, plus softmax) -- a ~19 GB +transient that is the pipeline's memory peak. transformers has no SDPA path +for this architecture (the decomposed bias must be added to the raw scores). + +Computing the same attention in query-row chunks is mathematically exact: +each output row depends on its own scores row only, and the decomposed bias +factorizes over query rows (``rel_h`` is indexed by query row, ``rel_w`` by +query column), so both slice cleanly. Peak drops from O(N^2) to O(N * chunk). +""" + + +def patch_vitdet_attention(model, chunk: int) -> int: + """Replace VitDetAttention.forward with a query-chunked equivalent. + + ``chunk`` is the number of query tokens per block; sequences no longer + than ``chunk`` (e.g. the 14x14 window blocks) keep the original one-shot + path. Returns the number of attention modules patched. + """ + import types + + import torch + from transformers.models.vitdet.modeling_vitdet import VitDetAttention, get_rel_pos + + def chunked_forward(self, hidden_state, output_attentions=False): + batch_size, height, width, _ = hidden_state.shape + tokens = height * width + if output_attentions or tokens <= chunk: + return VitDetAttention.forward(self, hidden_state, output_attentions) + + # Identical prologue to the stock forward. + qkv = ( + self.qkv(hidden_state) + .reshape(batch_size, tokens, 3, self.num_heads, -1) + .permute(2, 0, 3, 1, 4) + ) + queries, keys, values = qkv.reshape(3, batch_size * self.num_heads, tokens, -1).unbind(0) + + if self.use_relative_position_embeddings: + rel_h = get_rel_pos(height, height, self.rel_pos_h) + rel_w = get_rel_pos(width, width, self.rel_pos_w) + r_q = queries.reshape(batch_size * self.num_heads, height, width, -1) + + keys_t = keys.transpose(-2, -1) + rows = max(1, chunk // width) + out = torch.empty_like(queries) + for h0 in range(0, height, rows): + h1 = min(h0 + rows, height) + q_chunk = queries[:, h0 * width : h1 * width, :] + scores = (q_chunk * self.scale) @ keys_t + if self.use_relative_position_embeddings: + rq = r_q[:, h0:h1] + rel_height = torch.einsum("bhwc,hkc->bhwk", rq, rel_h[h0:h1]) + rel_width = torch.einsum("bhwc,wkc->bhwk", rq, rel_w) + scores = ( + scores.view(batch_size * self.num_heads, h1 - h0, width, height, width) + + rel_height[:, :, :, :, None] + + rel_width[:, :, :, None, :] + ).view(batch_size * self.num_heads, (h1 - h0) * width, tokens) + out[:, h0 * width : h1 * width, :] = scores.softmax(dim=-1) @ values + + hidden = out.view(batch_size, self.num_heads, height, width, -1) + hidden = hidden.permute(0, 2, 3, 1, 4).reshape(batch_size, height, width, -1) + return (self.proj(hidden),) + + count = 0 + for module in model.modules(): + if isinstance(module, VitDetAttention): + module.forward = types.MethodType(chunked_forward, module) + count += 1 + return count diff --git a/bgfilter/cli.py b/bgfilter/cli.py index 9d1ac93..bcf6ee9 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="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"), @@ -30,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: @@ -38,6 +40,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, @@ -47,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 f395c49..688aa1d 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 = { @@ -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 ce27d01..f1906aa 100644 --- a/bgfilter/pipeline.py +++ b/bgfilter/pipeline.py @@ -85,6 +85,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) @@ -135,8 +136,24 @@ def _process_rgb(rgb: np.ndarray, pipeline: MattingPipeline) -> MattingResult: 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) + reuse = ( + settings.cross_check.enabled + and settings.cross_check.reuse_as_seg + # Reuse swaps one BiRefNet-family model for another (validated + # equivalent topology). A different backend chosen explicitly + # (anime-seg) keeps its own forward; cross-check still runs + # independently on top of it. + and settings.segmentation.backend == "birefnet" + ) + if reuse: + # 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) @@ -163,10 +180,10 @@ def _process_rgb(rgb: np.ndarray, pipeline: MattingPipeline) -> MattingResult: 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 dc54240..1d8b58d 100644 --- a/bgfilter/segmentation.py +++ b/bgfilter/segmentation.py @@ -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]: diff --git a/bgfilter/service.py b/bgfilter/service.py index 532d28f..20c5364 100644 --- a/bgfilter/service.py +++ b/bgfilter/service.py @@ -105,6 +105,7 @@ class PipelineManager: model_name=check.model_name, device=self._base.segmentation.device, input_size=check.input_size, + precision=check.precision, ) ) return self._cross_checker diff --git a/bgfilter/settings.py b/bgfilter/settings.py index 0274f32..ea72a88 100644 --- a/bgfilter/settings.py +++ b/bgfilter/settings.py @@ -74,6 +74,17 @@ 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 + # applies to the birefnet segmentation backend (a same-family swap); with + # anime-seg the chosen segmenter keeps its own forward and the cross-check + # runs independently. Needs cross_check and segmentation both 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 @@ -119,6 +130,18 @@ 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 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. The BiRefNet models have + # their own precision fields; --precision sets all of them at once. + precision: str = "fp32" + # Query-chunked global attention (exact math, bitwise-identical output): + # the VitDet backbone's 4 global blocks materialize an N^2 attention map — + # ~19 GB at 2048x2048. Blocks seeing more than this many tokens compute it + # in query chunks of this size instead, capping the transient at + # O(N * chunk): 19.8 -> 4.0 GB for ~15% more ViTMatte time. 0 disables. + attn_query_chunk: int = 2048 @dataclass(frozen=True) @@ -128,6 +151,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 48a5004..2982a23 100644 --- a/bgfilter/vitmatte_infer.py +++ b/bgfilter/vitmatte_infer.py @@ -3,6 +3,7 @@ from __future__ import annotations import numpy as np from PIL import Image +from .precision import resolve_dtype from .settings import ModelSettings from .weights import resolve_model_source @@ -20,14 +21,16 @@ class ViTMatteRunner: self.torch = torch self.device = self._resolve_device(settings.device) - # fp16 only pays off on CUDA (tensor cores); CPU has no native fp16 conv/gemm - # kernels and degrades ~100x, so the dtype follows the device, never a param. - self.dtype = torch.float16 if self.device.type == "cuda" else torch.float32 + self.dtype = resolve_dtype(self.torch, self.device, settings.precision) source = resolve_model_source(settings.model_name) self.processor = VitMatteImageProcessor.from_pretrained(source) - self.model = VitMatteForImageMatting.from_pretrained(source, dtype=self.dtype) - self.model.to(self.device) + self.model = VitMatteForImageMatting.from_pretrained(source) + self.model.to(device=self.device, dtype=self.dtype) self.model.eval() + if settings.attn_query_chunk > 0: + from .attn_chunk import patch_vitdet_attention + + patch_vitdet_attention(self.model, settings.attn_query_chunk) def _resolve_device(self, requested: str): if requested == "cuda" and not self.torch.cuda.is_available(): @@ -41,9 +44,7 @@ class ViTMatteRunner: 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, dtype=self.dtype) - if value.is_floating_point() - else value.to(self.device) + 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(): diff --git a/configs/default.yaml b/configs/default.yaml index 48fbe9a..69a71e0 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -8,6 +8,15 @@ model: device: cpu matting_method: vitmatte fallback_to_chroma_alpha: false + # "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 + # Query-chunked global attention (exact, bitwise-identical output): caps the + # VitDet global blocks' N^2 attention transient at O(N * chunk) — 19.8 -> 4.0 + # GB at 2048x2048 for ~15% more ViTMatte time. 0 = stock one-shot attention. + attn_query_chunk: 2048 segmentation: # enabled: true = single-segmenter pipeline (one seg model drives the trimap) @@ -17,6 +26,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 @@ -56,6 +68,17 @@ 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). + # Only applies to the birefnet backend (same-family swap); anime-seg keeps + # its own forward, with the cross-check running independently on top. + # 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 diff --git a/docs/green_screen_matting_workflow.md b/docs/green_screen_matting_workflow.md index 5278cda..7cda5c7 100644 --- a/docs/green_screen_matting_workflow.md +++ b/docs/green_screen_matting_workflow.md @@ -11,6 +11,7 @@ - **两条管线**（`segmentation.enabled` 切换）：无分割模型的 Chroma + ViTMatte；或单个语义分割模型（默认 birefnet，可切 anime-seg，用 `--seg-backend`）定 trimap 拓扑，ViTMatte 细化。 - **背景色泛化**：`screen_color` 先验（如 `#CFEFFF`）；`null` 时自动从图像边框探测平背景色（探测不到干净平背景会报错），不再写死绿幕。 - **已移除**：染绿修复（recolor）、发丝补救（dropout-fill）、BiRefNet+anime-seg 结合方案——它们只是给“生图即染色/断裂”的坏素材擦屁股；源头改用干净背景生图后不再需要。 +- **新增（2026-07）**：交叉模型否决（`cross_check`，BiRefNet_HR-matting 第二意见 min 融合，清发丝缝隙残蓝，见 `docs/hair_gap_artifacts.md`）；其前向默认兼任 seg mask（`reuse_as_seg`，净模型数回到 2）；`--precision bf16` 一个旋钮加速全部三类模型；ViTMatte 全局 attention 默认按 query 分块（输出逐位不变，2048² 内存峰值 25→11 GB）。 下文其余章节描述绿幕基线流程，仍然有效。 diff --git a/docs/hair_gap_artifacts.md b/docs/hair_gap_artifacts.md index 50a2dce..1f85511 100644 --- a/docs/hair_gap_artifacts.md +++ b/docs/hair_gap_artifacts.md @@ -89,6 +89,14 @@ min 融合;B2 的意见锐化(remap 0.5→0.95)被否——数字更低(0.673 vs **已落地(2026-07-04,默认开启)。** 新增 `cross_check` 设置块 + `--cross-check/--no-cross-check`;pipeline 在 clean_alpha 后插入 `cross_check_alpha`(alpha_post.py),第二意见加载复用 BiRefNetSegmenter。 + +**成本已消化(2026-07-06)。** 当年"+1 个模型 +1 次推理"的账已改写: +`cross_check.reuse_as_seg`(默认开)让 HR-matting 前向兼任 seg mask, +**净模型数回到 2 个**(trimap 99.8% 不变,先导 Outputs/segreuse/); +`--precision bf16` 经 autocast + deform_conv2d shim 覆盖 BiRefNet +(35s→26s @2048);ViTMatte 全局 attention 分块(逐位等价)把内存峰值 +从 25.2GB 压到 11.4GB(2048×2048 bf16,9700X)。CPU 上开 cross-check +的常规图约 33s/张(bf16)。 验证:开启时 TestImage3 输出与评审原型逐字节一致;`--no-cross-check` 与 旧 baseline 逐字节一致;FixImage1(粉底)脸部无损、发缘粉残留被清, 唯一副作用是羽化权重外溢到暗豁免区边缘(弓弦零星像素最多 -0.5,中位