



'''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      UpperCamelCase_					(				__magic_name__ ,	__magic_name__ ,	unittest.TestCase     ):
	lowercase		=					UNetaDModel
	lowercase		=					'sample'





	@property
	def 					_lowercase(    self       )     ->     Tuple:
								UpperCAmelCase	:						int					=						4
								UpperCAmelCase	:						Any					=						14
								UpperCAmelCase	:						Union[str, Any]					=						16

								UpperCAmelCase	:						List[str]					=						floats_tensor((batch_size, num_features, seq_len)       ).to(A       )
								UpperCAmelCase	:						List[str]					=						torch.tensor([10] * batch_size       ).to(A       )

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





	@property
	def 					_lowercase(    self       )     ->     Any:
								return (4, 14, 16)





	@property
	def 					_lowercase(    self       )     ->     Tuple:
								return (4, 14, 16)





	def 					_lowercase(    self       )     ->     List[str]:
								pass





	def 					_lowercase(    self       )     ->     Any:
								pass





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





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





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





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





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





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





	def 					_lowercase(    self       )     ->     List[str]:
								UpperCAmelCase	:						Tuple					=						{
								    """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""",
								}
								UpperCAmelCase	:						Optional[int]					=						self.dummy_input
								return init_dict, inputs_dict





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

								model.to(A       )
								UpperCAmelCase	:						str					=						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 					_lowercase(    self       )     ->     int:
								UpperCAmelCase	:						List[Any]					=						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       )

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

								with torch.no_grad():
															UpperCAmelCase	:						int					=						model(A						,						A       ).sample.permute(0						,						2						,						1       )

								UpperCAmelCase	:						str					=						output[0, -3:, -3:].flatten()
								# fmt: off
								UpperCAmelCase	:						Optional[int]					=						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(A						,						A						,						rtol=1e-3       )       )





	def 					_lowercase(    self       )     ->     Optional[Any]:
								# Not implemented yet for this UNet
								pass






	@slow
	def 					_lowercase(    self       )     ->     List[Any]:
								UpperCAmelCase	:						Any					=						"""harmonai/maestro-150k"""
								UpperCAmelCase	:						Optional[int]					=						UNetaDModel.from_pretrained(A						,						subfolder="""unet"""       )
								model.to(A       )

								UpperCAmelCase	:						Optional[int]					=						65536
								UpperCAmelCase	:						Dict					=						torch.sin(torch.arange(A       )[None, None, :].repeat(1						,						2						,						1       )       ).to(A       )
								UpperCAmelCase	:						str					=						torch.tensor([1]       ).to(A       )

								with torch.no_grad():
															UpperCAmelCase	:						Union[str, Any]					=						model(A						,						A       ).sample

								UpperCAmelCase	:						Union[str, Any]					=						output.abs().sum()
								UpperCAmelCase	:						Any					=						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      UpperCamelCase_					(				__magic_name__ ,	__magic_name__ ,	unittest.TestCase     ):
	lowercase		=					UNetaDModel
	lowercase		=					'sample'





	@property
	def 					_lowercase(    self       )     ->     Dict:
								UpperCAmelCase	:						int					=						4
								UpperCAmelCase	:						Dict					=						14
								UpperCAmelCase	:						Optional[Any]					=						16

								UpperCAmelCase	:						Tuple					=						floats_tensor((batch_size, num_features, seq_len)       ).to(A       )
								UpperCAmelCase	:						Dict					=						torch.tensor([10] * batch_size       ).to(A       )

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





	@property
	def 					_lowercase(    self       )     ->     Optional[Any]:
								return (4, 14, 16)





	@property
	def 					_lowercase(    self       )     ->     Tuple:
								return (4, 14, 1)





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





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





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





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





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





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

								with torch.no_grad():
															UpperCAmelCase	:						Dict					=						model(**A       )

															if isinstance(A						,						A       ):
																						UpperCAmelCase	:						Optional[int]					=						output.sample

								self.assertIsNotNone(A       )
								UpperCAmelCase	:						List[str]					=						torch.Size((inputs_dict["""sample"""].shape[0], 1)       )
								self.assertEqual(output.shape						,						A						,						"""Input and output shapes do not match"""       )





	def 					_lowercase(    self       )     ->     str:
								pass





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





	def 					_lowercase(    self       )     ->     int:
								UpperCAmelCase	:						Optional[int]					=						{
								    """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""",
								}
								UpperCAmelCase	:						Optional[int]					=						self.dummy_input
								return init_dict, inputs_dict





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

								value_function.to(A       )
								UpperCAmelCase	:						Optional[Any]					=						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 					_lowercase(    self       )     ->     List[Any]:
								UpperCAmelCase	,	UpperCAmelCase	:						str					=						UNetaDModel.from_pretrained(
								    """bglick13/hopper-medium-v2-value-function-hor32"""						,						output_loading_info=A						,						subfolder="""value_function"""       )
								torch.manual_seed(0       )
								if torch.cuda.is_available():
															torch.cuda.manual_seed_all(0       )

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

								with torch.no_grad():
															UpperCAmelCase	:						Optional[Any]					=						value_function(A						,						A       ).sample

								# fmt: off
								UpperCAmelCase	:						Union[str, Any]					=						torch.tensor([1_6_5.2_5] * seq_len       )
								# fmt: on
								self.assertTrue(torch.allclose(A						,						A						,						rtol=1e-3       )       )






	def 					_lowercase(    self       )     ->     int:
								# Not implemented yet for this UNet
								pass






'''simple docstring'''




# Copyright 2023 The HuggingFace Inc. 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,
)


a  :     Optional[Any]       =       {
    """agents""": ["""Agent""", """AzureOpenAiAgent""", """HfAgent""", """LocalAgent""", """OpenAiAgent"""],
    """base""": ["""PipelineTool""", """RemoteTool""", """Tool""", """launch_gradio_demo""", """load_tool"""],
}

try:
							if not is_torch_available():
														raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
							pass
else:
							a  :     Dict       =       ["""DocumentQuestionAnsweringTool"""]
							a  :     List[str]       =       ["""ImageCaptioningTool"""]
							a  :     Tuple       =       ["""ImageQuestionAnsweringTool"""]
							a  :     Tuple       =       ["""ImageSegmentationTool"""]
							a  :     List[str]       =       ["""SpeechToTextTool"""]
							a  :     Dict       =       ["""TextClassificationTool"""]
							a  :     Union[str, Any]       =       ["""TextQuestionAnsweringTool"""]
							a  :     List[str]       =       ["""TextSummarizationTool"""]
							a  :     Optional[Any]       =       ["""TextToSpeechTool"""]
							a  :     str       =       ["""TranslationTool"""]

if TYPE_CHECKING:
							from .agents import Agent, AzureOpenAiAgent, HfAgent, LocalAgent, OpenAiAgent
							from .base import PipelineTool, RemoteTool, Tool, launch_gradio_demo, load_tool

							try:
														if not is_torch_available():
																					raise OptionalDependencyNotAvailable()
							except OptionalDependencyNotAvailable:
														pass
							else:
														from .document_question_answering import DocumentQuestionAnsweringTool
														from .image_captioning import ImageCaptioningTool
														from .image_question_answering import ImageQuestionAnsweringTool
														from .image_segmentation import ImageSegmentationTool
														from .speech_to_text import SpeechToTextTool
														from .text_classification import TextClassificationTool
														from .text_question_answering import TextQuestionAnsweringTool
														from .text_summarization import TextSummarizationTool
														from .text_to_speech import TextToSpeechTool
														from .translation import TranslationTool
else:
							import sys

							a  :     Dict       =       _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)

