From d5aa8772888d896ac64f59cc10756c290ac05fea Mon Sep 17 00:00:00 2001
From: Carter Richard <181430635+strbck@users.noreply.github.com>
Date: Wed, 13 May 2026 15:48:54 -0400
Subject: [PATCH] fix: reduce peak memory usage for large images and batch runs
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Three changes targeting the OOM crash reported in issue #13:

1. Reduce InsightFace det_size from (1024, 1024) to (640, 640).
   The detection grid was 2.6× larger than needed — 640 is the
   standard size for this model family and matches the value already
   used in the Dockerfile pre-warm step.

2. Add REFACER_MAX_DIMENSION cap (default 4096px longest edge).
   Images larger than this are downscaled with INTER_AREA before
   any inference runs. Prevents unbounded memory growth from very
   high-resolution inputs. Override via the REFACER_MAX_DIMENSION
   env var.

3. Call gc.collect() after each image in the batch loop.
   Ensures large intermediate arrays (full-resolution image buffers,
   ONNX activations) are released to the OS between images rather
   than accumulating across the run.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 models.py   |  2 +-
 pipeline.py | 17 +++++++++++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/models.py b/models.py
index ffcc5e9..5b04b73 100644
--- a/models.py
+++ b/models.py
@@ -59,7 +59,7 @@ def load_models(models_dir: str) -> ModelBundle:
         raise ImportError("insightface is not installed. Run: pip install insightface") from e
 
     app = FaceAnalysis(name="buffalo_l", providers=["CPUExecutionProvider"])
-    app.prepare(ctx_id=0, det_size=(1024, 1024))
+    app.prepare(ctx_id=0, det_size=(640, 640))
     logger.info("FaceAnalysis ready.")
 
     logger.info("Loading inswapper_128…")
diff --git a/pipeline.py b/pipeline.py
index 745f9a0..fef7a25 100644
--- a/pipeline.py
+++ b/pipeline.py
@@ -31,6 +31,7 @@
     print(stats)
 """
 
+import gc
 import logging
 import os
 import shutil
@@ -47,6 +48,10 @@
 
 SUPPORTED_EXTENSIONS = (".jpg", ".jpeg", ".png", ".tiff", ".webp")
 
+# Maximum pixel dimension (longest edge) before the image is downscaled.
+# Override with the REFACER_MAX_DIMENSION environment variable.
+MAX_DIMENSION = int(os.environ.get("REFACER_MAX_DIMENSION", "4096"))
+
 
 # ---------------------------------------------------------------------------
 # Result types
@@ -216,6 +221,17 @@ def _process_image(
         logger.error("Could not read %s, skipping.", filename)
         return result
 
+    # --- Downscale if needed ---
+    h, w = img.shape[:2]
+    if max(h, w) > MAX_DIMENSION:
+        scale = MAX_DIMENSION / max(h, w)
+        new_w, new_h = int(w * scale), int(h * scale)
+        logger.warning(
+            "%s — %dx%d exceeds MAX_DIMENSION=%d, downscaling to %dx%d",
+            filename, w, h, MAX_DIMENSION, new_w, new_h,
+        )
+        img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
+
     # --- Detect faces ---
     try:
         faces = models.app.get(img)
@@ -333,6 +349,7 @@ def run(
     for filename in filenames:
         logger.info("── Processing: %s", filename)
         image_result = _process_image(filename, input_dir, output_dir, models)
+        gc.collect()
         stats.image_results.append(image_result)
 
         stats.total_faces += image_result.faces_detected