




import unittest
from unittest import skip

from transformers import is_torch_available
from transformers.testing_utils import require_torch, slow
from transformers.trainer_utils import set_seed


if is_torch_available():
				import torch

				from transformers import JukeboxModel, JukeboxPrior, JukeboxTokenizer






@require_torch
class 					__lowerCAmelCase				( unittest.TestCase				):
			snake_case_   :				Optional[int]				  =  (JukeboxModel,) if is_torch_available() else ()
			snake_case_   :				int				  =  "openai/jukebox-1b-lyrics"
			snake_case_   :				Dict				  =  {
			    "artist": "Zac Brown Band",
			    "genres": "Country",
			    "lyrics": "I met a traveller from an antique land,\n    Who said \"Two vast and trunkless legs of stone\n    Stand in the desert. . . . Near them, on the sand,\n    Half sunk a shattered visage lies, whose frown,\n    And wrinkled lip, and sneer of cold command,\n    Tell that its sculptor well those passions read\n    Which yet survive, stamped on these lifeless things,\n    The hand that mocked them, and the heart that fed;\n    And on the pedestal, these words appear:\n    My name is Ozymandias, King of Kings;\n    Look on my Works, ye Mighty, and despair!\n    Nothing beside remains. Round the decay\n    Of that colossal Wreck, boundless and bare\n    The lone and level sands stretch far away\n    ",
			}
			# fmt: off
			snake_case_   :				Union[str, Any]				  =  [
			    1_864, 1_536, 1_213, 1_870, 1_357, 1_536, 519, 880, 1_323, 789, 1_082, 534,
			    1_000, 1_445, 1_105, 1_130, 967, 515, 1_434, 1_620, 534, 1_495, 283, 1_445,
			    333, 1_307, 539, 1_631, 1_528, 375, 1_434, 673, 627, 710, 778, 1_883,
			    1_405, 1_276, 1_455, 1_228
			]

			snake_case_   :				Union[str, Any]				  =  [
			    1_489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
			    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
			    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
			    653, 653, 653, 653
			]

			snake_case_   :				int				  =  [
			    1_125, 1_751, 697, 1_776, 1_141, 1_476, 391, 697, 1_125, 684, 867, 416,
			    844, 1_372, 1_274, 717, 1_274, 844, 1_299, 1_419, 697, 1_370, 317, 1_125,
			    191, 1_440, 1_370, 1_440, 1_370, 282, 1_621, 1_370, 368, 349, 867, 1_872,
			    1_262, 869, 1_728, 747
			]
			snake_case_   :				str				  =  [
			    416, 416, 1_125, 1_125, 416, 416, 416, 416, 416, 416, 416, 416,
			    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
			    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
			    416, 416, 416, 416
			]

			snake_case_   :				int				  =  [
			    1_755, 842, 307, 1_843, 1_022, 1_395, 234, 1_554, 806, 739, 1_022, 442,
			    616, 556, 268, 1_499, 933, 457, 1_440, 1_837, 755, 985, 308, 902,
			    293, 1_443, 1_671, 1_141, 1_533, 555, 1_562, 1_061, 287, 417, 1_022, 2_008,
			    1_186, 1_015, 1_777, 268
			]
			snake_case_   :				Union[str, Any]				  =  [
			    854, 842, 1_353, 114, 1_353, 842, 185, 842, 185, 114, 591, 842,
			    185, 417, 185, 842, 307, 842, 591, 842, 185, 842, 307, 842,
			    591, 842, 1_353, 842, 185, 842, 591, 842, 591, 114, 591, 842,
			    185, 842, 591, 89
			]

			snake_case_   :				Optional[int]				  =  [1_058_304, 0, 786_432, 7_169, 507, 76, 27, 40, 30, 76]

			snake_case_   :				List[str]				  =  [
			    390, 1_160, 1_002, 1_907, 1_788, 1_788, 1_788, 1_907, 1_002, 1_002, 1_854, 1_002,
			    1_002, 1_002, 1_002, 1_002, 1_002, 1_160, 1_160, 1_606, 596, 596, 1_160, 1_002,
			    1_516, 596, 1_002, 1_002, 1_002, 1_907, 1_788, 1_788, 1_788, 1_854, 1_788, 1_907,
			    1_907, 1_788, 596, 1_626
			]
			snake_case_   :				Optional[int]				  =  [
			    1_236, 1_668, 1_484, 1_920, 1_848, 1_409, 139, 864, 1_828, 1_272, 1_599, 824,
			    1_672, 139, 555, 1_484, 824, 1_920, 555, 596, 1_579, 1_599, 1_231, 1_599,
			    1_637, 1_407, 212, 824, 1_599, 116, 1_433, 824, 258, 1_599, 1_433, 1_895,
			    1_063, 1_433, 1_433, 1_599
			]
			snake_case_   :				Union[str, Any]				  =  [
			    1_684, 1_873, 1_119, 1_189, 395, 611, 1_901, 972, 890, 1_337, 1_392, 1_927,
			    96, 972, 672, 780, 1_119, 890, 158, 771, 1_073, 1_927, 353, 1_331,
			    1_269, 1_459, 1_333, 1_645, 812, 1_577, 1_337, 606, 353, 981, 1_466, 619,
			    197, 391, 302, 1_930
			]
			snake_case_   :				List[Any]				  =  [
			    390, 1_160, 1_002, 1_907, 1_788, 1_788, 1_788, 1_907, 1_002, 1_002, 1_854, 1_002,
			    1_002, 1_002, 1_002, 1_002, 1_002, 1_160, 1_160, 1_606, 596, 596, 1_160, 1_002,
			    1_516, 596, 1_002, 1_002, 1_002, 1_907, 1_788, 1_788, 1_788, 1_854, 1_788, 1_907,
			    1_907, 1_788, 596, 1_626
			]
			snake_case_   :				Tuple				  =  [
			    -0.0_492, -0.0_524, -0.0_565, -0.0_640, -0.0_686, -0.0_684, -0.0_677, -0.0_664,
			    -0.0_605, -0.0_490, -0.0_330, -0.0_168, -0.0_083, -0.0_075, -0.0_051, 0.0_025,
			    0.0_136, 0.0_261, 0.0_386, 0.0_497, 0.0_580, 0.0_599, 0.0_583, 0.0_614,
			    0.0_740, 0.0_889, 0.1_023, 0.1_162, 0.1_211, 0.1_212, 0.1_251, 0.1_336,
			    0.1_502, 0.1_686, 0.1_883, 0.2_148, 0.2_363, 0.2_458, 0.2_507, 0.2_531
			]
			snake_case_   :				Optional[int]				  =  [
			    0.0_256, -0.0_544, 0.1_600, -0.0_032, 0.1_066, 0.0_825, -0.0_013, 0.3_440,
			    0.0_210, 0.0_412, -0.1_777, -0.0_892, -0.0_164, 0.0_285, -0.0_613, -0.0_617,
			    -0.0_137, -0.0_201, -0.0_175, 0.0_215, -0.0_627, 0.0_520, -0.0_730, 0.0_970,
			    -0.0_100, 0.0_442, -0.0_586, 0.0_207, -0.0_015, -0.0_082
			]
			snake_case_   :				List[Any]				  =  [
			    0.0_415, 0.0_877, 0.0_022, -0.0_055, 0.0_751, 0.0_334, 0.0_324, -0.0_068,
			    0.0_011, 0.0_017, -0.0_676, 0.0_655, -0.0_143, 0.0_399, 0.0_303, 0.0_743,
			    -0.0_168, -0.0_394, -0.1_113, 0.0_124, 0.0_442, 0.0_267, -0.0_003, -0.1_536,
			    -0.0_116, -0.1_837, -0.0_180, -0.1_026, -0.0_777, -0.0_456
			]
			snake_case_   :				Optional[Any]				  =  [
			    76, 27, 40, 30, 76, 46, 44, 47, 40, 37, 38, 31, 45, 45, 76, 38, 31, 33,
			    45, 76, 41, 32, 76, 45, 46, 41, 40, 31, 78, 76
			]







			def UpperCamelCase    ( self					:    str   ):




									"""simple docstring"""


									_UpperCAmelCase									=  JukeboxTokenizer.from_pretrained(self.model_id   )
									_UpperCAmelCase									=  tokenizer(**self.metas   )["input_ids"]
									return tokens







			@slow
			def UpperCamelCase    ( self					:    Any   ):




									"""simple docstring"""


									_UpperCAmelCase									=  JukeboxModel.from_pretrained(self.model_id					,					min_duration=0   ).eval()
									_UpperCAmelCase									=  self.prepare_inputs()

									set_seed(0   )
									_UpperCAmelCase									=  [torch.zeros(1					,					0					,					dtype=torch.long   ).cpu() for _ in range(3   )]
									_UpperCAmelCase									=  model._sample(snake_case__					,					snake_case__					,					[0]					,					sample_length=40 * model.priors[0].raw_to_tokens					,					save_results=snake_case__   )
									self.assertIn(zs[0][0].detach().cpu().tolist()					,					[self.EXPECTED_OUTPUT_2, self.EXPECTED_OUTPUT_2_PT_2]   )

									set_seed(0   )
									_UpperCAmelCase									=  model._sample(snake_case__					,					snake_case__					,					[1]					,					sample_length=40 * model.priors[1].raw_to_tokens					,					save_results=snake_case__   )
									self.assertIn(zs[1][0].detach().cpu().tolist()					,					[self.EXPECTED_OUTPUT_1, self.EXPECTED_OUTPUT_1_PT_2]   )

									set_seed(0   )
									_UpperCAmelCase									=  model._sample(snake_case__					,					snake_case__					,					[2]					,					sample_length=40 * model.priors[2].raw_to_tokens					,					save_results=snake_case__   )
									self.assertIn(zs[2][0].detach().cpu().tolist()					,					[self.EXPECTED_OUTPUT_0, self.EXPECTED_OUTPUT_0_PT_2]   )







			@slow
			def UpperCamelCase    ( self					:    List[Any]   ):




									"""simple docstring"""


									_UpperCAmelCase									=  False
									_UpperCAmelCase									=  JukeboxModel.from_pretrained(self.model_id					,					min_duration=0   ).eval()

									_UpperCAmelCase									=  self.prepare_inputs()
									set_seed(0   )
									_UpperCAmelCase									=  [torch.zeros(1					,					0					,					dtype=torch.long   ) for _ in range(3   )]

									_UpperCAmelCase									=  model.priors[0]
									_UpperCAmelCase									=  0
									_UpperCAmelCase									=  top_prior.get_music_tokens_conds(snake_case__					,					start=snake_case__					,					end=start + top_prior.n_ctx   )
									_UpperCAmelCase									=  top_prior.get_metadata(labels[0].clone()					,					snake_case__					,					1_058_304					,					0   )

									self.assertIsNone(snake_case__   )
									self.assertListEqual(metadata.numpy()[0][:10].tolist()					,					self.EXPECTED_Y_COND   )

									_UpperCAmelCase							,				_UpperCAmelCase							,				_UpperCAmelCase									=  top_prior.get_cond(snake_case__					,					snake_case__   )
									torch.testing.assert_allclose(
									    audio_conditioning[0][0][:30].detach()					,					torch.tensor(self.EXPECTED_AUDIO_COND   )					,					atol=1e-4					,					rtol=1e-4   )
									torch.testing.assert_allclose(
									    metadata_conditioning[0][0][:30].detach()					,					torch.tensor(self.EXPECTED_META_COND   )					,					atol=1e-4					,					rtol=1e-4   )
									torch.testing.assert_allclose(
									    lyric_tokens[0, :30].detach()					,					torch.tensor(self.EXPECTED_LYRIC_COND   )					,					atol=1e-4					,					rtol=1e-4   )







			@slow
			def UpperCamelCase    ( self					:    Optional[int]   ):




									"""simple docstring"""


									_UpperCAmelCase									=  False

									_UpperCAmelCase									=  JukeboxModel.from_pretrained(self.model_id					,					min_duration=0   ).eval()
									set_seed(0   )
									_UpperCAmelCase									=  torch.rand((1, 5_120, 1)   )
									_UpperCAmelCase									=  list(self.prepare_inputs()   )

									_UpperCAmelCase									=  [model.vqvae.encode(snake_case__					,					start_level=2					,					bs_chunks=waveform.shape[0]   )[0], None, None]
									_UpperCAmelCase									=  model._sample(
									    snake_case__					,					snake_case__					,					sample_levels=[0]					,					save_results=snake_case__					,					sample_length=40 * model.priors[0].raw_to_tokens   )
									torch.testing.assert_allclose(zs[0][0][:40]					,					torch.tensor(self.EXPECTED_PRIMED_0   )   )

									_UpperCAmelCase									=  torch.cat((zs[0], torch.zeros(1					,					2_048 - zs[0].shape[-1]   ))					,					dim=-1   ).long()
									_UpperCAmelCase									=  [upper_a, model.vqvae.encode(snake_case__					,					start_level=1					,					bs_chunks=waveform.shape[0]   )[0], None]
									_UpperCAmelCase									=  model._sample(
									    snake_case__					,					snake_case__					,					sample_levels=[1]					,					save_results=snake_case__					,					sample_length=40 * model.priors[1].raw_to_tokens   )
									torch.testing.assert_allclose(zs[1][0][:40]					,					torch.tensor(self.EXPECTED_PRIMED_1   )   )

									_UpperCAmelCase									=  torch.cat((zs[1], torch.zeros(1					,					2_048 - zs[1].shape[-1]   ))					,					dim=-1   ).long()
									_UpperCAmelCase									=  [upper_a, upper_a, model.vqvae.encode(snake_case__					,					start_level=0					,					bs_chunks=waveform.shape[0]   )[0]]
									_UpperCAmelCase									=  model._sample(
									    snake_case__					,					snake_case__					,					sample_levels=[2]					,					save_results=snake_case__					,					sample_length=40 * model.priors[2].raw_to_tokens   )
									torch.testing.assert_allclose(zs[2][0][:40].cpu()					,					torch.tensor(self.EXPECTED_PRIMED_2   )   )







			@slow
			def UpperCamelCase    ( self					:    Tuple   ):




									"""simple docstring"""


									_UpperCAmelCase									=  JukeboxModel.from_pretrained(self.model_id					,					min_duration=0   ).eval()
									set_seed(0   )
									_UpperCAmelCase									=  torch.rand((1, 5_120, 1)   )
									with torch.no_grad():
															_UpperCAmelCase									=  model.vqvae.encode(snake_case__					,					start_level=2					,					bs_chunks=x.shape[0]   )
									torch.testing.assert_allclose(zs[0][0]					,					torch.tensor(self.EXPECTED_VQVAE_ENCODE   )   )

									with torch.no_grad():
															_UpperCAmelCase									=  model.vqvae.decode(snake_case__					,					start_level=2					,					bs_chunks=x.shape[0]   )
									torch.testing.assert_allclose(x[0, :40, 0]					,					torch.tensor(self.EXPECTED_VQVAE_DECODE   )					,					atol=1e-4					,					rtol=1e-4   )






