
import unittest

import torch

from diffusers import UNetaDModel
from diffusers.utils import floats_tensor, slow, torch_device

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin

class __snake_case	(   lowerCAmelCase      ,			lowerCAmelCase      ,			unittest.TestCase  ):
				_a	:		Any=       UNetaDModel
				_a	:		Any=       "sample"
				@property
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											lowercase							:  Tuple 			=   4
											lowercase							:  Any 			=   14
											lowercase							:  List[Any] 			=   16

											lowercase							:  Optional[int] 			=   floats_tensor((batch_size, num_features, seq_len)   ).to(snake_case   )
											lowercase							:  Union[str, Any] 			=   torch.tensor([10] * batch_size   ).to(snake_case   )

											return {"sample": noise, "timestep": time_step}
				@property
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											return (4, 14, 16)
				@property
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											return (4, 14, 16)
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											pass
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											pass
				@unittest.skipIf(torch_device == """mps"""	,"""mish op not supported in MPS"""   )
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											super().test_determinism()
				@unittest.skipIf(torch_device == """mps"""	,"""mish op not supported in MPS"""   )
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											super().test_outputs_equivalence()
				@unittest.skipIf(torch_device == """mps"""	,"""mish op not supported in MPS"""   )
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											super().test_from_save_pretrained()
				@unittest.skipIf(torch_device == """mps"""	,"""mish op not supported in MPS"""   )
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											super().test_from_save_pretrained_variant()
				@unittest.skipIf(torch_device == """mps"""	,"""mish op not supported in MPS"""   )
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											super().test_model_from_pretrained()
				@unittest.skipIf(torch_device == """mps"""	,"""mish op not supported in MPS"""   )
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											super().test_output()
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											lowercase							:  List[str] 			=   {
											    """block_out_channels""": (32, 64, 128, 256),
											    """in_channels""": 14,
											    """out_channels""": 14,
											    """time_embedding_type""": """positional""",
											    """use_timestep_embedding""": True,
											    """flip_sin_to_cos""": False,
											    """freq_shift""": 1.0,
											    """out_block_type""": """OutConv1DBlock""",
											    """mid_block_type""": """MidResTemporalBlock1D""",
											    """down_block_types""": ("""DownResnetBlock1D""", """DownResnetBlock1D""", """DownResnetBlock1D""", """DownResnetBlock1D"""),
											    """up_block_types""": ("""UpResnetBlock1D""", """UpResnetBlock1D""", """UpResnetBlock1D"""),
											    """act_fn""": """mish""",
											}
											lowercase							:  List[str] 			=   self.dummy_input
											return init_dict, inputs_dict
				@unittest.skipIf(torch_device == """mps"""	,"""mish op not supported in MPS"""   )
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											lowercase							:  str 			=   UNetaDModel.from_pretrained(
											    """bglick13/hopper-medium-v2-value-function-hor32"""	,output_loading_info=snake_case	,subfolder="""unet"""   )
											self.assertIsNotNone(snake_case   )
											self.assertEqual(len(loading_info["""missing_keys"""]   )	,0   )

											model.to(snake_case   )
											lowercase							:  Dict 			=   model(**self.dummy_input   )

											assert image is not None, "Make sure output is not None"
				@unittest.skipIf(torch_device == """mps"""	,"""mish op not supported in MPS"""   )
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											lowercase							:  Optional[int] 			=   UNetaDModel.from_pretrained("""bglick13/hopper-medium-v2-value-function-hor32"""	,subfolder="""unet"""   )
											torch.manual_seed(0   )
											if torch.cuda.is_available():
																		torch.cuda.manual_seed_all(0   )

											lowercase							:  List[Any] 			=   model.config.in_channels
											lowercase							:  int 			=   16
											lowercase							:  List[Any] 			=   torch.randn((1, seq_len, num_features)   ).permute(
											    0	,2	,1   )  # match original, we can update values and remove
											lowercase							:  List[str] 			=   torch.full((num_features,)	,0   )

											with torch.no_grad():
																		lowercase							:  str 			=   model(snake_case	,snake_case   ).sample.permute(0	,2	,1   )

											lowercase							:  int 			=   output[0, -3:, -3:].flatten()
											# fmt: off
											lowercase							:  Tuple 			=   torch.tensor([-2.137_172, 1.1_426_016, 0.3_688_687, -0.766_922, 0.7_303_146, 0.11_038_864, -0.4_760_633, 0.13_270_172, 0.02_591_348]   )
											# fmt: on
											self.assertTrue(torch.allclose(snake_case	,snake_case	,rtol=1e-3   )   )
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											pass
				@slow
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											lowercase							:  Tuple 			=   """harmonai/maestro-150k"""
											lowercase							:  List[Any] 			=   UNetaDModel.from_pretrained(snake_case	,subfolder="""unet"""   )
											model.to(snake_case   )

											lowercase							:  List[str] 			=   65536
											lowercase							:  int 			=   torch.sin(torch.arange(snake_case   )[None, None, :].repeat(1	,2	,1   )   ).to(snake_case   )
											lowercase							:  Dict 			=   torch.tensor([1]   ).to(snake_case   )

											with torch.no_grad():
																		lowercase							:  Any 			=   model(snake_case	,snake_case   ).sample

											lowercase							:  Optional[Any] 			=   output.abs().sum()
											lowercase							:  List[str] 			=   output.abs().max()

											assert (output_sum - 224.0_896).abs() < 0.5
											assert (output_max - 0.0_607).abs() < 4e-4






