mem: free intermediate arrays during YoloX inference

[KRRT7]

Free origin_img, img/ort_inputs, output, and the PIL pixel buffer at the points where they become dead in image_processing(), instead of letting them linger until function return.

The two main wins:

1. origin_img — the full-resolution numpy copy of the input PIL image stays alive through the entire ONNX session.run() call. del origin_img frees it before inference.
2. The PIL image itself — after np.array(image) copies the pixel data, the PIL buffer is no longer needed. image.close() frees it immediately while preserving PIL metadata (.width, .height, .format, .size).

Savings are proportional to image size: larger pages (higher DPI renders) carry bigger unused buffers through inference.

Benchmark

Azure Standard_D8s_v5 — 8 vCPU Intel Xeon Platinum 8473C, 32 GiB RAM, Python 3.12.12

Simulated ONNX session (35 MiB workspace), 1700x2200 letter-size image at 200 DPI.

Memory (peak per test)

Ref	Peak Memory	Allocations	Delta
b851f094d133 (base)	72.0 MiB	125
333c6d728837 (head)	47.0 MiB	119	-35%

Peak memory drops 25 MiB (-35%) by freeing dead buffers before ONNX inference. Timing is neutral (within noise).

Timing — image_processing (1700x2200 image, fake ONNX session)

Ref	Min	Median	Mean	OPS	Rounds
b851f094d133 (base)	388.2ms	392.4ms	392.8ms	2.55 ops/s	5
333c6d728837 (head)	399.2ms	406.3ms	412.6ms	2.42 ops/s	5
Speedup	0.97x	0.97x	0.95x	0.95x

Timing is within noise — the del statements and image.close() add negligible overhead vs. the ONNX inference wall time.

---

Generated by codeflash optimization agent

Reproduce the benchmark locally

# Full comparison (timing + memory):
uv run codeflash compare b851f094d133 333c6d728837 --memory \
  --functions 'unstructured_inference/models/yolox.py::image_processing' \
  --inject benchmarks/test_benchmark_yolox.py \
  --inject benchmarks/__init__.py \
  --inject pyproject.toml

Benchmark test source

"""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)

Test plan

• codeflash compare confirms -35% peak memory (72.0 MiB -> 47.0 MiB)
• Timing neutral (0.95-0.97x, within noise)
• Allocations reduced from 125 to 119 (-6)
• Existing unit tests pass