


'''simple docstring'''







import torch

from diffusers import KDPMaDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest
class 	UpperCamelCase__					( SCREAMING_SNAKE_CASE_	):



     """simple docstring"""






     __magic_name__    			=    (KDPMaDiscreteScheduler,)
     __magic_name__    			=    1_0







     def a	( self      ,	**snake_case__   ):

            '''simple docstring'''

            _lowerCAmelCase    :							Union[str, Any]   			=       {
                'num_train_timesteps': 1100,
                'beta_start': 0.0001,
                'beta_end': 0.02,
                'beta_schedule': 'linear',
            }

            config.update(**snake_case__   )
            return config







     def a	( self   ):

            '''simple docstring'''

            for timesteps in [10, 50, 100, 1000]:
                   self.check_over_configs(num_train_timesteps=snake_case__   )







     def a	( self   ):

            '''simple docstring'''

            for beta_start, beta_end in zip([0.0_0001, 0.0001, 0.001]      ,	[0.0002, 0.002, 0.02]   ):
                   self.check_over_configs(beta_start=snake_case__      ,	beta_end=snake_case__   )







     def a	( self   ):

            '''simple docstring'''

            for schedule in ["linear", "scaled_linear"]:
                   self.check_over_configs(beta_schedule=snake_case__   )







     def a	( self   ):

            '''simple docstring'''

            for prediction_type in ["epsilon", "v_prediction"]:
                   self.check_over_configs(prediction_type=snake_case__   )







     def a	( self   ):

            '''simple docstring'''

            _lowerCAmelCase    :							Union[str, Any]   			=       self.scheduler_classes[0]
            _lowerCAmelCase    :							Optional[Any]   			=       self.get_scheduler_config(prediction_type='v_prediction'   )
            _lowerCAmelCase    :							Any   			=       scheduler_class(**snake_case__   )

            scheduler.set_timesteps(self.num_inference_steps   )

            _lowerCAmelCase    :							Any   			=       self.dummy_model()
            _lowerCAmelCase    :							str   			=       self.dummy_sample_deter * scheduler.init_noise_sigma
            _lowerCAmelCase    :							List[Any]   			=       sample.to(snake_case__   )

            for i, t in enumerate(scheduler.timesteps   ):
                   _lowerCAmelCase    :							List[Any]   			=       scheduler.scale_model_input(snake_case__      ,	snake_case__   )

                   _lowerCAmelCase    :							List[Any]   			=       model(snake_case__      ,	snake_case__   )

                   _lowerCAmelCase    :							List[str]   			=       scheduler.step(snake_case__      ,	snake_case__      ,	snake_case__   )
                   _lowerCAmelCase    :							List[str]   			=       output.prev_sample

            _lowerCAmelCase    :							Dict   			=       torch.sum(torch.abs(snake_case__   )   )
            _lowerCAmelCase    :							Tuple   			=       torch.mean(torch.abs(snake_case__   )   )

            if torch_device in ["cpu", "mps"]:
                   assert abs(result_sum.item() - 4.6934E-07   ) < 1E-2
                   assert abs(result_mean.item() - 6.1112E-10   ) < 1E-3
            else:
                   # CUDA
                   assert abs(result_sum.item() - 4.693428650170972E-07   ) < 1E-2
                   assert abs(result_mean.item() - 0.0002   ) < 1E-3







     def a	( self   ):

            '''simple docstring'''

            if torch_device == "mps":
                   return
            _lowerCAmelCase    :							List[str]   			=       self.scheduler_classes[0]
            _lowerCAmelCase    :							Dict   			=       self.get_scheduler_config()
            _lowerCAmelCase    :							Optional[Any]   			=       scheduler_class(**snake_case__   )

            scheduler.set_timesteps(self.num_inference_steps   )

            _lowerCAmelCase    :							Tuple   			=       self.dummy_model()
            _lowerCAmelCase    :							List[str]   			=       self.dummy_sample_deter * scheduler.init_noise_sigma
            _lowerCAmelCase    :							List[Any]   			=       sample.to(snake_case__   )

            for i, t in enumerate(scheduler.timesteps   ):
                   _lowerCAmelCase    :							Union[str, Any]   			=       scheduler.scale_model_input(snake_case__      ,	snake_case__   )

                   _lowerCAmelCase    :							List[Any]   			=       model(snake_case__      ,	snake_case__   )

                   _lowerCAmelCase    :							Optional[Any]   			=       scheduler.step(snake_case__      ,	snake_case__      ,	snake_case__   )
                   _lowerCAmelCase    :							str   			=       output.prev_sample

            _lowerCAmelCase    :							Any   			=       torch.sum(torch.abs(snake_case__   )   )
            _lowerCAmelCase    :							Tuple   			=       torch.mean(torch.abs(snake_case__   )   )

            if torch_device in ["cpu", "mps"]:
                   assert abs(result_sum.item() - 20.4125   ) < 1E-2
                   assert abs(result_mean.item() - 0.0266   ) < 1E-3
            else:
                   # CUDA
                   assert abs(result_sum.item() - 20.4125   ) < 1E-2
                   assert abs(result_mean.item() - 0.0266   ) < 1E-3







     def a	( self   ):

            '''simple docstring'''

            if torch_device == "mps":
                   return
            _lowerCAmelCase    :							int   			=       self.scheduler_classes[0]
            _lowerCAmelCase    :							Optional[int]   			=       self.get_scheduler_config()
            _lowerCAmelCase    :							Union[str, Any]   			=       scheduler_class(**snake_case__   )

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

            _lowerCAmelCase    :							Dict   			=       self.dummy_model()
            _lowerCAmelCase    :							List[Any]   			=       self.dummy_sample_deter.to(snake_case__   ) * scheduler.init_noise_sigma

            for t in scheduler.timesteps:
                   _lowerCAmelCase    :							Optional[Any]   			=       scheduler.scale_model_input(snake_case__      ,	snake_case__   )

                   _lowerCAmelCase    :							Dict   			=       model(snake_case__      ,	snake_case__   )

                   _lowerCAmelCase    :							Optional[Any]   			=       scheduler.step(snake_case__      ,	snake_case__      ,	snake_case__   )
                   _lowerCAmelCase    :							Any   			=       output.prev_sample

            _lowerCAmelCase    :							Dict   			=       torch.sum(torch.abs(snake_case__   )   )
            _lowerCAmelCase    :							Union[str, Any]   			=       torch.mean(torch.abs(snake_case__   )   )

            if str(snake_case__   ).startswith('cpu'   ):
                   # The following sum varies between 148 and 156 on mps. Why?
                   assert abs(result_sum.item() - 20.4125   ) < 1E-2
                   assert abs(result_mean.item() - 0.0266   ) < 1E-3
            else:
                   # CUDA
                   assert abs(result_sum.item() - 20.4125   ) < 1E-2
                   assert abs(result_mean.item() - 0.0266   ) < 1E-3





'''simple docstring'''







from typing import TYPE_CHECKING

from ....utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


lowerCAmelCase    :     Optional[int]					=						{
    """configuration_mctct""": ["""MCTCT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """MCTCTConfig"""],
    """feature_extraction_mctct""": ["""MCTCTFeatureExtractor"""],
    """processing_mctct""": ["""MCTCTProcessor"""],
}


try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       lowerCAmelCase    :     List[str]					=						[
           """MCTCT_PRETRAINED_MODEL_ARCHIVE_LIST""",
           """MCTCTForCTC""",
           """MCTCTModel""",
           """MCTCTPreTrainedModel""",
       ]


if TYPE_CHECKING:
       from .configuration_mctct import MCTCT_PRETRAINED_CONFIG_ARCHIVE_MAP, MCTCTConfig
       from .feature_extraction_mctct import MCTCTFeatureExtractor
       from .processing_mctct import MCTCTProcessor

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_mctct import MCTCT_PRETRAINED_MODEL_ARCHIVE_LIST, MCTCTForCTC, MCTCTModel, MCTCTPreTrainedModel

else:
       import sys

       lowerCAmelCase    :     Tuple					=						_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

