





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 		__a							(	__UpperCamelCase			,__UpperCamelCase			,unittest.TestCase      ):
       __snake_case						:   Tuple      =		UNetaDModel
       __snake_case						:   List[Any]      =		"""sample"""



       @property
       def  A							(     self	:   List[str]			):
           lowerCAmelCase_					:	Union[str, Any]						=   4
           lowerCAmelCase_					:	List[Any]						=   14
           lowerCAmelCase_					:	int						=   16

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

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



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



       @property
       def  A							(     self	:   Union[str, Any]			):
           return (4, 14, 16)



       def  A							(     self	:   Any			):
           pass



       def  A							(     self	:   Optional[Any]			):
           pass



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



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



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



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



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



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



       def  A							(     self	:   List[str]			):
           lowerCAmelCase_					:	Optional[Any]						=   {
               """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_					:	Union[str, Any]						=   self.dummy_input
           return init_dict, inputs_dict



       @unittest.skipIf(torch_device == """mps"""					,			"""mish op not supported in MPS"""			)
       def  A							(     self	:   Optional[int]			):
           lowerCAmelCase_					:	Optional[int]						=   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_					:	int						=   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	:   List[str]			):
           lowerCAmelCase_					:	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			)

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

           with torch.no_grad():
               lowerCAmelCase_					:	Optional[Any]						=   model(UpperCAmelCase					,			UpperCAmelCase			).sample.permute(0					,			2					,			1			)

           lowerCAmelCase_					:	List[str]						=   output[0, -3:, -3:].flatten()
           # fmt: off
           lowerCAmelCase_					:	Optional[int]						=   torch.tensor([-2.13_7172, 1.142_6016, 0.368_8687, -0.76_6922, 0.730_3146, 0.1103_8864, -0.476_0633, 0.1327_0172, 0.0259_1348]			)
           # fmt: on
           self.assertTrue(torch.allclose(UpperCAmelCase					,			UpperCAmelCase					,			rtol=1e-3			)			)



       def  A							(     self	:   Tuple			):
           # Not implemented yet for this UNet
           pass



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

           lowerCAmelCase_					:	Any						=   6_55_36
           lowerCAmelCase_					:	str						=   torch.sin(torch.arange(UpperCAmelCase			)[None, None, :].repeat(1					,			2					,			1			)			).to(UpperCAmelCase			)
           lowerCAmelCase_					:	Dict						=   torch.tensor([1]			).to(UpperCAmelCase			)

           with torch.no_grad():
               lowerCAmelCase_					:	Dict						=   model(UpperCAmelCase					,			UpperCAmelCase			).sample

           lowerCAmelCase_					:	List[str]						=   output.abs().sum()
           lowerCAmelCase_					:	int						=   output.abs().max()

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






class 		__a							(	__UpperCamelCase			,__UpperCamelCase			,unittest.TestCase      ):
       __snake_case						:   List[str]      =		UNetaDModel
       __snake_case						:   Optional[Any]      =		"""sample"""



       @property
       def  A							(     self	:   Dict			):
           lowerCAmelCase_					:	List[Any]						=   4
           lowerCAmelCase_					:	Dict						=   14
           lowerCAmelCase_					:	Dict						=   16

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

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



       @property
       def  A							(     self	:   Union[str, Any]			):
           return (4, 14, 16)



       @property
       def  A							(     self	:   Optional[Any]			):
           return (4, 14, 1)



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



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



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



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



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



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

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

               if isinstance(UpperCAmelCase					,			UpperCAmelCase			):
                   lowerCAmelCase_					:	List[str]						=   output.sample

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



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



       def  A							(     self	:   Optional[Any]			):
           pass



       def  A							(     self	:   List[str]			):
           lowerCAmelCase_					:	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, 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_					:	Any						=   self.dummy_input
           return init_dict, inputs_dict



       @unittest.skipIf(torch_device == """mps"""					,			"""mish op not supported in MPS"""			)
       def  A							(     self	:   Optional[int]			):
           lowerCAmelCase_					:	int						=   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_					:	int						=   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	:   Dict			):
           lowerCAmelCase_					:	List[Any]						=   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_					:	int						=   value_function.config.in_channels
           lowerCAmelCase_					:	List[Any]						=   14
           lowerCAmelCase_					:	Dict						=   torch.randn((1, seq_len, num_features)			).permute(
               0					,			2					,			1			)  # match original, we can update values and remove
           lowerCAmelCase_					:	Optional[Any]						=   torch.full((num_features,)					,			0			)

           with torch.no_grad():
               lowerCAmelCase_					:	List[str]						=   value_function(UpperCAmelCase					,			UpperCAmelCase			).sample

           # fmt: off
           lowerCAmelCase_					:	List[str]						=   torch.tensor([165.25] * seq_len			)
           # fmt: on
           self.assertTrue(torch.allclose(UpperCAmelCase					,			UpperCAmelCase					,			rtol=1e-3			)			)



       def  A							(     self	:   List[str]			):
           # Not implemented yet for this UNet
           pass









import unittest

from transformers import BlenderbotTokenizer, BlenderbotTokenizerFast
from transformers.utils import cached_property

class 		__a							(	unittest.TestCase      ):



       @cached_property
       def  A							(     self	:   Optional[int]			):
           return BlenderbotTokenizer.from_pretrained("""facebook/blenderbot-3B"""			)



       @cached_property
       def  A							(     self	:   Optional[int]			):
           return BlenderbotTokenizerFast.from_pretrained("""facebook/blenderbot-3B"""			)



       def  A							(     self	:   Optional[Any]			):
           lowerCAmelCase_					:	List[str]						=   self.tokenizer_ab
           lowerCAmelCase_					:	Any						=   """ I am a small frog."""
           lowerCAmelCase_					:	Union[str, Any]						=   tok([src_text]					,			padding=UpperCAmelCase					,			truncation=UpperCAmelCase			)["""input_ids"""]
           lowerCAmelCase_					:	int						=   tok.batch_decode(UpperCAmelCase					,			skip_special_tokens=UpperCAmelCase					,			clean_up_tokenization_spaces=UpperCAmelCase			)[0]
           assert src_text == decoded



       def  A							(     self	:   List[Any]			):
           lowerCAmelCase_					:	Dict						=   self.rust_tokenizer_ab
           lowerCAmelCase_					:	str						=   """ I am a small frog."""
           lowerCAmelCase_					:	Tuple						=   tok([src_text]					,			padding=UpperCAmelCase					,			truncation=UpperCAmelCase			)["""input_ids"""]
           lowerCAmelCase_					:	int						=   tok.batch_decode(UpperCAmelCase					,			skip_special_tokens=UpperCAmelCase					,			clean_up_tokenization_spaces=UpperCAmelCase			)[0]
           assert src_text == decoded



       def  A							(     self	:   str			):
           assert self.tokenizer_ab.add_prefix_space
           assert self.tokenizer_ab([""" Sam""", """Sam"""]			).input_ids == [[55_02, 2], [55_02, 2]]



       def  A							(     self	:   List[str]			):
           assert self.rust_tokenizer_ab.add_prefix_space
           assert self.rust_tokenizer_ab([""" Sam""", """Sam"""]			).input_ids == [[55_02, 2], [55_02, 2]]


