diff --git a/DEPLOY.md b/DEPLOY.md index 2462dfb..f78405f 100644 --- a/DEPLOY.md +++ b/DEPLOY.md @@ -217,6 +217,8 @@ redirect back to `huggingface.co`, which recent `huggingface_hub` rejects with | `BGFILTER_CONFIG` | `configs/default.yaml` | Pipeline config file | | `BGFILTER_DEVICE` | `cpu` | Overrides **both** model and segmentation device (`cuda` for GPU) | | `BGFILTER_MAX_IMAGE_PIXELS` | `4194304` | Reject larger inputs with `413` (~4 MP) | +| `BGFILTER_FETCH_MAX_BYTES` | `26214400` | `image_url` mode: cap the downloaded body (~25 MB) → `413`. No SSRF filtering — keep auth on / callers trusted. | +| `BGFILTER_FETCH_TIMEOUT` | `15` | `image_url` mode: read timeout (s) for the fetch (connect fixed at 5 s). Outbound fetch concurrency is fixed at `2 × cpu_workers` (hard cap 64). | | `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` | config `server.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). Set it in the config (`server:` section) or override per box with this env var. | diff --git a/DEPLOY_ZH.md b/DEPLOY_ZH.md index 5d7cb49..eaef015 100644 --- a/DEPLOY_ZH.md +++ b/DEPLOY_ZH.md @@ -200,6 +200,8 @@ export HF_HUB_DISABLE_XET=1 # 这些 repo 是 Xet 存储;强 | `BGFILTER_CONFIG` | `configs/default.yaml` | 管线配置文件 | | `BGFILTER_DEVICE` | `cpu` | 同时覆盖 model 和 segmentation 的 device(GPU 用 `cuda`)| | `BGFILTER_MAX_IMAGE_PIXELS` | `4194304` | 超过则返回 `413`(约 4 MP)| +| `BGFILTER_FETCH_MAX_BYTES` | `26214400` | `image_url` 模式:下载体上限(约 25 MB)→ `413`。无 SSRF 过滤——务必开鉴权 / 调用方可信。 | +| `BGFILTER_FETCH_TIMEOUT` | `15` | `image_url` 模式:fetch 读超时(秒;连接超时固定 5 秒)。出站 fetch 并发固定为 `2 × cpu_workers`(硬顶 64)。 | | `BGFILTER_PRELOAD` | `1` | 启动时预加载默认模型(首个请求不用冷加载)| | `BGFILTER_CROSS_CHECK` | 配置(`on`)| 强制开/关 cross-check 否决(`0` 关闭;需 `BiRefNet_HR-matting`)| | `BGFILTER_CPU_WORKERS` | 配置 `server.cpu_workers`(`1`)| 请求流水线的线程池大小。`1` = 严格串行(一次一个请求)。`>1` 允许这么多请求并发,重叠各自的 CPU 段,同时用一把 GPU 互斥锁串行化模型前向——让 GPU 别闲着。CPU 为瓶颈时(如关 cross-check)收益最大。可写进配置的 `server:` 段,或用此 env 变量按机器覆盖。 | diff --git a/README.md b/README.md index e6e3279..189d3a5 100644 --- a/README.md +++ b/README.md @@ -207,7 +207,8 @@ is expected to authenticate). `/healthz` is always open. See [DEPLOY.md](DEPLOY. | Field | Required | Values | Meaning | | --- | --- | --- | --- | -| `file` | yes | image file | Source image (field name fixed as `file`). | +| `file` | one of `file`/`image_url` | image file | Source image as an upload (field name fixed as `file`). | +| `image_url` | one of `file`/`image_url` | `http(s)://…` | Source image by URL — the service fetches it (bounded by size/time; follows ≤3 redirects). Provide exactly one of `file` or `image_url`. No SSRF filtering: the endpoint trusts authenticated callers, so keep auth on and callers trusted. | | `background_mode` | no | `flat` (default) / `complex` | `flat` = solid-colour background (colour-keyed pipeline). `complex` = non-flat/scene background: no colour key, the segmenter alone drives the matte. `screen_color` is ignored in `complex`; expect flat backgrounds to stay stronger on hair detail. | | `screen_color` | no | `#RRGGBB`, or omit/empty/`auto` | Background-colour prior; omit to auto-detect from the border. Ignored when `background_mode=complex`. | | `seg_model` | no | `birefnet` (default) / `anime-seg` | Segmentation backend. | @@ -228,9 +229,9 @@ is expected to authenticate). `/healthz` is always open. See [DEPLOY.md](DEPLOY. | Code | When | | --- | --- | | `200` | Success — body is the RGBA PNG. | -| `400` | Bad input: unreadable image, invalid `screen_color`/`seg_model`/`cross_check`/`background_mode`, or auto-detect failed (pass `screen_color`). | +| `400` | Bad input: unreadable image, not exactly one of `file`/`image_url`, bad `image_url` / fetch failed, invalid `screen_color`/`seg_model`/`cross_check`/`background_mode`, or auto-detect failed (pass `screen_color`). | | `401` | Auth enabled and the token is missing or wrong. | -| `413` | Image exceeds `BGFILTER_MAX_IMAGE_PIXELS` (~4 MP by default). | +| `413` | Image exceeds `BGFILTER_MAX_IMAGE_PIXELS` (~4 MP by default), or an `image_url` body exceeds `BGFILTER_FETCH_MAX_BYTES` (~25 MB). | | `500` | Inference failed. | | `503` | Service still loading models (not ready). | @@ -256,6 +257,13 @@ curl -sS \ -F "file=@scene.jpg" \ -F "background_mode=complex" \ http://127.0.0.1:18083/remove-background -o output.png + +# by URL instead of an upload (service fetches it; use image_url OR file, not both) +curl -sS \ + -F "image_url=https://example.com/input.png" \ + -F "screen_color=#CFEFFF" \ + -H "X-Genarrative-Image-Token: " \ + http://127.0.0.1:18083/remove-background -o output.png ``` Config comes from `BGFILTER_CONFIG` (defaults to `configs/default.yaml` when present); diff --git a/README_ZH.md b/README_ZH.md index 1452df0..2fa9db0 100644 --- a/README_ZH.md +++ b/README_ZH.md @@ -186,7 +186,8 @@ python -m uvicorn app:app --host 127.0.0.1 --port 18083 --workers 1 | 字段 | 必填 | 取值 | 含义 | | --- | --- | --- | --- | -| `file` | 是 | 图片文件 | 源图(字段名固定为 `file`)。 | +| `file` | `file`/`image_url` 二选一 | 图片文件 | 以上传方式提供源图(字段名固定为 `file`)。 | +| `image_url` | `file`/`image_url` 二选一 | `http(s)://…` | 以 URL 提供源图,服务端自行 fetch(受大小/超时限制,跟随 ≤3 跳重定向)。`file` 与 `image_url` 必须恰好提供一个。无 SSRF 过滤:端点信任已鉴权的调用方,请保持鉴权开启且调用方可信。 | | `background_mode` | 否 | `flat`(默认)/ `complex` | `flat` = 纯色背景(走色键管线)。`complex` = 非纯色/实景背景:无色键,完全由分割器驱动抠图;此模式下 `screen_color` 被忽略,发丝细节通常弱于纯色模式。 | | `screen_color` | 否 | `#RRGGBB`,或留空/`auto` | 背景色先验;留空则从边框自动探测。`background_mode=complex` 时忽略。 | | `seg_model` | 否 | `birefnet`(默认)/ `anime-seg` | 分割后端。 | diff --git a/app.py b/app.py index 5e27064..deb7133 100644 --- a/app.py +++ b/app.py @@ -16,7 +16,9 @@ import time from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager from pathlib import Path +from urllib.parse import urlparse +import httpx import numpy as np from fastapi import Depends, FastAPI, File, Form, Header, HTTPException, Response, UploadFile from PIL import Image, ImageOps, UnidentifiedImageError @@ -46,6 +48,16 @@ def _int_env(name: str, default: int) -> int: MAX_IMAGE_PIXELS = _int_env("BGFILTER_MAX_IMAGE_PIXELS", 4_194_304) +# /remove-background accepts either an uploaded `file` or an `image_url` the service +# fetches itself. These bound the fetch so a huge or slow URL cannot hang a worker or +# blow RAM. NOTE: no SSRF hardening here (no host/IP allowlist) -- the endpoint is +# guarded by token auth and assumes trusted callers; do NOT expose it to untrusted +# clients without adding private-address / redirect filtering to _fetch_image. +FETCH_MAX_BYTES = _int_env("BGFILTER_FETCH_MAX_BYTES", 25 * 1024 * 1024) +FETCH_TIMEOUT_S = float(os.environ.get("BGFILTER_FETCH_TIMEOUT", "15")) +# Outbound fetch concurrency = 2 x cpu_workers, hard-capped here so a burst can never +# exhaust file descriptors / ephemeral ports (each fetch = one socket). +FETCH_CONCURRENCY_CAP = 64 # We enforce our own pixel budget from the image header before decoding, so disable # Pillow's separate decompression-bomb ceiling (which would otherwise raise its own error). Image.MAX_IMAGE_PIXELS = None @@ -81,6 +93,11 @@ _manager: PipelineManager | None = None # drains. _executor: ThreadPoolExecutor | None = None _sem: asyncio.Semaphore | None = None +# Outbound image_url fetches: a FIFO admission bound (restores the throttle a direct +# upload has implicitly) and a shared client for connection reuse. Both built in +# lifespan once the pool size is known. See _fetch_image / §image_url mode. +_fetch_sem: asyncio.Semaphore | None = None +_http_client: httpx.AsyncClient | None = None def _build_settings(): @@ -106,13 +123,24 @@ def _build_settings(): @asynccontextmanager async def lifespan(app: FastAPI): - global _manager, _executor, _sem + global _manager, _executor, _sem, _fetch_sem, _http_client settings = _build_settings() preload = os.environ.get("BGFILTER_PRELOAD", "1").strip().lower() not in ("0", "false", "no", "") # Pool size: config default (server.cpu_workers), env override wins. cpu_workers = max(1, _int_env("BGFILTER_CPU_WORKERS", settings.server.cpu_workers)) _executor = ThreadPoolExecutor(max_workers=cpu_workers, thread_name_prefix="bgf") _sem = asyncio.Semaphore(cpu_workers) + # Outbound-fetch bound: 2x the pool (a small prefetch buffer), hard-capped. The + # shared client's max_connections matches so httpx never opens more sockets than + # the semaphore admits, and reuses connections across fetches. + fetch_conc = min(2 * cpu_workers, FETCH_CONCURRENCY_CAP) + _fetch_sem = asyncio.Semaphore(fetch_conc) + _http_client = httpx.AsyncClient( + follow_redirects=True, + max_redirects=3, + timeout=httpx.Timeout(connect=5.0, read=FETCH_TIMEOUT_S, write=10.0, pool=None), + limits=httpx.Limits(max_connections=fetch_conc), + ) logger.info( "starting bgfilter service: device=%s default_seg=%s cross_check=%s preload=%s " "cpu_workers=%d gpu_concurrency=%d max_pixels=%d auth=%s", @@ -130,7 +158,10 @@ async def lifespan(app: FastAPI): "BGFILTER_AUTH_TOKEN not set: /remove-background is OPEN (no app-level auth)" ) _manager = PipelineManager(settings, preload=preload) - yield + try: + yield + finally: + await _http_client.aclose() app = FastAPI(title="BGfilter", version=VERSION, lifespan=lifespan) @@ -205,6 +236,34 @@ def _normalize_background_mode(value: str | None) -> str | None: raise HTTPException(status_code=400, detail="invalid background_mode") +async def _fetch_image(url: str) -> bytes: + """Download image bytes from a URL, bounded by size and time. + + Robustness only -- no SSRF filtering (see FETCH_* note above): callers are + trusted via the auth token. Streams the body so an over-large response is + rejected without being fully buffered. Uses the shared client (connection reuse); + the caller holds _fetch_sem so outbound concurrency stays bounded. + """ + parsed = urlparse(url) + if parsed.scheme not in ("http", "https") or not parsed.hostname: + raise HTTPException(status_code=400, detail="image_url must be an http(s) URL") + try: + async with _http_client.stream("GET", url) as resp: + if resp.status_code != 200: + raise HTTPException( + status_code=400, + detail=f"image_url fetch failed (status {resp.status_code})", + ) + buf = bytearray() + async for chunk in resp.aiter_bytes(): + buf += chunk + if len(buf) > FETCH_MAX_BYTES: + raise HTTPException(status_code=413, detail="image_url body too large") + return bytes(buf) + except httpx.HTTPError: + raise HTTPException(status_code=400, detail="image_url fetch failed") + + def _decode_image(data: bytes) -> np.ndarray: if not data: raise HTTPException(status_code=400, detail="empty image") @@ -235,7 +294,8 @@ def _encode_png(rgb: np.ndarray, alpha: np.ndarray) -> bytes: @app.post("/remove-background") async def remove_background( - file: UploadFile = File(...), + file: UploadFile | None = File(None), + image_url: str | None = Form(None), screen_color: str | None = Form(None), seg_model: str | None = Form(None), cross_check: str | None = Form(None), @@ -245,6 +305,14 @@ async def remove_background( if _manager is None: raise HTTPException(status_code=503, detail="service not ready") + # Exactly one image source. (Auth already ran, so only authenticated callers can + # trigger an outbound fetch.) + url = image_url.strip() if image_url else None + if (file is None) == (not url): + raise HTTPException( + status_code=400, detail="provide exactly one of 'file' or 'image_url'" + ) + normalized_color = _normalize_screen_color(screen_color) normalized_seg = _normalize_seg_model(seg_model) or _manager.default_seg_model normalized_cc = _normalize_cross_check(cross_check) @@ -260,11 +328,20 @@ 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) + # Acquire the raw bytes BEFORE the compute gate: a slow fetch must not hold a + # pool/compute slot. The upload body is already buffered by Starlette (fast); + # a URL is fetched under its own looser bound (_fetch_sem). + if url is not None: + async with _fetch_sem: + raw = await _fetch_image(url) + else: + raw = await file.read() + # 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()) + rgb = _decode_image(raw) start = time.perf_counter() loop = asyncio.get_running_loop() try: diff --git a/requirements.txt b/requirements.txt index 566302f..3b9a161 100644 --- a/requirements.txt +++ b/requirements.txt @@ -21,3 +21,4 @@ PyYAML fastapi uvicorn[standard] python-multipart +httpx