From 7af1a1e18505380abdc5b07292a48b18e6e4e4ea Mon Sep 17 00:00:00 2001 From: tecnomanu Date: Mon, 22 Jun 2026 09:30:14 -0300 Subject: [PATCH] feat(server): evict cached model sessions to keep RAM bounded MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Electron app was growing past 20 GB because every model the user touched stayed loaded in `_SESSIONS` forever. Three knobs to fix that: - Idle TTL: a background sweeper unloads sessions that have not been used in `RBL_MODEL_IDLE_TTL` seconds (default 600). The pinned model is exempt — it is the one the user is actively working with. - /set_default_model: when the UI changes the global model dropdown, the previous one is evicted from RAM immediately so we don't pile up BiRefNet + ISNet + U2Net at once. - /remove?transient=true: per-image overrides (reprocess) load their model but do not steal the pinned default. The idle TTL reclaims them later. Plus /unload_model for manual freeing without touching the on-disk cache, and /health now reports the pinned model + TTL. Co-Authored-By: Claude Opus 4.7 --- server.py | 150 ++++++++++++++++++++++++++++++++++++++++++- static/index.html | 12 +++- tests/test_server.py | 150 +++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 309 insertions(+), 3 deletions(-) diff --git a/server.py b/server.py index 68fa4a6..c485f26 100644 --- a/server.py +++ b/server.py @@ -36,6 +36,7 @@ import io import os +import gc import glob import json import time @@ -177,6 +178,12 @@ def _pkg_version() -> str: MAX_UPLOAD_MB = int(os.environ.get("MAX_UPLOAD_MB", "30")) MAX_UPLOAD_BYTES = MAX_UPLOAD_MB * 1024 * 1024 +# How long an idle (unused) model session stays in RAM before being evicted. +# Set to 0 to disable the background evictor entirely. +MODEL_IDLE_TTL = int(os.environ.get("RBL_MODEL_IDLE_TTL", "600")) +# How often the evictor wakes up to check for idle models. +MODEL_EVICTOR_INTERVAL = int(os.environ.get("RBL_MODEL_EVICTOR_INTERVAL", "30")) + # Execution providers for onnxruntime. CPU is the reliable default (CoreML hangs # on some models on Apple Silicon). Override with REMBG_PROVIDERS if you want to # experiment, e.g. "CoreMLExecutionProvider,CPUExecutionProvider". @@ -225,13 +232,68 @@ def download_progress(name: str): _SESSIONS: dict[str, object] = {} _MODEL_STATE: dict[str, str] = {} # name -> "loading" | "ready" | "error" _MODEL_ERROR: dict[str, str] = {} +_LAST_USED: dict[str, float] = {} # name -> monotonic timestamp of last touch _LOAD_LOCKS: dict[str, asyncio.Lock] = {} +# The "pinned" model is the one the UI currently shows as the active default. +# It is exempt from idle-TTL eviction. /set_default_model rotates this; an +# explicit per-image override via /remove?transient=true does NOT. +_PINNED_MODEL: str = DEFAULT_MODEL # Serializes inference (processing queue). Created lazily inside the running # event loop: on Python 3.9 an asyncio.Lock() built at import time binds to the # wrong loop and raises "got Future attached to a different loop" under uvicorn. _INFER_LOCK: asyncio.Lock | None = None +def _touch(model_name: str) -> None: + _LAST_USED[model_name] = time.monotonic() + + +def _evict(model_name: str) -> bool: + """Drop a loaded session from RAM. Does NOT touch the on-disk .onnx file.""" + if model_name not in _SESSIONS: + return False + _SESSIONS.pop(model_name, None) + _LAST_USED.pop(model_name, None) + _MODEL_STATE.pop(model_name, None) + _MODEL_ERROR.pop(model_name, None) + gc.collect() + log.info("Unloaded model %s from RAM", model_name) + return True + + +def _evictor_sweep() -> list[str]: + """One pass of the idle-TTL eviction logic. Returns the names evicted. + + Pure (no awaits) so it can be unit-tested without driving the event loop. + """ + if MODEL_IDLE_TTL <= 0: + return [] + now = time.monotonic() + evicted: list[str] = [] + for name in [n for n in list(_SESSIONS) if n != _PINNED_MODEL]: + last = _LAST_USED.get(name, now) + if now - last > MODEL_IDLE_TTL and _evict(name): + evicted.append(name) + return evicted + + +async def _evictor_loop() -> None: + """Background task: drops loaded models that have been idle past the TTL. + + The pinned model is never evicted by the timer — it's the one the user is + actively working with. Everything else (overrides, leftovers from a model + switch) is fair game. + """ + while True: + try: + await asyncio.sleep(MODEL_EVICTOR_INTERVAL) + _evictor_sweep() + except asyncio.CancelledError: + raise + except Exception: # noqa: BLE001 + log.exception("Model evictor loop error (continuing)") + + def get_infer_lock() -> asyncio.Lock: global _INFER_LOCK if _INFER_LOCK is None: @@ -262,11 +324,13 @@ async def ensure_session(model_name: str): """ _check_model(model_name) if model_name in _SESSIONS: + _touch(model_name) return _SESSIONS[model_name] lock = _LOAD_LOCKS.setdefault(model_name, asyncio.Lock()) async with lock: if model_name in _SESSIONS: + _touch(model_name) return _SESSIONS[model_name] _MODEL_STATE[model_name] = "loading" _MODEL_ERROR.pop(model_name, None) @@ -282,6 +346,7 @@ async def ensure_session(model_name: str): raise _SESSIONS[model_name] = session _MODEL_STATE[model_name] = "ready" + _touch(model_name) log.info("Model %s ready in %.1fs", model_name, time.time() - t0) return session @@ -292,6 +357,32 @@ async def ensure_session(model_name: str): app = FastAPI(title="remove-background-local", version=APP_VERSION) +_EVICTOR_TASK: asyncio.Task | None = None + + +@app.on_event("startup") +async def _start_evictor() -> None: + global _EVICTOR_TASK + if MODEL_IDLE_TTL > 0 and _EVICTOR_TASK is None: + _EVICTOR_TASK = asyncio.create_task(_evictor_loop()) + log.info( + "Model evictor running: idle TTL %ds, check every %ds (pinned=%s)", + MODEL_IDLE_TTL, MODEL_EVICTOR_INTERVAL, _PINNED_MODEL, + ) + + +@app.on_event("shutdown") +async def _stop_evictor() -> None: + global _EVICTOR_TASK + if _EVICTOR_TASK is not None: + _EVICTOR_TASK.cancel() + try: + await _EVICTOR_TASK + except (asyncio.CancelledError, Exception): # noqa: BLE001 + pass + _EVICTOR_TASK = None + + if STATIC_DIR.exists(): app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static") @@ -311,7 +402,9 @@ async def health(): "ok": True, "version": APP_VERSION, "default_model": DEFAULT_MODEL, + "pinned_model": _PINNED_MODEL, "loaded_models": list(_SESSIONS.keys()), + "idle_ttl_seconds": MODEL_IDLE_TTL, } @@ -364,11 +457,55 @@ async def _bg(): return {"model": model, "state": state_of(model), "size_mb": MODEL_SIZES_MB.get(model)} +@app.post("/unload_model") +async def unload_model(model: str = Form(...)): + """Drop a loaded model from RAM without touching its on-disk .onnx cache. + + Useful for freeing memory when a model is no longer needed. The next + request that uses it will reload it from the local cache (no re-download). + """ + _check_model(model) + unloaded = _evict(model) + return {"model": model, "unloaded": unloaded, "state": state_of(model)} + + +@app.post("/set_default_model") +async def set_default_model(model: str = Form(...), warmup: bool = Form(False)): + """Pin a new active default model and evict every other loaded model. + + The UI calls this when the user changes the model in the main dropdown so + that switching from BiRefNet to ISNet (or vice-versa) doesn't leave both + sitting in RAM. Per-image overrides (/remove with transient=true) do NOT + use this — they go through the idle TTL instead. + """ + global _PINNED_MODEL + _check_model(model) + _PINNED_MODEL = model + evicted = [] + for name in list(_SESSIONS): + if name != model and _evict(name): + evicted.append(name) + if warmup and model not in _SESSIONS and state_of(model) != "loading": + async def _bg(): + try: + await ensure_session(model) + except Exception: # noqa: BLE001 + pass + asyncio.create_task(_bg()) + return { + "pinned": _PINNED_MODEL, + "evicted": evicted, + "loaded": list(_SESSIONS.keys()), + "state": state_of(model), + } + + @app.post("/delete_model") async def delete_model(model: str = Form(...)): """Delete a model's cached .onnx file from disk and unload it from memory.""" _check_model(model) _SESSIONS.pop(model, None) + _LAST_USED.pop(model, None) _MODEL_STATE.pop(model, None) _MODEL_ERROR.pop(model, None) path = model_file(model) @@ -397,8 +534,15 @@ async def remove_background( alpha_matting_foreground_threshold: int = Form(240), alpha_matting_background_threshold: int = Form(10), alpha_matting_erode_size: int = Form(10), + transient: bool = Form(False), ): - """Process an image and return a PNG with a transparent background.""" + """Process an image and return a PNG with a transparent background. + + When ``transient`` is true the request is treated as a one-off override: + the model is loaded if needed but the pinned default is left alone, so the + idle TTL will reclaim this model later. Used for per-image reprocessing + with a different model from the UI. + """ # Validate size raw = await image.read() if len(raw) == 0: @@ -418,8 +562,9 @@ async def remove_background( raise HTTPException(status_code=400, detail=f"Invalid image: {exc}") log.info( - "Processing %s (%dx%d, %.1f KB) with model=%s", + "Processing %s (%dx%d, %.1f KB) with model=%s%s", image.filename, img.width, img.height, len(raw) / 1024, model, + " (transient)" if transient else "", ) # Loads/downloads the model if needed (in a worker thread). @@ -444,6 +589,7 @@ async def remove_background( elapsed = time.time() - t0 png_bytes = await run_in_threadpool(_encode_png, out) + _touch(model) log.info("Done in %.2fs", elapsed) headers = { diff --git a/static/index.html b/static/index.html index 28d2d1a..e1f36ee 100644 --- a/static/index.html +++ b/static/index.html @@ -380,6 +380,12 @@