@require_torch
class 					__lowerCAmelCase				( unittest.TestCase				):
			snake_case_   :				str				  =  (JukeboxModel,) if is_torch_available() else ()
			snake_case_   :				Any				  =  "openai/jukebox-5b-lyrics"
			snake_case_   :				Optional[Any]				  =  {
			    "artist": "Zac Brown Band",
			    "genres": "Country",
			    "lyrics": "I met a traveller from an antique land,\n    Who said \"Two vast and trunkless legs of stone\n    Stand in the desert. . . . Near them, on the sand,\n    Half sunk a shattered visage lies, whose frown,\n    And wrinkled lip, and sneer of cold command,\n    Tell that its sculptor well those passions read\n    Which yet survive, stamped on these lifeless things,\n    The hand that mocked them, and the heart that fed;\n    And on the pedestal, these words appear:\n    My name is Ozymandias, King of Kings;\n    Look on my Works, ye Mighty, and despair!\n    Nothing beside remains. Round the decay\n    Of that colossal Wreck, boundless and bare\n    The lone and level sands stretch far away\n    ",
			}

			# fmt: off
			snake_case_   :				Dict				  =  [
			    1_489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
			    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
			    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
			    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
			    1_489, 1_489, 1_489, 1_489, 1_150, 1_853, 1_509, 1_150, 1_357, 1_509, 6, 1_272
			]
			snake_case_   :				Tuple				  =  [
			    1_489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
			    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
			    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
			    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
			    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653
			]

			snake_case_   :				Dict				  =  [
			    1_125, 416, 1_125, 1_125, 1_125, 1_125, 1_125, 416, 416, 416, 416, 416,
			    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
			    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
			    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
			    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416
			]
			snake_case_   :				int				  =  [
			    416, 416, 1_125, 1_125, 416, 416, 416, 416, 416, 416, 416, 416, 416,
			    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
			    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
			    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
			    416, 416, 416, 416, 416, 416, 416, 416
			]

			snake_case_   :				Dict				  =  [
			    1_755, 1_061, 234, 1_755, 1_061, 1_755, 185, 290, 307, 307, 616, 616,
			    616, 616, 616, 616, 307, 290, 417, 1_755, 234, 1_755, 185, 290,
			    290, 290, 307, 616, 616, 616, 616, 616, 290, 234, 234, 1_755,
			    234, 234, 1_755, 234, 185, 185, 307, 616, 616, 616, 616, 290,
			    1_755, 1_755, 1_755, 234, 234, 1_755, 1_572, 290, 307, 616, 34, 616
			]
			snake_case_   :				Dict				  =  [
			    854, 842, 1_353, 114, 1_353, 842, 185, 842, 185, 114, 591, 842, 185,
			    417, 185, 842, 307, 842, 591, 842, 185, 842, 185, 842, 591, 842,
			    1_353, 842, 185, 842, 591, 842, 591, 114, 591, 842, 185, 842, 591,
			    89, 591, 842, 591, 842, 591, 417, 1_372, 842, 1_372, 842, 34, 842,
			    185, 89, 591, 842, 185, 842, 591, 632
			]

			snake_case_   :				Optional[int]				  =  [
			    1_489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
			    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
			    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
			    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
			    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653
			]
			snake_case_   :				List[Any]				  =  [
			    1_489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
			    653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
			    653, 653, 653, 653, 653, 653, 653, 1_853, 1_177, 1_536, 1_228,
			    710, 475, 1_489, 1_229, 1_224, 231, 1_224, 252, 1_434, 653, 475,
			    1_106, 1_877, 1_599, 1_228, 1_600, 1_683, 1_182, 1_853, 475, 1_864,
			    252, 1_229, 1_434, 2_001
			]

			snake_case_   :				Optional[Any]				  =  [
			    1_125, 1_125, 416, 1_125, 1_125, 416, 1_125, 1_125, 416, 416, 1_125, 416,
			    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
			    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
			    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
			    416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416
			]
			snake_case_   :				Tuple				  =  [
			    491, 1_755, 34, 1_613, 1_755, 417, 992, 1_613, 222, 842, 1_353, 1_613,
			    844, 632, 185, 1_613, 844, 632, 185, 1_613, 185, 842, 677, 1_613,
			    185, 114, 1_353, 1_613, 307, 89, 844, 1_613, 307, 1_332, 234, 1_979,
			    307, 89, 1_353, 616, 34, 842, 185, 842, 34, 842, 185, 842,
			    307, 114, 185, 89, 34, 1_268, 185, 89, 34, 842, 185, 89
			]







			def UpperCamelCase    ( self					:    str					,					snake_case__					:    Union[str, Any]   ):




									"""simple docstring"""


									_UpperCAmelCase									=  JukeboxTokenizer.from_pretrained(snake_case__   )
									_UpperCAmelCase									=  tokenizer(**self.metas   )["input_ids"]
									return tokens







			@slow
			def UpperCamelCase    ( self					:    Union[str, Any]   ):




									"""simple docstring"""


									_UpperCAmelCase									=  JukeboxModel.from_pretrained(self.model_id					,					min_duration=0   ).eval()
									_UpperCAmelCase									=  self.prepare_inputs(self.model_id   )

									set_seed(0   )
									_UpperCAmelCase									=  [torch.zeros(1					,					0					,					dtype=torch.long   ).cpu() for _ in range(3   )]
									_UpperCAmelCase									=  model._sample(snake_case__					,					snake_case__					,					[0]					,					sample_length=60 * model.priors[0].raw_to_tokens					,					save_results=snake_case__   )
									self.assertIn(zs[0][0].detach().cpu().tolist()					,					[self.EXPECTED_OUTPUT_2, self.EXPECTED_OUTPUT_2_PT_2]   )

									set_seed(0   )
									_UpperCAmelCase									=  model._sample(snake_case__					,					snake_case__					,					[1]					,					sample_length=60 * model.priors[1].raw_to_tokens					,					save_results=snake_case__   )
									self.assertIn(zs[1][0].detach().cpu().tolist()					,					[self.EXPECTED_OUTPUT_1, self.EXPECTED_OUTPUT_1_PT_2]   )

									set_seed(0   )
									_UpperCAmelCase									=  model._sample(snake_case__					,					snake_case__					,					[2]					,					sample_length=60 * model.priors[2].raw_to_tokens					,					save_results=snake_case__   )
									self.assertIn(zs[2][0].detach().cpu().tolist()					,					[self.EXPECTED_OUTPUT_0, self.EXPECTED_OUTPUT_0_PT_2]   )







			@slow
			@skip("Not enough GPU memory on CI runners"   )
			def UpperCamelCase    ( self					:    Any   ):




									"""simple docstring"""


									_UpperCAmelCase									=  JukeboxModel.from_pretrained(self.model_id					,					min_duration=0   ).eval()
									_UpperCAmelCase									=  [i.cuda() for i in self.prepare_inputs(self.model_id   )]

									set_seed(0   )
									model.priors[0].cuda()
									_UpperCAmelCase									=  [torch.zeros(1					,					0					,					dtype=torch.long   ).cuda() for _ in range(3   )]
									_UpperCAmelCase									=  model._sample(snake_case__					,					snake_case__					,					[0]					,					sample_length=60 * model.priors[0].raw_to_tokens					,					save_results=snake_case__   )
									torch.testing.assert_allclose(zs[0][0].cpu()					,					torch.tensor(self.EXPECTED_GPU_OUTPUTS_2   )   )
									model.priors[0].cpu()

									set_seed(0   )
									model.priors[1].cuda()
									_UpperCAmelCase									=  model._sample(snake_case__					,					snake_case__					,					[1]					,					sample_length=60 * model.priors[1].raw_to_tokens					,					save_results=snake_case__   )
									torch.testing.assert_allclose(zs[1][0].cpu()					,					torch.tensor(self.EXPECTED_GPU_OUTPUTS_1   )   )
									model.priors[1].cpu()

									set_seed(0   )
									model.priors[2].cuda()
									_UpperCAmelCase									=  model._sample(snake_case__					,					snake_case__					,					[2]					,					sample_length=60 * model.priors[2].raw_to_tokens					,					save_results=snake_case__   )
									torch.testing.assert_allclose(zs[2][0].cpu()					,					torch.tensor(self.EXPECTED_GPU_OUTPUTS_0   )   )







			@slow
			def UpperCamelCase    ( self					:    Union[str, Any]   ):




									"""simple docstring"""


									_UpperCAmelCase									=  "ArthurZ/jukebox_prior_0"
									_UpperCAmelCase									=  JukeboxPrior.from_pretrained(snake_case__					,					min_duration=0   ).eval().half().to("cuda"   )

									_UpperCAmelCase									=  self.prepare_inputs(snake_case__   )[0].cuda()
									_UpperCAmelCase									=  model.get_metadata(snake_case__					,					0					,					7_680					,					0   )
									set_seed(0   )
									_UpperCAmelCase									=  model.sample(1					,					metadata=snake_case__					,					sample_tokens=60   )
									self.assertIn(outputs[0].cpu().tolist()					,					[self.EXPECTED_GPU_OUTPUTS_2, self.EXPECTED_GPU_OUTPUTS_2_PT_2]   )










