
"""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 	lowerCamelCase__      (      __magic_name__				,			__magic_name__				,			unittest.TestCase       ):





			'''simple docstring'''
			lowerCamelCase						      =	UNetaDModel
			lowerCamelCase						      =	'''sample'''



			@property
			def       _lowerCAmelCase	(							self						)				->  int:
				_lowerCAmelCase			=4
				_lowerCAmelCase			=14
				_lowerCAmelCase			=16

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

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



			@property
			def       _lowerCAmelCase	(							self						)				->  List[str]:
				return (4, 14, 16)



			@property
			def       _lowerCAmelCase	(							self						)				->  int:
				return (4, 14, 16)



			def       _lowerCAmelCase	(							self						)				->  str:
				pass



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



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



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



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



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



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



			@unittest.skipIf(torch_device == """mps"""  ,  """mish op not supported in MPS"""						)
			def       _lowerCAmelCase	(							self						)				->  Dict:
				super().test_output()



			def       _lowerCAmelCase	(							self						)				->  List[str]:
				_lowerCAmelCase			={
				    """block_out_channels""": (32, 64, 1_28, 2_56),
				    """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""",
				}
				_lowerCAmelCase			=self.dummy_input
				return init_dict, inputs_dict



			@unittest.skipIf(torch_device == """mps"""  ,  """mish op not supported in MPS"""						)
			def       _lowerCAmelCase	(							self						)				->  Optional[int]:
				_lowerCAmelCase		,							_lowerCAmelCase			=UNetaDModel.from_pretrained(
				    """bglick13/hopper-medium-v2-value-function-hor32"""  ,  output_loading_info=__UpperCAmelCase  ,  subfolder="""unet"""						)
				self.assertIsNotNone(__UpperCAmelCase						)
				self.assertEqual(len(loading_info["""missing_keys"""]						)  ,  0						)

				model.to(__UpperCAmelCase						)
				_lowerCAmelCase			=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       _lowerCAmelCase	(							self						)				->  Dict:
				_lowerCAmelCase			=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						)

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

				with torch.no_grad():
					_lowerCAmelCase			=model(__UpperCAmelCase  ,  __UpperCAmelCase						).sample.permute(0  ,  2  ,  1						)

				_lowerCAmelCase			=output[0, -3:, -3:].flatten()
				# fmt: off
				_lowerCAmelCase			=torch.tensor([-2.1_3_7_1_7_2, 1.1_4_2_6_0_1_6, 0.3_6_8_8_6_8_7, -0.7_6_6_9_2_2, 0.7_3_0_3_1_4_6, 0.1_1_0_3_8_8_6_4, -0.4_7_6_0_6_3_3, 0.1_3_2_7_0_1_7_2, 0.0_2_5_9_1_3_4_8]						)
				# fmt: on
				self.assertTrue(torch.allclose(__UpperCAmelCase  ,  __UpperCAmelCase  ,  rtol=1e-3						)						)



			def       _lowerCAmelCase	(							self						)				->  Any:
				# Not implemented yet for this UNet
				pass



			@slow
			def       _lowerCAmelCase	(							self						)				->  str:
				_lowerCAmelCase			="""harmonai/maestro-150k"""
				_lowerCAmelCase			=UNetaDModel.from_pretrained(__UpperCAmelCase  ,  subfolder="""unet"""						)
				model.to(__UpperCAmelCase						)

				_lowerCAmelCase			=6_55_36
				_lowerCAmelCase			=torch.sin(torch.arange(__UpperCAmelCase						)[None, None, :].repeat(1  ,  2  ,  1						)						).to(__UpperCAmelCase						)
				_lowerCAmelCase			=torch.tensor([1]						).to(__UpperCAmelCase						)

				with torch.no_grad():
					_lowerCAmelCase			=model(__UpperCAmelCase  ,  __UpperCAmelCase						).sample

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

				assert (output_sum - 2_2_4.0_8_9_6).abs() < 0.5
				assert (output_max - 0.0_6_0_7).abs() < 4e-4







class 	lowerCamelCase__      (      __magic_name__				,			__magic_name__				,			unittest.TestCase       ):





			'''simple docstring'''
			lowerCamelCase						      =	UNetaDModel
			lowerCamelCase						      =	'''sample'''



			@property
			def       _lowerCAmelCase	(							self						)				->  Dict:
				_lowerCAmelCase			=4
				_lowerCAmelCase			=14
				_lowerCAmelCase			=16

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

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



			@property
			def       _lowerCAmelCase	(							self						)				->  int:
				return (4, 14, 16)



			@property
			def       _lowerCAmelCase	(							self						)				->  Optional[int]:
				return (4, 14, 1)



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



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



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



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



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



			@unittest.skipIf(torch_device == """mps"""  ,  """mish op not supported in MPS"""						)
			def       _lowerCAmelCase	(							self						)				->  Optional[Any]:
				# UNetRL is a value-function is different output shape
				_lowerCAmelCase		,							_lowerCAmelCase			=self.prepare_init_args_and_inputs_for_common()
				_lowerCAmelCase			=self.model_class(**__UpperCAmelCase						)
				model.to(__UpperCAmelCase						)
				model.eval()

				with torch.no_grad():
					_lowerCAmelCase			=model(**__UpperCAmelCase						)

					if isinstance(__UpperCAmelCase  ,  __UpperCAmelCase						):
						_lowerCAmelCase			=output.sample

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



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



			def       _lowerCAmelCase	(							self						)				->  Optional[Any]:
				pass



			def       _lowerCAmelCase	(							self						)				->  Optional[Any]:
				_lowerCAmelCase			={
				    """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, 1_28, 2_56],
				    """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""",
				}
				_lowerCAmelCase			=self.dummy_input
				return init_dict, inputs_dict



			@unittest.skipIf(torch_device == """mps"""  ,  """mish op not supported in MPS"""						)
			def       _lowerCAmelCase	(							self						)				->  int:
				_lowerCAmelCase		,							_lowerCAmelCase			=UNetaDModel.from_pretrained(
				    """bglick13/hopper-medium-v2-value-function-hor32"""  ,  output_loading_info=__UpperCAmelCase  ,  subfolder="""value_function"""						)
				self.assertIsNotNone(__UpperCAmelCase						)
				self.assertEqual(len(vf_loading_info["""missing_keys"""]						)  ,  0						)

				value_function.to(__UpperCAmelCase						)
				_lowerCAmelCase			=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       _lowerCAmelCase	(							self						)				->  str:
				_lowerCAmelCase		,							_lowerCAmelCase			=UNetaDModel.from_pretrained(
				    """bglick13/hopper-medium-v2-value-function-hor32"""  ,  output_loading_info=__UpperCAmelCase  ,  subfolder="""value_function"""						)
				torch.manual_seed(0						)
				if torch.cuda.is_available():
					torch.cuda.manual_seed_all(0						)

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

				with torch.no_grad():
					_lowerCAmelCase			=value_function(__UpperCAmelCase  ,  __UpperCAmelCase						).sample

				# fmt: off
				_lowerCAmelCase			=torch.tensor([1_6_5.2_5] * seq_len						)
				# fmt: on
				self.assertTrue(torch.allclose(__UpperCAmelCase  ,  __UpperCAmelCase  ,  rtol=1e-3						)						)



			def       _lowerCAmelCase	(							self						)				->  str:
				# Not implemented yet for this UNet
				pass





"""simple docstring"""


# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

# rely on isort to merge the imports
from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


__A 		= {
    'configuration_convnextv2': [
        'CONVNEXTV2_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'ConvNextV2Config',
    ]
}

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			__A 		= [
			    'CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST',
			    'ConvNextV2ForImageClassification',
			    'ConvNextV2Model',
			    'ConvNextV2PreTrainedModel',
			    'ConvNextV2Backbone',
			]


if TYPE_CHECKING:
			from .configuration_convnextva import (
			    CONVNEXTV2_PRETRAINED_CONFIG_ARCHIVE_MAP,
			    ConvNextVaConfig,
			)

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_convnextva import (
						    CONVNEXTV2_PRETRAINED_MODEL_ARCHIVE_LIST,
						    ConvNextVaBackbone,
						    ConvNextVaForImageClassification,
						    ConvNextVaModel,
						    ConvNextVaPreTrainedModel,
						)

else:
			import sys

			__A 		= _LazyModule(__name__, globals()['__file__'], _import_structure)



