




'''simple docstring'''
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 	SCREAMING_SNAKE_CASE(    __A       ,				__A       ,				unittest.TestCase					):
       snake_case_			:					Union[str, Any]   = UNetaDModel
       snake_case_			:					str   = """sample"""

       @property
       def    snake_case__      (   self					)   -> List[str]:

           """simple docstring"""



           __lowercase          =						4
           __lowercase          =						14
           __lowercase          =						16

           __lowercase          =						floats_tensor((batch_size, num_features, seq_len)					).to(lowerCamelCase__					)
           __lowercase          =						torch.tensor([10] * batch_size					).to(lowerCamelCase__					)

           return {"sample": noise, "timestep": time_step}

       @property
       def    snake_case__      (   self					)   -> Tuple:

           """simple docstring"""



           return (4, 14, 16)

       @property
       def    snake_case__      (   self					)   -> Optional[int]:

           """simple docstring"""



           return (4, 14, 16)

       def    snake_case__      (   self					)   -> str:

           """simple docstring"""



           pass

       def    snake_case__      (   self					)   -> Union[str, Any]:

           """simple docstring"""



           pass

       @unittest.skipIf(torch_device == """mps""" ,		"""mish op not supported in MPS"""					)
       def    snake_case__      (   self					)   -> Union[str, Any]:

           """simple docstring"""



           super().test_determinism()

       @unittest.skipIf(torch_device == """mps""" ,		"""mish op not supported in MPS"""					)
       def    snake_case__      (   self					)   -> int:

           """simple docstring"""



           super().test_outputs_equivalence()

       @unittest.skipIf(torch_device == """mps""" ,		"""mish op not supported in MPS"""					)
       def    snake_case__      (   self					)   -> List[str]:

           """simple docstring"""



           super().test_from_save_pretrained()

       @unittest.skipIf(torch_device == """mps""" ,		"""mish op not supported in MPS"""					)
       def    snake_case__      (   self					)   -> Union[str, Any]:

           """simple docstring"""



           super().test_from_save_pretrained_variant()

       @unittest.skipIf(torch_device == """mps""" ,		"""mish op not supported in MPS"""					)
       def    snake_case__      (   self					)   -> Optional[Any]:

           """simple docstring"""



           super().test_model_from_pretrained()

       @unittest.skipIf(torch_device == """mps""" ,		"""mish op not supported in MPS"""					)
       def    snake_case__      (   self					)   -> str:

           """simple docstring"""



           super().test_output()

       def    snake_case__      (   self					)   -> List[Any]:

           """simple docstring"""



           __lowercase          =						{
               """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          =						self.dummy_input
           return init_dict, inputs_dict

       @unittest.skipIf(torch_device == """mps""" ,		"""mish op not supported in MPS"""					)
       def    snake_case__      (   self					)   -> Tuple:

           """simple docstring"""



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

           model.to(lowerCamelCase__					)
           __lowercase          =						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    snake_case__      (   self					)   -> Union[str, Any]:

           """simple docstring"""



           __lowercase          =						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          =						model.config.in_channels
           __lowercase          =						16
           __lowercase          =						torch.randn((1, seq_len, num_features)					).permute(
               0 ,		2 ,		1					)  # match original, we can update values and remove
           __lowercase          =						torch.full((num_features,) ,		0					)

           with torch.no_grad():
               __lowercase          =						model(lowerCamelCase__ ,		lowerCamelCase__					).sample.permute(0 ,		2 ,		1					)

           __lowercase          =						output[0, -3:, -3:].flatten()
           # fmt: off
           __lowercase          =						torch.tensor([-2.13_71_72, 1.1_42_60_16, 0.3_68_86_87, -0.76_69_22, 0.7_30_31_46, 0.11_03_88_64, -0.4_76_06_33, 0.13_27_01_72, 0.02_59_13_48]					)
           # fmt: on
           self.assertTrue(torch.allclose(lowerCamelCase__ ,		lowerCamelCase__ ,		rtol=1E-3					)					)

       def    snake_case__      (   self					)   -> Optional[int]:

           """simple docstring"""



           pass




       @slow
       def    snake_case__      (   self					)   -> List[str]:

           """simple docstring"""



           __lowercase          =						"""harmonai/maestro-150k"""
           __lowercase          =						UNetaDModel.from_pretrained(lowerCamelCase__ ,		subfolder="""unet"""					)
           model.to(lowerCamelCase__					)

           __lowercase          =						6_5536
           __lowercase          =						torch.sin(torch.arange(lowerCamelCase__					)[None, None, :].repeat(1 ,		2 ,		1					)					).to(lowerCamelCase__					)
           __lowercase          =						torch.tensor([1]					).to(lowerCamelCase__					)

           with torch.no_grad():
               __lowercase          =						model(lowerCamelCase__ ,		lowerCamelCase__					).sample

           __lowercase          =						output.abs().sum()
           __lowercase          =						output.abs().max()

           assert (output_sum - 2_24.08_96).abs() < 0.5
           assert (output_max - 0.06_07).abs() < 4E-4






class 	SCREAMING_SNAKE_CASE(    __A       ,				__A       ,				unittest.TestCase					):
       snake_case_			:					int   = UNetaDModel
       snake_case_			:					List[str]   = """sample"""

       @property
       def    snake_case__      (   self					)   -> Tuple:

           """simple docstring"""



           __lowercase          =						4
           __lowercase          =						14
           __lowercase          =						16

           __lowercase          =						floats_tensor((batch_size, num_features, seq_len)					).to(lowerCamelCase__					)
           __lowercase          =						torch.tensor([10] * batch_size					).to(lowerCamelCase__					)

           return {"sample": noise, "timestep": time_step}

       @property
       def    snake_case__      (   self					)   -> Any:

           """simple docstring"""



           return (4, 14, 16)

       @property
       def    snake_case__      (   self					)   -> Optional[Any]:

           """simple docstring"""



           return (4, 14, 1)

       @unittest.skipIf(torch_device == """mps""" ,		"""mish op not supported in MPS"""					)
       def    snake_case__      (   self					)   -> List[str]:

           """simple docstring"""



           super().test_determinism()

       @unittest.skipIf(torch_device == """mps""" ,		"""mish op not supported in MPS"""					)
       def    snake_case__      (   self					)   -> int:

           """simple docstring"""



           super().test_outputs_equivalence()

       @unittest.skipIf(torch_device == """mps""" ,		"""mish op not supported in MPS"""					)
       def    snake_case__      (   self					)   -> List[str]:

           """simple docstring"""



           super().test_from_save_pretrained()

       @unittest.skipIf(torch_device == """mps""" ,		"""mish op not supported in MPS"""					)
       def    snake_case__      (   self					)   -> Union[str, Any]:

           """simple docstring"""



           super().test_from_save_pretrained_variant()

       @unittest.skipIf(torch_device == """mps""" ,		"""mish op not supported in MPS"""					)
       def    snake_case__      (   self					)   -> List[Any]:

           """simple docstring"""



           super().test_model_from_pretrained()

       @unittest.skipIf(torch_device == """mps""" ,		"""mish op not supported in MPS"""					)
       def    snake_case__      (   self					)   -> str:

           """simple docstring"""



           __lowercase    ,__lowercase          =						self.prepare_init_args_and_inputs_for_common()
           __lowercase          =						self.model_class(**lowerCamelCase__					)
           model.to(lowerCamelCase__					)
           model.eval()

           with torch.no_grad():
               __lowercase          =						model(**lowerCamelCase__					)

               if isinstance(lowerCamelCase__ ,		lowerCamelCase__					):
                   __lowercase          =						output.sample

           self.assertIsNotNone(lowerCamelCase__					)
           __lowercase          =						torch.Size((inputs_dict["""sample"""].shape[0], 1)					)
           self.assertEqual(output.shape ,		lowerCamelCase__ ,		"""Input and output shapes do not match"""					)

       def    snake_case__      (   self					)   -> List[str]:

           """simple docstring"""



           pass

       def    snake_case__      (   self					)   -> str:

           """simple docstring"""



           pass

       def    snake_case__      (   self					)   -> Tuple:

           """simple docstring"""



           __lowercase          =						{
               """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          =						self.dummy_input
           return init_dict, inputs_dict

       @unittest.skipIf(torch_device == """mps""" ,		"""mish op not supported in MPS"""					)
       def    snake_case__      (   self					)   -> List[str]:

           """simple docstring"""



           __lowercase    ,__lowercase          =						UNetaDModel.from_pretrained(
               """bglick13/hopper-medium-v2-value-function-hor32""" ,		output_loading_info=lowerCamelCase__ ,		subfolder="""value_function"""					)
           self.assertIsNotNone(lowerCamelCase__					)
           self.assertEqual(len(vf_loading_info["""missing_keys"""]					) ,		0					)

           value_function.to(lowerCamelCase__					)
           __lowercase          =						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    snake_case__      (   self					)   -> Union[str, Any]:

           """simple docstring"""



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

           __lowercase          =						value_function.config.in_channels
           __lowercase          =						14
           __lowercase          =						torch.randn((1, seq_len, num_features)					).permute(
               0 ,		2 ,		1					)  # match original, we can update values and remove
           __lowercase          =						torch.full((num_features,) ,		0					)

           with torch.no_grad():
               __lowercase          =						value_function(lowerCamelCase__ ,		lowerCamelCase__					).sample

           # fmt: off
           __lowercase          =						torch.tensor([1_65.25] * seq_len					)
           # fmt: on
           self.assertTrue(torch.allclose(lowerCamelCase__ ,		lowerCamelCase__ ,		rtol=1E-3					)					)




       def    snake_case__      (   self					)   -> List[str]:

           """simple docstring"""



           pass










'''simple docstring'''
from typing import TYPE_CHECKING

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


A							:				Tuple			     =	{
    """configuration_dpr""": ["""DPR_PRETRAINED_CONFIG_ARCHIVE_MAP""", """DPRConfig"""],
    """tokenization_dpr""": [
        """DPRContextEncoderTokenizer""",
        """DPRQuestionEncoderTokenizer""",
        """DPRReaderOutput""",
        """DPRReaderTokenizer""",
    ],
}


try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  A							:				str			     =	[
      """DPRContextEncoderTokenizerFast""",
      """DPRQuestionEncoderTokenizerFast""",
      """DPRReaderTokenizerFast""",
  ]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  A							:				Tuple			     =	[
      """DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """DPRContextEncoder""",
      """DPRPretrainedContextEncoder""",
      """DPRPreTrainedModel""",
      """DPRPretrainedQuestionEncoder""",
      """DPRPretrainedReader""",
      """DPRQuestionEncoder""",
      """DPRReader""",
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  A							:				Tuple			     =	[
      """TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST""",
      """TFDPRContextEncoder""",
      """TFDPRPretrainedContextEncoder""",
      """TFDPRPretrainedQuestionEncoder""",
      """TFDPRPretrainedReader""",
      """TFDPRQuestionEncoder""",
      """TFDPRReader""",
  ]


if TYPE_CHECKING:
  from .configuration_dpr import DPR_PRETRAINED_CONFIG_ARCHIVE_MAP, DPRConfig
  from .tokenization_dpr import (
      DPRContextEncoderTokenizer,
      DPRQuestionEncoderTokenizer,
      DPRReaderOutput,
      DPRReaderTokenizer,
  )

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_dpr_fast import (
        DPRContextEncoderTokenizerFast,
        DPRQuestionEncoderTokenizerFast,
        DPRReaderTokenizerFast,
    )

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_dpr import (
        DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
        DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
        DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
        DPRContextEncoder,
        DPRPretrainedContextEncoder,
        DPRPreTrainedModel,
        DPRPretrainedQuestionEncoder,
        DPRPretrainedReader,
        DPRQuestionEncoder,
        DPRReader,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_dpr import (
        TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
        TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
        TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFDPRContextEncoder,
        TFDPRPretrainedContextEncoder,
        TFDPRPretrainedQuestionEncoder,
        TFDPRPretrainedReader,
        TFDPRQuestionEncoder,
        TFDPRReader,
    )

else:
  import sys

  A							:				int			     =	_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)