import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowercase_			:       Tuple						       =				logging.get_logger(__name__)

lowercase_			:       List[Any]						       =				{
    'facebook/sam-vit-huge': 'https://huggingface.co/facebook/sam-vit-huge/resolve/main/config.json',
    'facebook/sam-vit-large': 'https://huggingface.co/facebook/sam-vit-large/resolve/main/config.json',
    'facebook/sam-vit-base': 'https://huggingface.co/facebook/sam-vit-base/resolve/main/config.json',
}






class 					__lowerCAmelCase				( UpperCAmelCase__				):







			def __init__( self					:    Optional[int]					,					snake_case__					:    int=256					,					snake_case__					:    Dict=1_024					,					snake_case__					:    Dict=16					,					snake_case__					:    Union[str, Any]=16					,					snake_case__					:    Any=4					,					snake_case__					:    Tuple="gelu"					,					snake_case__					:    Tuple=1e-6					,					**snake_case__					:    str					,					):




									"""simple docstring"""


									super().__init__(**snake_case__   )
									_UpperCAmelCase									=  hidden_size
									_UpperCAmelCase									=  image_size
									_UpperCAmelCase									=  patch_size
									_UpperCAmelCase									=  image_size // patch_size
									_UpperCAmelCase									=  mask_input_channels
									_UpperCAmelCase									=  num_point_embeddings
									_UpperCAmelCase									=  hidden_act
									_UpperCAmelCase									=  layer_norm_eps






