diff --git a/DEPLOY.md b/DEPLOY.md index 1a1717d..cd44a4e 100644 --- a/DEPLOY.md +++ b/DEPLOY.md @@ -219,6 +219,8 @@ redirect back to `huggingface.co`, which recent `huggingface_hub` rejects with | `BGFILTER_MAX_IMAGE_PIXELS` | `4194304` | Reject larger inputs with `413` (~4 MP) | | `BGFILTER_PRELOAD` | `1` | Load default models at startup (first request isn't cold) | | `BGFILTER_CROSS_CHECK` | config (`on`) | Force the cross-check veto on/off (`0` disables; needs `BiRefNet_HR-matting`) | +| `BGFILTER_CPU_WORKERS` | `1` | Thread-pool size for request pipelining. `1` = strictly serial (one request at a time). `>1` lets that many requests run concurrently, overlapping their CPU stages while a GPU mutex serializes the model forwards — keeps the GPU fed instead of idle. Biggest win when the pipeline is CPU-bound (e.g. cross-check off). | +| `BGFILTER_GPU_CONCURRENCY` | `1` | How many requests may be inside a GPU forward at once. `1` = single-stream (safe). `2` lets two forwards co-schedule on the SMs (batch-1 underfills them), raising GPU-bound throughput — but concurrent activations multiply VRAM, so size it as `weights + N × peak_forward ≤ VRAM` (on a 32 GB card, `2` peaks ~19 GB at 2048 cross-check). Needs `BGFILTER_CPU_WORKERS>1` to matter. | | `BGFILTER_AUTH_TOKEN` | — | If set, `/remove-background` requires header `X-Genarrative-Image-Token` to equal it (constant-time; `401` otherwise). Unset = **open** (auth left to a fronting proxy). `/healthz` stays open. See §10. | | `BGFILTER_WEIGHTS_DIR` | `models` | Where local weight folders are looked up (Option B) | | `HF_HOME` | — | HF cache location (Option A) | diff --git a/DEPLOY_ZH.md b/DEPLOY_ZH.md index 745d6c8..79a93e6 100644 --- a/DEPLOY_ZH.md +++ b/DEPLOY_ZH.md @@ -202,6 +202,8 @@ export HF_HUB_DISABLE_XET=1 # 这些 repo 是 Xet 存储；强 | `BGFILTER_MAX_IMAGE_PIXELS` | `4194304` | 超过则返回 `413`（约 4 MP）| | `BGFILTER_PRELOAD` | `1` | 启动时预加载默认模型（首个请求不用冷加载）| | `BGFILTER_CROSS_CHECK` | 配置（`on`）| 强制开/关 cross-check 否决（`0` 关闭；需 `BiRefNet_HR-matting`）| +| `BGFILTER_CPU_WORKERS` | `1` | 请求流水线的线程池大小。`1` = 严格串行（一次一个请求）。`>1` 允许这么多请求并发,重叠各自的 CPU 段,同时用一把 GPU 互斥锁串行化模型前向——让 GPU 别闲着。CPU 为瓶颈时(如关 cross-check)收益最大。 | +| `BGFILTER_GPU_CONCURRENCY` | `1` | 允许几个请求同时进入 GPU 前向。`1` = 单流(安全)。`2` 让两次前向在 SM 上共同调度(batch-1 喂不满 SM),提升 GPU 瓶颈下的吞吐——但并发激活会成倍占显存,按 `权重 + N × 单次前向峰值 ≤ 显存` 来定(32GB 卡上 `2` 在 2048 cross-check 时峰值约 19GB)。需 `BGFILTER_CPU_WORKERS>1` 才有意义。 | | `BGFILTER_AUTH_TOKEN` | — | 若设置，`/remove-background` 要求 header `X-Genarrative-Image-Token` 等于它（常量时间比较；否则 `401`）。不设 = **开放**（鉴权交给前置代理）。`/healthz` 始终开放。见 §10。 | | `BGFILTER_WEIGHTS_DIR` | `models` | 本地权重文件夹的查找目录（方案 B）| | `HF_HOME` | — | HF 缓存位置（方案 A）| diff --git a/app.py b/app.py index 5c16494..8c24c81 100644 --- a/app.py +++ b/app.py @@ -13,6 +13,7 @@ import logging import os import re import time +from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager from pathlib import Path @@ -67,7 +68,18 @@ async def require_auth(x_genarrative_image_token: str | None = Header(default=No _manager: PipelineManager | None = None -_lock = asyncio.Lock() + +# Request pipelining. Each request's process() runs on a small thread pool so the +# CPU stages (chroma / trimap / alpha_post / foreground) of concurrent requests +# overlap, while a GPU mutex inside the pipeline keeps the model forwards mutually +# exclusive on the single device -- i.e. when one request steps off the GPU to do +# CPU work, another steps on, keeping the GPU fed instead of idle. +# BGFILTER_CPU_WORKERS=1 (default) gives a 1-thread pool == the old strictly serial +# behaviour. A semaphore bounds in-flight requests (hence decoded-image RAM) to the +# pool size so a burst cannot pile unbounded decoded images before the pool drains. +CPU_WORKERS = max(1, _int_env("BGFILTER_CPU_WORKERS", 1)) +_executor = ThreadPoolExecutor(max_workers=CPU_WORKERS, thread_name_prefix="bgf") +_sem = asyncio.Semaphore(CPU_WORKERS) def _build_settings(): @@ -240,14 +252,16 @@ async def remove_background( # Honors the lane lock (a locked cpu-fast box reports/runs "off" even on cross_check=on). effective_cc = _manager.effective_cross_check(normalized_cc) - rgb = _decode_image(await file.read()) - - start = time.perf_counter() - async with _lock: # one inference at a time per process + # Bound in-flight requests to the pool size, then run on the pool. The GPU + # sections are serialized inside the pipeline (a shared GPU mutex); the CPU + # sections of concurrent requests overlap across the pool threads. + async with _sem: + rgb = _decode_image(await file.read()) + start = time.perf_counter() loop = asyncio.get_running_loop() try: result = await loop.run_in_executor( - None, + _executor, lambda: _manager.process( rgb, screen_color=normalized_color, @@ -268,7 +282,7 @@ async def remove_background( except Exception: logger.exception("inference failed") raise HTTPException(status_code=500, detail="inference failed") - elapsed_ms = int((time.perf_counter() - start) * 1000) + elapsed_ms = int((time.perf_counter() - start) * 1000) png = _encode_png(result.rgb, result.alpha) headers = { diff --git a/bgfilter/pipeline.py b/bgfilter/pipeline.py index b0aeaad..c89ab45 100644 --- a/bgfilter/pipeline.py +++ b/bgfilter/pipeline.py @@ -1,5 +1,6 @@ from __future__ import annotations +import contextlib import json from dataclasses import asdict, dataclass from pathlib import Path @@ -60,6 +61,7 @@ class MattingPipeline: runner: ViTMatteRunner | None = None, segmenter=None, cross_checker=None, + gpu_lock=None, ): self.settings = settings # Pre-loaded models can be injected so a long-lived service shares them @@ -67,6 +69,10 @@ class MattingPipeline: self._runner: ViTMatteRunner | None = runner self._segmenter = segmenter self._cross_checker = cross_checker + # Held around every GPU model forward so concurrent requests never drive the + # device at once (models are shared, single GPU). None = no contention + # (CLI / single-thread), so a no-op context keeps that path lock-free. + self._gpu_lock = gpu_lock if gpu_lock is not None else contextlib.nullcontext() def _segment(self, rgb: np.ndarray) -> np.ndarray: if self._segmenter is None: @@ -166,10 +172,12 @@ def _process_rgb(rgb: np.ndarray, pipeline: MattingPipeline) -> MattingResult: 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) + with pipeline._gpu_lock: + second_alpha = pipeline._second_opinion(rgb) seg_mask = second_alpha else: - seg_mask = pipeline._segment(rgb) + with pipeline._gpu_lock: + seg_mask = pipeline._segment(rgb) # The directional modes are colour tests against the key colour; without # one, the seg-topology trimap is the only meaningful choice. mode = settings.trimap.mode if settings.chroma.enabled else "seg" @@ -188,7 +196,8 @@ def _process_rgb(rgb: np.ndarray, pipeline: MattingPipeline) -> MattingResult: else: seg_mask = None trimap, trimap_stats = generate_trimap(bg_confidence, settings.trimap) - raw_alpha, alpha_source = pipeline._predict_alpha(rgb, trimap, bg_confidence) + with pipeline._gpu_lock: + raw_alpha, alpha_source = pipeline._predict_alpha(rgb, trimap, bg_confidence) alpha = enforce_trimap(raw_alpha, trimap) if settings.chroma.enabled: @@ -202,7 +211,8 @@ def _process_rgb(rgb: np.ndarray, pipeline: MattingPipeline) -> MattingResult: if settings.cross_check.enabled: lab = convert_color_spaces(rgb)[1] if second_alpha is None: - second_alpha = pipeline._second_opinion(rgb) + with pipeline._gpu_lock: + second_alpha = pipeline._second_opinion(rgb) alpha = cross_check_alpha( alpha, second_alpha, @@ -213,9 +223,12 @@ def _process_rgb(rgb: np.ndarray, pipeline: MattingPipeline) -> MattingResult: trimap, settings.cross_check, ) - foreground = estimate_foreground_rgb( - rgb, alpha, bg_confidence, model, settings.foreground - ) + # Foreground colour estimation runs on the GPU (cupy) when enabled, so it shares + # the mutex; the CPU fallback path is short and simply runs while holding it. + with pipeline._gpu_lock: + foreground = estimate_foreground_rgb( + rgb, alpha, bg_confidence, model, settings.foreground + ) corrected_rgb, color_mask = despill(foreground, alpha, model, settings.despill) metadata = { diff --git a/bgfilter/service.py b/bgfilter/service.py index a315ed9..ec1178a 100644 --- a/bgfilter/service.py +++ b/bgfilter/service.py @@ -1,6 +1,7 @@ from __future__ import annotations import logging +import os import threading from dataclasses import replace @@ -37,6 +38,14 @@ class PipelineManager: 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() @@ -169,11 +178,18 @@ class PipelineManager: self._get_cross_checker() if settings.cross_check.enabled else None ) pipeline = MattingPipeline( - settings, runner=runner, segmenter=segmenter, cross_checker=cross_checker + 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. - release_freed_memory() + # 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()