Setting up rm.background local

// ===================== Model status API ================================== async function getStatus(model) { const r = await fetch("/model_status?model=" + encodeURIComponent(model)); if (!r.ok) throw new Error("status check failed"); return r.json(); } async function warmup(model) { const fd = new FormData(); fd.append("model", model); await fetch("/warmup", { method: "POST", body: fd }); } +async function pinDefaultModel(model, warm) { + // Tells the server this is the active default; the server evicts any other + // loaded model from RAM so we don't pile up BiRefNet + ISNet + U2Net at once. + const fd = new FormData(); fd.append("model", model); if (warm) fd.append("warmup", "true"); + try { await fetch("/set_default_model", { method: "POST", body: fd }); } catch {} +} async function waitForModel(model, onTick) { let s = await getStatus(model); if (s.state === "ready") return s; await warmup(model); @@ -404,7 +410,7 @@

Setting up rm.background local

opt.className = "dd-opt" + (key === selectedModel ? " sel" : "") + (avail ? "" : " disabled"); const na = avail ? "" : `Not downloaded`; opt.innerHTML = `${info2.title}~${SIZES[key]} MB${na}${info2.tagline || ""}${ICON_CHECK}`; - if (avail) opt.addEventListener("click", () => { selectedModel = key; closeDropdown(); renderModelDropdown(); }); + if (avail) opt.addEventListener("click", () => { selectedModel = key; closeDropdown(); renderModelDropdown(); pinDefaultModel(key, false); }); else opt.title = "Download it on the Models page first"; menu.appendChild(opt); } @@ -518,6 +524,8 @@

