


'''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 					lowercase__						(		_snake_case  ,  _snake_case  ,  unittest.TestCase      ):





							'''simple docstring'''


							A_							:   List[Any]   			=						UNetaDModel
							A_							:   str   			=						"""sample"""






							@property
							def UpperCAmelCase_			(					self						):
											_SCREAMING_SNAKE_CASE				: str  	=   4
											_SCREAMING_SNAKE_CASE				: int  	=   14
											_SCREAMING_SNAKE_CASE				: List[Any]  	=   16

											_SCREAMING_SNAKE_CASE				: List[str]  	=   floats_tensor((batch_size, num_features, seq_len)						).to(__snake_case						)
											_SCREAMING_SNAKE_CASE				: Optional[int]  	=   torch.tensor([10] * batch_size						).to(__snake_case						)

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






							@property
							def UpperCAmelCase_			(					self						):
											return (4, 14, 16)






							@property
							def UpperCAmelCase_			(					self						):
											return (4, 14, 16)






							def UpperCAmelCase_			(					self						):
											pass






							def UpperCAmelCase_			(					self						):
											pass






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






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






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






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






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






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






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






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

											model.to(__snake_case						)
											_SCREAMING_SNAKE_CASE				: 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 UpperCAmelCase_			(					self						):
											_SCREAMING_SNAKE_CASE				: Optional[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						)

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

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

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






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






							@slow
							def UpperCAmelCase_			(					self						):
											_SCREAMING_SNAKE_CASE				: List[str]  	=   """harmonai/maestro-150k"""
											_SCREAMING_SNAKE_CASE				: str  	=   UNetaDModel.from_pretrained(__snake_case					,      subfolder="""unet"""						)
											model.to(__snake_case						)

											_SCREAMING_SNAKE_CASE				: List[Any]  	=   6_5536
											_SCREAMING_SNAKE_CASE				: Optional[int]  	=   torch.sin(torch.arange(__snake_case						)[None, None, :].repeat(1					,      2					,      1						)						).to(__snake_case						)
											_SCREAMING_SNAKE_CASE				: Tuple  	=   torch.tensor([1]						).to(__snake_case						)

											with torch.no_grad():
															_SCREAMING_SNAKE_CASE				: Dict  	=   model(__snake_case					,      __snake_case						).sample

											_SCREAMING_SNAKE_CASE				: List[Any]  	=   output.abs().sum()
											_SCREAMING_SNAKE_CASE				: List[str]  	=   output.abs().max()

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





class 					lowercase__						(		_snake_case  ,  _snake_case  ,  unittest.TestCase      ):





							'''simple docstring'''


							A_							:   Optional[Any]   			=						UNetaDModel
							A_							:   Any   			=						"""sample"""






							@property
							def UpperCAmelCase_			(					self						):
											_SCREAMING_SNAKE_CASE				: List[Any]  	=   4
											_SCREAMING_SNAKE_CASE				: Any  	=   14
											_SCREAMING_SNAKE_CASE				: Any  	=   16

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

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






							@property
							def UpperCAmelCase_			(					self						):
											return (4, 14, 16)






							@property
							def UpperCAmelCase_			(					self						):
											return (4, 14, 1)






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






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






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






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






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






							@unittest.skipIf(torch_device == """mps"""					,      """mish op not supported in MPS"""						)
							def UpperCAmelCase_			(					self						):
											# UNetRL is a value-function is different output shape
											_SCREAMING_SNAKE_CASE	, _SCREAMING_SNAKE_CASE				: Tuple  	=   self.prepare_init_args_and_inputs_for_common()
											_SCREAMING_SNAKE_CASE				: Tuple  	=   self.model_class(**__snake_case						)
											model.to(__snake_case						)
											model.eval()

											with torch.no_grad():
															_SCREAMING_SNAKE_CASE				: int  	=   model(**__snake_case						)

															if isinstance(__snake_case					,      __snake_case						):
																			_SCREAMING_SNAKE_CASE				: List[str]  	=   output.sample

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






							def UpperCAmelCase_			(					self						):
											pass






							def UpperCAmelCase_			(					self						):
											pass






							def UpperCAmelCase_			(					self						):
											_SCREAMING_SNAKE_CASE				: List[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, 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""",
											}
											_SCREAMING_SNAKE_CASE				: List[str]  	=   self.dummy_input
											return init_dict, inputs_dict






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

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

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

											with torch.no_grad():
															_SCREAMING_SNAKE_CASE				: Dict  	=   value_function(__snake_case					,      __snake_case						).sample

											# fmt: off
											_SCREAMING_SNAKE_CASE				: Dict  	=   torch.tensor([165.25] * seq_len						)
											# fmt: on
											self.assertTrue(torch.allclose(__snake_case					,      __snake_case					,      rtol=1e-3						)						)






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








'''simple docstring'''


def     snake_case_     (						SCREAMING_SNAKE_CASE__							,   SCREAMING_SNAKE_CASE__  ):







				"""simple docstring"""





				return (input_a, input_a).count(0  ) % 2







def     snake_case_     (						):







				"""simple docstring"""





				assert xor_gate(0							,   0  ) == 0
				assert xor_gate(0							,   1  ) == 1
				assert xor_gate(1							,   0  ) == 1
				assert xor_gate(1							,   1  ) == 0


if __name__ == "__main__":
				print(xor_gate(0, 0))
				print(xor_gate(0, 1))




