mirror of
https://github.com/invoke-ai/InvokeAI
synced 2024-08-30 20:32:17 +00:00
Tidy SilenceWarnings context manager:
- Fix type errors - Enable SilenceWarnings to be used as both a context manager and a decorator - Remove duplicate implementation - Check the initial verbosity on __enter__() rather than __init__()
This commit is contained in:
@ -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")
|
||||
|
Reference in New Issue
Block a user