



import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin


if is_torch_available():
						import torch

if is_vision_available():
						from PIL import Image

						from transformers import SwinaSRImageProcessor
						from transformers.image_transforms import get_image_size



class     __snake_case						(					unittest.TestCase						):




							def __init__(	self					:	List[Any]    ,    A_					:	str    ,    A_					:	Optional[int]=7    ,    A_					:	int=3    ,    A_					:	Union[str, Any]=1_8    ,    A_					:	str=3_0    ,    A_					:	List[Any]=4_0_0    ,    A_					:	Optional[Any]=True    ,    A_					:	List[Any]=1 / 2_5_5    ,    A_					:	int=True    ,    A_					:	Dict=8    ,    ):
									lowerCAmelCase_      :					int     =			parent
									lowerCAmelCase_      :					Union[str, Any]     =			batch_size
									lowerCAmelCase_      :					Optional[int]     =			num_channels
									lowerCAmelCase_      :					Tuple     =			image_size
									lowerCAmelCase_      :					int     =			min_resolution
									lowerCAmelCase_      :					List[str]     =			max_resolution
									lowerCAmelCase_      :					Union[str, Any]     =			do_rescale
									lowerCAmelCase_      :					Dict     =			rescale_factor
									lowerCAmelCase_      :					int     =			do_pad
									lowerCAmelCase_      :					int     =			pad_size




							def        UpperCAmelCase__      (	self					:	Any):
									return {
									    "do_rescale": self.do_rescale,
									    "rescale_factor": self.rescale_factor,
									    "do_pad": self.do_pad,
									    "pad_size": self.pad_size,
									}







							def        UpperCAmelCase__      (	self					:	Dict    ,    A_					:	Dict=False    ,    A_					:	int=False    ,    A_					:	int=False):
									assert not (numpify and torchify), "You cannot specify both numpy and PyTorch tensors at the same time"

									if equal_resolution:
											lowerCAmelCase_      :					Optional[int]     =			[]
											for i in range(self.batch_size):
													image_inputs.append(
													    np.random.randint(
													        2_5_5    ,    size=(self.num_channels, self.max_resolution, self.max_resolution)    ,    dtype=np.uinta))
									else:
											lowerCAmelCase_      :					Optional[Any]     =			[]
											for i in range(self.batch_size):
													lowerCAmelCase_ ,  lowerCAmelCase_      :					Optional[Any]     =			np.random.choice(np.arange(self.min_resolution    ,    self.max_resolution)    ,    2)
													image_inputs.append(np.random.randint(2_5_5    ,    size=(self.num_channels, width, height)    ,    dtype=np.uinta))

									if not numpify and not torchify:
											# PIL expects the channel dimension as last dimension
											lowerCAmelCase_      :					Dict     =			[Image.fromarray(np.moveaxis(A_    ,    0    ,    -1)) for x in image_inputs]

									if torchify:
											lowerCAmelCase_      :					Union[str, Any]     =			[torch.from_numpy(A_) for x in image_inputs]

									return image_inputs





@require_torch
@require_vision
class     __snake_case						(					UpperCamelCase_						,unittest.TestCase						):
							_a      				=			SwinaSRImageProcessor if is_vision_available() else None




							def        UpperCAmelCase__      (	self					:	int):
									lowerCAmelCase_      :					Optional[int]     =			SwinaSRImageProcessingTester(self)




							@property
							def        UpperCAmelCase__      (	self					:	List[Any]):
									return self.image_processor_tester.prepare_image_processor_dict()




							def        UpperCAmelCase__      (	self					:	Optional[Any]):
									lowerCAmelCase_      :					Tuple     =			self.image_processing_class(**self.image_processor_dict)
									self.assertTrue(hasattr(A_    ,    '''do_rescale'''))
									self.assertTrue(hasattr(A_    ,    '''rescale_factor'''))
									self.assertTrue(hasattr(A_    ,    '''do_pad'''))
									self.assertTrue(hasattr(A_    ,    '''pad_size'''))




							def        UpperCAmelCase__      (	self					:	List[str]):
									pass




							def        UpperCAmelCase__      (	self					:	Optional[int]    ,    A_					:	List[Any]):
									lowerCAmelCase_ ,  lowerCAmelCase_      :					Any     =			get_image_size(A_)
									lowerCAmelCase_      :					Any     =			self.image_processor_tester.pad_size

									lowerCAmelCase_      :					str     =			(old_height // size + 1) * size - old_height
									lowerCAmelCase_      :					Union[str, Any]     =			(old_width // size + 1) * size - old_width
									return old_height + pad_height, old_width + pad_width




							def        UpperCAmelCase__      (	self					:	Dict):
									# Initialize image_processor
									lowerCAmelCase_      :					Any     =			self.image_processing_class(**self.image_processor_dict)
									# create random PIL images
									lowerCAmelCase_      :					Any     =			self.image_processor_tester.prepare_inputs(equal_resolution=A_)
									for image in image_inputs:
											self.assertIsInstance(A_    ,    Image.Image)

									# Test not batched input
									lowerCAmelCase_      :					List[str]     =			image_processor(image_inputs[0]    ,    return_tensors='''pt''').pixel_values
									lowerCAmelCase_ ,  lowerCAmelCase_      :					Optional[Any]     =			self.calculate_expected_size(np.array(image_inputs[0]))
									self.assertEqual(
									    encoded_images.shape    ,    (
									        1,
									        self.image_processor_tester.num_channels,
									        expected_height,
									        expected_width,
									    )    ,    )




							def        UpperCAmelCase__      (	self					:	Dict):
									# Initialize image_processor
									lowerCAmelCase_      :					Dict     =			self.image_processing_class(**self.image_processor_dict)
									# create random numpy tensors
									lowerCAmelCase_      :					Union[str, Any]     =			self.image_processor_tester.prepare_inputs(equal_resolution=A_    ,    numpify=A_)
									for image in image_inputs:
											self.assertIsInstance(A_    ,    np.ndarray)

									# Test not batched input
									lowerCAmelCase_      :					Union[str, Any]     =			image_processor(image_inputs[0]    ,    return_tensors='''pt''').pixel_values
									lowerCAmelCase_ ,  lowerCAmelCase_      :					Any     =			self.calculate_expected_size(image_inputs[0])
									self.assertEqual(
									    encoded_images.shape    ,    (
									        1,
									        self.image_processor_tester.num_channels,
									        expected_height,
									        expected_width,
									    )    ,    )







							def        UpperCAmelCase__      (	self					:	List[Any]):
									# Initialize image_processor
									lowerCAmelCase_      :					int     =			self.image_processing_class(**self.image_processor_dict)
									# create random PyTorch tensors
									lowerCAmelCase_      :					Union[str, Any]     =			self.image_processor_tester.prepare_inputs(equal_resolution=A_    ,    torchify=A_)
									for image in image_inputs:
											self.assertIsInstance(A_    ,    torch.Tensor)

									# Test not batched input
									lowerCAmelCase_      :					List[Any]     =			image_processor(image_inputs[0]    ,    return_tensors='''pt''').pixel_values
									lowerCAmelCase_ ,  lowerCAmelCase_      :					Optional[Any]     =			self.calculate_expected_size(image_inputs[0])
									self.assertEqual(
									    encoded_images.shape    ,    (
									        1,
									        self.image_processor_tester.num_channels,
									        expected_height,
									        expected_width,
									    )    ,    )










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])





