fix: stop RAM growth during full pipeline runs

Processing all pages (detect/OCR/translate/inpaint/render) grew RAM
monotonically and never released it. Three independent leaks:

- Frontend: useBlobImage created object URLs that were never revoked, so
  every sprite/inpaint/render blob pinned a full decoded image in the
  webview for the whole run. Revoke URLs on query-cache eviction and drop
  the blob gcTime from 10min to 1min so inactive pages release promptly.

- Metal: candle's Metal ops and the MPSGraph FFT allocate autoreleased
  Objective-C objects (command buffers, MPS intermediates), and the
  pipeline runs on tokio worker threads with no autorelease pool draining,
  so they accumulate for the entire run. Wrap the LaMa per-crop forward
  and the FFT calls in objc2 autorelease pools (no-op off Metal).

- Unbounded caches: the LaMa FFT-plan caches (keyed by crop shape) and the
  flux2 Qwen prompt cache lived on whole-run engine instances and were
  never evicted. Bound them with LruCache.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: simon-mds

Cargo.lock

@@ -22,6 +22,7 @@ candle-transformers = { workspace = true }
 candle-nn = { workspace = true }
 candle-flash-attn = { workspace = true, optional = true }
 tokenizers = { workspace = true }
+lru = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 tracing = { workspace = true }

crates/koharu-ml/Cargo.toml

@@ -1,7 +1,9 @@
+use lru::LruCache;
 use std::{
     collections::HashMap,
     fs::File,
     io::{Read, Seek},
+    num::NonZeroUsize,
     path::Path,
     sync::{Arc, Mutex},
 };
@@ -399,11 +401,18 @@ fn causal_padding_mask(
     Tensor::from_vec(values, (batch, 1, seq_len, seq_len), device)?.to_dtype(dtype)
 }
 