class 					__lowerCAmelCase				( UpperCAmelCase__				):







			def __init__( self					:    str					,					snake_case__					:    List[str]=256					,					snake_case__					:    Tuple="relu"					,					snake_case__					:    int=2_048					,					snake_case__					:    Optional[int]=2					,					snake_case__					:    Dict=8					,					snake_case__					:    List[Any]=2					,					snake_case__					:    List[Any]=3					,					snake_case__					:    List[Any]=3					,					snake_case__					:    Optional[int]=256					,					snake_case__					:    Tuple=1e-6					,					**snake_case__					:    Any					,					):




									"""simple docstring"""


									super().__init__(**snake_case__   )
									_UpperCAmelCase									=  hidden_size
									_UpperCAmelCase									=  hidden_act
									_UpperCAmelCase									=  mlp_dim
									_UpperCAmelCase									=  num_hidden_layers
									_UpperCAmelCase									=  num_attention_heads
									_UpperCAmelCase									=  attention_downsample_rate
									_UpperCAmelCase									=  num_multimask_outputs
									_UpperCAmelCase									=  iou_head_depth
									_UpperCAmelCase									=  iou_head_hidden_dim
									_UpperCAmelCase									=  layer_norm_eps






class 					__lowerCAmelCase				( UpperCAmelCase__				):







			def __init__( self					:    List[str]					,					snake_case__					:    str=768					,					snake_case__					:    Dict=256					,					snake_case__					:    Any=12					,					snake_case__					:    Union[str, Any]=12					,					snake_case__					:    Any=3					,					snake_case__					:    List[str]=1_024					,					snake_case__					:    Union[str, Any]=16					,					snake_case__					:    Any="gelu"					,					snake_case__					:    List[str]=1e-06					,					snake_case__					:    Union[str, Any]=0.0					,					snake_case__					:    Optional[int]=1e-10					,					snake_case__					:    int=True					,					snake_case__					:    Dict=4.0					,					snake_case__					:    Optional[Any]=True					,					snake_case__					:    Tuple=True					,					snake_case__					:    List[Any]=14					,					snake_case__					:    List[str]=[2, 5, 8, 11]					,					snake_case__					:    List[Any]=128					,					snake_case__					:    List[Any]=None					,					**snake_case__					:    Any					,					):




									"""simple docstring"""


									super().__init__(**snake_case__   )

									_UpperCAmelCase									=  hidden_size
									_UpperCAmelCase									=  output_channels
									_UpperCAmelCase									=  num_hidden_layers
									_UpperCAmelCase									=  num_attention_heads
									_UpperCAmelCase									=  num_channels
									_UpperCAmelCase									=  image_size
									_UpperCAmelCase									=  patch_size
									_UpperCAmelCase									=  hidden_act
									_UpperCAmelCase									=  layer_norm_eps
									_UpperCAmelCase									=  attention_dropout
									_UpperCAmelCase									=  initializer_range
									_UpperCAmelCase									=  qkv_bias
									_UpperCAmelCase									=  mlp_ratio
									_UpperCAmelCase									=  use_abs_pos
									_UpperCAmelCase									=  use_rel_pos
									_UpperCAmelCase									=  window_size
									_UpperCAmelCase									=  global_attn_indexes
									_UpperCAmelCase									=  num_pos_feats
									_UpperCAmelCase									=  int(hidden_size * mlp_ratio   ) if mlp_dim is None else mlp_dim






