
import functools
import logging
import os
import sys
import threading
from logging import (
    CRITICAL,  # NOQA
    DEBUG,  # NOQA
    ERROR,  # NOQA
    FATAL,  # NOQA
    INFO,  # NOQA
    NOTSET,  # NOQA
    WARN,  # NOQA
    WARNING,  # NOQA
)
from typing import Optional

import huggingface_hub.utils as hf_hub_utils
from tqdm import auto as tqdm_lib


__magic_name__            =  threading.Lock()
__magic_name__            =  None

__magic_name__            =  {
    '''debug''': logging.DEBUG,
    '''info''': logging.INFO,
    '''warning''': logging.WARNING,
    '''error''': logging.ERROR,
    '''critical''': logging.CRITICAL,
}

__magic_name__            =  logging.WARNING

__magic_name__            =  True





def 	SCREAMING_SNAKE_CASE__    (      ):
    snake_case__						=				os.getenv("TRANSFORMERS_VERBOSITY"    ,		__lowerCAmelCase							)
    if env_level_str:
        if env_level_str in log_levels:
            return log_levels[env_level_str]
        else:
            logging.getLogger().warning(
                F"""Unknown option TRANSFORMERS_VERBOSITY={env_level_str}, """
                F"""has to be one of: { ", ".join(log_levels.keys()							) }"""							)
    return _default_log_level





def 	SCREAMING_SNAKE_CASE__    (      ):
    return __name__.split("."							)[0]





def 	SCREAMING_SNAKE_CASE__    (      ):
    return logging.getLogger(_get_library_name()							)





def 	SCREAMING_SNAKE_CASE__    (      ):
    global _default_handler

    with _lock:
        if _default_handler:
            # This library has already configured the library root logger.
            return
        snake_case__						=				logging.StreamHandler()  # Set sys.stderr as stream.
        snake_case__						=				sys.stderr.flush

        # Apply our default configuration to the library root logger.
        snake_case__						=				_get_library_root_logger()
        library_root_logger.addHandler(_default_handler							)
        library_root_logger.setLevel(_get_default_logging_level()							)
        snake_case__						=				False





def 	SCREAMING_SNAKE_CASE__    (      ):
    global _default_handler

    with _lock:
        if not _default_handler:
            return

        snake_case__						=				_get_library_root_logger()
        library_root_logger.removeHandler(_default_handler							)
        library_root_logger.setLevel(logging.NOTSET							)
        snake_case__						=				None





def 	SCREAMING_SNAKE_CASE__    (      ):
    return log_levels





def 	SCREAMING_SNAKE_CASE__    (      __lowerCAmelCase = None							):

    if name is None:
        snake_case__						=				_get_library_name()

    _configure_library_root_logger()
    return logging.getLogger(__lowerCAmelCase							)





def 	SCREAMING_SNAKE_CASE__    (      ):
    _configure_library_root_logger()
    return _get_library_root_logger().getEffectiveLevel()





def 	SCREAMING_SNAKE_CASE__    (      __lowerCAmelCase							):
    _configure_library_root_logger()
    _get_library_root_logger().setLevel(__lowerCAmelCase							)





def 	SCREAMING_SNAKE_CASE__    (      ):
    return set_verbosity(__lowerCAmelCase							)





def 	SCREAMING_SNAKE_CASE__    (      ):
    return set_verbosity(__lowerCAmelCase							)





def 	SCREAMING_SNAKE_CASE__    (      ):
    return set_verbosity(__lowerCAmelCase							)





def 	SCREAMING_SNAKE_CASE__    (      ):
    return set_verbosity(__lowerCAmelCase							)





def 	SCREAMING_SNAKE_CASE__    (      ):
    _configure_library_root_logger()

    assert _default_handler is not None
    _get_library_root_logger().removeHandler(_default_handler							)





def 	SCREAMING_SNAKE_CASE__    (      ):
    _configure_library_root_logger()

    assert _default_handler is not None
    _get_library_root_logger().addHandler(_default_handler							)





def 	SCREAMING_SNAKE_CASE__    (      __lowerCAmelCase							):
    _configure_library_root_logger()

    assert handler is not None
    _get_library_root_logger().addHandler(__lowerCAmelCase							)





def 	SCREAMING_SNAKE_CASE__    (      __lowerCAmelCase							):
    _configure_library_root_logger()

    assert handler is not None and handler not in _get_library_root_logger().handlers
    _get_library_root_logger().removeHandler(__lowerCAmelCase							)





def 	SCREAMING_SNAKE_CASE__    (      ):
    _configure_library_root_logger()
    snake_case__						=				False





