




from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


__a          =			logging.get_logger(__name__)

__a          =			{
    'microsoft/beit-base-patch16-224-pt22k': (
        'https://huggingface.co/microsoft/beit-base-patch16-224-pt22k/resolve/main/config.json'
    ),
    # See all BEiT models at https://huggingface.co/models?filter=beit
}

class   __a(					_a       ):






    """simple docstring"""





    lowerCAmelCase        =				'''beit'''





    def __init__(				self       ,_SCREAMING_SNAKE_CASE=8_192       ,_SCREAMING_SNAKE_CASE=768       ,_SCREAMING_SNAKE_CASE=12       ,_SCREAMING_SNAKE_CASE=12       ,_SCREAMING_SNAKE_CASE=3_072       ,_SCREAMING_SNAKE_CASE="gelu"       ,_SCREAMING_SNAKE_CASE=0.0       ,_SCREAMING_SNAKE_CASE=0.0       ,_SCREAMING_SNAKE_CASE=0.02       ,_SCREAMING_SNAKE_CASE=1e-12       ,_SCREAMING_SNAKE_CASE=224       ,_SCREAMING_SNAKE_CASE=16       ,_SCREAMING_SNAKE_CASE=3       ,_SCREAMING_SNAKE_CASE=False       ,_SCREAMING_SNAKE_CASE=False       ,_SCREAMING_SNAKE_CASE=False       ,_SCREAMING_SNAKE_CASE=False       ,_SCREAMING_SNAKE_CASE=0.1       ,_SCREAMING_SNAKE_CASE=0.1       ,_SCREAMING_SNAKE_CASE=True       ,_SCREAMING_SNAKE_CASE=[3, 5, 7, 11]       ,_SCREAMING_SNAKE_CASE=[1, 2, 3, 6]       ,_SCREAMING_SNAKE_CASE=True       ,_SCREAMING_SNAKE_CASE=0.4       ,_SCREAMING_SNAKE_CASE=256       ,_SCREAMING_SNAKE_CASE=1       ,_SCREAMING_SNAKE_CASE=False       ,_SCREAMING_SNAKE_CASE=255       ,**_SCREAMING_SNAKE_CASE       ,)			->		Union[str, Any]:
         super().__init__(**_SCREAMING_SNAKE_CASE  )

         UpperCAmelCase_      :							Optional[int]       =      vocab_size
         UpperCAmelCase_      :							Union[str, Any]       =      hidden_size
         UpperCAmelCase_      :							Any       =      num_hidden_layers
         UpperCAmelCase_      :							Union[str, Any]       =      num_attention_heads
         UpperCAmelCase_      :							Optional[Any]       =      intermediate_size
         UpperCAmelCase_      :							List[str]       =      hidden_act
         UpperCAmelCase_      :							Any       =      hidden_dropout_prob
         UpperCAmelCase_      :							int       =      attention_probs_dropout_prob
         UpperCAmelCase_      :							List[Any]       =      initializer_range
         UpperCAmelCase_      :							Optional[int]       =      layer_norm_eps

         UpperCAmelCase_      :							Dict       =      image_size
         UpperCAmelCase_      :							Optional[Any]       =      patch_size
         UpperCAmelCase_      :							Tuple       =      num_channels
         UpperCAmelCase_      :							Tuple       =      use_mask_token
         UpperCAmelCase_      :							List[Any]       =      use_absolute_position_embeddings
         UpperCAmelCase_      :							List[Any]       =      use_relative_position_bias
         UpperCAmelCase_      :							Tuple       =      use_shared_relative_position_bias
         UpperCAmelCase_      :							Union[str, Any]       =      layer_scale_init_value
         UpperCAmelCase_      :							Optional[int]       =      drop_path_rate
         UpperCAmelCase_      :							List[Any]       =      use_mean_pooling
         # decode head attributes (semantic segmentation)
         UpperCAmelCase_      :							Union[str, Any]       =      out_indices
         UpperCAmelCase_      :							Optional[int]       =      pool_scales
         # auxiliary head attributes (semantic segmentation)
         UpperCAmelCase_      :							Any       =      use_auxiliary_head
         UpperCAmelCase_      :							List[str]       =      auxiliary_loss_weight
         UpperCAmelCase_      :							Dict       =      auxiliary_channels
         UpperCAmelCase_      :							Any       =      auxiliary_num_convs
         UpperCAmelCase_      :							Optional[Any]       =      auxiliary_concat_input
         UpperCAmelCase_      :							Dict       =      semantic_loss_ignore_index


class   __a(					_a       ):






    """simple docstring"""





    lowerCAmelCase        =				version.parse('''1.11'''       )





    @property
    def 							a__				(				self  )			->		Mapping[str, Mapping[int, str]]:
         return OrderedDict(
             [
                 ('''pixel_values''', {0: '''batch''', 1: '''num_channels''', 2: '''height''', 3: '''width'''}),
             ]  )





    @property
    def 							a__				(				self  )			->		float:
         return 1e-4





