fix(annotations): stop changing the process warnings filters during capture

header_dimensions wrapped Image.open in warnings.catch_warnings with an
error filter for DecompressionBombWarning. catch_warnings is not
thread-safe, so two overlapping captures could leave that error filter
installed for every later Pillow open in the ComfyUI process. Only
Pillow's DecompressionBombError is now mapped to HeaderTooLarge; a size
that only warns is refused by the existing inference pixel cap.
This commit is contained in:
Vito Sansevero
2026-10-22 09:49:30 -07:00
parent 4c8b24bac3
commit ba4dfa3681
2 changed files with 66 additions and 8 deletions
+57
View File
@@ -8,6 +8,7 @@ import tempfile
import threading
import time
import unittest
import warnings
import zlib
from pathlib import Path
from unittest.mock import patch
@@ -632,6 +633,62 @@ class CaptureStoreTestCase(unittest.TestCase):
)
assert_closed(self, self.parent)
def test_header_dimensions_leaves_the_warning_filters_alone(self):
# The capture workers run in the ComfyUI process: touching the
# process-global warnings filters from them races (catch_warnings is
# not thread-safe) and can leave an "error" filter installed for every
# later Pillow open in the host.
from PIL import Image
path = str(self.write("dims.png", png_bytes(7, 5)))
before = list(warnings.filters)
seen_during = []
real_open = Image.open
def recording_open(*args, **kwargs):
seen_during.append(list(warnings.filters))
return real_open(*args, **kwargs)
with patch.object(Image, "open", recording_open):
self.assertEqual(capture_module.header_dimensions(path), (7, 5))
self.assertEqual(seen_during, [before])
self.assertEqual(list(warnings.filters), before)
barrier = threading.Barrier(8)
errors = []
def worker():
try:
barrier.wait()
for _ in range(50):
capture_module.header_dimensions(path)
except Exception as exc: # pragma: no cover - reported below
errors.append(exc)
threads = [threading.Thread(target=worker) for _ in range(8)]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
self.assertEqual(errors, [])
self.assertEqual(list(warnings.filters), before)
assert_closed(self, self.parent)
def test_header_over_the_warning_size_is_refused_by_the_pixel_cap(self):
# 100 MP: past Pillow's warning size, under its error size. The
# header is read, then the inference pixel cap refuses it.
from PIL import Image
path = str(self.write("big.png", png_bytes(10_000, 10_000)))
with self.assertWarns(Image.DecompressionBombWarning):
self.assertEqual(capture_module.header_dimensions(path), (10_000, 10_000))
store = self.make_store()
with self.assertWarns(Image.DecompressionBombWarning):
cap = self.capture(store, "big.png")
self.assertFalse(cap.inferable)
self.assertIn("4096", cap.inferable_reason)
assert_closed(self, self.parent)
def test_video_capture_is_not_inferable(self):
data = b"\x00\x00\x00\x18ftypmp42" + os.urandom(32)
self.write("clip.mp4", data)
+9 -8
View File
@@ -26,7 +26,6 @@ import stat as stat_module
import tempfile
import threading
import time
import warnings
from dataclasses import dataclass
from typing import Callable
@@ -607,15 +606,17 @@ class HeaderTooLarge(HeaderUnreadable):
def header_dimensions(path: str) -> tuple[int, int]:
"""``(width, height)`` from the image header alone; the file is closed.
``Image.open`` reads only the header; a decompression-bomb warning is an
error here, so an absurd size is refused before anything is decoded.
``Image.open`` reads only the header. Pillow's decompression-bomb error
(twice its warning size) is ``HeaderTooLarge``; a size between the two
only warns, and the caller's inference pixel cap, far below Pillow's
warning size, refuses it. The process-global warnings filters are never
touched: this runs on capture workers inside the host process, and
``warnings.catch_warnings`` is not thread-safe.
"""
try:
with warnings.catch_warnings():
warnings.simplefilter("error", Image.DecompressionBombWarning)
with Image.open(path) as image:
width, height = image.size
except (Image.DecompressionBombError, Image.DecompressionBombWarning) as exc:
with Image.open(path) as image:
width, height = image.size
except Image.DecompressionBombError as exc:
raise HeaderTooLarge(str(exc)) from None
except Exception as exc: # any decoder failure means "no usable header"
raise HeaderUnreadable(str(exc)) from None