From b9fe63a72e307d573c8bd4d4fb5ca426ee953f13 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jos=C3=A9=20Luis=20C=C3=B3rdova?= <107780279+Damkohler@users.noreply.github.com> Date: Thu, 17 Sep 2026 22:09:22 -0400 Subject: [PATCH] Handle Joy cast diagnostics emitted through bitsandbytes logging --- docs/joy-8bit-installation.md | 12 ++++-- engines/captionforge_joy_warnings.py | 34 ++++++++++++++- tests/test_joy_hardening.py | 64 ++++++++++++++++++++++++++++ 3 files changed, 105 insertions(+), 5 deletions(-) diff --git a/docs/joy-8bit-installation.md b/docs/joy-8bit-installation.md index e360103..74160e8 100644 --- a/docs/joy-8bit-installation.md +++ b/docs/joy-8bit-installation.md @@ -27,7 +27,7 @@ References: ## Warning scope The old import-time filters were removed. During Balanced (8-bit) generation -only, CaptionForge ignores this exact `UserWarning` from +only, CaptionForge ignores this exact `UserWarning` or WARNING log record from `bitsandbytes.autograd._functions`: ```text @@ -36,8 +36,12 @@ MatMul8bitLt: inputs will be cast from torch.bfloat16 to float16 during quantiza FP32 cast warnings, different messages, other modules, other warning categories, processor/loading/cleanup diagnostics and Default-mode warnings remain visible. -The caller's filters are restored even if generation fails. This does not alter -tensors, quantization parameters, model outputs or logging handlers. +The caller's filters are restored even if generation fails. Recent bitsandbytes versions emit this through `logger.warning` instead of +`warnings.warn`; both routes are covered. The logging filter matches the fully +formatted message and originating logger, and applies only to the inference +thread. It is removed on normal exit or failure, without changing logger levels, +handlers, propagation, or pre-existing filters. This does not alter tensors, +quantization parameters or model outputs. Python 3.10-3.12 warning filters are process-wide while a `catch_warnings` scope is active. An identical warning from concurrent bitsandbytes work could therefore @@ -79,7 +83,7 @@ CPU-only focused checks (Python 3.11/3.12; Python 3.10 also needs `tomli`): python -m unittest discover -s tests -p test_joy_hardening.py -v ``` -They cover six repeated caption warning bursts, near-match diagnostics, +They cover six repeated caption bursts through both warnings and logging, near-match diagnostics, Default mode, filter restoration, exceptions, and matching dependency lists. Generation tests execute the real `caption_pil` method with test doubles for processor, model and torch; they do not test numerical inference. diff --git a/engines/captionforge_joy_warnings.py b/engines/captionforge_joy_warnings.py index 5d24ece..692678d 100644 --- a/engines/captionforge_joy_warnings.py +++ b/engines/captionforge_joy_warnings.py @@ -1,9 +1,32 @@ """Scoped handling of Joy's expected bitsandbytes activation-cast warning.""" from contextlib import contextmanager +import logging +import threading import warnings +_CAST_MESSAGE = ( + "MatMul8bitLt: inputs will be cast from torch.bfloat16 " + "to float16 during quantization" +) +_BNB_LOGGER = "bitsandbytes.autograd._functions" + + +class _JoyCastLogFilter(logging.Filter): + def __init__(self): + super().__init__() + self.thread_id = threading.get_ident() + + def filter(self, record): + return not ( + record.name == _BNB_LOGGER + and record.levelno == logging.WARNING + and record.thread == self.thread_id + and record.getMessage() == _CAST_MESSAGE + ) + + @contextmanager def joy_8bit_inference_warnings(enabled: bool): """Preserve all diagnostics except the known BF16 cast during Joy 8-bit inference. @@ -27,4 +50,13 @@ def joy_8bit_inference_warnings(enabled: bool): category=UserWarning, module=r"\Abitsandbytes\.autograd\._functions\Z", ) - yield + # Recent bitsandbytes versions use logging rather than warnings.warn. + # Attach to the emitting logger so propagation/ComfyUI handlers remain + # untouched. Limit this route to the current inference thread as well. + logger = logging.getLogger(_BNB_LOGGER) + log_filter = _JoyCastLogFilter() + logger.addFilter(log_filter) + try: + yield + finally: + logger.removeFilter(log_filter) diff --git a/tests/test_joy_hardening.py b/tests/test_joy_hardening.py index 89a77ba..ceeeacc 100644 --- a/tests/test_joy_hardening.py +++ b/tests/test_joy_hardening.py @@ -2,6 +2,8 @@ import ast from contextlib import nullcontext import importlib.util +import logging +import threading from pathlib import Path import types import unittest @@ -85,6 +87,64 @@ class WarningScopeTests(unittest.TestCase): emit('unrelated') +class LoggingScopeTests(unittest.TestCase): + def setUp(self): + self.logger = logging.getLogger(MODULE) + self.before = (list(self.logger.filters), self.logger.level, list(self.logger.handlers), self.logger.propagate) + + def tearDown(self): + self.assertEqual((list(self.logger.filters), self.logger.level, list(self.logger.handlers), self.logger.propagate), self.before) + + def test_six_logging_bursts_keep_other_diagnostics(self): + with self.assertLogs(MODULE, level='WARNING') as caught: + for _ in range(6): + with scope(True): + for _ in range(100): + self.logger.warning('MatMul8bitLt: inputs will be cast from %s to float16 during quantization', 'torch.bfloat16') + self.logger.warning('Useful diagnostic') + self.assertEqual([r.getMessage() for r in caught.records], ['Useful diagnostic'] * 6) + + def test_near_matches_levels_child_logger_and_other_thread(self): + with self.assertLogs(MODULE, level='INFO') as caught: + with scope(True): + self.logger.warning(MESSAGE.replace('bfloat16', 'float32')) + self.logger.warning(MESSAGE + ' extra') + self.logger.error(MESSAGE) + self.logger.info(MESSAGE) + logging.getLogger(MODULE + '.other').warning(MESSAGE) + thread = threading.Thread(target=lambda: self.logger.warning(MESSAGE)) + thread.start() + thread.join() + self.assertEqual(len(caught.records), 6) + + def test_default_after_scope_and_failure(self): + with self.assertLogs(MODULE, level='WARNING') as caught: + with scope(False): + self.logger.warning(MESSAGE) + with self.assertRaises(ValueError): + with scope(True): + self.logger.warning(MESSAGE) + raise ValueError('generation failed') + self.logger.warning(MESSAGE) + self.assertEqual(len(caught.records), 2) + + def test_existing_filters_preserved_in_nested_scope(self): + existing = logging.Filter() + self.logger.addFilter(existing) + try: + with self.assertLogs(MODULE, level='WARNING') as caught: + with scope(True): + with scope(True): + self.logger.warning(MESSAGE) + self.logger.warning(MESSAGE) + self.assertIn(existing, self.logger.filters) + self.logger.warning(MESSAGE) + self.assertEqual(len(caught.records), 1) + self.assertEqual(self.logger.filters, self.before[0] + [existing]) + finally: + self.logger.removeFilter(existing) + + class GenerationScopeTests(unittest.TestCase): """Execute the real caption_pil method with processor/torch/model test doubles.""" def run_generation(self, mode, fail=False): @@ -115,6 +175,10 @@ class GenerationScopeTests(unittest.TestCase): def generate(**kwargs): emit() + logging.getLogger(MODULE).warning( + 'MatMul8bitLt: inputs will be cast from %s to float16 during quantization', + 'torch.bfloat16', + ) emit('Useful generation diagnostic') if fail: raise ValueError('generation failed')