def 	SCREAMING_SNAKE_CASE__    (      ):
    _configure_library_root_logger()
    snake_case__						=				True





def 	SCREAMING_SNAKE_CASE__    (      ):
    snake_case__						=				_get_library_root_logger().handlers

    for handler in handlers:
        snake_case__						=				logging.Formatter("[%(levelname)s|%(filename)s:%(lineno)s] %(asctime)s >> %(message)s"							)
        handler.setFormatter(__lowerCAmelCase							)





def 	SCREAMING_SNAKE_CASE__    (      ):
    snake_case__						=				_get_library_root_logger().handlers

    for handler in handlers:
        handler.setFormatter(__lowerCAmelCase							)





def 	SCREAMING_SNAKE_CASE__    (      self    ,		*__lowerCAmelCase    ,		**__lowerCAmelCase							):
    snake_case__						=				os.getenv("TRANSFORMERS_NO_ADVISORY_WARNINGS"    ,		__lowerCAmelCase							)
    if no_advisory_warnings:
        return
    self.warning(*__lowerCAmelCase    ,		**__lowerCAmelCase							)


__magic_name__            =  warning_advice





@functools.lru_cache(__lowerCAmelCase							)
def 	SCREAMING_SNAKE_CASE__    (      self    ,		*__lowerCAmelCase    ,		**__lowerCAmelCase							):
    self.warning(*__lowerCAmelCase    ,		**__lowerCAmelCase							)


__magic_name__            =  warning_once




class 					_SCREAMING_SNAKE_CASE  :
      def __init__(			self		,	*lowerCamelCase		,	**lowerCamelCase ):  # pylint: disable=unused-argument
          snake_case__						=				args[0] if args else None
      def __iter__(			self ):
          return iter(self._iterator )
      def __getattr__(			self		,	lowerCamelCase ):

          def empty_fn(*lowerCamelCase		,	**lowerCamelCase ):  # pylint: disable=unused-argument
              return

          return empty_fn
      def __enter__(			self ):
          return self



      def __exit__(			self		,	lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase ):
          return





class 					_SCREAMING_SNAKE_CASE  :
      def __call__(			self		,	*lowerCamelCase		,	**lowerCamelCase ):
          if _tqdm_active:
              return tqdm_lib.tqdm(*lowerCamelCase		,	**lowerCamelCase )
          else:
              return EmptyTqdm(*lowerCamelCase		,	**lowerCamelCase )
      def A_			(			self		,	*lowerCamelCase		,	**lowerCamelCase ):
          snake_case__						=				None
          if _tqdm_active:
              return tqdm_lib.tqdm.set_lock(*lowerCamelCase		,	**lowerCamelCase )



      def A_			(			self ):
          if _tqdm_active:
              return tqdm_lib.tqdm.get_lock()





__magic_name__            =  _tqdm_cls()





def 	SCREAMING_SNAKE_CASE__    (      ):
    global _tqdm_active
    return bool(_tqdm_active							)





def 	SCREAMING_SNAKE_CASE__    (      ):
    global _tqdm_active
    snake_case__						=				True
    hf_hub_utils.enable_progress_bars()





def 	SCREAMING_SNAKE_CASE__    (      ):
    global _tqdm_active
    snake_case__						=				False
    hf_hub_utils.disable_progress_bars()





import torch

from diffusers import KDPMaDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest




