diff --git a/README.md b/README.md index 54b3bd4..07c3e80 100644 --- a/README.md +++ b/README.md @@ -99,7 +99,9 @@ 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` 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` +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`. 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/settings.py b/bgfilter/settings.py index 72063b6..b5aaf23 100644 --- a/bgfilter/settings.py +++ b/bgfilter/settings.py @@ -134,6 +134,12 @@ class ModelSettings: # 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" + # 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) diff --git a/bgfilter/vitmatte_infer.py b/bgfilter/vitmatte_infer.py index 70e2a20..8b3265a 100644 --- a/bgfilter/vitmatte_infer.py +++ b/bgfilter/vitmatte_infer.py @@ -25,6 +25,10 @@ class ViTMatteRunner: self.dtype = resolve_dtype(self.torch, self.device, settings.precision) 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(): diff --git a/configs/default.yaml b/configs/default.yaml index 8f3b3eb..5cf00de 100644 --- a/configs/default.yaml +++ b/configs/default.yaml @@ -13,6 +13,10 @@ model: # 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)