




'''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 				_snake_case	(			lowercase_		,		lowercase_		,		unittest.TestCase							):
			lowerCAmelCase_  :   str		 =					UNetaDModel
			lowerCAmelCase_  :   List[str]		 =					"sample"




			@property
			def      lowerCAmelCase__       (      self							)       ->       int:

										'''simple docstring'''




										snake_case_			      =   4
										snake_case_			      =   14
										snake_case_			      =   16

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

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




			@property
			def      lowerCAmelCase__       (      self							)       ->       Dict:

										'''simple docstring'''




										return (4, 14, 16)




			@property
			def      lowerCAmelCase__       (      self							)       ->       Dict:

										'''simple docstring'''




										return (4, 14, 16)




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

										'''simple docstring'''




										pass




			def      lowerCAmelCase__       (      self							)       ->       str:

										'''simple docstring'''




										pass




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

										'''simple docstring'''




										super().test_determinism()




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

										'''simple docstring'''




										super().test_outputs_equivalence()




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

										'''simple docstring'''




										super().test_from_save_pretrained()




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

										'''simple docstring'''




										super().test_from_save_pretrained_variant()




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

										'''simple docstring'''




										super().test_model_from_pretrained()




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

										'''simple docstring'''




										super().test_output()




			def      lowerCAmelCase__       (      self							)       ->       str:

										'''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      lowerCAmelCase__       (      self							)       ->       Any:

										'''simple docstring'''




										snake_case_      ,   snake_case_			      =   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__							)
										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      lowerCAmelCase__       (      self							)       ->       Any:

										'''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(a__       ,							a__							).sample.permute(0       ,							2       ,							1							)

										snake_case_			      =   output[0, -3:, -3:].flatten()
										# fmt: off
										snake_case_			      =   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      lowerCAmelCase__       (      self							)       ->       Dict:

										'''simple docstring'''




										pass




			@slow
			def      lowerCAmelCase__       (      self							)       ->       int:

										'''simple docstring'''




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

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

										with torch.no_grad():
																	snake_case_			      =   model(a__       ,							a__							).sample

										snake_case_			      =   output.abs().sum()
										snake_case_			      =   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 				_snake_case	(			lowercase_		,		lowercase_		,		unittest.TestCase							):
			lowerCAmelCase_  :   Union[str, Any]		 =					UNetaDModel
			lowerCAmelCase_  :   Union[str, Any]		 =					"sample"




			@property
			def      lowerCAmelCase__       (      self							)       ->       str:

										'''simple docstring'''




										snake_case_			      =   4
										snake_case_			      =   14
										snake_case_			      =   16

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

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




			@property
			def      lowerCAmelCase__       (      self							)       ->       Any:

										'''simple docstring'''




										return (4, 14, 16)




			@property
			def      lowerCAmelCase__       (      self							)       ->       Dict:

										'''simple docstring'''




										return (4, 14, 1)




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

										'''simple docstring'''




										super().test_determinism()




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

										'''simple docstring'''




										super().test_outputs_equivalence()




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

										'''simple docstring'''




										super().test_from_save_pretrained()




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

										'''simple docstring'''




										super().test_from_save_pretrained_variant()




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

										'''simple docstring'''




										super().test_model_from_pretrained()




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

										'''simple docstring'''




										snake_case_      ,   snake_case_			      =   self.prepare_init_args_and_inputs_for_common()
										snake_case_			      =   self.model_class(**a__							)
										model.to(a__							)
										model.eval()

										with torch.no_grad():
																	snake_case_			      =   model(**a__							)

																	if isinstance(a__       ,							a__							):
																								snake_case_			      =   output.sample

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




			def      lowerCAmelCase__       (      self							)       ->       int:

										'''simple docstring'''




										pass




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

										'''simple docstring'''




										pass




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

										'''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      lowerCAmelCase__       (      self							)       ->       Union[str, Any]:

										'''simple docstring'''




										snake_case_      ,   snake_case_			      =   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__							)
										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      lowerCAmelCase__       (      self							)       ->       Tuple:

										'''simple docstring'''




										snake_case_      ,   snake_case_			      =   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							)

										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(a__       ,							a__							).sample

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




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

										'''simple docstring'''




										pass










'''simple docstring'''





import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
		import torch

		from transformers import CamembertModel




@require_torch
@require_sentencepiece
@require_tokenizers
class 				_snake_case	(			unittest.TestCase							):




			@slow
			def      lowerCAmelCase__       (      self							)       ->       List[Any]:

										'''simple docstring'''




										snake_case_			      =   CamembertModel.from_pretrained("camembert-base"							)
										model.to(a__							)

										snake_case_			      =   torch.tensor(
										    [[5, 121, 11, 660, 16, 730, 25_543, 110, 83, 6]]       ,							device=a__       ,							dtype=torch.long       ,							)  # J'aime le camembert !
										with torch.no_grad():
																	snake_case_			      =   model(a__							)["last_hidden_state"]
										snake_case_			      =   torch.Size((1, 10, 768)							)
										self.assertEqual(output.shape       ,							a__							)
										# compare the actual values for a slice.
										snake_case_			      =   torch.tensor(
										    [[[-0.0_2_5_4, 0.0_2_3_5, 0.1_0_2_7], [0.0_6_0_6, -0.1_8_1_1, -0.0_4_1_8], [-0.1_5_6_1, -0.1_1_2_7, 0.2_6_8_7]]]       ,							device=a__       ,							dtype=torch.float       ,							)
										# camembert = torch.hub.load('pytorch/fairseq', 'camembert.v0')
										# camembert.eval()
										# expected_slice = roberta.model.forward(input_ids)[0][:, :3, :3].detach()

										self.assertTrue(torch.allclose(output[:, :3, :3]       ,							a__       ,							atol=1e-4							)							)