class __snake_case	(   lowerCAmelCase      ,			lowerCAmelCase      ,			unittest.TestCase  ):
				_a	:		int=       UNetaDModel
				_a	:		Union[str, Any]=       "sample"
				@property
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											lowercase							:  Tuple 			=   4
											lowercase							:  Any 			=   14
											lowercase							:  List[str] 			=   16

											lowercase							:  Optional[int] 			=   floats_tensor((batch_size, num_features, seq_len)   ).to(snake_case   )
											lowercase							:  str 			=   torch.tensor([10] * batch_size   ).to(snake_case   )

											return {"sample": noise, "timestep": time_step}
				@property
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											return (4, 14, 16)
				@property
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											return (4, 14, 1)
				@unittest.skipIf(torch_device == """mps"""	,"""mish op not supported in MPS"""   )
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											super().test_determinism()
				@unittest.skipIf(torch_device == """mps"""	,"""mish op not supported in MPS"""   )
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											super().test_outputs_equivalence()
				@unittest.skipIf(torch_device == """mps"""	,"""mish op not supported in MPS"""   )
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											super().test_from_save_pretrained()
				@unittest.skipIf(torch_device == """mps"""	,"""mish op not supported in MPS"""   )
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											super().test_from_save_pretrained_variant()
				@unittest.skipIf(torch_device == """mps"""	,"""mish op not supported in MPS"""   )
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											super().test_model_from_pretrained()
				@unittest.skipIf(torch_device == """mps"""	,"""mish op not supported in MPS"""   )
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											lowercase							:  Optional[Any] 			=   self.prepare_init_args_and_inputs_for_common()
											lowercase							:  Any 			=   self.model_class(**snake_case   )
											model.to(snake_case   )
											model.eval()

											with torch.no_grad():
																		lowercase							:  Dict 			=   model(**snake_case   )

																		if isinstance(snake_case	,snake_case   ):
																									lowercase							:  str 			=   output.sample

											self.assertIsNotNone(snake_case   )
											lowercase							:  Optional[int] 			=   torch.Size((inputs_dict["""sample"""].shape[0], 1)   )
											self.assertEqual(output.shape	,snake_case	,"""Input and output shapes do not match"""   )
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											pass
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											pass
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											lowercase							:  List[Any] 			=   {
											    """in_channels""": 14,
											    """out_channels""": 14,
											    """down_block_types""": ["""DownResnetBlock1D""", """DownResnetBlock1D""", """DownResnetBlock1D""", """DownResnetBlock1D"""],
											    """up_block_types""": [],
											    """out_block_type""": """ValueFunction""",
											    """mid_block_type""": """ValueFunctionMidBlock1D""",
											    """block_out_channels""": [32, 64, 128, 256],
											    """layers_per_block""": 1,
											    """downsample_each_block""": True,
											    """use_timestep_embedding""": True,
											    """freq_shift""": 1.0,
											    """flip_sin_to_cos""": False,
											    """time_embedding_type""": """positional""",
											    """act_fn""": """mish""",
											}
											lowercase							:  Optional[Any] 			=   self.dummy_input
											return init_dict, inputs_dict
				@unittest.skipIf(torch_device == """mps"""	,"""mish op not supported in MPS"""   )
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											lowercase							:  Union[str, Any] 			=   UNetaDModel.from_pretrained(
											    """bglick13/hopper-medium-v2-value-function-hor32"""	,output_loading_info=snake_case	,subfolder="""value_function"""   )
											self.assertIsNotNone(snake_case   )
											self.assertEqual(len(vf_loading_info["""missing_keys"""]   )	,0   )

											value_function.to(snake_case   )
											lowercase							:  Dict 			=   value_function(**self.dummy_input   )

											assert image is not None, "Make sure output is not None"
				@unittest.skipIf(torch_device == """mps"""	,"""mish op not supported in MPS"""   )
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											lowercase							:  Tuple 			=   UNetaDModel.from_pretrained(
											    """bglick13/hopper-medium-v2-value-function-hor32"""	,output_loading_info=snake_case	,subfolder="""value_function"""   )
											torch.manual_seed(0   )
											if torch.cuda.is_available():
																		torch.cuda.manual_seed_all(0   )

											lowercase							:  Tuple 			=   value_function.config.in_channels
											lowercase							:  Tuple 			=   14
											lowercase							:  Tuple 			=   torch.randn((1, seq_len, num_features)   ).permute(
											    0	,2	,1   )  # match original, we can update values and remove
											lowercase							:  List[Any] 			=   torch.full((num_features,)	,0   )

											with torch.no_grad():
																		lowercase							:  List[str] 			=   value_function(snake_case	,snake_case   ).sample

											# fmt: off
											lowercase							:  Tuple 			=   torch.tensor([165.25] * seq_len   )
											# fmt: on
											self.assertTrue(torch.allclose(snake_case	,snake_case	,rtol=1e-3   )   )
				def    _SCREAMING_SNAKE_CASE (						self   ):



											'''simple docstring'''







											pass


