BGfilter/bgfilter/service.py

from __future__ import annotations

import logging
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``) never trigger a model reload -- ``screen_color``
    only affects this request's chroma/despill, and ``seg_model``/``cross_check``
    select an already-cached (or first-time-loaded, then cached) model.
    """

    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()
        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 _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,
    ) -> 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 = keep the service default; True/False = per-request override.
            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
        )
        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.
            release_freed_memory()