
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__  (			UpperCAmelCase							,UpperCAmelCase							,unittest.TestCase      ):





	'''simple docstring'''



	_UpperCAmelCase :	List[Any]					     =      UNetaDModel
	_UpperCAmelCase :	Optional[Any]					     =      "sample"



	@property
	def 				A							(						self			:					Dict			):


				'''simple docstring'''


				_snake_case						=    4
				_snake_case						=    14
				_snake_case						=    16

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

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



	@property
	def 				A							(						self			:					List[str]			):


				'''simple docstring'''


				return (4, 14, 16)



	@property
	def 				A							(						self			:					Any			):


				'''simple docstring'''


				return (4, 14, 16)



	def 				A							(						self			:					Any			):


				'''simple docstring'''


				pass



	def 				A							(						self			:					Tuple			):


				'''simple docstring'''


				pass



	@unittest.skipIf(torch_device == 'mps'    ,						'mish op not supported in MPS'			)
	def 				A							(						self			:					List[str]			):


				'''simple docstring'''


				super().test_determinism()



	@unittest.skipIf(torch_device == 'mps'    ,						'mish op not supported in MPS'			)
	def 				A							(						self			:					Optional[Any]			):


				'''simple docstring'''


				super().test_outputs_equivalence()



	@unittest.skipIf(torch_device == 'mps'    ,						'mish op not supported in MPS'			)
	def 				A							(						self			:					List[str]			):


				'''simple docstring'''


				super().test_from_save_pretrained()



	@unittest.skipIf(torch_device == 'mps'    ,						'mish op not supported in MPS'			)
	def 				A							(						self			:					str			):


				'''simple docstring'''


				super().test_from_save_pretrained_variant()



	@unittest.skipIf(torch_device == 'mps'    ,						'mish op not supported in MPS'			)
	def 				A							(						self			:					Dict			):


				'''simple docstring'''


				super().test_model_from_pretrained()



	@unittest.skipIf(torch_device == 'mps'    ,						'mish op not supported in MPS'			)
	def 				A							(						self			:					Optional[int]			):


				'''simple docstring'''


				super().test_output()



	def 				A							(						self			:					Union[str, Any]			):


				'''simple docstring'''


				_snake_case						=    {
				    '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',
				}
				_snake_case						=    self.dummy_input
				return init_dict, inputs_dict



	@unittest.skipIf(torch_device == 'mps'    ,						'mish op not supported in MPS'			)
	def 				A							(						self			:					int			):


				'''simple docstring'''


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

				model.to(lowercase			)
				_snake_case						=    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 				A							(						self			:					str			):


				'''simple docstring'''


				_snake_case						=    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			)

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

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

				_snake_case						=    output[0, -3:, -3:].flatten()
				# fmt: off
				_snake_case						=    torch.tensor([-2.137172, 1.1426016, 0.3688687, -0.766922, 0.7303146, 0.11038864, -0.4760633, 0.13270172, 0.02591348]			)
				# fmt: on
				self.assertTrue(torch.allclose(lowercase    ,						lowercase    ,						rtol=1E-3			)			)



	def 				A							(						self			:					Optional[Any]			):


				'''simple docstring'''


				pass





	@slow
	def 				A							(						self			:					int			):


				'''simple docstring'''


				_snake_case						=    'harmonai/maestro-150k'
				_snake_case						=    UNetaDModel.from_pretrained(lowercase    ,						subfolder='unet'			)
				model.to(lowercase			)

				_snake_case						=    65_536
				_snake_case						=    torch.sin(torch.arange(lowercase			)[None, None, :].repeat(1    ,						2    ,						1			)			).to(lowercase			)
				_snake_case						=    torch.tensor([1]			).to(lowercase			)

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

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

				assert (output_sum - 224.0896).abs() < 0.5
				assert (output_max - 0.0607).abs() < 4E-4