from typing import TYPE_CHECKING

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


lowercase						:      Any							=							{
    """configuration_gpt2""": ["""GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP""", """GPT2Config""", """GPT2OnnxConfig"""],
    """tokenization_gpt2""": ["""GPT2Tokenizer"""],
}

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowercase						:      List[Any]							=							["""GPT2TokenizerFast"""]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowercase						:      str							=							[
      """GPT2_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """GPT2DoubleHeadsModel""",
      """GPT2ForQuestionAnswering""",
      """GPT2ForSequenceClassification""",
      """GPT2ForTokenClassification""",
      """GPT2LMHeadModel""",
      """GPT2Model""",
      """GPT2PreTrainedModel""",
      """load_tf_weights_in_gpt2""",
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowercase						:      Tuple							=							[
      """TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """TFGPT2DoubleHeadsModel""",
      """TFGPT2ForSequenceClassification""",
      """TFGPT2LMHeadModel""",
      """TFGPT2MainLayer""",
      """TFGPT2Model""",
      """TFGPT2PreTrainedModel""",
  ]

try:
  if not is_keras_nlp_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowercase						:      List[str]							=							["""TFGPT2Tokenizer"""]

try:
  if not is_flax_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  lowercase						:      str							=							["""FlaxGPT2LMHeadModel""", """FlaxGPT2Model""", """FlaxGPT2PreTrainedModel"""]

if TYPE_CHECKING:
  from .configuration_gpta import GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTaConfig, GPTaOnnxConfig
  from .tokenization_gpta import GPTaTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_gpta_fast import GPTaTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_gpta import (
        GPT2_PRETRAINED_MODEL_ARCHIVE_LIST,
        GPTaDoubleHeadsModel,
        GPTaForQuestionAnswering,
        GPTaForSequenceClassification,
        GPTaForTokenClassification,
        GPTaLMHeadModel,
        GPTaModel,
        GPTaPreTrainedModel,
        load_tf_weights_in_gpta,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_gpta import (
        TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFGPTaDoubleHeadsModel,
        TFGPTaForSequenceClassification,
        TFGPTaLMHeadModel,
        TFGPTaMainLayer,
        TFGPTaModel,
        TFGPTaPreTrainedModel,
    )

  try:
    if not is_keras_nlp_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_gpta_tf import TFGPTaTokenizer

  try:
    if not is_flax_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_flax_gpta import FlaxGPTaLMHeadModel, FlaxGPTaModel, FlaxGPTaPreTrainedModel

else:
  import sys

  lowercase						:      Union[str, Any]							=							_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
