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:
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user