class 					__lowerCAmelCase				( UpperCAmelCase__				):
			snake_case_   :				Tuple				  =  "sam"
			snake_case_   :				Tuple				  =  True







			def __init__( self					:    Optional[int]					,					snake_case__					:    Optional[Any]=None					,					snake_case__					:    Any=None					,					snake_case__					:    int=None					,					snake_case__					:    List[Any]=0.02					,					**snake_case__					:    str					,					):




									"""simple docstring"""


									super().__init__(**snake_case__   )
									_UpperCAmelCase									=  vision_config if vision_config is not None else {}
									_UpperCAmelCase									=  prompt_encoder_config if prompt_encoder_config is not None else {}
									_UpperCAmelCase									=  mask_decoder_config if mask_decoder_config is not None else {}

									if isinstance(snake_case__					,					snake_case__   ):
															_UpperCAmelCase									=  vision_config.to_dict()
									if isinstance(snake_case__					,					snake_case__   ):
															_UpperCAmelCase									=  prompt_encoder_config.to_dict()
									if isinstance(snake_case__					,					snake_case__   ):
															_UpperCAmelCase									=  mask_decoder_config.to_dict()

									_UpperCAmelCase									=  SamVisionConfig(**snake_case__   )
									_UpperCAmelCase									=  SamPromptEncoderConfig(**snake_case__   )
									_UpperCAmelCase									=  SamMaskDecoderConfig(**snake_case__   )
									_UpperCAmelCase									=  initializer_range







			def UpperCamelCase    ( self					:    Any   ):




									"""simple docstring"""


									_UpperCAmelCase									=  copy.deepcopy(self.__dict__   )
									_UpperCAmelCase									=  self.vision_config.to_dict()
									_UpperCAmelCase									=  self.prompt_encoder_config.to_dict()
									_UpperCAmelCase									=  self.mask_decoder_config.to_dict()
									_UpperCAmelCase									=  self.__class__.model_type
									return output




