diff --git a/invokeai/backend/model_manager/probe.py b/invokeai/backend/model_manager/probe.py index 8f33e4b49f..4aef281bac 100644 --- a/invokeai/backend/model_manager/probe.py +++ b/invokeai/backend/model_manager/probe.py @@ -10,7 +10,7 @@ from picklescan.scanner import scan_file_path import invokeai.backend.util.logging as logger from invokeai.app.util.misc import uuid_string from invokeai.backend.model_hash.model_hash import HASHING_ALGORITHMS, ModelHash -from invokeai.backend.util.util import SilenceWarnings +from invokeai.backend.util.silence_warnings import SilenceWarnings from .config import ( AnyModelConfig, diff --git a/invokeai/backend/util/silence_warnings.py b/invokeai/backend/util/silence_warnings.py index 4c566ba759..0cd6d0738d 100644 --- a/invokeai/backend/util/silence_warnings.py +++ b/invokeai/backend/util/silence_warnings.py @@ -1,29 +1,36 @@ -"""Context class to silence transformers and diffusers warnings.""" - import warnings -from typing import Any +from contextlib import ContextDecorator -from diffusers import logging as diffusers_logging +from diffusers.utils import logging as diffusers_logging from transformers import logging as transformers_logging -class SilenceWarnings(object): - """Use in context to temporarily turn off warnings from transformers & diffusers modules. +# Inherit from ContextDecorator to allow using SilenceWarnings as both a context manager and a decorator. +class SilenceWarnings(ContextDecorator): + """A context manager that disables warnings from transformers & diffusers modules while active. + As context manager: + ``` with SilenceWarnings(): # do something + ``` + + As decorator: + ``` + @SilenceWarnings() + def some_function(): + # do something + ``` """ - def __init__(self) -> None: - self.transformers_verbosity = transformers_logging.get_verbosity() - self.diffusers_verbosity = diffusers_logging.get_verbosity() - def __enter__(self) -> None: + self._transformers_verbosity = transformers_logging.get_verbosity() + self._diffusers_verbosity = diffusers_logging.get_verbosity() transformers_logging.set_verbosity_error() diffusers_logging.set_verbosity_error() warnings.simplefilter("ignore") - def __exit__(self, *args: Any) -> None: - transformers_logging.set_verbosity(self.transformers_verbosity) - diffusers_logging.set_verbosity(self.diffusers_verbosity) + def __exit__(self, *args) -> None: + transformers_logging.set_verbosity(self._transformers_verbosity) + diffusers_logging.set_verbosity(self._diffusers_verbosity) warnings.simplefilter("default") diff --git a/invokeai/backend/util/util.py b/invokeai/backend/util/util.py index 1ee89dcc66..b3466ddba9 100644 --- a/invokeai/backend/util/util.py +++ b/invokeai/backend/util/util.py @@ -3,12 +3,9 @@ import io import os import re import unicodedata -import warnings from pathlib import Path -from diffusers import logging as diffusers_logging from PIL import Image -from transformers import logging as transformers_logging # actual size of a gig GIG = 1073741824 @@ -80,21 +77,3 @@ class Chdir(object): def __exit__(self, *args): os.chdir(self.original) - - -class SilenceWarnings(object): - """Context manager to temporarily lower verbosity of diffusers & transformers warning messages.""" - - def __enter__(self): - """Set verbosity to error.""" - self.transformers_verbosity = transformers_logging.get_verbosity() - self.diffusers_verbosity = diffusers_logging.get_verbosity() - transformers_logging.set_verbosity_error() - diffusers_logging.set_verbosity_error() - warnings.simplefilter("ignore") - - def __exit__(self, type, value, traceback): - """Restore logger verbosity to state before context was entered.""" - transformers_logging.set_verbosity(self.transformers_verbosity) - diffusers_logging.set_verbosity(self.diffusers_verbosity) - warnings.simplefilter("default")