class      SCREAMING_SNAKE_CASE__  (			UpperCAmelCase							,UpperCAmelCase							,unittest.TestCase      ):





	'''simple docstring'''



	_UpperCAmelCase :	Optional[Any]					     =      UNetaDModel
	_UpperCAmelCase :	Dict					     =      "sample"



	@property
	def 				A							(						self			:					Optional[Any]			):


				'''simple docstring'''


				_snake_case						=    4
				_snake_case						=    14
				_snake_case						=    16

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

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



	@property
	def 				A							(						self			:					Tuple			):


				'''simple docstring'''


				return (4, 14, 16)



	@property
	def 				A							(						self			:					Union[str, Any]			):


				'''simple docstring'''


				return (4, 14, 1)



	@unittest.skipIf(torch_device == 'mps'    ,						'mish op not supported in MPS'			)
	def 				A							(						self			:					Tuple			):


				'''simple docstring'''


				super().test_determinism()



	@unittest.skipIf(torch_device == 'mps'    ,						'mish op not supported in MPS'			)
	def 				A							(						self			:					List[str]			):


				'''simple docstring'''


				super().test_outputs_equivalence()



	@unittest.skipIf(torch_device == 'mps'    ,						'mish op not supported in MPS'			)
	def 				A							(						self			:					Any			):


				'''simple docstring'''


				super().test_from_save_pretrained()



	@unittest.skipIf(torch_device == 'mps'    ,						'mish op not supported in MPS'			)
	def 				A							(						self			:					List[Any]			):


				'''simple docstring'''


				super().test_from_save_pretrained_variant()



	@unittest.skipIf(torch_device == 'mps'    ,						'mish op not supported in MPS'			)
	def 				A							(						self			:					List[str]			):


				'''simple docstring'''


				super().test_model_from_pretrained()



	@unittest.skipIf(torch_device == 'mps'    ,						'mish op not supported in MPS'			)
	def 				A							(						self			:					Optional[Any]			):


				'''simple docstring'''


				_snake_case			,      _snake_case						=    self.prepare_init_args_and_inputs_for_common()
				_snake_case						=    self.model_class(**lowercase			)
				model.to(lowercase			)
				model.eval()

				with torch.no_grad():
							_snake_case						=    model(**lowercase			)

							if isinstance(lowercase    ,						lowercase			):
										_snake_case						=    output.sample

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



	def 				A							(						self			:					List[str]			):


				'''simple docstring'''


				pass



	def 				A							(						self			:					int			):


				'''simple docstring'''


				pass



	def 				A							(						self			:					Optional[int]			):


				'''simple docstring'''


				_snake_case						=    {
				    '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',
				}
				_snake_case						=    self.dummy_input
				return init_dict, inputs_dict



	@unittest.skipIf(torch_device == 'mps'    ,						'mish op not supported in MPS'			)
	def 				A							(						self			:					List[str]			):


				'''simple docstring'''


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

				value_function.to(lowercase			)
				_snake_case						=    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 				A							(						self			:					int			):


				'''simple docstring'''


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

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

				with torch.no_grad():
							_snake_case						=    value_function(lowercase    ,						lowercase			).sample

				# fmt: off
				_snake_case						=    torch.tensor([165.25] * seq_len			)
				# fmt: on
				self.assertTrue(torch.allclose(lowercase    ,						lowercase    ,						rtol=1E-3			)			)





	def 				A							(						self			:					Optional[int]			):


				'''simple docstring'''


				pass

# 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

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
    is_vision_available,
)


_lowerCamelCase :       int     					=    {
    '''configuration_tvlt''': ['''TVLT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''TvltConfig'''],
    '''feature_extraction_tvlt''': ['''TvltFeatureExtractor'''],
    '''processing_tvlt''': ['''TvltProcessor'''],
}

try:
			if not is_torch_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_lowerCamelCase :       str     					=    [
			    '''TVLT_PRETRAINED_MODEL_ARCHIVE_LIST''',
			    '''TvltModel''',
			    '''TvltForPreTraining''',
			    '''TvltForAudioVisualClassification''',
			    '''TvltPreTrainedModel''',
			]

try:
			if not is_vision_available():
						raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
			pass
else:
			_lowerCamelCase :       str     					=    ['''TvltImageProcessor''']


if TYPE_CHECKING:
			from .configuration_tvlt import TVLT_PRETRAINED_CONFIG_ARCHIVE_MAP, TvltConfig
			from .processing_tvlt import TvltProcessor
			from .feature_extraction_tvlt import TvltFeatureExtractor

			try:
						if not is_torch_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .modeling_tvlt import (
						    TVLT_PRETRAINED_MODEL_ARCHIVE_LIST,
						    TvltForAudioVisualClassification,
						    TvltForPreTraining,
						    TvltModel,
						    TvltPreTrainedModel,
						)

			try:
						if not is_vision_available():
									raise OptionalDependencyNotAvailable()
			except OptionalDependencyNotAvailable:
						pass
			else:
						from .image_processing_tvlt import TvltImageProcessor


else:
			import sys

			_lowerCamelCase :       Dict     					=    _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)