fix(vram): reclaim idle pool after session teardown + offload upload encoder

Session start/close cycles left the pod sitting on the session's
transient allocation peak: measured driver-level, ~3.6-4.7 GB idle
after a plain 60s session and ~6 GB after one that swapped to the
240s decoder profile, on a server whose true idle floor is ~1 GB.
No objects leak — teardown frees its last tensors asynchronously
(recv thread joined with timeout, polygraphy/TRT finalizer chains),
AFTER every in-band empty_cache() has already run, so the freed
blocks stay reserved in PyTorch's caching pool indefinitely. That
pool is invisible to TensorRT's cudaMalloc workspaces, so it eats
exactly the headroom the next session's engine loads and swap-time
stem extraction need — the slow 'OOM over time' pattern on fleet
pods.

- ws_adapter: idle VRAM janitor — when no session is registered and
  the pool holds >512 MiB of freed blocks, gc.collect() +
  empty_cache(). Idle-only by construction; never competes with the
  realtime loop. Cycle-tested: idle returns to ~1 GB after every
  session (was 3.6-6 GB).
- ws_adapter: final trim in handle_client after the body frame (the
  last session-ref holder) is gone.
- ws_adapter: the upload-encoder Session now offloads to CPU
  (offload_to_cpu + offload_dit_to_cpu). It previously pinned a full
  eager DiT+VAE+text-encoder copy (~6 GB) on the GPU permanently
  from the first upload onward, next to the streaming TRT engines.
  prepare_source hops weights per call; uploads measured 14-22s and
  settle at ~1.4 GB.
- Session: expose offload_dit_to_cpu (passthrough to ModelContext);
  without it offload_to_cpu lets the DiT go GPU-resident on first
  use.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: leszko

acestep/engine/session.py

