





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_			(       lowercase_				,							lowercase_				,							unittest.TestCase ):





					"""simple docstring"""


					A__								=     UNetaDModel
					A__								=     "sample"





					@property
					def        __magic_name__			(			self						):
										lowerCamelCase__      =							4
										lowerCamelCase__      =							14
										lowerCamelCase__      =							16

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

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





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





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





					def        __magic_name__			(			self						):
										pass





					def        __magic_name__			(			self						):
										pass





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





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





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





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





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





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





					def        __magic_name__			(			self						):
										lowerCamelCase__      =							{
										    "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",
										}
										lowerCamelCase__      =							self.dummy_input
										return init_dict, inputs_dict





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

										model.to(_lowerCAmelCase						)
										lowerCamelCase__      =							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        __magic_name__			(			self						):
										lowerCamelCase__      =							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__      =							model.config.in_channels
										lowerCamelCase__      =							16
										lowerCamelCase__      =							torch.randn((1, seq_len, num_features)						).permute(
										    0						,	2						,	1						)  # match original, we can update values and remove
										lowerCamelCase__      =							torch.full((num_features,)						,	0						)

										with torch.no_grad():
															lowerCamelCase__      =							model(_lowerCAmelCase						,	_lowerCAmelCase						).sample.permute(0						,	2						,	1						)

										lowerCamelCase__      =							output[0, -3:, -3:].flatten()
										# fmt: off
										lowerCamelCase__      =							torch.tensor([-2.137_172, 1.1_426_016, 0.3_688_687, -0.766_922, 0.7_303_146, 0.11_038_864, -0.4_760_633, 0.13_270_172, 0.02_591_348]						)
										# fmt: on
										self.assertTrue(torch.allclose(_lowerCAmelCase						,	_lowerCAmelCase						,	rtol=1E-3						)						)





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






					@slow
					def        __magic_name__			(			self						):
										lowerCamelCase__      =							"harmonai/maestro-150k"
										lowerCamelCase__      =							UNetaDModel.from_pretrained(_lowerCAmelCase						,	subfolder="unet"						)
										model.to(_lowerCAmelCase						)

										lowerCamelCase__      =							6_5536
										lowerCamelCase__      =							torch.sin(torch.arange(_lowerCAmelCase						)[None, None, :].repeat(1						,	2						,	1						)						).to(_lowerCAmelCase						)
										lowerCamelCase__      =							torch.tensor([1]						).to(_lowerCAmelCase						)

										with torch.no_grad():
															lowerCamelCase__      =							model(_lowerCAmelCase						,	_lowerCAmelCase						).sample

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

										assert (output_sum - 224.0_896).abs() < 0.5
										assert (output_max - 0.0_607).abs() < 4E-4

class SCREAMING_SNAKE_CASE_			(       lowercase_				,							lowercase_				,							unittest.TestCase ):





					"""simple docstring"""


					A__								=     UNetaDModel
					A__								=     "sample"





					@property
					def        __magic_name__			(			self						):
										lowerCamelCase__      =							4
										lowerCamelCase__      =							14
										lowerCamelCase__      =							16

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

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





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





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





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





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





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





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





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





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

										with torch.no_grad():
															lowerCamelCase__      =							model(**_lowerCAmelCase						)

															if isinstance(_lowerCAmelCase						,	_lowerCAmelCase						):
																				lowerCamelCase__      =							output.sample

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





					def        __magic_name__			(			self						):
										pass





					def        __magic_name__			(			self						):
										pass





					def        __magic_name__			(			self						):
										lowerCamelCase__      =							{
										    "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",
										}
										lowerCamelCase__      =							self.dummy_input
										return init_dict, inputs_dict





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

										value_function.to(_lowerCAmelCase						)
										lowerCamelCase__      =							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        __magic_name__			(			self						):
										lowerCamelCase__     ,		lowerCamelCase__      =							UNetaDModel.from_pretrained(
										    "bglick13/hopper-medium-v2-value-function-hor32"						,	output_loading_info=_lowerCAmelCase						,	subfolder="value_function"						)
										torch.manual_seed(0						)
										if torch.cuda.is_available():
															torch.cuda.manual_seed_all(0						)

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

										with torch.no_grad():
															lowerCamelCase__      =							value_function(_lowerCAmelCase						,	_lowerCAmelCase						).sample

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






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








# this script builds a small sample spm file tests/fixtures/test_sentencepiece_no_bos.model, with features needed by pegasus

# 1. pip install sentencepiece
#
# 2. wget https://raw.githubusercontent.com/google/sentencepiece/master/data/botchan.txt

# 3. build
import sentencepiece as spm

# pegasus:
# 1. no bos
# 2. eos_id is 1
# 3. unk_id is 2
# build a sample spm file accordingly
spm.SentencePieceTrainer.train(
    "--input=botchan.txt --model_prefix=test_sentencepiece_no_bos --bos_id=-1 --unk_id=2  --eos_id=1  --vocab_size=1000"
)

# 4. now update the fixture
# mv test_sentencepiece_no_bos.model ../../tests/fixtures/