+/// Upper bound on cached prompt embeddings. Each entry retains a
+/// `(1, max_sequence_len, hidden)` tensor; keys are full prompt strings, which
+/// are effectively unique per translated block, so an unbounded map grew one
+/// large tensor per page for the whole run. Reuse across pages is rare, so a
+/// small LRU is plenty (it still absorbs the repeated empty/negative prompt).
+const PROMPT_CACHE_CAP: usize = 8;
+
 #[derive(Debug)]
 pub struct PromptEmbedder {
     tokenizer: Tokenizer,
     text_encoder: QwenTextEncoder,
-    cache: Mutex<HashMap<String, Arc<QwenPromptEmbeddings>>>,
+    cache: Mutex<LruCache<String, Arc<QwenPromptEmbeddings>>>,
     pad_token_id: u32,
     max_sequence_len: usize,
     hidden_layer_indices: Vec<usize>,
@@ -425,7 +434,9 @@ impl PromptEmbedder {
         Ok(Self {
             tokenizer,
             text_encoder,
-            cache: Mutex::new(HashMap::new()),
+            cache: Mutex::new(LruCache::new(
+                NonZeroUsize::new(PROMPT_CACHE_CAP).expect("cache capacity is non-zero"),
+            )),
             pad_token_id,
             max_sequence_len: DEFAULT_MAX_PROMPT_SEQUENCE_LEN,
             hidden_layer_indices: DEFAULT_HIDDEN_LAYERS.to_vec(),
@@ -471,7 +482,7 @@ impl PromptEmbedder {
         self.cache
             .lock()
             .expect("prompt cache poisoned")
-            .insert(key, embeddings.clone());
+            .put(key, embeddings.clone());
         Ok(embeddings)
     }
 }

crates/koharu-ml/src/flux2_klein/qwen.rs

@@ -8,11 +8,19 @@ use cudarc::{
     cufft::{result as cufft, sys},
     driver::{DevicePtr, DevicePtrMut},
 };
+use lru::LruCache;
 use std::{
-    collections::HashMap,
+    num::NonZeroUsize,
     sync::{Arc, Mutex, OnceLock},
 };
 
+/// Upper bound on distinct cached cuFFT plans. Keyed by FFT shape, which
+/// varies per LaMa crop resolution, so an unbounded map retains one
+/// `cufftHandle` per distinct crop size for the whole run. Cap it so stale
+/// handles are evicted (and destroyed via `CachedPlan::drop`) as new shapes
+/// appear. Mirrors the Metal path's `FFT_PLAN_CACHE_CAP`.
+const FFT_PLAN_CACHE_CAP: usize = 64;
+
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
 enum PlanKind {
     R2C,
@@ -41,9 +49,13 @@ impl Drop for CachedPlan {
     }
 }
 
-fn plan_cache() -> &'static Mutex<HashMap<PlanKey, Arc<CachedPlan>>> {
-    static CACHE: OnceLock<Mutex<HashMap<PlanKey, Arc<CachedPlan>>>> = OnceLock::new();
-    CACHE.get_or_init(|| Mutex::new(HashMap::new()))
+fn plan_cache() -> &'static Mutex<LruCache<PlanKey, Arc<CachedPlan>>> {
+    static CACHE: OnceLock<Mutex<LruCache<PlanKey, Arc<CachedPlan>>>> = OnceLock::new();
+    CACHE.get_or_init(|| {
+        Mutex::new(LruCache::new(
+            NonZeroUsize::new(FFT_PLAN_CACHE_CAP).expect("cache capacity is non-zero"),
+        ))
+    })
 }
 
 fn get_or_create_plan(
@@ -104,7 +116,7 @@ fn get_or_create_plan(
         handle,
         lock: Mutex::new(()),
     });
-    cache.insert(key, plan.clone());
+    cache.put(key, plan.clone());
     Ok(plan)
 }
 

crates/koharu-ml/src/lama/fft/cuda.rs

@@ -4,14 +4,28 @@ use candle_core::{
     bail,
     metal_backend::{DeviceId, MetalError, MetalStorage},
 };
-use objc2::{AnyThread, rc::Retained, runtime::ProtocolObject};
+use objc2::{
+    AnyThread,
+    rc::{Retained, autoreleasepool},
+    runtime::ProtocolObject,
+};
 use objc2_foundation::{NSArray, NSCopying, NSDictionary, NSNumber};
 use objc2_metal_performance_shaders::MPSDataType;
 use objc2_metal_performance_shaders_graph::{
     MPSGraph, MPSGraphFFTDescriptor, MPSGraphFFTScalingMode, MPSGraphTensor, MPSGraphTensorData,
     MPSGraphTensorDataDictionary,
 };
-use std::{cell::RefCell, collections::HashMap, ptr::NonNull};
+use lru::LruCache;
+use std::{cell::RefCell, num::NonZeroUsize, ptr::NonNull};
+
+/// Upper bound on distinct cached FFT plans (MPSGraphs). Each LaMa crop runs
+/// at its native resolution, and crop sizes vary per bubble/page, so this map
+/// is keyed by a continuously-varying shape. Without a cap it grows one
+/// retained `MPSGraph` per distinct crop size for the whole run — the source
+/// of the steady RAM climb when processing many pages. The working set of
+/// shapes touched by a single inpaint pass is small, so an LRU comfortably
+/// keeps hot plans while evicting (and freeing) stale ones.
+const FFT_PLAN_CACHE_CAP: usize = 64;
 
 fn nsarray_from_usize(values: &[usize]) -> Result<Retained<objc2_foundation::NSArray<NSNumber>>> {
     let nums: Vec<Retained<NSNumber>> = values
@@ -78,8 +92,10 @@ struct FftPlan {
 }
 
 thread_local! {
-    static FFT_PLANS: RefCell<HashMap<FftKey, FftPlan>> = RefCell::new(HashMap::new());
-    static COMMAND_QUEUES: RefCell<HashMap<DeviceId, Retained<ProtocolObject<dyn objc2_metal::MTLCommandQueue>>>> = RefCell::new(HashMap::new());
+    static FFT_PLANS: RefCell<LruCache<FftKey, FftPlan>> = RefCell::new(LruCache::new(
+        NonZeroUsize::new(FFT_PLAN_CACHE_CAP).expect("cache capacity is non-zero"),
+    ));
+    static COMMAND_QUEUES: RefCell<std::collections::HashMap<DeviceId, Retained<ProtocolObject<dyn objc2_metal::MTLCommandQueue>>>> = RefCell::new(std::collections::HashMap::new());
 }
 
 fn shared_command_queue(
@@ -101,7 +117,7 @@ fn shared_command_queue(
 
 fn fft_plan(key: FftKey) -> Result<FftPlan> {
     FFT_PLANS.with(|plans| {
-        if let Some(plan) = plans.borrow().get(&key) {
+        if let Some(plan) = plans.borrow_mut().get(&key) {
             return Ok(plan.clone());
         }
 
@@ -164,12 +180,25 @@ fn fft_plan(key: FftKey) -> Result<FftPlan> {
             output_shape,
         };
 
-        plans.borrow_mut().insert(key, plan.clone());
+        plans.borrow_mut().put(key, plan.clone());
         Ok(plan)
     })
 }
 
+/// Each MPSGraph run and the Cocoa factory calls below (`NSArray`,
+/// `NSDictionary`, `MPSGraphTensorData`, and the command buffers allocated
+/// inside `runWithMTLCommandQueue`) return autoreleased objects. The pipeline
+/// runs on tokio worker threads that have no autorelease pool draining, so
+/// without an explicit pool these temporaries — including sizeable Metal
+/// command buffers and MPS intermediates — accumulate for the entire run and
+/// are the dominant per-page RAM climb on the LaMa+Metal path. Wrapping each
+/// call drains them immediately; the returned `output_buffer` is candle-owned
+/// (held by an `Arc`), so it safely outlives the pool.
 pub fn rfft2(storage: &MetalStorage, layout: &Layout) -> Result<(MetalStorage, Shape)> {
+    autoreleasepool(|_| rfft2_impl(storage, layout))
+}
+
+fn rfft2_impl(storage: &MetalStorage, layout: &Layout) -> Result<(MetalStorage, Shape)> {
     let dims = layout.dims();
     if dims.len() != 4 {
         bail!("rfft2 expects rank-4 input, got {:?}", dims)
@@ -250,6 +279,15 @@ pub fn irfft2(
     storage: &MetalStorage,
     layout: &Layout,
     width: usize,
+) -> Result<(MetalStorage, Shape)> {
+    // See `rfft2` — drain autoreleased Metal/MPS temporaries per call.
+    autoreleasepool(|_| irfft2_impl(storage, layout, width))
+}
+
+fn irfft2_impl(
+    storage: &MetalStorage,
+    layout: &Layout,
+    width: usize,
 ) -> Result<(MetalStorage, Shape)> {
     let dims = layout.dims();
     if dims.len() != 5 || dims[4] != 2 {

crates/koharu-ml/src/lama/fft/metal.rs

@@ -21,6 +21,18 @@ const HF_REPO: &str = "mayocream/lama-manga";
 const BLOCK_WINDOW_RATIO: f64 = 1.7;
 const BLOCK_WINDOW_ASPECT_RATIO: f64 = 1.0;
 
+/// Run `f` inside an Objective-C autorelease pool on Metal builds so the Metal
+/// temporaries it creates are freed immediately; a plain call elsewhere.
+#[cfg(feature = "metal")]
+fn drain_autorelease_pool<T>(f: impl FnOnce() -> T) -> T {
+    objc2::rc::autoreleasepool(|_| f())
+}
+
+#[cfg(not(feature = "metal"))]
+fn drain_autorelease_pool<T>(f: impl FnOnce() -> T) -> T {
+    f()
+}
+
 type Xyxy = [u32; 4];
 
 koharu_runtime::declare_hf_model_package!(
@@ -164,9 +176,18 @@ impl Lama {
 
     #[instrument(level = "debug", skip_all)]
     fn inference_model(&self, image: &RgbImage, mask: &GrayImage) -> Result<RgbImage> {
-        let (image_tensor, mask_tensor) = self.preprocess(image, mask)?;
-        let output = self.forward(&image_tensor, &mask_tensor)?;
-        self.postprocess(&output)
+        // Drain autoreleased Metal temporaries per crop. candle's Metal ops
+        // (and the FFC layers' MPSGraph FFTs) allocate command buffers and MPS
+        // intermediates that are autoreleased; on the pool-less pipeline worker
+        // threads a leaked command buffer also pins every GPU buffer it
+        // referenced, so without a pool GPU/unified memory climbs crop after
+        // crop for the whole run. The returned `RgbImage` is heap-owned and
+        // safely outlives the pool.
+        drain_autorelease_pool(|| {
+            let (image_tensor, mask_tensor) = self.preprocess(image, mask)?;
+            let output = self.forward(&image_tensor, &mask_tensor)?;
+            self.postprocess(&output)
+        })
     }
 
     #[instrument(level = "debug", skip_all)]

crates/koharu-ml/src/lama/mod.rs

@@ -5,6 +5,14 @@ import { keepPreviousData, useQuery } from '@tanstack/react-query'
 import { getBlob } from '@/lib/api/default/default'
 import { convertToBlob } from '@/lib/io/blobConvert'
 
+// Each cached entry holds a full-resolution decoded image (raw bytes, or a
+// Blob pinned by an object URL). Batch runs touch a new set of hashes per page,
+// so a long gcTime keeps a large sliding window of full images resident at
+// once. Keep just enough to make navigating a few pages back instant; inactive
+// entries are evicted quickly (and their object URLs revoked — see
+// `queryClient.ts`). Actively-observed queries are never evicted regardless.
+const BLOB_GC_TIME = 60 * 1000
+
 const blobQueryOptions = (hash: string) => ({
   queryKey: ['blob', hash] as const,
   queryFn: async () => {
@@ -13,7 +21,7 @@ const blobQueryOptions = (hash: string) => ({
     return new Uint8Array(buf)
   },
   staleTime: Infinity,
-  gcTime: 10 * 60 * 1000,
+  gcTime: BLOB_GC_TIME,
   structuralSharing: false as const,
 })
 
@@ -44,7 +52,7 @@ const blobImageQueryOptions = (hash: string) => ({
     return url
   },
   staleTime: Infinity,
-  gcTime: 10 * 60 * 1000,
+  gcTime: BLOB_GC_TIME,
   structuralSharing: false as const,
 })
 

ui/hooks/useBlobData.ts

@@ -11,3 +11,20 @@ import { QueryClient } from '@tanstack/react-query'
  * calls and React Query hooks share the exact same cache.
  */
 export const queryClient = new QueryClient()
+
+// `useBlobImage` caches ready-to-paint object URLs (`blob:` URLs) as query
+// data. Object URLs pin their backing `Blob` (a full decoded image) in memory
+// until explicitly revoked, so a cached URL that is merely garbage-collected
+// still leaks the image. During a "process all pages" run every page mints new
+// sprite/inpaint/render blobs, so without this the webview's memory climbs one
+// full image per hash for the whole run. Revoke the URL when its query leaves
+// the cache, tying the Blob's lifetime to the cache entry's.
+queryClient.getQueryCache().subscribe((event) => {
+  if (event.type !== 'removed') return
+  const [kind] = event.query.queryKey as [unknown]
+  if (kind !== 'blobImage') return
+  const url = event.query.state.data
+  if (typeof url === 'string' && url.startsWith('blob:')) {
+    URL.revokeObjectURL(url)
+  }
+})