

'''simple docstring'''

import torch

from diffusers import HeunDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest


class 					lowerCAmelCase       ( A				):
       lowerCAmelCase_				=	(HeunDiscreteScheduler,)
       lowerCAmelCase_				=	1_0







       def      snake_case				(							self     :	Optional[Any]  , **__lowercase     :	List[Any] ):







            """simple docstring"""







            __lowercase     						={
                'num_train_timesteps': 1100,
                'beta_start': 0.0_0_0_1,
                'beta_end': 0.0_2,
                'beta_schedule': 'linear',
            }

            config.update(**__lowercase )
            return config







       def      snake_case				(							self     :	str ):







            """simple docstring"""







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







       def      snake_case				(							self     :	List[Any] ):







            """simple docstring"""







            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=__lowercase  , beta_end=__lowercase )







       def      snake_case				(							self     :	Union[str, Any] ):







            """simple docstring"""







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







       def      snake_case				(							self     :	Tuple ):







            """simple docstring"""







            for clip_sample_range in [1.0, 2.0, 3.0]:
                 self.check_over_configs(clip_sample_range=__lowercase  , clip_sample=__lowercase )







       def      snake_case				(							self     :	Any ):







            """simple docstring"""







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







       def      snake_case				(							self     :	str ):







            """simple docstring"""







            __lowercase     						=self.scheduler_classes[0]
            __lowercase     						=self.get_scheduler_config()
            __lowercase     						=scheduler_class(**__lowercase )

            scheduler.set_timesteps(self.num_inference_steps )

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

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

                 __lowercase     						=model(__lowercase  , __lowercase )

                 __lowercase     						=scheduler.step(__lowercase  , __lowercase  , __lowercase )
                 __lowercase     						=output.prev_sample

            __lowercase     						=torch.sum(torch.abs(__lowercase ) )
            __lowercase     						=torch.mean(torch.abs(__lowercase ) )

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







       def      snake_case				(							self     :	List[str] ):







            """simple docstring"""







            __lowercase     						=self.scheduler_classes[0]
            __lowercase     						=self.get_scheduler_config(prediction_type='v_prediction' )
            __lowercase     						=scheduler_class(**__lowercase )

            scheduler.set_timesteps(self.num_inference_steps )

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

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

                 __lowercase     						=model(__lowercase  , __lowercase )

                 __lowercase     						=scheduler.step(__lowercase  , __lowercase  , __lowercase )
                 __lowercase     						=output.prev_sample

            __lowercase     						=torch.sum(torch.abs(__lowercase ) )
            __lowercase     						=torch.mean(torch.abs(__lowercase ) )

            if torch_device in ["cpu", "mps"]:
                 assert abs(result_sum.item() - 4.6_934E-07 ) < 1E-2
                 assert abs(result_mean.item() - 6.1_112E-10 ) < 1E-3
            else:
                 # CUDA
                 assert abs(result_sum.item() - 4.693_428_650_170_972E-07 ) < 1E-2
                 assert abs(result_mean.item() - 0.0_0_0_2 ) < 1E-3







       def      snake_case				(							self     :	Tuple ):







            """simple docstring"""







            __lowercase     						=self.scheduler_classes[0]
            __lowercase     						=self.get_scheduler_config()
            __lowercase     						=scheduler_class(**__lowercase )

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

            __lowercase     						=self.dummy_model()
            __lowercase     						=self.dummy_sample_deter.to(__lowercase ) * scheduler.init_noise_sigma

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

                 __lowercase     						=model(__lowercase  , __lowercase )

                 __lowercase     						=scheduler.step(__lowercase  , __lowercase  , __lowercase )
                 __lowercase     						=output.prev_sample

            __lowercase     						=torch.sum(torch.abs(__lowercase ) )
            __lowercase     						=torch.mean(torch.abs(__lowercase ) )

            if str(__lowercase ).startswith('cpu' ):
                 # The following sum varies between 148 and 156 on mps. Why?
                 assert abs(result_sum.item() - 0.1_2_3_3 ) < 1E-2
                 assert abs(result_mean.item() - 0.0_0_0_2 ) < 1E-3
            elif str(__lowercase ).startswith('mps' ):
                 # Larger tolerance on mps
                 assert abs(result_mean.item() - 0.0_0_0_2 ) < 1E-2
            else:
                 # CUDA
                 assert abs(result_sum.item() - 0.1_2_3_3 ) < 1E-2
                 assert abs(result_mean.item() - 0.0_0_0_2 ) < 1E-3







       def      snake_case				(							self     :	Optional[Any] ):







            """simple docstring"""







            __lowercase     						=self.scheduler_classes[0]
            __lowercase     						=self.get_scheduler_config()
            __lowercase     						=scheduler_class(**__lowercase  , use_karras_sigmas=__lowercase )

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

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

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

                 __lowercase     						=model(__lowercase  , __lowercase )

                 __lowercase     						=scheduler.step(__lowercase  , __lowercase  , __lowercase )
                 __lowercase     						=output.prev_sample

            __lowercase     						=torch.sum(torch.abs(__lowercase ) )
            __lowercase     						=torch.mean(torch.abs(__lowercase ) )

            assert abs(result_sum.item() - 0.0_0_0_1_5 ) < 1E-2
            assert abs(result_mean.item() - 1.9_869_554_535_034_695E-07 ) < 1E-2





'''simple docstring'''

from typing import TYPE_CHECKING

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


UpperCAmelCase							=					{
    '''configuration_videomae''': ['''VIDEOMAE_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''VideoMAEConfig'''],
}

try:
       if not is_torch_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       UpperCAmelCase							=					[
           '''VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST''',
           '''VideoMAEForPreTraining''',
           '''VideoMAEModel''',
           '''VideoMAEPreTrainedModel''',
           '''VideoMAEForVideoClassification''',
       ]

try:
       if not is_vision_available():
              raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
       pass
else:
       UpperCAmelCase							=					['''VideoMAEFeatureExtractor''']
       UpperCAmelCase							=					['''VideoMAEImageProcessor''']

if TYPE_CHECKING:
       from .configuration_videomae import VIDEOMAE_PRETRAINED_CONFIG_ARCHIVE_MAP, VideoMAEConfig

       try:
              if not is_torch_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .modeling_videomae import (
                  VIDEOMAE_PRETRAINED_MODEL_ARCHIVE_LIST,
                  VideoMAEForPreTraining,
                  VideoMAEForVideoClassification,
                  VideoMAEModel,
                  VideoMAEPreTrainedModel,
              )

       try:
              if not is_vision_available():
                     raise OptionalDependencyNotAvailable()
       except OptionalDependencyNotAvailable:
              pass
       else:
              from .feature_extraction_videomae import VideoMAEFeatureExtractor
              from .image_processing_videomae import VideoMAEImageProcessor

else:
       import sys

       UpperCAmelCase							=					_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)


