



from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_speech_available,
    is_torch_available,
)


A__  :  Union[str, Any]       = {
    '''configuration_speech_to_text_2''': ['''SPEECH_TO_TEXT_2_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''Speech2Text2Config'''],
    '''processing_speech_to_text_2''': ['''Speech2Text2Processor'''],
    '''tokenization_speech_to_text_2''': ['''Speech2Text2Tokenizer'''],
}


try:
						if not is_torch_available():
												raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
						pass
else:
						A__  :  Any       = [
						    '''SPEECH_TO_TEXT_2_PRETRAINED_MODEL_ARCHIVE_LIST''',
						    '''Speech2Text2ForCausalLM''',
						    '''Speech2Text2PreTrainedModel''',
						]


if TYPE_CHECKING:
						from .configuration_speech_to_text_a import SPEECH_TO_TEXT_2_PRETRAINED_CONFIG_ARCHIVE_MAP, SpeechaTextaConfig
						from .processing_speech_to_text_a import SpeechaTextaProcessor
						from .tokenization_speech_to_text_a import SpeechaTextaTokenizer

						try:
												if not is_torch_available():
																		raise OptionalDependencyNotAvailable()
						except OptionalDependencyNotAvailable:
												pass
						else:
												from .modeling_speech_to_text_a import (
												    SPEECH_TO_TEXT_2_PRETRAINED_MODEL_ARCHIVE_LIST,
												    SpeechaTextaForCausalLM,
												    SpeechaTextaPreTrainedModel,
												)

