Cache rembg sessions between requests — remove-background spends ~11s loading the model and ~1s inferring

[photoevents]

Problem or use case

On a GPU host, remove-background takes ~12 s per image end-to-end, but the actual inference is under 1 s. The remaining ~11 s is the ONNX model being loaded into VRAM again on every request. Caching the session between requests would turn a 12 s call into roughly 1 s.

Proposed solution

TITLE:
Cache rembg sessions between requests — remove-background spends ~11s loading the model and ~1s inferring

BODY:

Summary

On a GPU host, remove-background takes ~12 s per image end-to-end, but the actual inference is under 1 s. The remaining ~11 s is the ONNX model being loaded into VRAM again on every request. Caching the session between requests would turn a 12 s call into roughly 1 s.

Measurements (SnapOtter 2.2.0, RTX 3090)

Same 1200×1800 JPEG each time, backgroundType: transparent, via POST /api/v1/tools/image/remove-background:

Model	.onnx size	End-to-end
isnet-general-use	170 MB	4.6 s
u2net	168 MB	5.0 s
birefnet-general-lite	214 MB	9.5 s
birefnet-matting	928 MB	11.4 s
birefnet-portrait	928 MB	11.5 s
birefnet-general	928 MB	14.0 s

The time tracks the model file size almost linearly, which points at load time rather than compute.

Confirming that, with a warm session inside /data/ai/venv (same image, same models):

birefnet-portrait   session create 10.3 s → 1.71 s / 0.83 s / 0.58 s
birefnet-general    session create  9.1 s → 0.95 s / 0.96 s / 0.58 s

Other checks:

• The GPU is genuinely used: 83 % utilisation and 9.4 GB VRAM during an API call.
• Timing is identical over localhost:1349 and through a reverse proxy, so it is not network or proxy overhead.
• onnxruntime-gpu 1.20.1 reports ['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']; TensorRT libs are absent so it falls back to CUDA, which is fine.

Where it comes from

The Python side already has the hard part solved — dispatcher.py is a persistent process, so imports and the CUDA context stay warm. But each job runs the tool script like this:

module_globals = {"__name__": "__main__", "__file__": script_path}
with open(script_path) as f:
    code = compile(f.read(), script_path, "exec")

A fresh globals dict per call, so nothing the script defines survives to the next request. remove_bg.py then calls new_session(model, providers=providers) inside that run, and rembg builds a new InferenceSession — re-reading the .onnx and re-uploading the weights to VRAM every time.

Suggested fix

