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). Config value # (``server.gpu_concurrency``) is the default; the env var overrides it. gpu_conc = max(1, int(os.environ.get("BGFILTER_GPU_CONCURRENCY", base_settings.server.gpu_concurrency))) 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()