else:
						import sys

						A__  :  Tuple       = _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)










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

							_a      				=			[
							    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
							]

							_a      				=			[
							    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
							]
							_a      				=			[
							    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
							]

							_a      				=			[
							    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
							]
							_a      				=			[
							    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
							]

							_a      				=			[1_058_304, 0, 786_432, 7_169, 507, 76, 27, 40, 30, 76]

							_a      				=			[
							    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
							]
							_a      				=			[
							    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
							]
							_a      				=			[
							    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
							]
							_a      				=			[
							    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
							]
							_a      				=			[
							    -0.0_4_9_2, -0.0_5_2_4, -0.0_5_6_5, -0.0_6_4_0, -0.0_6_8_6, -0.0_6_8_4, -0.0_6_7_7, -0.0_6_6_4,
							    -0.0_6_0_5, -0.0_4_9_0, -0.0_3_3_0, -0.0_1_6_8, -0.0_0_8_3, -0.0_0_7_5, -0.0_0_5_1, 0.0_0_2_5,
							    0.0_1_3_6, 0.0_2_6_1, 0.0_3_8_6, 0.0_4_9_7, 0.0_5_8_0, 0.0_5_9_9, 0.0_5_8_3, 0.0_6_1_4,
							    0.0_7_4_0, 0.0_8_8_9, 0.1_0_2_3, 0.1_1_6_2, 0.1_2_1_1, 0.1_2_1_2, 0.1_2_5_1, 0.1_3_3_6,
							    0.1_5_0_2, 0.1_6_8_6, 0.1_8_8_3, 0.2_1_4_8, 0.2_3_6_3, 0.2_4_5_8, 0.2_5_0_7, 0.2_5_3_1
							]
							_a      				=			[
							    0.0_2_5_6, -0.0_5_4_4, 0.1_6_0_0, -0.0_0_3_2, 0.1_0_6_6, 0.0_8_2_5, -0.0_0_1_3, 0.3_4_4_0,
							    0.0_2_1_0, 0.0_4_1_2, -0.1_7_7_7, -0.0_8_9_2, -0.0_1_6_4, 0.0_2_8_5, -0.0_6_1_3, -0.0_6_1_7,
							    -0.0_1_3_7, -0.0_2_0_1, -0.0_1_7_5, 0.0_2_1_5, -0.0_6_2_7, 0.0_5_2_0, -0.0_7_3_0, 0.0_9_7_0,
							    -0.0_1_0_0, 0.0_4_4_2, -0.0_5_8_6, 0.0_2_0_7, -0.0_0_1_5, -0.0_0_8_2
							]
							_a      				=			[
							    0.0_4_1_5, 0.0_8_7_7, 0.0_0_2_2, -0.0_0_5_5, 0.0_7_5_1, 0.0_3_3_4, 0.0_3_2_4, -0.0_0_6_8,
							    0.0_0_1_1, 0.0_0_1_7, -0.0_6_7_6, 0.0_6_5_5, -0.0_1_4_3, 0.0_3_9_9, 0.0_3_0_3, 0.0_7_4_3,
							    -0.0_1_6_8, -0.0_3_9_4, -0.1_1_1_3, 0.0_1_2_4, 0.0_4_4_2, 0.0_2_6_7, -0.0_0_0_3, -0.1_5_3_6,
							    -0.0_1_1_6, -0.1_8_3_7, -0.0_1_8_0, -0.1_0_2_6, -0.0_7_7_7, -0.0_4_5_6
							]
							_a      				=			[
							    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					:	Tuple):
									lowerCAmelCase_      :					Any     =			JukeboxTokenizer.from_pretrained(self.model_id)
									lowerCAmelCase_      :					Union[str, Any]     =			tokenizer(**self.metas)['''input_ids''']
									return tokens




							@slow
							def        UpperCAmelCase__      (	self					:	Any):
									lowerCAmelCase_      :					Tuple     =			JukeboxModel.from_pretrained(self.model_id    ,    min_duration=0).eval()
									lowerCAmelCase_      :					Tuple     =			self.prepare_inputs()

									set_seed(0)
									lowerCAmelCase_      :					List[str]     =			[torch.zeros(1    ,    0    ,    dtype=torch.long).cpu() for _ in range(3)]
									lowerCAmelCase_      :					int     =			model._sample(A_    ,    A_    ,    [0]    ,    sample_length=4_0 * model.priors[0].raw_to_tokens    ,    save_results=A_)
									self.assertIn(zs[0][0].detach().cpu().tolist()    ,    [self.EXPECTED_OUTPUT_2, self.EXPECTED_OUTPUT_2_PT_2])

									set_seed(0)
									lowerCAmelCase_      :					str     =			model._sample(A_    ,    A_    ,    [1]    ,    sample_length=4_0 * model.priors[1].raw_to_tokens    ,    save_results=A_)
									self.assertIn(zs[1][0].detach().cpu().tolist()    ,    [self.EXPECTED_OUTPUT_1, self.EXPECTED_OUTPUT_1_PT_2])

									set_seed(0)
									lowerCAmelCase_      :					str     =			model._sample(A_    ,    A_    ,    [2]    ,    sample_length=4_0 * model.priors[2].raw_to_tokens    ,    save_results=A_)
									self.assertIn(zs[2][0].detach().cpu().tolist()    ,    [self.EXPECTED_OUTPUT_0, self.EXPECTED_OUTPUT_0_PT_2])




							@slow
							def        UpperCAmelCase__      (	self					:	List[Any]):
									lowerCAmelCase_      :					int     =			False
									lowerCAmelCase_      :					Any     =			JukeboxModel.from_pretrained(self.model_id    ,    min_duration=0).eval()

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

									lowerCAmelCase_      :					List[str]     =			model.priors[0]
									lowerCAmelCase_      :					str     =			0
									lowerCAmelCase_      :					Dict     =			top_prior.get_music_tokens_conds(A_    ,    start=A_    ,    end=start + top_prior.n_ctx)
									lowerCAmelCase_      :					Dict     =			top_prior.get_metadata(labels[0].clone()    ,    A_    ,    1_0_5_8_3_0_4    ,    0)

									self.assertIsNone(A_)
									self.assertListEqual(metadata.numpy()[0][:1_0].tolist()    ,    self.EXPECTED_Y_COND)

									lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_      :					List[str]     =			top_prior.get_cond(A_    ,    A_)
									torch.testing.assert_allclose(
									    audio_conditioning[0][0][:3_0].detach()    ,    torch.tensor(self.EXPECTED_AUDIO_COND)    ,    atol=1e-4    ,    rtol=1e-4)
									torch.testing.assert_allclose(
									    metadata_conditioning[0][0][:3_0].detach()    ,    torch.tensor(self.EXPECTED_META_COND)    ,    atol=1e-4    ,    rtol=1e-4)
									torch.testing.assert_allclose(
									    lyric_tokens[0, :3_0].detach()    ,    torch.tensor(self.EXPECTED_LYRIC_COND)    ,    atol=1e-4    ,    rtol=1e-4)




							@slow
							def        UpperCAmelCase__      (	self					:	Union[str, Any]):
									lowerCAmelCase_      :					int     =			False

									lowerCAmelCase_      :					Optional[Any]     =			JukeboxModel.from_pretrained(self.model_id    ,    min_duration=0).eval()
									set_seed(0)
									lowerCAmelCase_      :					Dict     =			torch.rand((1, 5_1_2_0, 1))
									lowerCAmelCase_      :					List[str]     =			list(self.prepare_inputs())

									lowerCAmelCase_      :					List[str]     =			[model.vqvae.encode(A_    ,    start_level=2    ,    bs_chunks=waveform.shape[0])[0], None, None]
									lowerCAmelCase_      :					Tuple     =			model._sample(
									    A_    ,    A_    ,    sample_levels=[0]    ,    save_results=A_    ,    sample_length=4_0 * model.priors[0].raw_to_tokens)
									torch.testing.assert_allclose(zs[0][0][:4_0]    ,    torch.tensor(self.EXPECTED_PRIMED_0))

									lowerCAmelCase_      :					Dict     =			torch.cat((zs[0], torch.zeros(1    ,    2_0_4_8 - zs[0].shape[-1]))    ,    dim=-1).long()
									lowerCAmelCase_      :					Optional[int]     =			[upper_a, model.vqvae.encode(A_    ,    start_level=1    ,    bs_chunks=waveform.shape[0])[0], None]
									lowerCAmelCase_      :					List[Any]     =			model._sample(
									    A_    ,    A_    ,    sample_levels=[1]    ,    save_results=A_    ,    sample_length=4_0 * model.priors[1].raw_to_tokens)
									torch.testing.assert_allclose(zs[1][0][:4_0]    ,    torch.tensor(self.EXPECTED_PRIMED_1))

									lowerCAmelCase_      :					Dict     =			torch.cat((zs[1], torch.zeros(1    ,    2_0_4_8 - zs[1].shape[-1]))    ,    dim=-1).long()
									lowerCAmelCase_      :					Tuple     =			[upper_a, upper_a, model.vqvae.encode(A_    ,    start_level=0    ,    bs_chunks=waveform.shape[0])[0]]
									lowerCAmelCase_      :					Optional[Any]     =			model._sample(
									    A_    ,    A_    ,    sample_levels=[2]    ,    save_results=A_    ,    sample_length=4_0 * model.priors[2].raw_to_tokens)
									torch.testing.assert_allclose(zs[2][0][:4_0].cpu()    ,    torch.tensor(self.EXPECTED_PRIMED_2))







							@slow
							def        UpperCAmelCase__      (	self					:	List[Any]):
									lowerCAmelCase_      :					str     =			JukeboxModel.from_pretrained(self.model_id    ,    min_duration=0).eval()
									set_seed(0)
									lowerCAmelCase_      :					List[Any]     =			torch.rand((1, 5_1_2_0, 1))
									with torch.no_grad():
											lowerCAmelCase_      :					Dict     =			model.vqvae.encode(A_    ,    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():
											lowerCAmelCase_      :					Dict     =			model.vqvae.decode(A_    ,    start_level=2    ,    bs_chunks=x.shape[0])
									torch.testing.assert_allclose(x[0, :4_0, 0]    ,    torch.tensor(self.EXPECTED_VQVAE_DECODE)    ,    atol=1e-4    ,    rtol=1e-4)





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

							# fmt: off
							_a      				=			[
							    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
							]
							_a      				=			[
							    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
							]

							_a      				=			[
							    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
							]
							_a      				=			[
							    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
							]

							_a      				=			[
							    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
							]
							_a      				=			[
							    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
							]

							_a      				=			[
							    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
							]
							_a      				=			[
							    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
							]

							_a      				=			[
							    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
							]
							_a      				=			[
							    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					:	int    ,    A_					:	Union[str, Any]):
									lowerCAmelCase_      :					Tuple     =			JukeboxTokenizer.from_pretrained(A_)
									lowerCAmelCase_      :					Any     =			tokenizer(**self.metas)['''input_ids''']
									return tokens




							@slow
							def        UpperCAmelCase__      (	self					:	Dict):
									lowerCAmelCase_      :					Optional[Any]     =			JukeboxModel.from_pretrained(self.model_id    ,    min_duration=0).eval()
									lowerCAmelCase_      :					Union[str, Any]     =			self.prepare_inputs(self.model_id)

									set_seed(0)
									lowerCAmelCase_      :					Optional[int]     =			[torch.zeros(1    ,    0    ,    dtype=torch.long).cpu() for _ in range(3)]
									lowerCAmelCase_      :					Dict     =			model._sample(A_    ,    A_    ,    [0]    ,    sample_length=6_0 * model.priors[0].raw_to_tokens    ,    save_results=A_)
									self.assertIn(zs[0][0].detach().cpu().tolist()    ,    [self.EXPECTED_OUTPUT_2, self.EXPECTED_OUTPUT_2_PT_2])

									set_seed(0)
									lowerCAmelCase_      :					Optional[int]     =			model._sample(A_    ,    A_    ,    [1]    ,    sample_length=6_0 * model.priors[1].raw_to_tokens    ,    save_results=A_)
									self.assertIn(zs[1][0].detach().cpu().tolist()    ,    [self.EXPECTED_OUTPUT_1, self.EXPECTED_OUTPUT_1_PT_2])

									set_seed(0)
									lowerCAmelCase_      :					Tuple     =			model._sample(A_    ,    A_    ,    [2]    ,    sample_length=6_0 * model.priors[2].raw_to_tokens    ,    save_results=A_)
									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					:	List[str]):
									lowerCAmelCase_      :					str     =			JukeboxModel.from_pretrained(self.model_id    ,    min_duration=0).eval()
									lowerCAmelCase_      :					Optional[Any]     =			[i.cuda() for i in self.prepare_inputs(self.model_id)]

									set_seed(0)
									model.priors[0].cuda()
									lowerCAmelCase_      :					Tuple     =			[torch.zeros(1    ,    0    ,    dtype=torch.long).cuda() for _ in range(3)]
									lowerCAmelCase_      :					Any     =			model._sample(A_    ,    A_    ,    [0]    ,    sample_length=6_0 * model.priors[0].raw_to_tokens    ,    save_results=A_)
									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()
									lowerCAmelCase_      :					Any     =			model._sample(A_    ,    A_    ,    [1]    ,    sample_length=6_0 * model.priors[1].raw_to_tokens    ,    save_results=A_)
									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()
									lowerCAmelCase_      :					Union[str, Any]     =			model._sample(A_    ,    A_    ,    [2]    ,    sample_length=6_0 * model.priors[2].raw_to_tokens    ,    save_results=A_)
									torch.testing.assert_allclose(zs[2][0].cpu()    ,    torch.tensor(self.EXPECTED_GPU_OUTPUTS_0))







							@slow
							def        UpperCAmelCase__      (	self					:	Optional[int]):
									lowerCAmelCase_      :					Tuple     =			'''ArthurZ/jukebox_prior_0'''
									lowerCAmelCase_      :					Any     =			JukeboxPrior.from_pretrained(A_    ,    min_duration=0).eval().half().to('''cuda''')

									lowerCAmelCase_      :					Any     =			self.prepare_inputs(A_)[0].cuda()
									lowerCAmelCase_      :					Any     =			model.get_metadata(A_    ,    0    ,    7_6_8_0    ,    0)
									set_seed(0)
									lowerCAmelCase_      :					Optional[int]     =			model.sample(1    ,    metadata=A_    ,    sample_tokens=6_0)
									self.assertIn(outputs[0].cpu().tolist()    ,    [self.EXPECTED_GPU_OUTPUTS_2, self.EXPECTED_GPU_OUTPUTS_2_PT_2])





