




'''simple docstring'''




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


a_   :	Union[str, Any]							   = {
    """configuration_openai""": ["""OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP""", """OpenAIGPTConfig"""],
    """tokenization_openai""": ["""OpenAIGPTTokenizer"""],
}

try:
     if not is_tokenizers_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     a_   :	Optional[int]							   = ["""OpenAIGPTTokenizerFast"""]

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     a_   :	int							   = [
         """OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST""",
         """OpenAIGPTDoubleHeadsModel""",
         """OpenAIGPTForSequenceClassification""",
         """OpenAIGPTLMHeadModel""",
         """OpenAIGPTModel""",
         """OpenAIGPTPreTrainedModel""",
         """load_tf_weights_in_openai_gpt""",
     ]

try:
     if not is_tf_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     a_   :	Dict							   = [
         """TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST""",
         """TFOpenAIGPTDoubleHeadsModel""",
         """TFOpenAIGPTForSequenceClassification""",
         """TFOpenAIGPTLMHeadModel""",
         """TFOpenAIGPTMainLayer""",
         """TFOpenAIGPTModel""",
         """TFOpenAIGPTPreTrainedModel""",
     ]


if TYPE_CHECKING:
     from .configuration_openai import OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP, OpenAIGPTConfig
     from .tokenization_openai import OpenAIGPTTokenizer

     try:
          if not is_tokenizers_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .tokenization_openai_fast import OpenAIGPTTokenizerFast

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_openai import (
              OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
              OpenAIGPTDoubleHeadsModel,
              OpenAIGPTForSequenceClassification,
              OpenAIGPTLMHeadModel,
              OpenAIGPTModel,
              OpenAIGPTPreTrainedModel,
              load_tf_weights_in_openai_gpt,
          )

     try:
          if not is_tf_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_tf_openai import (
              TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
              TFOpenAIGPTDoubleHeadsModel,
              TFOpenAIGPTForSequenceClassification,
              TFOpenAIGPTLMHeadModel,
              TFOpenAIGPTMainLayer,
              TFOpenAIGPTModel,
              TFOpenAIGPTPreTrainedModel,
          )

else:
     import sys

     a_   :	List[Any]							   = _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)










'''simple docstring'''




import torch

from diffusers import HeunDiscreteScheduler
from diffusers.utils import torch_device

from .test_schedulers import SchedulerCommonTest


class 						__UpperCamelCase (				lowerCamelCase__       ):
       lowercase			:					Union[str, Any]   		=(HeunDiscreteScheduler,)
       lowercase			:					int   		=10






       def   lowercase__  (							self,	**lowerCAmelCase						):



              """simple docstring"""







              lowerCamelCase_							={
                  '''num_train_timesteps''': 1_100,
                  '''beta_start''': 0.0_0_0_1,
                  '''beta_end''': 0.0_2,
                  '''beta_schedule''': '''linear''',
              }

              config.update(**lowerCAmelCase						)
              return config






       def   lowercase__  (							self						):



              """simple docstring"""







              for timesteps in [10, 50, 100, 1_000]:
                     self.check_over_configs(num_train_timesteps=lowerCAmelCase						)






       def   lowercase__  (							self						):



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






       def   lowercase__  (							self						):



              """simple docstring"""







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






       def   lowercase__  (							self						):



              """simple docstring"""







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






       def   lowercase__  (							self						):



              """simple docstring"""







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






       def   lowercase__  (							self						):



              """simple docstring"""







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

              scheduler.set_timesteps(self.num_inference_steps						)

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

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

                     lowerCamelCase_							=model(lowerCAmelCase,	lowerCAmelCase						)

                     lowerCamelCase_							=scheduler.step(lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase						)
                     lowerCamelCase_							=output.prev_sample

              lowerCamelCase_							=torch.sum(torch.abs(lowerCAmelCase						)						)
              lowerCamelCase_							=torch.mean(torch.abs(lowerCAmelCase						)						)

              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   lowercase__  (							self						):



              """simple docstring"""







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

              scheduler.set_timesteps(self.num_inference_steps						)

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

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

                     lowerCamelCase_							=model(lowerCAmelCase,	lowerCAmelCase						)

                     lowerCamelCase_							=scheduler.step(lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase						)
                     lowerCamelCase_							=output.prev_sample

              lowerCamelCase_							=torch.sum(torch.abs(lowerCAmelCase						)						)
              lowerCamelCase_							=torch.mean(torch.abs(lowerCAmelCase						)						)

              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   lowercase__  (							self						):



              """simple docstring"""







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

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

              lowerCamelCase_							=self.dummy_model()
              lowerCamelCase_							=self.dummy_sample_deter.to(lowerCAmelCase						) * scheduler.init_noise_sigma

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

                     lowerCamelCase_							=model(lowerCAmelCase,	lowerCAmelCase						)

                     lowerCamelCase_							=scheduler.step(lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase						)
                     lowerCamelCase_							=output.prev_sample

              lowerCamelCase_							=torch.sum(torch.abs(lowerCAmelCase						)						)
              lowerCamelCase_							=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() - 0.1_2_3_3						) < 1e-2
                     assert abs(result_mean.item() - 0.0_0_0_2						) < 1e-3
              elif str(lowerCAmelCase						).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   lowercase__  (							self						):



              """simple docstring"""







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

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

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

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

                     lowerCamelCase_							=model(lowerCAmelCase,	lowerCAmelCase						)

                     lowerCamelCase_							=scheduler.step(lowerCAmelCase,	lowerCAmelCase,	lowerCAmelCase						)
                     lowerCamelCase_							=output.prev_sample

              lowerCamelCase_							=torch.sum(torch.abs(lowerCAmelCase						)						)
              lowerCamelCase_							=torch.mean(torch.abs(lowerCAmelCase						)						)

              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




