declare-lab/tango2

1# coding=utf-82# Copyright 2023 Optuna, Hugging Face3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15""" Logging utilities."""16 17import logging18import os19import sys20import threading21from logging import (22    CRITICAL,  # NOQA23    DEBUG,  # NOQA24    ERROR,  # NOQA25    FATAL,  # NOQA26    INFO,  # NOQA27    NOTSET,  # NOQA28    WARN,  # NOQA29    WARNING,  # NOQA30)31from typing import Optional32 33from tqdm import auto as tqdm_lib34 35 36_lock = threading.Lock()37_default_handler: Optional[logging.Handler] = None38 39log_levels = {40    "debug": logging.DEBUG,41    "info": logging.INFO,42    "warning": logging.WARNING,43    "error": logging.ERROR,44    "critical": logging.CRITICAL,45}46 47_default_log_level = logging.WARNING48 49_tqdm_active = True50 51 52def _get_default_logging_level():53    """54    If DIFFUSERS_VERBOSITY env var is set to one of the valid choices return that as the new default level. If it is55    not - fall back to `_default_log_level`56    """57    env_level_str = os.getenv("DIFFUSERS_VERBOSITY", None)58    if env_level_str:59        if env_level_str in log_levels:60            return log_levels[env_level_str]61        else:62            logging.getLogger().warning(63                f"Unknown option DIFFUSERS_VERBOSITY={env_level_str}, "64                f"has to be one of: { ', '.join(log_levels.keys()) }"65            )66    return _default_log_level67 68 69def _get_library_name() -> str:70    return __name__.split(".")[0]71 72 73def _get_library_root_logger() -> logging.Logger:74    return logging.getLogger(_get_library_name())75 76 77def _configure_library_root_logger() -> None:78    global _default_handler79 80    with _lock:81        if _default_handler:82            # This library has already configured the library root logger.83            return84        _default_handler = logging.StreamHandler()  # Set sys.stderr as stream.85        _default_handler.flush = sys.stderr.flush86 87        # Apply our default configuration to the library root logger.88        library_root_logger = _get_library_root_logger()89        library_root_logger.addHandler(_default_handler)90        library_root_logger.setLevel(_get_default_logging_level())91        library_root_logger.propagate = False92 93 94def _reset_library_root_logger() -> None:95    global _default_handler96 97    with _lock:98        if not _default_handler:99            return100 101        library_root_logger = _get_library_root_logger()102        library_root_logger.removeHandler(_default_handler)103        library_root_logger.setLevel(logging.NOTSET)104        _default_handler = None105 106 107def get_log_levels_dict():108    return log_levels109 110 111def get_logger(name: Optional[str] = None) -> logging.Logger:112    """113    Return a logger with the specified name.114 115    This function is not supposed to be directly accessed unless you are writing a custom diffusers module.116    """117 118    if name is None:119        name = _get_library_name()120 121    _configure_library_root_logger()122    return logging.getLogger(name)123 124 125def get_verbosity() -> int:126    """127    Return the current level for the 🤗 Diffusers' root logger as an int.128 129    Returns:130        `int`: The logging level.131 132    <Tip>133 134    🤗 Diffusers has following logging levels:135 136    - 50: `diffusers.logging.CRITICAL` or `diffusers.logging.FATAL`137    - 40: `diffusers.logging.ERROR`138    - 30: `diffusers.logging.WARNING` or `diffusers.logging.WARN`139    - 20: `diffusers.logging.INFO`140    - 10: `diffusers.logging.DEBUG`141 142    </Tip>"""143 144    _configure_library_root_logger()145    return _get_library_root_logger().getEffectiveLevel()146 147 148def set_verbosity(verbosity: int) -> None:149    """150    Set the verbosity level for the 🤗 Diffusers' root logger.151 152    Args:153        verbosity (`int`):154            Logging level, e.g., one of:155 156            - `diffusers.logging.CRITICAL` or `diffusers.logging.FATAL`157            - `diffusers.logging.ERROR`158            - `diffusers.logging.WARNING` or `diffusers.logging.WARN`159            - `diffusers.logging.INFO`160            - `diffusers.logging.DEBUG`161    """162 163    _configure_library_root_logger()164    _get_library_root_logger().setLevel(verbosity)165 166 167def set_verbosity_info():168    """Set the verbosity to the `INFO` level."""169    return set_verbosity(INFO)170 171 172def set_verbosity_warning():173    """Set the verbosity to the `WARNING` level."""174    return set_verbosity(WARNING)175 176 177def set_verbosity_debug():178    """Set the verbosity to the `DEBUG` level."""179    return set_verbosity(DEBUG)180 181 182def set_verbosity_error():183    """Set the verbosity to the `ERROR` level."""184    return set_verbosity(ERROR)185 186 187def disable_default_handler() -> None:188    """Disable the default handler of the HuggingFace Diffusers' root logger."""189 190    _configure_library_root_logger()191 192    assert _default_handler is not None193    _get_library_root_logger().removeHandler(_default_handler)194 195 196def enable_default_handler() -> None:197    """Enable the default handler of the HuggingFace Diffusers' root logger."""198 199    _configure_library_root_logger()200 201    assert _default_handler is not None202    _get_library_root_logger().addHandler(_default_handler)203 204 205def add_handler(handler: logging.Handler) -> None:206    """adds a handler to the HuggingFace Diffusers' root logger."""207 208    _configure_library_root_logger()209 210    assert handler is not None211    _get_library_root_logger().addHandler(handler)212 213 214def remove_handler(handler: logging.Handler) -> None:215    """removes given handler from the HuggingFace Diffusers' root logger."""216 217    _configure_library_root_logger()218 219    assert handler is not None and handler not in _get_library_root_logger().handlers220    _get_library_root_logger().removeHandler(handler)221 222 223def disable_propagation() -> None:224    """225    Disable propagation of the library log outputs. Note that log propagation is disabled by default.226    """227 228    _configure_library_root_logger()229    _get_library_root_logger().propagate = False230 231 232def enable_propagation() -> None:233    """234    Enable propagation of the library log outputs. Please disable the HuggingFace Diffusers' default handler to prevent235    double logging if the root logger has been configured.236    """237 238    _configure_library_root_logger()239    _get_library_root_logger().propagate = True240 241 242def enable_explicit_format() -> None:243    """244    Enable explicit formatting for every HuggingFace Diffusers' logger. The explicit formatter is as follows:245    ```246        [LEVELNAME|FILENAME|LINE NUMBER] TIME >> MESSAGE247    ```248    All handlers currently bound to the root logger are affected by this method.249    """250    handlers = _get_library_root_logger().handlers251 252    for handler in handlers:253        formatter = logging.Formatter("[%(levelname)s|%(filename)s:%(lineno)s] %(asctime)s >> %(message)s")254        handler.setFormatter(formatter)255 256 257def reset_format() -> None:258    """259    Resets the formatting for HuggingFace Diffusers' loggers.260 261    All handlers currently bound to the root logger are affected by this method.262    """263    handlers = _get_library_root_logger().handlers264 265    for handler in handlers:266        handler.setFormatter(None)267 268 269def warning_advice(self, *args, **kwargs):270    """271    This method is identical to `logger.warning()`, but if env var DIFFUSERS_NO_ADVISORY_WARNINGS=1 is set, this272    warning will not be printed273    """274    no_advisory_warnings = os.getenv("DIFFUSERS_NO_ADVISORY_WARNINGS", False)275    if no_advisory_warnings:276        return277    self.warning(*args, **kwargs)278 279 280logging.Logger.warning_advice = warning_advice281 282 283class EmptyTqdm:284    """Dummy tqdm which doesn't do anything."""285 286    def __init__(self, *args, **kwargs):  # pylint: disable=unused-argument287        self._iterator = args[0] if args else None288 289    def __iter__(self):290        return iter(self._iterator)291 292    def __getattr__(self, _):293        """Return empty function."""294 295        def empty_fn(*args, **kwargs):  # pylint: disable=unused-argument296            return297 298        return empty_fn299 300    def __enter__(self):301        return self302 303    def __exit__(self, type_, value, traceback):304        return305 306 307class _tqdm_cls:308    def __call__(self, *args, **kwargs):309        if _tqdm_active:310            return tqdm_lib.tqdm(*args, **kwargs)311        else:312            return EmptyTqdm(*args, **kwargs)313 314    def set_lock(self, *args, **kwargs):315        self._lock = None316        if _tqdm_active:317            return tqdm_lib.tqdm.set_lock(*args, **kwargs)318 319    def get_lock(self):320        if _tqdm_active:321            return tqdm_lib.tqdm.get_lock()322 323 324tqdm = _tqdm_cls()325 326 327def is_progress_bar_enabled() -> bool:328    """Return a boolean indicating whether tqdm progress bars are enabled."""329    global _tqdm_active330    return bool(_tqdm_active)331 332 333def enable_progress_bar():334    """Enable tqdm progress bar."""335    global _tqdm_active336    _tqdm_active = True337 338 339def disable_progress_bar():340    """Disable tqdm progress bar."""341    global _tqdm_active342    _tqdm_active = False343