Keep the session somewhere that outlives a single exec, since the process itself already persists. For example a small module that the script imports (imported modules stay in sys.modules, unlike the script's globals):

# session_cache.py
_SESSIONS = {}

def get_session(model, providers):
    key = (model, tuple(providers))
    if key not in _SESSIONS:
        from rembg import new_session
        _SESSIONS[key] = new_session(model, providers=providers)
    return _SESSIONS[key]

with remove_bg.py calling get_session(...) instead of new_session(...).

Worth bounding so VRAM stays predictable: an LRU of 1–2 sessions, and/or eviction after an idle timeout, configurable through an env var (AI_SESSION_CACHE_SIZE / AI_SESSION_IDLE_TIMEOUT) for people who would rather keep VRAM free. The same pattern would help the other AI tools that build a session per call.

Why it matters here

I run event photography software that composites a cut-out subject into a scene while the guests wait at the booth. At ~12 s per photo the operator feels every second; at ~1 s it becomes invisible. It also decides whether the GPU is worth having at all — right now a 168 MB model on CPU-grade quality is faster end-to-end than a 928 MB BiRefNet on a 3090, purely because of reload overhead, which inverts the quality/speed trade-off the models are meant to offer.

Environment

• SnapOtter 2.2.0 (single container, embedded PostgreSQL 17 + Redis 8), snapotter/snapotter:2.2.0
• Docker Desktop 4.86 on Windows 10, WSL2, --gpus all
• NVIDIA RTX 3090 24 GB, background-removal bundle 2.0.0 (rembg 2.0.62, onnxruntime-gpu 1.20.1, numpy 1.26.4)

Alternatives considered

No response

Area

Existing tool improvement

Contribution

• I'm willing to submit a PR for this feature

[snapotter-hq]

Thanks @photoevents, excellent write-up. The measurements make the cause obvious: the dispatcher process persists but each script runs in a fresh globals dict (dispatcher.py:288), so remove_bg rebuilds the session every call. Tracked in #776, with your LRU + idle-timeout as the VRAM guard. The same per-call pattern hits the other AI tools, so we'll aim for a shared session cache rather than a rembg-only one.

[photoevents]

Glad it was useful — and a shared cache is the better call. I only scoped mine to rembg because that was the tool I could measure. Three things I hit that get harder once the cache is shared across tools, in case they save you time: Bound it by VRAM, not by session count. My SNAPOTTER_SESSION_CACHE=1 default was sized for exactly one BiRefNet: ~1.3 GB of weights plus the inference arena ONNX Runtime grows on first use — about 7.5 GB at 1200×1800, and it never shrinks. Two of those exhausted a 24 GB card and stalled requests for ten minutes. With one cache shared across tools that ceiling arrives sooner and from combinations you didn't test: two different tools each holding one session is the same failure with a count of 1 per tool. A byte budget, or a per-tool cap, fails predictably where a session count does not. Idle eviction only returns VRAM after a gc. Dropping the reference wasn't enough in my testing — I needed an explicit gc.collect() before the memory actually came back. Relevant if the timeout runs on a timer thread, which also wants to not race a request that is about to hit the entry it is evicting. Providers belong in the cache key. Mine was (model, tuple(providers)). That matters more across tools than it did for me: a tool that deliberately asks for CPU shouldn't be handed a CUDA session because something else warmed one first. I also passed any call with extra args straight through uncached, on the assumption that per-call options shouldn't be silently shared between jobs — worth re-checking that assumption against the tools I couldn't see. One interaction worth deciding early in #776: DISPATCHER_MAX_REQUESTS recycling throws the cache away with the process, so the recycle default effectively caps how much any cache can help. I measured the dispatcher at +2 MB RSS over 100 requests with the cache active, so 500 was comfortable here — but that was one tool, and a shared cache changes the memory profile you're recycling to protect against. Happy to test a pre-release on the 3090. I still have the harness these numbers came from (120-photo runs with VRAM sampling) and a real event workload to point at it.   "SnapOtter" ***@***.*** – 13 augustus 2026 om 11:23 

  Thanks @photoevents, excellent write-up. The measurements make the cause obvious: the dispatcher process persists but each script runs in a fresh globals dict (dispatcher.py:288), so remove_bg rebuilds the session every call. Tracked in #776, with your LRU + idle-timeout as the VRAM guard. The same per-call pattern hits the other AI tools, so we'll aim for a shared session cache rather than a rembg-only one. —Reply to this email directly, view it on GitHub, or unsubscribe.Triage notifications, keep track of coding agent tasks and review pull requests on the go with GitHub Mobile for iOS and Android. Download it today! You are receiving this because you were mentioned.Message ID:

***@***.***>

[snapotter-hq]

The byte budget is the right call and I've moved #776 to it. A session count was always going to break once the cache is shared, for exactly the reason you give: two tools holding one session each is your failure case with a count of 1 everywhere. Providers in the key too. That one would have bitten us the first time somebody pinned a tool to CPU.

The gc.collect() detail is the kind of thing we'd have shipped without and then chased for a week. Thanks.

DISPATCHER_MAX_REQUESTS is the open question. Recycling exists because we didn't trust the sidecar's memory over long runs, and a cache that dies every 50 requests isn't a cache. Your +2 MB over 100 gives us something to argue from, though a shared cache changes what we're measuring.

Yes to the 3090 when there's a build, if you're still up for it. A 24 GB card with your harness on it is the config most likely to catch the eviction path going wrong, and we don't have one.

[photoevents]

Thanks — the byte budget and providers in the key both sound right.

One data point on DISPATCHER_MAX_REQUESTS, from a live event on 11 September, that argues for recycling on health rather than on a count.

On the 3090 (2.2.0, my session-cache hook, DISPATCHER_MAX_REQUESTS=500), every remove-background job suddenly failed with [ONNXRuntimeError] : 1 : FAIL : CUDA failure 999: unknown error: 81 failures in about four hours. The container stayed "healthy" the whole time, because the healthcheck (curl /api/v1/health) only proves that the API answers. /health kept returning 200 in between the failed jobs.

What ended it was not a recycle. After a model switch, the dispatcher crashed on the first job with the new model, and the bridge restarted it ("Dispatcher crash #1, backing off 1000ms before restart"). The new process came up with a fresh CUDA context, and the errors were gone.

The trigger was almost certainly on my side: an NVIDIA driver update two days earlier, without a reboot. The failure mode is what matters for #776, though. A long-lived process holding GPU sessions can end up in a state that only a restart clears, and a count of 500 never came close to firing. A shared cache makes the dispatcher live longer by design, so this gets more likely, not less.

The good news is that the crash-restart path already does the right thing; it just needs a better trigger. So I'd treat recycling as a recovery mechanism more than as memory hygiene:

• Recycle on failures: after N consecutive failed jobs (a CUDA/ORT error being the obvious trigger), instead of or on top of a request count. Here that would have cleared it within a few jobs instead of four hours later.
• Recycle on a memory budget (RSS and/or VRAM): that is what the count stands in for. On my setup it never needed to fire: +2 MB RSS over 100 requests, and VRAM flat at 16.5 GB over a 120-photo run.
• Let /api/v1/health go unhealthy when the last few jobs failed, so whatever watches the container, or the person running it, can see it and restart.

If a count stays, a high default plus re-warming the cache after a recycle would keep most requests on the warm path.

And yes, still up for testing on the 3090. Ping me when there's a build.

[snapotter-hq]

That event failure makes the case for health-based recycling better than any benchmark could.

Right now dispatcher.py only breaks out of its loop when request_count >= MAX_REQUESTS, and on the Node side bridge.ts only replaces the child when the process itself exits or its pipe breaks. A script that catches a poisoned CUDA context and returns exitCode: 1 leaves the process alive to fail the next 80 requests in a row.

Filed #2709 specifically for recycling the dispatcher on CUDA/ORT context-poison errors (or after a short run of consecutive script failures) and surfacing that state in /api/v1/admin/health, and linked your numbers onto #776. One detail we noted on #2709: the recycle has to exit cleanly from Python rather than go through recordCrash() in bridge.ts, otherwise five corrupt user uploads in a minute would trip the 5-crashes-in-60s guard and disable the sidecar for the life of the container.

We will ping you here when the #776 + #2709 build is ready for the 3090.

[photoevents]

Thanks for filing #2709 — recycling on the poison error (or a short run of failures) is exactly what would have saved that afternoon. And agreed on exiting cleanly rather than through recordCrash(): losing the sidecar for the rest of the container's life mid-event would be worse than the original problem.

Two small notes for #2709:

• The plain /api/v1/health as well, not only the admin one. The container's healthcheck curls /api/v1/health without an admin key, and that is what Docker (and I) looked at during the incident. If a recycle doesn't clear the error — a driver that really needs a host reboot, say — that endpoint should stop reporting healthy, so the container shows up as unhealthy instead of quietly failing every job.
• Bad input shouldn't count towards the failure run. With perf: cache AI model sessions between requests so remove-background stops reloading the model each call #776 every recycle throws away the warm sessions (~11 s model load in my numbers), so a few corrupt or unsupported uploads in a row shouldn't trigger one; CUDA/ORT errors should.

Happy to test the #776 + #2709 build on the 3090 as soon as it's ready.