lhk229
c748fc22a4
The single asyncio lock serialized the whole request (chroma -> seg ->
trimap -> ViTMatte -> cross-check -> foreground -> despill), so the GPU
sat idle ~70-80% of each request while the CPU stages ran. Replace it
with a thread pool that runs several requests concurrently, and move the
mutex down to just the GPU forwards: when one request steps off the GPU
to do CPU work, another steps on, keeping the device fed.
Two opt-in env knobs, both defaulting to the old strictly-serial
behaviour:
* BGFILTER_CPU_WORKERS (pool size) attacks the CPU-bound regime -- with
cross-check off, throughput went 0.91 -> 1.67 req/s (1.84x) at pool=8
on a 4080S, vs a mere +6% from a second worker process, at half the
VRAM (one model copy).
* BGFILTER_GPU_CONCURRENCY caps how many forwards may co-schedule on the
SMs. batch-1 underfills the GPU, so =2 recovers the SM overlap a lock
forgoes: cross-check on went 0.83 -> 1.01 req/s, beating two worker
processes (0.92) with one model copy (~19 GB vs 24 GB).
The CLI/single-thread path passes no lock and uses a nullcontext, so it
stays allocation- and contention-free. Concurrent output verified
identical (10x same-frame -> one hash). A per-worker admission semaphore
bounds in-flight decoded images to the pool size.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
196 lines
8.2 KiB
Python
196 lines
8.2 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
import threading
|
|
from dataclasses import replace
|
|
|
|
import numpy as np
|
|
|
|
from .config import _SEG_MODELS, override_settings
|
|
from .memtune import release_freed_memory
|
|
from .pipeline import MattingPipeline, MattingResult
|
|
from .segmentation import BiRefNetSegmenter, make_segmenter
|
|
from .settings import PipelineSettings, SegmentationSettings
|
|
from .vitmatte_infer import ViTMatteRunner
|
|
|
|
logger = logging.getLogger("bgfilter.service")
|
|
|
|
# HTTP-facing seg_model values, in preference order (first = default).
|
|
SEG_MODELS: tuple[str, ...] = tuple(_SEG_MODELS)
|
|
|
|
|
|
class PipelineManager:
|
|
"""Long-lived holder of the matting models.
|
|
|
|
ViTMatte is the matting model shared by both segmentation backends, so it is
|
|
loaded once and reused. Each segmentation backend is lazily loaded once on
|
|
first use and cached by name. Per-request parameters (``screen_color``,
|
|
``seg_model``, ``cross_check``, ``chroma``) never trigger a model reload --
|
|
they only re-parametrise the pipeline over already-cached (or first-time-
|
|
loaded, then cached) models. ``chroma=False`` selects non-flat/complex
|
|
background mode (no colour key); ``screen_color`` is then unused.
|
|
"""
|
|
|
|
def __init__(self, base_settings: PipelineSettings, preload: bool = True):
|
|
self._base = base_settings
|
|
self._runner: ViTMatteRunner | None = None
|
|
self._segmenters: dict[str, object] = {}
|
|
self._cross_checker: BiRefNetSegmenter | None = None
|
|
self._load_lock = threading.Lock()
|
|
# Caps how many concurrent requests may be inside a GPU forward at once
|
|
# (the app may run several process() calls on a thread pool). Shared by every
|
|
# pipeline this manager builds. =1 (default) makes the device strictly
|
|
# single-stream; >1 lets that many forwards overlap so batch-1 kernels that
|
|
# underfill the SMs can co-schedule (higher GPU-bound throughput, but the
|
|
# concurrent activations multiply GPU memory). Uncontended when CPU_WORKERS=1.
|
|
gpu_conc = max(1, int(os.environ.get("BGFILTER_GPU_CONCURRENCY", "1")))
|
|
self._gpu_lock = threading.Lock() if gpu_conc == 1 else threading.BoundedSemaphore(gpu_conc)
|
|
if preload:
|
|
if base_settings.model.matting_method == "vitmatte":
|
|
self._get_runner()
|
|
# Cross-check is default behaviour, so its HR-matting model is part of
|
|
# the warm set. With reuse-as-seg on the birefnet backend that same
|
|
# forward also serves as the seg mask, so the primary segmenter is
|
|
# never used on the default path -- only preload it when reuse won't
|
|
# cover segmentation (reuse off, or a non-birefnet default backend).
|
|
if base_settings.cross_check.enabled:
|
|
self._get_cross_checker()
|
|
reuse_covers_seg = (
|
|
base_settings.cross_check.enabled
|
|
and base_settings.cross_check.reuse_as_seg
|
|
and base_settings.segmentation.backend == "birefnet"
|
|
)
|
|
if base_settings.segmentation.enabled and not reuse_covers_seg:
|
|
self._get_segmenter(base_settings.segmentation.backend)
|
|
|
|
@property
|
|
def default_seg_model(self) -> str:
|
|
return self._base.segmentation.backend
|
|
|
|
@property
|
|
def device(self) -> str:
|
|
return self._base.model.device
|
|
|
|
@property
|
|
def cross_check_enabled(self) -> bool:
|
|
return self._base.cross_check.enabled
|
|
|
|
def effective_cross_check(self, requested: bool | None) -> bool:
|
|
"""Resolve the cross-check state for a request, honoring the lane lock.
|
|
|
|
When ``cross_check.lock`` is set (e.g. the cpu-fast lane) the base default
|
|
is authoritative and a per-request override is ignored -- so a locked
|
|
off-lane never loads the HR-matting model or risks an OOM, and never errors
|
|
on a ``cross_check=on`` request. Otherwise the request value wins when
|
|
given, else the base default applies.
|
|
"""
|
|
base = self._base.cross_check
|
|
if base.lock:
|
|
return base.enabled
|
|
return requested if requested is not None else base.enabled
|
|
|
|
def _get_runner(self) -> ViTMatteRunner:
|
|
if self._runner is None:
|
|
with self._load_lock:
|
|
if self._runner is None:
|
|
logger.info(
|
|
"loading ViTMatte: model=%s device=%s",
|
|
self._base.model.model_name,
|
|
self._base.model.device,
|
|
)
|
|
self._runner = ViTMatteRunner(self._base.model)
|
|
return self._runner
|
|
|
|
def _get_segmenter(self, backend: str):
|
|
seg = self._segmenters.get(backend)
|
|
if seg is None:
|
|
with self._load_lock:
|
|
seg = self._segmenters.get(backend)
|
|
if seg is None:
|
|
seg_settings = replace(
|
|
self._base.segmentation,
|
|
backend=backend,
|
|
model_name=_SEG_MODELS[backend],
|
|
)
|
|
logger.info(
|
|
"loading segmenter: backend=%s model=%s device=%s",
|
|
backend,
|
|
seg_settings.model_name,
|
|
seg_settings.device,
|
|
)
|
|
seg = make_segmenter(seg_settings)
|
|
self._segmenters[backend] = seg
|
|
return seg
|
|
|
|
def _get_cross_checker(self) -> BiRefNetSegmenter:
|
|
# The cross-check model depends only on base settings (never on per-request
|
|
# params), so one shared instance serves all requests.
|
|
if self._cross_checker is None:
|
|
with self._load_lock:
|
|
if self._cross_checker is None:
|
|
check = self._base.cross_check
|
|
logger.info(
|
|
"loading cross-check model: model=%s device=%s",
|
|
check.model_name,
|
|
self._base.segmentation.device,
|
|
)
|
|
self._cross_checker = BiRefNetSegmenter(
|
|
SegmentationSettings(
|
|
model_name=check.model_name,
|
|
device=self._base.segmentation.device,
|
|
input_size=check.input_size,
|
|
precision=check.precision,
|
|
)
|
|
)
|
|
return self._cross_checker
|
|
|
|
def process(
|
|
self,
|
|
rgb: np.ndarray,
|
|
*,
|
|
screen_color: str | None,
|
|
seg_model: str,
|
|
cross_check: bool | None = None,
|
|
chroma: bool | None = None,
|
|
) -> MattingResult:
|
|
if seg_model not in _SEG_MODELS:
|
|
raise ValueError(f"invalid seg_model '{seg_model}'")
|
|
settings = override_settings(
|
|
self._base,
|
|
screen_color=screen_color,
|
|
seg_backend=seg_model,
|
|
# None leaves the base default; False selects non-flat/complex mode.
|
|
chroma=chroma,
|
|
# Honor the lane lock: a locked lane ignores the per-request override
|
|
# (a cpu-fast box never runs cross-check, even on cross_check=on).
|
|
cross_check=self.effective_cross_check(cross_check),
|
|
)
|
|
runner = (
|
|
self._get_runner()
|
|
if settings.model.matting_method == "vitmatte"
|
|
else None
|
|
)
|
|
segmenter = (
|
|
self._get_segmenter(seg_model) if settings.segmentation.enabled else None
|
|
)
|
|
cross_checker = (
|
|
self._get_cross_checker() if settings.cross_check.enabled else None
|
|
)
|
|
pipeline = MattingPipeline(
|
|
settings,
|
|
runner=runner,
|
|
segmenter=segmenter,
|
|
cross_checker=cross_checker,
|
|
gpu_lock=self._gpu_lock,
|
|
)
|
|
try:
|
|
return pipeline.run_rgb(rgb)
|
|
finally:
|
|
# Return this request's large freed activations to the OS so RSS does
|
|
# not ratchet up across requests in the long-lived service process.
|
|
# Hold the GPU lock: empty_cache must not race a concurrent request's
|
|
# in-flight allocations on the same device.
|
|
with self._gpu_lock:
|
|
release_freed_memory()
|