





"""simple docstring"""





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       ):
      lowerCamelCase__      =			(JukeboxModel,) if is_torch_available() else ()
      lowerCamelCase__      =			'''openai/jukebox-1b-lyrics'''
      lowerCamelCase__      =			{
          '''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
      lowerCamelCase__      =			[
          1864, 1536, 1213, 1870, 1357, 1536, 519, 880, 1323, 789, 1082, 534,
          1000, 1445, 1105, 1130, 967, 515, 1434, 1620, 534, 1495, 283, 1445,
          333, 1307, 539, 1631, 1528, 375, 1434, 673, 627, 710, 778, 1883,
          1405, 1276, 1455, 1228
      ]

      lowerCamelCase__      =			[
          1489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653
      ]

      lowerCamelCase__      =			[
          1125, 1751, 697, 1776, 1141, 1476, 391, 697, 1125, 684, 867, 416,
          844, 1372, 1274, 717, 1274, 844, 1299, 1419, 697, 1370, 317, 1125,
          191, 1440, 1370, 1440, 1370, 282, 1621, 1370, 368, 349, 867, 1872,
          1262, 869, 1728, 747
      ]
      lowerCamelCase__      =			[
          416, 416, 1125, 1125, 416, 416, 416, 416, 416, 416, 416, 416,
          416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
          416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
          416, 416, 416, 416
      ]

      lowerCamelCase__      =			[
          1755, 842, 307, 1843, 1022, 1395, 234, 1554, 806, 739, 1022, 442,
          616, 556, 268, 1499, 933, 457, 1440, 1837, 755, 985, 308, 902,
          293, 1443, 1671, 1141, 1533, 555, 1562, 1061, 287, 417, 1022, 2008,
          1186, 1015, 1777, 268
      ]
      lowerCamelCase__      =			[
          854, 842, 1353, 114, 1353, 842, 185, 842, 185, 114, 591, 842,
          185, 417, 185, 842, 307, 842, 591, 842, 185, 842, 307, 842,
          591, 842, 1353, 842, 185, 842, 591, 842, 591, 114, 591, 842,
          185, 842, 591, 89
      ]

      lowerCamelCase__      =			[1058304, 0, 786432, 7169, 507, 76, 27, 40, 30, 76]

      lowerCamelCase__      =			[
          390, 1160, 1002, 1907, 1788, 1788, 1788, 1907, 1002, 1002, 1854, 1002,
          1002, 1002, 1002, 1002, 1002, 1160, 1160, 1606, 596, 596, 1160, 1002,
          1516, 596, 1002, 1002, 1002, 1907, 1788, 1788, 1788, 1854, 1788, 1907,
          1907, 1788, 596, 1626
      ]
      lowerCamelCase__      =			[
          1236, 1668, 1484, 1920, 1848, 1409, 139, 864, 1828, 1272, 1599, 824,
          1672, 139, 555, 1484, 824, 1920, 555, 596, 1579, 1599, 1231, 1599,
          1637, 1407, 212, 824, 1599, 116, 1433, 824, 258, 1599, 1433, 1895,
          1063, 1433, 1433, 1599
      ]
      lowerCamelCase__      =			[
          1684, 1873, 1119, 1189, 395, 611, 1901, 972, 890, 1337, 1392, 1927,
          96, 972, 672, 780, 1119, 890, 158, 771, 1073, 1927, 353, 1331,
          1269, 1459, 1333, 1645, 812, 1577, 1337, 606, 353, 981, 1466, 619,
          197, 391, 302, 1930
      ]
      lowerCamelCase__      =			[
          390, 1160, 1002, 1907, 1788, 1788, 1788, 1907, 1002, 1002, 1854, 1002,
          1002, 1002, 1002, 1002, 1002, 1160, 1160, 1606, 596, 596, 1160, 1002,
          1516, 596, 1002, 1002, 1002, 1907, 1788, 1788, 1788, 1854, 1788, 1907,
          1907, 1788, 596, 1626
      ]
      lowerCamelCase__      =			[
          -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
      ]
      lowerCamelCase__      =			[
          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
      ]
      lowerCamelCase__      =			[
          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
      ]
      lowerCamelCase__      =			[
          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 						SCREAMING_SNAKE_CASE_      (      self       :int						):
       __SCREAMING_SNAKE_CASE			:      Any 							=	JukeboxTokenizer.from_pretrained(self.model_id						)
       __SCREAMING_SNAKE_CASE			:      List[str] 							=	tokenizer(**self.metas						)['''input_ids''']
       return tokens

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

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

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

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

      @slow
      def 						SCREAMING_SNAKE_CASE_      (      self       :Any						):
       __SCREAMING_SNAKE_CASE			:      Tuple 							=	False
       __SCREAMING_SNAKE_CASE			:      Union[str, Any] 							=	JukeboxModel.from_pretrained(self.model_id						,							min_duration=0						).eval()

       __SCREAMING_SNAKE_CASE			:      Dict 							=	self.prepare_inputs()
       set_seed(0						)
       __SCREAMING_SNAKE_CASE			:      List[Any] 							=	[torch.zeros(1						,							0						,							dtype=torch.long						) for _ in range(3						)]

       __SCREAMING_SNAKE_CASE			:      str 							=	model.priors[0]
       __SCREAMING_SNAKE_CASE			:      Union[str, Any] 							=	0
       __SCREAMING_SNAKE_CASE			:      List[Any] 							=	top_prior.get_music_tokens_conds(_lowerCamelCase						,							start=_lowerCamelCase						,							end=start + top_prior.n_ctx						)
       __SCREAMING_SNAKE_CASE			:      str 							=	top_prior.get_metadata(labels[0].clone()						,							_lowerCamelCase						,							1_0_5_8_3_0_4						,							0						)

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

       __SCREAMING_SNAKE_CASE			:      str 							=	top_prior.get_cond(_lowerCamelCase						,							_lowerCamelCase						)
       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 						SCREAMING_SNAKE_CASE_      (      self       :Optional[int]						):
       __SCREAMING_SNAKE_CASE			:      Dict 							=	False

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

       __SCREAMING_SNAKE_CASE			:      Any 							=	[model.vqvae.encode(_lowerCamelCase						,							start_level=2						,							bs_chunks=waveform.shape[0]						)[0], None, None]
       __SCREAMING_SNAKE_CASE			:      List[str] 							=	model._sample(
           _lowerCamelCase						,							_lowerCamelCase						,							sample_levels=[0]						,							save_results=_lowerCamelCase						,							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						)						)

       __SCREAMING_SNAKE_CASE			:      int 							=	torch.cat((zs[0], torch.zeros(1						,							2_0_4_8 - zs[0].shape[-1]						))						,							dim=-1						).long()
       __SCREAMING_SNAKE_CASE			:      List[str] 							=	[upper_a, model.vqvae.encode(_lowerCamelCase						,							start_level=1						,							bs_chunks=waveform.shape[0]						)[0], None]
       __SCREAMING_SNAKE_CASE			:      str 							=	model._sample(
           _lowerCamelCase						,							_lowerCamelCase						,							sample_levels=[1]						,							save_results=_lowerCamelCase						,							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						)						)

       __SCREAMING_SNAKE_CASE			:      List[Any] 							=	torch.cat((zs[1], torch.zeros(1						,							2_0_4_8 - zs[1].shape[-1]						))						,							dim=-1						).long()
       __SCREAMING_SNAKE_CASE			:      Optional[Any] 							=	[upper_a, upper_a, model.vqvae.encode(_lowerCamelCase						,							start_level=0						,							bs_chunks=waveform.shape[0]						)[0]]
       __SCREAMING_SNAKE_CASE			:      List[Any] 							=	model._sample(
           _lowerCamelCase						,							_lowerCamelCase						,							sample_levels=[2]						,							save_results=_lowerCamelCase						,							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 						SCREAMING_SNAKE_CASE_      (      self       :Tuple						):
       __SCREAMING_SNAKE_CASE			:      int 							=	JukeboxModel.from_pretrained(self.model_id						,							min_duration=0						).eval()
       set_seed(0						)
       __SCREAMING_SNAKE_CASE			:      Dict 							=	torch.rand((1, 5_1_2_0, 1)						)
       with torch.no_grad():
        __SCREAMING_SNAKE_CASE			:      Tuple 							=	model.vqvae.encode(_lowerCamelCase						,							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():
        __SCREAMING_SNAKE_CASE			:      str 							=	model.vqvae.decode(_lowerCamelCase						,							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       ):
      lowerCamelCase__      =			(JukeboxModel,) if is_torch_available() else ()
      lowerCamelCase__      =			'''openai/jukebox-5b-lyrics'''
      lowerCamelCase__      =			{
          '''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
      lowerCamelCase__      =			[
          1489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          1489, 1489, 1489, 1489, 1150, 1853, 1509, 1150, 1357, 1509, 6, 1272
      ]
      lowerCamelCase__      =			[
          1489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653
      ]

      lowerCamelCase__      =			[
          1125, 416, 1125, 1125, 1125, 1125, 1125, 416, 416, 416, 416, 416,
          416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
          416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
          416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
          416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416
      ]
      lowerCamelCase__      =			[
          416, 416, 1125, 1125, 416, 416, 416, 416, 416, 416, 416, 416, 416,
          416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
          416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
          416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
          416, 416, 416, 416, 416, 416, 416, 416
      ]

      lowerCamelCase__      =			[
          1755, 1061, 234, 1755, 1061, 1755, 185, 290, 307, 307, 616, 616,
          616, 616, 616, 616, 307, 290, 417, 1755, 234, 1755, 185, 290,
          290, 290, 307, 616, 616, 616, 616, 616, 290, 234, 234, 1755,
          234, 234, 1755, 234, 185, 185, 307, 616, 616, 616, 616, 290,
          1755, 1755, 1755, 234, 234, 1755, 1572, 290, 307, 616, 34, 616
      ]
      lowerCamelCase__      =			[
          854, 842, 1353, 114, 1353, 842, 185, 842, 185, 114, 591, 842, 185,
          417, 185, 842, 307, 842, 591, 842, 185, 842, 185, 842, 591, 842,
          1353, 842, 185, 842, 591, 842, 591, 114, 591, 842, 185, 842, 591,
          89, 591, 842, 591, 842, 591, 417, 1372, 842, 1372, 842, 34, 842,
          185, 89, 591, 842, 185, 842, 591, 632
      ]

      lowerCamelCase__      =			[
          1489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653
      ]
      lowerCamelCase__      =			[
          1489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 1853, 1177, 1536, 1228,
          710, 475, 1489, 1229, 1224, 231, 1224, 252, 1434, 653, 475,
          1106, 1877, 1599, 1228, 1600, 1683, 1182, 1853, 475, 1864,
          252, 1229, 1434, 2001
      ]

      lowerCamelCase__      =			[
          1125, 1125, 416, 1125, 1125, 416, 1125, 1125, 416, 416, 1125, 416,
          416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
          416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
          416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
          416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416
      ]
      lowerCamelCase__      =			[
          491, 1755, 34, 1613, 1755, 417, 992, 1613, 222, 842, 1353, 1613,
          844, 632, 185, 1613, 844, 632, 185, 1613, 185, 842, 677, 1613,
          185, 114, 1353, 1613, 307, 89, 844, 1613, 307, 1332, 234, 1979,
          307, 89, 1353, 616, 34, 842, 185, 842, 34, 842, 185, 842,
          307, 114, 185, 89, 34, 1268, 185, 89, 34, 842, 185, 89
      ]

      def 						SCREAMING_SNAKE_CASE_      (      self       :Optional[Any]						,							_lowerCamelCase       :Union[str, Any]						):
       __SCREAMING_SNAKE_CASE			:      Tuple 							=	JukeboxTokenizer.from_pretrained(_lowerCamelCase						)
       __SCREAMING_SNAKE_CASE			:      Dict 							=	tokenizer(**self.metas						)['''input_ids''']
       return tokens

      @slow
      def 						SCREAMING_SNAKE_CASE_      (      self       :Optional[Any]						):
       __SCREAMING_SNAKE_CASE			:      Optional[int] 							=	JukeboxModel.from_pretrained(self.model_id						,							min_duration=0						).eval()
       __SCREAMING_SNAKE_CASE			:      Union[str, Any] 							=	self.prepare_inputs(self.model_id						)

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

       set_seed(0						)
       __SCREAMING_SNAKE_CASE			:      List[Any] 							=	model._sample(_lowerCamelCase						,							_lowerCamelCase						,							[1]						,							sample_length=6_0 * model.priors[1].raw_to_tokens						,							save_results=_lowerCamelCase						)
       self.assertIn(zs[1][0].detach().cpu().tolist()						,							[self.EXPECTED_OUTPUT_1, self.EXPECTED_OUTPUT_1_PT_2]						)

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

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



      @slow
      def 						SCREAMING_SNAKE_CASE_      (      self       :List[Any]						):
       __SCREAMING_SNAKE_CASE			:      Dict 							=	'''ArthurZ/jukebox_prior_0'''
       __SCREAMING_SNAKE_CASE			:      int 							=	JukeboxPrior.from_pretrained(_lowerCamelCase						,							min_duration=0						).eval().half().to('''cuda'''						)

       __SCREAMING_SNAKE_CASE			:      Union[str, Any] 							=	self.prepare_inputs(_lowerCamelCase						)[0].cuda()
       __SCREAMING_SNAKE_CASE			:      Union[str, Any] 							=	model.get_metadata(_lowerCamelCase						,							0						,							7_6_8_0						,							0						)
       set_seed(0						)
       __SCREAMING_SNAKE_CASE			:      str 							=	model.sample(1						,							metadata=_lowerCamelCase						,							sample_tokens=6_0						)
       self.assertIn(outputs[0].cpu().tolist()						,							[self.EXPECTED_GPU_OUTPUTS_2, self.EXPECTED_GPU_OUTPUTS_2_PT_2]						)








"""simple docstring"""





from ..utils import DummyObject, requires_backends





class       snake_case		(       metaclass=__UpperCAmelCase       ):
      lowerCamelCase__      =			['''torch''', '''transformers''', '''invisible_watermark''']

      def __init__(      self       :Optional[int]						,							*_lowerCamelCase       :int						,							**_lowerCamelCase       :Any						):
       requires_backends(self						,							['''torch''', '''transformers''', '''invisible_watermark''']						)

      @classmethod
      def 						SCREAMING_SNAKE_CASE_      (      cls       :Dict						,							*_lowerCamelCase       :Optional[Any]						,							**_lowerCamelCase       :int						):
       requires_backends(cls						,							['''torch''', '''transformers''', '''invisible_watermark''']						)



      @classmethod
      def 						SCREAMING_SNAKE_CASE_      (      cls       :Optional[int]						,							*_lowerCamelCase       :Optional[Any]						,							**_lowerCamelCase       :Union[str, Any]						):
       requires_backends(cls						,							['''torch''', '''transformers''', '''invisible_watermark''']						)







class       snake_case		(       metaclass=__UpperCAmelCase       ):
      lowerCamelCase__      =			['''torch''', '''transformers''', '''invisible_watermark''']

      def __init__(      self       :Dict						,							*_lowerCamelCase       :Dict						,							**_lowerCamelCase       :Optional[Any]						):
       requires_backends(self						,							['''torch''', '''transformers''', '''invisible_watermark''']						)

      @classmethod
      def 						SCREAMING_SNAKE_CASE_      (      cls       :Optional[int]						,							*_lowerCamelCase       :int						,							**_lowerCamelCase       :Optional[int]						):
       requires_backends(cls						,							['''torch''', '''transformers''', '''invisible_watermark''']						)



      @classmethod
      def 						SCREAMING_SNAKE_CASE_      (      cls       :List[Any]						,							*_lowerCamelCase       :Dict						,							**_lowerCamelCase       :List[str]						):
       requires_backends(cls						,							['''torch''', '''transformers''', '''invisible_watermark''']						)

