From ba4dfa3681278e618391ee024aa3e93a3b000dbb Mon Sep 17 00:00:00 2001
From: Vito Sansevero <vito.sansevero@gmail.com>
Date: Thu, 22 Oct 2026 09:49:30 -0700
Subject: [PATCH] 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.
---
 tests/test_media_capture.py | 57 +++++++++++++++++++++++++++++++++++++
 utils/media_capture.py      | 17 +++++------
 2 files changed, 66 insertions(+), 8 deletions(-)

diff --git a/tests/test_media_capture.py b/tests/test_media_capture.py
index 69e43fa5..a183ab36 100644
--- a/tests/test_media_capture.py
+++ b/tests/test_media_capture.py
@@ -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)
diff --git a/utils/media_capture.py b/utils/media_capture.py
index f9a54153..e8a49d6e 100644
--- a/utils/media_capture.py
+++ b/utils/media_capture.py
@@ -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