class 					_SCREAMING_SNAKE_CASE  (      __UpperCamelCase ):
      _A						:		Any    	=     (KDPMaDiscreteScheduler,)
      _A						:		Dict    	=     10
      def A_			(			self		,	**lowerCamelCase ):
          snake_case__						=				{
              "num_train_timesteps": 11_00,
              "beta_start": 0.0_0_0_1,
              "beta_end": 0.0_2,
              "beta_schedule": "linear",
          }

          config.update(**lowerCamelCase )
          return config
      def A_			(			self ):
          for timesteps in [10, 50, 1_00, 10_00]:
              self.check_over_configs(num_train_timesteps=lowerCamelCase )
      def A_			(			self ):
          for beta_start, beta_end in zip([0.0_0_0_0_1, 0.0_0_0_1, 0.0_0_1]		,	[0.0_0_0_2, 0.0_0_2, 0.0_2] ):
              self.check_over_configs(beta_start=lowerCamelCase		,	beta_end=lowerCamelCase )
      def A_			(			self ):
          for schedule in ["linear", "scaled_linear"]:
              self.check_over_configs(beta_schedule=lowerCamelCase )
      def A_			(			self ):
          for prediction_type in ["epsilon", "v_prediction"]:
              self.check_over_configs(prediction_type=lowerCamelCase )
      def A_			(			self ):
          snake_case__						=				self.scheduler_classes[0]
          snake_case__						=				self.get_scheduler_config(prediction_type="v_prediction" )
          snake_case__						=				scheduler_class(**lowerCamelCase )

          scheduler.set_timesteps(self.num_inference_steps )

          snake_case__						=				self.dummy_model()
          snake_case__						=				self.dummy_sample_deter * scheduler.init_noise_sigma
          snake_case__						=				sample.to(lowerCamelCase )

          for i, t in enumerate(scheduler.timesteps ):
              snake_case__						=				scheduler.scale_model_input(lowerCamelCase		,	lowerCamelCase )

              snake_case__						=				model(lowerCamelCase		,	lowerCamelCase )

              snake_case__						=				scheduler.step(lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase )
              snake_case__						=				output.prev_sample

          snake_case__						=				torch.sum(torch.abs(lowerCamelCase ) )
          snake_case__						=				torch.mean(torch.abs(lowerCamelCase ) )

          if torch_device in ["cpu", "mps"]:
              assert abs(result_sum.item() - 4.69_34e-07 ) < 1e-2
              assert abs(result_mean.item() - 6.11_12e-10 ) < 1e-3
          else:
              # CUDA
              assert abs(result_sum.item() - 4.6_93_42_86_50_17_09_72e-07 ) < 1e-2
              assert abs(result_mean.item() - 0.0_0_0_2 ) < 1e-3
      def A_			(			self ):
          if torch_device == "mps":
              return
          snake_case__						=				self.scheduler_classes[0]
          snake_case__						=				self.get_scheduler_config()
          snake_case__						=				scheduler_class(**lowerCamelCase )

          scheduler.set_timesteps(self.num_inference_steps )

          snake_case__						=				self.dummy_model()
          snake_case__						=				self.dummy_sample_deter * scheduler.init_noise_sigma
          snake_case__						=				sample.to(lowerCamelCase )

          for i, t in enumerate(scheduler.timesteps ):
              snake_case__						=				scheduler.scale_model_input(lowerCamelCase		,	lowerCamelCase )

              snake_case__						=				model(lowerCamelCase		,	lowerCamelCase )

              snake_case__						=				scheduler.step(lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase )
              snake_case__						=				output.prev_sample

          snake_case__						=				torch.sum(torch.abs(lowerCamelCase ) )
          snake_case__						=				torch.mean(torch.abs(lowerCamelCase ) )

          if torch_device in ["cpu", "mps"]:
              assert abs(result_sum.item() - 2_0.4_1_2_5 ) < 1e-2
              assert abs(result_mean.item() - 0.0_2_6_6 ) < 1e-3
          else:
              # CUDA
              assert abs(result_sum.item() - 2_0.4_1_2_5 ) < 1e-2
              assert abs(result_mean.item() - 0.0_2_6_6 ) < 1e-3



      def A_			(			self ):
          if torch_device == "mps":
              return
          snake_case__						=				self.scheduler_classes[0]
          snake_case__						=				self.get_scheduler_config()
          snake_case__						=				scheduler_class(**lowerCamelCase )

          scheduler.set_timesteps(self.num_inference_steps		,	device=lowerCamelCase )

          snake_case__						=				self.dummy_model()
          snake_case__						=				self.dummy_sample_deter.to(lowerCamelCase ) * scheduler.init_noise_sigma

          for t in scheduler.timesteps:
              snake_case__						=				scheduler.scale_model_input(lowerCamelCase		,	lowerCamelCase )

              snake_case__						=				model(lowerCamelCase		,	lowerCamelCase )

              snake_case__						=				scheduler.step(lowerCamelCase		,	lowerCamelCase		,	lowerCamelCase )
              snake_case__						=				output.prev_sample

          snake_case__						=				torch.sum(torch.abs(lowerCamelCase ) )
          snake_case__						=				torch.mean(torch.abs(lowerCamelCase ) )

          if str(lowerCamelCase ).startswith("cpu" ):
              # The following sum varies between 148 and 156 on mps. Why?
              assert abs(result_sum.item() - 2_0.4_1_2_5 ) < 1e-2
              assert abs(result_mean.item() - 0.0_2_6_6 ) < 1e-3
          else:
              # CUDA
              assert abs(result_sum.item() - 2_0.4_1_2_5 ) < 1e-2
              assert abs(result_mean.item() - 0.0_2_6_6 ) < 1e-3



