mem: free intermediate arrays during YoloX inference

CHANGELOG.md

@@ -1,3 +1,8 @@
+## 1.6.1
+
+### Enhancement
+- Free intermediate arrays (`origin_img`, `img`, `ort_inputs`, `output`) and PIL pixel buffer at dead points during YoloX `image_processing()` to reduce peak memory during inference
+
 ## 1.6.0
 
 ### Fix

benchmarks/test_benchmark_yolox.py

@@ -0,0 +1,63 @@
+"""Benchmark for YoloX image_processing() memory optimization.
+
+Uses a fake ONNX session to isolate the memory behavior of image_processing()
+without requiring the real model weights. The fake session allocates a realistic
+35 MiB workspace to simulate ONNX inference memory pressure.
+"""
+
+import numpy as np
+from PIL import Image as PILImage
+
+from unstructured_inference.models.yolox import UnstructuredYoloXModel
+
+
+class _FakeInput:
+    def __init__(self) -> None:
+        self.name = "input"
+
+
+class _FakeSession:
+    """Simulates an ONNX inference session with realistic memory allocation."""
+
+    def get_inputs(self):
+        return [_FakeInput()]
+
+    def run(self, _names, _inputs):
+        workspace = np.empty((35 * 1024 * 1024,), dtype=np.uint8)  # 35 MiB  # noqa: F841
+        # input_shape (1024,768), strides [8,16,32] → 128*96 + 64*48 + 32*24 = 16128
+        return [np.random.randn(1, 16128, 16).astype(np.float32)]
+
+
+def make_model() -> UnstructuredYoloXModel:
+    model = object.__new__(UnstructuredYoloXModel)
+    model.model = _FakeSession()
+    model.model_path = "yolox_fake"
+    model.layout_classes = {
+        0: "Caption",
+        1: "Footnote",
+        2: "Formula",
+        3: "List-item",
+        4: "Page-footer",
+        5: "Page-header",
+        6: "Picture",
+        7: "Section-header",
+        8: "Table",
+        9: "Text",
+        10: "Title",
+    }
+    return model
+
+
+# Letter-size page at 200 DPI — the default render resolution
+def make_letter_200dpi() -> PILImage.Image:
+    return PILImage.fromarray(np.random.randint(0, 255, (2200, 1700, 3), dtype=np.uint8))
+
+
+def run_image_processing():
+    model = make_model()
+    img = make_letter_200dpi()
+    return model.image_processing(img)
+
+
+def test_benchmark_yolox_image_processing(benchmark):
+    benchmark(run_image_processing)

pyproject.toml

@@ -128,5 +128,8 @@ filterwarnings = [
     "ignore::DeprecationWarning",
 ]
 
+[tool.codeflash]
+benchmarks-root = "benchmarks"
+
 [tool.coverage.report]
 fail_under = 90

unstructured_inference/__version__.py

@@ -1 +1 @@
-__version__ = "1.6.0"  # pragma: no cover
+__version__ = "1.6.1"  # pragma: no cover

unstructured_inference/models/yolox.py

@@ -107,13 +107,17 @@ def image_processing(
         # TODO (benjamin): check other shapes for inference
         input_shape = (1024, 768)
         origin_img = np.array(image)
+        image.close()
         img, ratio = preprocess(origin_img, input_shape)
+        del origin_img  # Free full-size image array before ONNX inference
         session = self.model
 
         ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
         output = session.run(None, ort_inputs)
+        del img, ort_inputs  # Free preprocessed inputs after inference
         # TODO(benjamin): check for p6
         predictions = demo_postprocess(output[0], input_shape, p6=False)[0]
+        del output
 
         boxes = predictions[:, :4]
         scores = predictions[:, 4:5] * predictions[:, 5:]