import tempfile

import torch

from diffusers import IPNDMScheduler

from .test_schedulers import SchedulerCommonTest

class   __a(					_a       ):






    """simple docstring"""





    lowerCAmelCase        =				(IPNDMScheduler,)
    lowerCAmelCase        =				(('''num_inference_steps''', 50),)





    def 							a__				(				self       ,**_SCREAMING_SNAKE_CASE  )			->		int:
         UpperCAmelCase_      :							Tuple       =      {'''num_train_timesteps''': 1_000}
         config.update(**_SCREAMING_SNAKE_CASE  )
         return config





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE=0       ,**_SCREAMING_SNAKE_CASE  )			->		Union[str, Any]:
         UpperCAmelCase_      :							Optional[int]       =      dict(self.forward_default_kwargs  )
         UpperCAmelCase_      :							int       =      kwargs.pop('''num_inference_steps'''       ,_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							str       =      self.dummy_sample
         UpperCAmelCase_      :							Optional[Any]       =      0.1 * sample
         UpperCAmelCase_      :							Optional[int]       =      [residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

         for scheduler_class in self.scheduler_classes:
              UpperCAmelCase_      :							Tuple       =      self.get_scheduler_config(**_SCREAMING_SNAKE_CASE  )
              UpperCAmelCase_      :							Optional[int]       =      scheduler_class(**_SCREAMING_SNAKE_CASE  )
              scheduler.set_timesteps(_SCREAMING_SNAKE_CASE  )
              # copy over dummy past residuals
              UpperCAmelCase_      :							Optional[Any]       =      dummy_past_residuals[:]

              if time_step is None:
                   UpperCAmelCase_      :							Tuple       =      scheduler.timesteps[len(scheduler.timesteps  ) // 2]

              with tempfile.TemporaryDirectory() as tmpdirname:
                   scheduler.save_config(_SCREAMING_SNAKE_CASE  )
                   UpperCAmelCase_      :							List[Any]       =      scheduler_class.from_pretrained(_SCREAMING_SNAKE_CASE  )
                   new_scheduler.set_timesteps(_SCREAMING_SNAKE_CASE  )
                   # copy over dummy past residuals
                   UpperCAmelCase_      :							Optional[int]       =      dummy_past_residuals[:]

              UpperCAmelCase_      :							List[Any]       =      scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample
              UpperCAmelCase_      :							str       =      new_scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample

              assert torch.sum(torch.abs(output - new_output  )  ) < 1e-5, "Scheduler outputs are not identical"

              UpperCAmelCase_      :							Union[str, Any]       =      scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample
              UpperCAmelCase_      :							Optional[int]       =      new_scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample

              assert torch.sum(torch.abs(output - new_output  )  ) < 1e-5, "Scheduler outputs are not identical"





    def 							a__				(				self  )			->		List[Any]:
         pass





    def 							a__				(				self       ,_SCREAMING_SNAKE_CASE=0       ,**_SCREAMING_SNAKE_CASE  )			->		List[Any]:
         UpperCAmelCase_      :							List[str]       =      dict(self.forward_default_kwargs  )
         UpperCAmelCase_      :							Tuple       =      kwargs.pop('''num_inference_steps'''       ,_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							int       =      self.dummy_sample
         UpperCAmelCase_      :							Tuple       =      0.1 * sample
         UpperCAmelCase_      :							int       =      [residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

         for scheduler_class in self.scheduler_classes:
              UpperCAmelCase_      :							str       =      self.get_scheduler_config()
              UpperCAmelCase_      :							str       =      scheduler_class(**_SCREAMING_SNAKE_CASE  )
              scheduler.set_timesteps(_SCREAMING_SNAKE_CASE  )

              # copy over dummy past residuals (must be after setting timesteps)
              UpperCAmelCase_      :							Tuple       =      dummy_past_residuals[:]

              if time_step is None:
                   UpperCAmelCase_      :							Dict       =      scheduler.timesteps[len(scheduler.timesteps  ) // 2]

              with tempfile.TemporaryDirectory() as tmpdirname:
                   scheduler.save_config(_SCREAMING_SNAKE_CASE  )
                   UpperCAmelCase_      :							Optional[int]       =      scheduler_class.from_pretrained(_SCREAMING_SNAKE_CASE  )
                   # copy over dummy past residuals
                   new_scheduler.set_timesteps(_SCREAMING_SNAKE_CASE  )

                   # copy over dummy past residual (must be after setting timesteps)
                   UpperCAmelCase_      :							str       =      dummy_past_residuals[:]

              UpperCAmelCase_      :							Tuple       =      scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample
              UpperCAmelCase_      :							Optional[int]       =      new_scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample

              assert torch.sum(torch.abs(output - new_output  )  ) < 1e-5, "Scheduler outputs are not identical"

              UpperCAmelCase_      :							str       =      scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample
              UpperCAmelCase_      :							List[str]       =      new_scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample

              assert torch.sum(torch.abs(output - new_output  )  ) < 1e-5, "Scheduler outputs are not identical"





    def 							a__				(				self       ,**_SCREAMING_SNAKE_CASE  )			->		List[Any]:
         UpperCAmelCase_      :							Tuple       =      self.scheduler_classes[0]
         UpperCAmelCase_      :							List[Any]       =      self.get_scheduler_config(**_SCREAMING_SNAKE_CASE  )
         UpperCAmelCase_      :							Tuple       =      scheduler_class(**_SCREAMING_SNAKE_CASE  )

         UpperCAmelCase_      :							Any       =      10
         UpperCAmelCase_      :							Tuple       =      self.dummy_model()
         UpperCAmelCase_      :							List[str]       =      self.dummy_sample_deter
         scheduler.set_timesteps(_SCREAMING_SNAKE_CASE  )

         for i, t in enumerate(scheduler.timesteps  ):
              UpperCAmelCase_      :							int       =      model(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )
              UpperCAmelCase_      :							Union[str, Any]       =      scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  ).prev_sample

         for i, t in enumerate(scheduler.timesteps  ):
              UpperCAmelCase_      :							Optional[int]       =      model(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  )
              UpperCAmelCase_      :							Dict       =      scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE  ).prev_sample

         return sample





    def 							a__				(				self  )			->		str:
         UpperCAmelCase_      :							List[Any]       =      dict(self.forward_default_kwargs  )

         UpperCAmelCase_      :							Any       =      kwargs.pop('''num_inference_steps'''       ,_SCREAMING_SNAKE_CASE  )

         for scheduler_class in self.scheduler_classes:
              UpperCAmelCase_      :							int       =      self.get_scheduler_config()
              UpperCAmelCase_      :							List[Any]       =      scheduler_class(**_SCREAMING_SNAKE_CASE  )

              UpperCAmelCase_      :							Optional[int]       =      self.dummy_sample
              UpperCAmelCase_      :							List[Any]       =      0.1 * sample

              if num_inference_steps is not None and hasattr(_SCREAMING_SNAKE_CASE       ,'''set_timesteps'''  ):
                   scheduler.set_timesteps(_SCREAMING_SNAKE_CASE  )
              elif num_inference_steps is not None and not hasattr(_SCREAMING_SNAKE_CASE       ,'''set_timesteps'''  ):
                   UpperCAmelCase_      :							Union[str, Any]       =      num_inference_steps

              # copy over dummy past residuals (must be done after set_timesteps)
              UpperCAmelCase_      :							str       =      [residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]
              UpperCAmelCase_      :							int       =      dummy_past_residuals[:]

              UpperCAmelCase_      :							List[Any]       =      scheduler.timesteps[5]
              UpperCAmelCase_      :							Dict       =      scheduler.timesteps[6]

              UpperCAmelCase_      :							List[Any]       =      scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample
              UpperCAmelCase_      :							Dict       =      scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample

              self.assertEqual(output_a.shape       ,sample.shape  )
              self.assertEqual(output_a.shape       ,output_a.shape  )

              UpperCAmelCase_      :							List[Any]       =      scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample
              UpperCAmelCase_      :							Dict       =      scheduler.step(_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,_SCREAMING_SNAKE_CASE       ,**_SCREAMING_SNAKE_CASE  ).prev_sample

              self.assertEqual(output_a.shape       ,sample.shape  )
              self.assertEqual(output_a.shape       ,output_a.shape  )





    def 							a__				(				self  )			->		Any:
         for timesteps in [100, 1_000]:
              self.check_over_configs(num_train_timesteps=_SCREAMING_SNAKE_CASE       ,time_step=_SCREAMING_SNAKE_CASE  )





    def 							a__				(				self  )			->		Union[str, Any]:
         for t, num_inference_steps in zip([1, 5, 10]       ,[10, 50, 100]  ):
              self.check_over_forward(num_inference_steps=_SCREAMING_SNAKE_CASE       ,time_step=_SCREAMING_SNAKE_CASE  )





    def 							a__				(				self  )			->		Union[str, Any]:
         UpperCAmelCase_      :							Optional[int]       =      self.full_loop()
         UpperCAmelCase_      :							str       =      torch.mean(torch.abs(_SCREAMING_SNAKE_CASE  )  )

         assert abs(result_mean.item() - 2_540_529  ) < 10