Setting up rm.background local

if (st.state !== "ready") { job.state = "loading-model"; renderAll(); setStatus(`downloading ${job.model}…`, true); await waitForModel(job.model); } job.state = "processing"; renderAll(); setStatus("processing…", true); const fd = new FormData(); fd.append("image", job.file); fd.append("model", job.model); + // A one-off reprocess with a non-default model: keep the global pinned model alive. + if (job.transient || job.model !== selectedModel) fd.append("transient", "true"); if ($("am-enabled").checked) { fd.append("alpha_matting", "true"); fd.append("alpha_matting_foreground_threshold", $("am-fg").value); fd.append("alpha_matting_background_threshold", $("am-bg").value); fd.append("alpha_matting_erode_size", $("am-erode").value); } const r = await fetch("/remove", { method: "POST", body: fd }); if (!r.ok) { let d = "Server error"; try { d = (await r.json()).detail || d; } catch {} throw new Error(d); } @@ -531,6 +539,8 @@

Setting up rm.background local

function reprocessJob(job, newModel) { if (!confirm(`Reprocess "${job.name}" with ${INFO[newModel]?.title || newModel}?\n\nThis replaces the current result for this image.`)) return; job.model = newModel; job.state = "queued"; job.err = null; + // Per-image override: load the model but do not steal the pinned default. + job.transient = newModel !== selectedModel; if (job.outUrl) { URL.revokeObjectURL(job.outUrl); job.outUrl = null; } job.outBlob = null; job.ms = null; job.outKB = null; renderAll(); pump(); } diff --git a/tests/test_server.py b/tests/test_server.py index 9c3a23a..1600b1f 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -30,6 +30,9 @@ def test_health(client): body = r.json() assert body["ok"] is True assert body["default_model"] == server.DEFAULT_MODEL + assert body["pinned_model"] == server._PINNED_MODEL + assert body["idle_ttl_seconds"] == server.MODEL_IDLE_TTL + assert isinstance(body["loaded_models"], list) def test_models_payload(client): @@ -147,6 +150,153 @@ def test_delete_model_not_downloaded_ok(client, tmp_path, monkeypatch): assert r.json()["deleted"] is False +@pytest.fixture +def clean_sessions(): + """Reset the in-RAM session bookkeeping before and after a test.""" + saved_pinned = server._PINNED_MODEL + server._SESSIONS.clear() + server._LAST_USED.clear() + server._MODEL_STATE.clear() + server._MODEL_ERROR.clear() + yield + server._SESSIONS.clear() + server._LAST_USED.clear() + server._MODEL_STATE.clear() + server._MODEL_ERROR.clear() + server._PINNED_MODEL = saved_pinned + + +def test_unload_model_unknown_is_400(client): + r = client.post("/unload_model", data={"model": "does-not-exist"}) + assert r.status_code == 400 + + +def test_unload_model_drops_session(client, clean_sessions): + name = server.DEFAULT_MODEL + server._SESSIONS[name] = object() + server._LAST_USED[name] = 123.0 + server._MODEL_STATE[name] = "ready" + r = client.post("/unload_model", data={"model": name}) + assert r.status_code == 200 + body = r.json() + assert body["unloaded"] is True + assert body["state"] == "idle" + assert name not in server._SESSIONS + assert name not in server._LAST_USED + + +def test_unload_model_when_not_loaded(client, clean_sessions): + r = client.post("/unload_model", data={"model": server.DEFAULT_MODEL}) + assert r.status_code == 200 + assert r.json()["unloaded"] is False + + +def test_set_default_model_unknown_is_400(client): + r = client.post("/set_default_model", data={"model": "does-not-exist"}) + assert r.status_code == 400 + + +def test_set_default_model_evicts_others(client, clean_sessions): + keep = "isnet-general-use" + drop = "u2net" + server._SESSIONS[keep] = object() + server._SESSIONS[drop] = object() + server._LAST_USED[keep] = 1.0 + server._LAST_USED[drop] = 1.0 + r = client.post("/set_default_model", data={"model": keep}) + assert r.status_code == 200 + body = r.json() + assert body["pinned"] == keep + assert body["evicted"] == [drop] + assert keep in server._SESSIONS + assert drop not in server._SESSIONS + assert server._PINNED_MODEL == keep + + +def test_set_default_model_no_warmup_by_default(client, monkeypatch, clean_sessions): + called = {"n": 0} + + async def fake_ensure_session(model): + called["n"] += 1 + return object() + + monkeypatch.setattr(server, "ensure_session", fake_ensure_session) + r = client.post("/set_default_model", data={"model": server.DEFAULT_MODEL}) + assert r.status_code == 200 + # warmup defaults to false → no eager load. + assert called["n"] == 0 + + +def test_remove_transient_flag_is_accepted(client, monkeypatch, clean_sessions): + """The transient flag must round-trip and not break the happy path.""" + seen = {} + + async def fake_ensure_session(model): + seen["model"] = model + return object() + + def fake_remove(img, session=None, **kwargs): + return img.convert("RGBA") + + monkeypatch.setattr(server, "ensure_session", fake_ensure_session) + monkeypatch.setattr(server, "remove", fake_remove) + + r = client.post( + "/remove", + files={"image": ("photo.jpg", png_bytes(), "image/png")}, + data={"model": "u2net", "transient": "true"}, + ) + assert r.status_code == 200 + assert seen["model"] == "u2net" + # Pinned default must NOT be hijacked by a transient request. + assert server._PINNED_MODEL == server.DEFAULT_MODEL + + +def test_evict_helper_clears_state(clean_sessions): + name = server.DEFAULT_MODEL + server._SESSIONS[name] = object() + server._LAST_USED[name] = 7.0 + server._MODEL_STATE[name] = "ready" + assert server._evict(name) is True + assert name not in server._SESSIONS + assert name not in server._LAST_USED + assert name not in server._MODEL_STATE + assert server._evict(name) is False # idempotent + + +def test_evictor_sweep_drops_idle_non_pinned(monkeypatch, clean_sessions): + """A single sweep must drop idle non-pinned models past the TTL.""" + pinned = "isnet-general-use" + stale = "u2net" + fresh = "u2net_human_seg" + server._PINNED_MODEL = pinned + server._SESSIONS[pinned] = object() + server._SESSIONS[stale] = object() + server._SESSIONS[fresh] = object() + + monkeypatch.setattr(server, "MODEL_IDLE_TTL", 100) + server._LAST_USED[pinned] = 0.0 # would be evicted if it weren't pinned + server._LAST_USED[stale] = 0.0 + server._LAST_USED[fresh] = 1_000_000.0 # in the future + + monkeypatch.setattr(server.time, "monotonic", lambda: 1_000.0) + + evicted = server._evictor_sweep() + + assert evicted == [stale] + assert pinned in server._SESSIONS, "pinned must survive" + assert fresh in server._SESSIONS, "fresh must survive" + assert stale not in server._SESSIONS, "stale must be evicted" + + +def test_evictor_sweep_disabled_when_ttl_zero(monkeypatch, clean_sessions): + server._SESSIONS["u2net"] = object() + server._LAST_USED["u2net"] = 0.0 + monkeypatch.setattr(server, "MODEL_IDLE_TTL", 0) + assert server._evictor_sweep() == [] + assert "u2net" in server._SESSIONS + + def test_download_progress_helpers(monkeypatch): # A downloaded model reports full progress; helpers must not raise. monkeypatch.setattr(server, "is_downloaded", lambda name: True)