@@ -160,6 +160,7 @@ def __init__(
         vae_backend: str = "eager",
         use_flash_attention: bool = True,
         offload_to_cpu: bool = False,
+        offload_dit_to_cpu: bool = False,
         offload_text_encoder: bool = False,
         quantization: Optional[str] = None,
         trt_engines: Optional[dict[str, str]] = None,
@@ -181,6 +182,13 @@ def __init__(
             trt_engines: Engine paths, used iff a *_backend is "tensorrt".
                 Keys: "decoder", "vae_encode", "vae_decode". Use
                 acestep.paths.default_trt_engines() for the canonical paths.
+            offload_to_cpu: Build every model on CPU and hop weights to
+                the GPU per call. NOTE: without ``offload_dit_to_cpu``
+                the DiT becomes GPU-resident on its FIRST use and stays
+                (see ``ModelContext._load_model_context``); set both for
+                a session whose steady-state VRAM must stay ~0.
+            offload_dit_to_cpu: Move the DiT back to CPU after each use
+                instead of letting it go resident on first touch.
             offload_text_encoder: Override text encoder placement policy.
                 Defaults to ``False`` so prompt edits do not pay CPU/GPU
                 transfer cost. Set ``True`` for lower steady VRAM usage.
@@ -249,6 +257,7 @@ def _discover_loras_safely():
                 device=device,
                 use_flash_attention=use_flash_attention,
                 offload_to_cpu=offload_to_cpu,
+                offload_dit_to_cpu=offload_dit_to_cpu,
                 offload_text_encoder=offload_text_encoder,
                 quantization=quantization,
                 **ctx_flags,

demos/realtime_motion_graph_web/server.py

@@ -590,7 +590,13 @@ def main():
         # Defer the heavy import until we know we need it. Pulling this in
         # loads torch + acestep + TRT machinery; in --no-backend we never
         # touch any of it.
-        from .ws_adapter import handle_client
+        from .ws_adapter import handle_client, start_idle_vram_janitor
+
+        # Session teardown frees its last GPU tensors asynchronously;
+        # the janitor returns the caching allocator's freed pool to the
+        # driver whenever the pod sits idle (see its docstring for the
+        # measured numbers).
+        start_idle_vram_janitor()
 
         def ws_handler(ws):
             handle_client(

demos/realtime_motion_graph_web/ws_adapter.py

@@ -98,6 +98,77 @@
 from .protocol import COMMAND_NAMES, SAMPLE_RATE, coerce_command_payload
 
 
+# ---------------------------------------------------------------------------
+# Idle VRAM janitor
+# ---------------------------------------------------------------------------
+
+_JANITOR_STARTED = False
+_JANITOR_LOCK = threading.Lock()
+
+
+def start_idle_vram_janitor(
+    *,
+    interval_s: float = 5.0,
+    min_trim_bytes: int = 512 * 2**20,
+) -> None:
+    """Trim the CUDA caching allocator while no session is live.
+
+    Session teardown frees its last tensors asynchronously: the recv
+    thread is joined with a timeout, and polygraphy/TRT finalizer
+    chains can release buffers seconds after the connection handler
+    returned — after every in-band ``empty_cache()`` call has already
+    run. Those late frees land in PyTorch's caching pool and stay
+    reserved against the driver for as long as the pod idles
+    (measured: ~3 GB after a plain 60 s session, ~5 GB after one that
+    swapped to the 240 s decoder profile). That reserved-but-unused
+    pool is invisible to TensorRT, whose workspace comes from
+    cudaMalloc, so it directly eats the headroom the next session's
+    engine loads and stem extraction need.
+
+    A point-in-time trim can't win that race, so this janitor owns it:
+    every ``interval_s`` it checks that no session is registered and,
+    when the pool holds more than ``min_trim_bytes`` of freed blocks,
+    runs ``gc.collect()`` + ``torch.cuda.empty_cache()``. Idle-only by
+    construction — it never competes with a live session's allocator.
+    """
+    global _JANITOR_STARTED
+    with _JANITOR_LOCK:
+        if _JANITOR_STARTED:
+            return
+        _JANITOR_STARTED = True
+
+    def _loop() -> None:
+        import gc
+
+        while True:
+            time.sleep(interval_s)
+            try:
+                if session_registry.list_handles():
+                    continue
+                if not torch.cuda.is_available():
+                    continue
+                reserved = torch.cuda.memory_reserved()
+                allocated = torch.cuda.memory_allocated()
+                if reserved - allocated < min_trim_bytes:
+                    continue
+                gc.collect()
+                torch.cuda.empty_cache()
+                freed = reserved - torch.cuda.memory_reserved()
+                logger.info(
+                    "idle_vram_trim freed_mib={:.0f} reserved_mib={:.0f} "
+                    "allocated_mib={:.0f}",
+                    freed / 2**20,
+                    torch.cuda.memory_reserved() / 2**20,
+                    torch.cuda.memory_allocated() / 2**20,
+                )
+            except Exception as exc:  # never kill the janitor
+                logger.warning("idle_vram_trim_failed error={}", exc)
+
+    threading.Thread(
+        target=_loop, name="idle-vram-janitor", daemon=True,
+    ).start()
+
+
 # ---------------------------------------------------------------------------
 # Canonical user-upload packet processing
 # ---------------------------------------------------------------------------
@@ -122,11 +193,22 @@ def _upload_encoder_session(checkpoint: str) -> Session:
         session = _UPLOAD_ENCODERS.get(checkpoint)
         if session is None:
             logger.info("upload_encoder_load_start checkpoint={}", checkpoint)
+            # offload_to_cpu: this Session lives for the process and only
+            # serves prepare_source (VAE encode + semantic extract) for
+            # uploads — both paths hop weights to the GPU per call via
+            # _load_model_context and hop them back after. Keeping the
+            # weights CPU-resident turns what used to be a PERMANENT
+            # ~7 GB GPU copy (full eager DiT + VAE + text encoder, pinned
+            # from the first upload onward, next to the streaming
+            # session's TRT engines) into a few seconds of H2D/D2H per
+            # upload.
             session = Session(
                 project_root=str(checkpoints_dir()),
                 config_path=checkpoint,
                 decoder_backend="eager",
                 vae_backend="eager",
+                offload_to_cpu=True,
+                offload_dit_to_cpu=True,
             )
             _UPLOAD_ENCODERS[checkpoint] = session
             logger.info("upload_encoder_loaded checkpoint={}", checkpoint)
@@ -278,14 +360,33 @@ def handle_client(
     this wrapper exists only to own a single ``ExitStack`` so the
     contextvar tokens bound for session / track unwind in reverse
     order on every exit path."""
-    with contextlib.ExitStack() as ctx_stack:
-        _handle_client_body(
-            ws, ctx_stack,
-            decoder_backend=decoder_backend,
-            vae_backend=vae_backend,
-            checkpoint=checkpoint,
-            offload_text_encoder=offload_text_encoder,
-        )
+    try:
+        with contextlib.ExitStack() as ctx_stack:
+            _handle_client_body(
+                ws, ctx_stack,
+                decoder_backend=decoder_backend,
+                vae_backend=vae_backend,
+                checkpoint=checkpoint,
+                offload_text_encoder=offload_text_encoder,
+            )
+    finally:
+        # Final allocator trim, AFTER the body frame (the last holder of
+        # the session, its state, codec, and recv-thread closures) is
+        # gone. ``Session.close()`` runs its own gc + empty_cache, but at
+        # that point this connection still references those objects, so
+        # their pool blocks are live and the trim can't return them.
+        # Once the body returns they are garbage — without this trim the
+        # caching allocator keeps the session's transient peak reserved
+        # (~3 GB after a 60s session, ~6 GB after a 240s-profile one,
+        # measured driver-level) for as long as the pod idles, which is
+        # exactly the headroom the next session's engine loads and stem
+        # extraction need. Reserved-but-unallocated VRAM also can't be
+        # used by TensorRT, whose workspace comes from cudaMalloc.
+        import gc
+
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
 
 def _handle_client_body(