refactor(word): lazy-load DOCX images to reduce peak memory without changing output (#13233)

**Summary**
This PR tackles a significant memory bottleneck when processing
image-heavy Word documents. Previously, our pipeline eagerly decoded
DOCX images into `PIL.Image` objects, which caused high peak memory
usage. To solve this, I've introduced a **lazy-loading approach**:
images are now stored as raw blobs and only decoded exactly when and
where they are consumed.

This successfully reduces the memory footprint while keeping the parsing
output completely identical to before.

**What's Changed**
Instead of a dry file-by-file list, here is the logical breakdown of the
updates:

* **The Core Abstraction (`lazy_image.py`)**: Introduced `LazyDocxImage`
along with helper APIs to handle lazy decoding, image-type checks, and
NumPy compatibility. It also supports `.close()` and detached PIL access
to ensure safe lifecycle management and prevent memory leaks.
* **Pipeline Integration (`naive.py`, `figure_parser.py`, etc.)**:
Updated the general DOCX picture extraction to return these new lazy
images. Downstream consumers (like the figure/VLM flow and base64
encoding paths) now decode images right at the use site using detached
PIL instances, avoiding shared-instance side effects.
* **Compatibility Hooks (`operators.py`, `book.py`, etc.)**: Added
necessary compatibility conversions so these lazy images flow smoothly
through existing merging, filtering, and presentation steps without
breaking.

**Scope & What is Intentionally Left Out**
To keep this PR focused, I have restricted these changes strictly to the
**general Word pipeline** and its downstream consumers.
The `QA` and `manual` Word parsing pipelines are explicitly **not
modified** in this PR. They can be safely migrated to this new lazy-load
model in a subsequent, standalone PR.

**Design Considerations**
I briefly considered adding image compression during processing, but
decided against it to avoid any potential quality degradation in the
derived outputs. I also held off on a massive pipeline re-architecture
to avoid overly invasive changes right now.

**Validation & Testing**
I've tested this to ensure no regressions:

* Compared identical DOCX inputs before and after this branch: chunk
counts, extracted text, table HTML, and image descriptions match
perfectly.
* **Confirmed a noticeable drop in peak memory usage when processing
image-dense documents.** For a 30MB Word document containing 243 1080p
screenshots, memory consumption is reduced by approximately 1.5GB.

**Breaking Changes**
None.
This commit is contained in:
eviaaaaa
2026-02-28 11:22:31 +08:00
committed by GitHub
parent 4f0c892b32
commit fa71f8d0c7
8 changed files with 195 additions and 38 deletions

View File

@@ -24,19 +24,24 @@ from common.connection_utils import timeout
from rag.app.picture import vision_llm_chunk as picture_vision_llm_chunk
from rag.prompts.generator import vision_llm_figure_describe_prompt, vision_llm_figure_describe_prompt_with_context
from rag.nlp import append_context2table_image4pdf
from rag.utils.lazy_image import ensure_pil_image, open_image_for_processing, is_image_like
# need to delete before pr
def vision_figure_parser_figure_data_wrapper(figures_data_without_positions):
if not figures_data_without_positions:
return []
return [
(
(figure_data[1], [figure_data[0]]),
[(0, 0, 0, 0, 0)],
res = []
for figure_data in figures_data_without_positions:
img = ensure_pil_image(figure_data[1])
if not isinstance(img, Image.Image):
continue
res.append(
(
(img, [figure_data[0]]),
[(0, 0, 0, 0, 0)],
)
)
for figure_data in figures_data_without_positions
if isinstance(figure_data[1], Image.Image)
]
return res
def vision_figure_parser_docx_wrapper(sections, tbls, callback=None,**kwargs):
if not sections:
@@ -96,7 +101,7 @@ def vision_figure_parser_pdf_wrapper(tbls, callback=None, **kwargs):
if vision_model:
def is_figure_item(item):
return isinstance(item[0][0], Image.Image) and isinstance(item[0][1], list)
return is_image_like(item[0][0]) and isinstance(item[0][1], list)
figures_data = [item for item in tbls if is_figure_item(item)]
figure_contexts = []
@@ -134,6 +139,9 @@ def vision_figure_parser_docx_wrapper_naive(chunks, idx_lst, callback=None, **kw
if vision_model:
@timeout(30, 3)
def worker(idx, ck):
img, close_after = open_image_for_processing(ck.get("image"), allow_bytes=True)
if not isinstance(img, Image.Image):
return idx, ""
context_above = ck.get("context_above", "")
context_below = ck.get("context_below", "")
if context_above or context_below:
@@ -149,13 +157,20 @@ def vision_figure_parser_docx_wrapper_naive(chunks, idx_lst, callback=None, **kw
prompt = vision_llm_figure_describe_prompt()
logging.info(f"[VisionFigureParser] figure={idx} context_len=0 prompt=default")
description_text = picture_vision_llm_chunk(
binary=ck.get("image"),
vision_model=vision_model,
prompt=prompt,
callback=callback,
)
return idx, description_text
try:
description_text = picture_vision_llm_chunk(
binary=img,
vision_model=vision_model,
prompt=prompt,
callback=callback,
)
return idx, description_text
finally:
if close_after and isinstance(img, Image.Image):
try:
img.close()
except Exception:
pass
with ThreadPoolExecutor(max_workers=10) as executor:
futures = [
@@ -187,13 +202,15 @@ class VisionFigureParser:
# position
if len(item) == 2 and isinstance(item[0], tuple) and len(item[0]) == 2 and isinstance(item[1], list) and isinstance(item[1][0], tuple) and len(item[1][0]) == 5:
img_desc = item[0]
assert len(img_desc) == 2 and isinstance(img_desc[0], Image.Image) and isinstance(img_desc[1], list), "Should be (figure, [description])"
self.figures.append(img_desc[0])
img = ensure_pil_image(img_desc[0])
assert len(img_desc) == 2 and isinstance(img, Image.Image) and isinstance(img_desc[1], list), "Should be (figure, [description])"
self.figures.append(img)
self.descriptions.append(img_desc[1])
self.positions.append(item[1])
else:
assert len(item) == 2 and isinstance(item[0], Image.Image) and isinstance(item[1], list), f"Unexpected form of figure data: get {len(item)=}, {item=}"
self.figures.append(item[0])
img = ensure_pil_image(item[0])
assert len(item) == 2 and isinstance(img, Image.Image) and isinstance(item[1], list), f"Unexpected form of figure data: get {len(item)=}, {item=}"
self.figures.append(img)
self.descriptions.append(item[1])
def _assemble(self):