



"""simple docstring"""







import tempfile

import torch

from diffusers import (
    DEISMultistepScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    UniPCMultistepScheduler,
)

from .test_schedulers import SchedulerCommonTest







class 				UpperCamelCase__       (     _lowerCamelCase  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE      =					(DPMSolverSinglestepScheduler,)
      _SCREAMING_SNAKE_CASE      =					(("""num_inference_steps""", 25),)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Dict   ,      **SCREAMING_SNAKE_CASE_    :				Tuple	):
        lowerCAmelCase_       :    Dict								=			{
            'num_train_timesteps': 1_0_0_0,
            'beta_start': 0.00_01,
            'beta_end': 0.02,
            'beta_schedule': 'linear',
            'solver_order': 2,
            'prediction_type': 'epsilon',
            'thresholding': False,
            'sample_max_value': 1.0,
            'algorithm_type': 'dpmsolver++',
            'solver_type': 'midpoint',
            'lambda_min_clipped': -float('inf'	),
            'variance_type': None,
        }

        config.update(**SCREAMING_SNAKE_CASE_	)
        return config



      def 					SCREAMING_SNAKE_CASE__			(				self    :				int   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]=0   ,      **SCREAMING_SNAKE_CASE_    :				Union[str, Any]	):
        lowerCAmelCase_       :    List[Any]								=			dict(self.forward_default_kwargs	)
        lowerCAmelCase_       :    str								=			kwargs.pop('num_inference_steps'   ,      SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    Optional[Any]								=			self.dummy_sample
        lowerCAmelCase_       :    Optional[Any]								=			0.1 * sample
        lowerCAmelCase_       :    Optional[int]								=			[residual + 0.2, residual + 0.15, residual + 0.10]

        for scheduler_class in self.scheduler_classes:
          lowerCAmelCase_       :    Tuple								=			self.get_scheduler_config(**SCREAMING_SNAKE_CASE_	)
          lowerCAmelCase_       :    List[Any]								=			scheduler_class(**SCREAMING_SNAKE_CASE_	)
          scheduler.set_timesteps(SCREAMING_SNAKE_CASE_	)
          # copy over dummy past residuals
          lowerCAmelCase_       :    Any								=			dummy_past_residuals[: scheduler.config.solver_order]

          with tempfile.TemporaryDirectory() as tmpdirname:
            scheduler.save_config(SCREAMING_SNAKE_CASE_	)
            lowerCAmelCase_       :    Dict								=			scheduler_class.from_pretrained(SCREAMING_SNAKE_CASE_	)
            new_scheduler.set_timesteps(SCREAMING_SNAKE_CASE_	)
            # copy over dummy past residuals
            lowerCAmelCase_       :    Any								=			dummy_past_residuals[: new_scheduler.config.solver_order]

          lowerCAmelCase_						,lowerCAmelCase_       :    Dict								=			sample, sample
          for t in range(SCREAMING_SNAKE_CASE_   ,      time_step + scheduler.config.solver_order + 1	):
            lowerCAmelCase_       :    Union[str, Any]								=			scheduler.step(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_   ,      **SCREAMING_SNAKE_CASE_	).prev_sample
            lowerCAmelCase_       :    Tuple								=			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 					SCREAMING_SNAKE_CASE__			(				self    :				str	):
        pass



      def 					SCREAMING_SNAKE_CASE__			(				self    :				int   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]=0   ,      **SCREAMING_SNAKE_CASE_    :				Optional[Any]	):
        lowerCAmelCase_       :    List[Any]								=			dict(self.forward_default_kwargs	)
        lowerCAmelCase_       :    Dict								=			kwargs.pop('num_inference_steps'   ,      SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    Optional[Any]								=			self.dummy_sample
        lowerCAmelCase_       :    List[str]								=			0.1 * sample
        lowerCAmelCase_       :    List[str]								=			[residual + 0.2, residual + 0.15, residual + 0.10]

        for scheduler_class in self.scheduler_classes:
          lowerCAmelCase_       :    List[Any]								=			self.get_scheduler_config()
          lowerCAmelCase_       :    Union[str, Any]								=			scheduler_class(**SCREAMING_SNAKE_CASE_	)
          scheduler.set_timesteps(SCREAMING_SNAKE_CASE_	)

          # copy over dummy past residuals (must be after setting timesteps)
          lowerCAmelCase_       :    int								=			dummy_past_residuals[: scheduler.config.solver_order]

          with tempfile.TemporaryDirectory() as tmpdirname:
            scheduler.save_config(SCREAMING_SNAKE_CASE_	)
            lowerCAmelCase_       :    List[Any]								=			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)
            lowerCAmelCase_       :    Dict								=			dummy_past_residuals[: new_scheduler.config.solver_order]

          lowerCAmelCase_       :    Dict								=			scheduler.step(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_   ,      **SCREAMING_SNAKE_CASE_	).prev_sample
          lowerCAmelCase_       :    List[Any]								=			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 					SCREAMING_SNAKE_CASE__			(				self    :				str   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]=None   ,      **SCREAMING_SNAKE_CASE_    :				int	):


        if scheduler is None:
          lowerCAmelCase_       :    Optional[int]								=			self.scheduler_classes[0]
          lowerCAmelCase_       :    Any								=			self.get_scheduler_config(**SCREAMING_SNAKE_CASE_	)
          lowerCAmelCase_       :    str								=			scheduler_class(**SCREAMING_SNAKE_CASE_	)

        lowerCAmelCase_       :    List[str]								=			self.scheduler_classes[0]
        lowerCAmelCase_       :    Tuple								=			self.get_scheduler_config(**SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    int								=			scheduler_class(**SCREAMING_SNAKE_CASE_	)

        lowerCAmelCase_       :    int								=			1_0
        lowerCAmelCase_       :    int								=			self.dummy_model()
        lowerCAmelCase_       :    List[Any]								=			self.dummy_sample_deter
        scheduler.set_timesteps(SCREAMING_SNAKE_CASE_	)

        for i, t in enumerate(scheduler.timesteps	):
          lowerCAmelCase_       :    Tuple								=			model(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	)
          lowerCAmelCase_       :    List[Any]								=			scheduler.step(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	).prev_sample

        return sample



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[Any]	):
        lowerCAmelCase_       :    Optional[Any]								=			DPMSolverSinglestepScheduler(**self.get_scheduler_config()	)
        lowerCAmelCase_       :    Any								=			5_0
        lowerCAmelCase_       :    int								=			self.dummy_model()
        lowerCAmelCase_       :    Union[str, Any]								=			self.dummy_sample_deter
        scheduler.set_timesteps(SCREAMING_SNAKE_CASE_	)

        # make sure that the first t is uneven
        for i, t in enumerate(scheduler.timesteps[3:]	):
          lowerCAmelCase_       :    Tuple								=			model(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	)
          lowerCAmelCase_       :    Dict								=			scheduler.step(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	).prev_sample

        lowerCAmelCase_       :    Optional[Any]								=			torch.mean(torch.abs(SCREAMING_SNAKE_CASE_	)	)

        assert abs(result_mean.item() - 0.25_74	) < 1E-3



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[Any]	):


        for timesteps in [2_5, 5_0, 1_0_0, 9_9_9, 1_0_0_0]:
          self.check_over_configs(num_train_timesteps=SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Dict	):
        lowerCAmelCase_       :    str								=			DPMSolverSinglestepScheduler(**self.get_scheduler_config()	)
        lowerCAmelCase_       :    int								=			self.full_loop(scheduler=SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    Optional[Any]								=			torch.mean(torch.abs(SCREAMING_SNAKE_CASE_	)	)

        assert abs(result_mean.item() - 0.27_91	) < 1E-3

        lowerCAmelCase_       :    Dict								=			DEISMultistepScheduler.from_config(scheduler.config	)
        lowerCAmelCase_       :    List[str]								=			DPMSolverMultistepScheduler.from_config(scheduler.config	)
        lowerCAmelCase_       :    Optional[Any]								=			UniPCMultistepScheduler.from_config(scheduler.config	)
        lowerCAmelCase_       :    Any								=			DPMSolverSinglestepScheduler.from_config(scheduler.config	)

        lowerCAmelCase_       :    List[str]								=			self.full_loop(scheduler=SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    Optional[Any]								=			torch.mean(torch.abs(SCREAMING_SNAKE_CASE_	)	)

        assert abs(result_mean.item() - 0.27_91	) < 1E-3



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Union[str, Any]	):
        self.check_over_configs(thresholding=SCREAMING_SNAKE_CASE_	)
        for order in [1, 2, 3]:
          for solver_type in ["midpoint", "heun"]:
            for threshold in [0.5, 1.0, 2.0]:
              for prediction_type in ["epsilon", "sample"]:
                self.check_over_configs(
                    thresholding=SCREAMING_SNAKE_CASE_   ,      prediction_type=SCREAMING_SNAKE_CASE_   ,      sample_max_value=SCREAMING_SNAKE_CASE_   ,      algorithm_type='dpmsolver++'   ,      solver_order=SCREAMING_SNAKE_CASE_   ,      solver_type=SCREAMING_SNAKE_CASE_   ,      )



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Any	):


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



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Tuple	):


        for algorithm_type in ["dpmsolver", "dpmsolver++"]:
          for solver_type in ["midpoint", "heun"]:
            for order in [1, 2, 3]:
              for prediction_type in ["epsilon", "sample"]:
                self.check_over_configs(
                    solver_order=SCREAMING_SNAKE_CASE_   ,      solver_type=SCREAMING_SNAKE_CASE_   ,      prediction_type=SCREAMING_SNAKE_CASE_   ,      algorithm_type=SCREAMING_SNAKE_CASE_   ,      )
                lowerCAmelCase_       :    List[str]								=			self.full_loop(
                    solver_order=SCREAMING_SNAKE_CASE_   ,      solver_type=SCREAMING_SNAKE_CASE_   ,      prediction_type=SCREAMING_SNAKE_CASE_   ,      algorithm_type=SCREAMING_SNAKE_CASE_   ,      )
                assert not torch.isnan(SCREAMING_SNAKE_CASE_	).any(), "Samples have nan numbers"



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Union[str, Any]	):
        self.check_over_configs(lower_order_final=SCREAMING_SNAKE_CASE_	)
        self.check_over_configs(lower_order_final=SCREAMING_SNAKE_CASE_	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Tuple	):
        self.check_over_configs(lambda_min_clipped=-float('inf'	)	)
        self.check_over_configs(lambda_min_clipped=-5.1	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Union[str, Any]	):
        self.check_over_configs(variance_type=SCREAMING_SNAKE_CASE_	)
        self.check_over_configs(variance_type='learned_range'	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[Any]	):


        for num_inference_steps in [1, 2, 3, 5, 1_0, 5_0, 1_0_0, 9_9_9, 1_0_0_0]:
          self.check_over_forward(num_inference_steps=SCREAMING_SNAKE_CASE_   ,      time_step=0	)



      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[str]	):
        lowerCAmelCase_       :    Any								=			self.full_loop()
        lowerCAmelCase_       :    List[str]								=			torch.mean(torch.abs(SCREAMING_SNAKE_CASE_	)	)

        assert abs(result_mean.item() - 0.27_91	) < 1E-3



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Optional[int]	):
        lowerCAmelCase_       :    Optional[int]								=			self.full_loop(use_karras_sigmas=SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    Tuple								=			torch.mean(torch.abs(SCREAMING_SNAKE_CASE_	)	)

        assert abs(result_mean.item() - 0.22_48	) < 1E-3



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Any	):
        lowerCAmelCase_       :    Dict								=			self.full_loop(prediction_type='v_prediction'	)
        lowerCAmelCase_       :    Tuple								=			torch.mean(torch.abs(SCREAMING_SNAKE_CASE_	)	)

        assert abs(result_mean.item() - 0.14_53	) < 1E-3



      def 					SCREAMING_SNAKE_CASE__			(				self    :				Tuple	):
        lowerCAmelCase_       :    Union[str, Any]								=			self.full_loop(prediction_type='v_prediction'   ,      use_karras_sigmas=SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    Dict								=			torch.mean(torch.abs(SCREAMING_SNAKE_CASE_	)	)

        assert abs(result_mean.item() - 0.06_49	) < 1E-3





      def 					SCREAMING_SNAKE_CASE__			(				self    :				Any	):
        lowerCAmelCase_       :    Optional[int]								=			self.scheduler_classes[0]
        lowerCAmelCase_       :    Any								=			self.get_scheduler_config(thresholding=SCREAMING_SNAKE_CASE_   ,      dynamic_thresholding_ratio=0	)
        lowerCAmelCase_       :    Union[str, Any]								=			scheduler_class(**SCREAMING_SNAKE_CASE_	)

        lowerCAmelCase_       :    Tuple								=			1_0
        lowerCAmelCase_       :    List[str]								=			self.dummy_model()
        lowerCAmelCase_       :    Optional[Any]								=			self.dummy_sample_deter.half()
        scheduler.set_timesteps(SCREAMING_SNAKE_CASE_	)

        for i, t in enumerate(scheduler.timesteps	):
          lowerCAmelCase_       :    Optional[Any]								=			model(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	)
          lowerCAmelCase_       :    Any								=			scheduler.step(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	).prev_sample

        assert sample.dtype == torch.floataa










"""simple docstring"""







from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowercase__		:    List[Any]				   =     logging.get_logger(__name__)

lowercase__		:    List[str]				   =     {
    """s-JoL/Open-Llama-V1""": """https://huggingface.co/s-JoL/Open-Llama-V1/blob/main/config.json""",
}







class 				UpperCamelCase__       (     lowercase_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE      =					"""open-llama"""



      def __init__(				self    :				Optional[int]   ,      SCREAMING_SNAKE_CASE_    :				Tuple=1_0_0_0_0_0   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]=4_0_9_6   ,      SCREAMING_SNAKE_CASE_    :				List[Any]=1_1_0_0_8   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]=3_2   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]=3_2   ,      SCREAMING_SNAKE_CASE_    :				str="silu"   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]=2_0_4_8   ,      SCREAMING_SNAKE_CASE_    :				Union[str, Any]=0.02   ,      SCREAMING_SNAKE_CASE_    :				int=1E-6   ,      SCREAMING_SNAKE_CASE_    :				List[str]=True   ,      SCREAMING_SNAKE_CASE_    :				Dict=0   ,      SCREAMING_SNAKE_CASE_    :				Dict=1   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]=2   ,      SCREAMING_SNAKE_CASE_    :				str=False   ,      SCREAMING_SNAKE_CASE_    :				Optional[int]=True   ,      SCREAMING_SNAKE_CASE_    :				List[Any]=0.1   ,      SCREAMING_SNAKE_CASE_    :				Optional[Any]=0.1   ,      SCREAMING_SNAKE_CASE_    :				Dict=True   ,      SCREAMING_SNAKE_CASE_    :				Any=True   ,      SCREAMING_SNAKE_CASE_    :				Tuple=None   ,      **SCREAMING_SNAKE_CASE_    :				Optional[int]   ,      ):
        lowerCAmelCase_       :    Union[str, Any]								=			vocab_size
        lowerCAmelCase_       :    int								=			max_position_embeddings
        lowerCAmelCase_       :    Union[str, Any]								=			hidden_size
        lowerCAmelCase_       :    Optional[Any]								=			intermediate_size
        lowerCAmelCase_       :    List[Any]								=			num_hidden_layers
        lowerCAmelCase_       :    Optional[int]								=			num_attention_heads
        lowerCAmelCase_       :    List[Any]								=			hidden_act
        lowerCAmelCase_       :    List[Any]								=			initializer_range
        lowerCAmelCase_       :    List[str]								=			rms_norm_eps
        lowerCAmelCase_       :    List[Any]								=			use_cache
        lowerCAmelCase_       :    Optional[int]								=			kwargs.pop(
            'use_memorry_efficient_attention'   ,      SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    List[Any]								=			hidden_dropout_prob
        lowerCAmelCase_       :    List[str]								=			attention_dropout_prob
        lowerCAmelCase_       :    Tuple								=			use_stable_embedding
        lowerCAmelCase_       :    Optional[Any]								=			shared_input_output_embedding
        lowerCAmelCase_       :    Optional[Any]								=			rope_scaling
        self._rope_scaling_validation()

        super().__init__(
            pad_token_id=SCREAMING_SNAKE_CASE_   ,      bos_token_id=SCREAMING_SNAKE_CASE_   ,      eos_token_id=SCREAMING_SNAKE_CASE_   ,      tie_word_embeddings=SCREAMING_SNAKE_CASE_   ,      **SCREAMING_SNAKE_CASE_   ,      )





      def 					SCREAMING_SNAKE_CASE__			(				self    :				List[str]	):
        if self.rope_scaling is None:
          return

        if not isinstance(self.rope_scaling   ,      SCREAMING_SNAKE_CASE_	) or len(self.rope_scaling	) != 2:
          raise ValueError(
              '`rope_scaling` must be a dictionary with with two fields, `name` and `factor`, '
              F"got {self.rope_scaling}"	)
        lowerCAmelCase_       :    int								=			self.rope_scaling.get('type'   ,      SCREAMING_SNAKE_CASE_	)
        lowerCAmelCase_       :    Union[str, Any]								=			self.rope_scaling.get('factor'   ,      SCREAMING_SNAKE_CASE_	)
        if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
          raise ValueError(
              F"`rope_scaling`'s name field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"	)
        if rope_scaling_factor is None or not isinstance(SCREAMING_SNAKE_CASE_   ,      SCREAMING_SNAKE_CASE_	) or rope_scaling_factor <= 1.0:
          raise ValueError(F"`rope_scaling`'s factor field must be an float > 1, got {rope_scaling_factor}"	)





