


'''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 					_A						(						unittest.TestCase							):



   '''simple docstring'''






   _lowercase    						= (JukeboxModel,) if is_torch_available() else ()
   _lowercase    						= """openai/jukebox-1b-lyrics"""
   _lowercase    						= {
       """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
   _lowercase    						= [
       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
   ]

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

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

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

   _lowercase    						= [105_8304, 0, 78_6432, 7169, 507, 76, 27, 40, 30, 76]

   _lowercase    						= [
       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
   ]
   _lowercase    						= [
       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
   ]
   _lowercase    						= [
       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
   ]
   _lowercase    						= [
       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
   ]
   _lowercase    						= [
       -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
   ]
   _lowercase    						= [
       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
   ]
   _lowercase    						= [
       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
   ]
   _lowercase    						= [
       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   __lowerCAmelCase						(	self	:					List[Any]						)->	Any:
          snake_case__					:	str          =					JukeboxTokenizer.from_pretrained(self.model_id						)
          snake_case__					:	List[str]          =					tokenizer(**self.metas						)['''input_ids''']
          return tokens




   @slow
   def   __lowerCAmelCase						(	self	:					Dict						)->	Tuple:
          snake_case__					:	List[str]          =					JukeboxModel.from_pretrained(self.model_id					, min_duration=0						).eval()
          snake_case__					:	str          =					self.prepare_inputs()

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

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

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




   @slow
   def   __lowerCAmelCase						(	self	:					Dict						)->	Optional[Any]:
          snake_case__					:	Tuple          =					False
          snake_case__					:	Tuple          =					JukeboxModel.from_pretrained(self.model_id					, min_duration=0						).eval()

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

          snake_case__					:	List[str]          =					model.priors[0]
          snake_case__					:	List[Any]          =					0
          snake_case__					:	Union[str, Any]          =					top_prior.get_music_tokens_conds(__lowercase					, start=__lowercase					, end=start + top_prior.n_ctx						)
          snake_case__					:	Optional[int]          =					top_prior.get_metadata(labels[0].clone()					, __lowercase					, 1_058_304					, 0						)

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

          snake_case__					:	str          =					top_prior.get_cond(__lowercase					, __lowercase						)
          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   __lowerCAmelCase						(	self	:					Dict						)->	Union[str, Any]:
          snake_case__					:	List[Any]          =					False

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

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

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

          snake_case__					:	Union[str, Any]          =					torch.cat((zs[1], torch.zeros(1					, 2_048 - zs[1].shape[-1]						))					, dim=-1						).long()
          snake_case__					:	Any          =					[upper_a, upper_a, model.vqvae.encode(__lowercase					, start_level=0					, bs_chunks=waveform.shape[0]						)[0]]
          snake_case__					:	Optional[int]          =					model._sample(
              __lowercase					, __lowercase					, sample_levels=[2]					, save_results=__lowercase					, 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   __lowerCAmelCase						(	self	:					List[str]						)->	Any:
          snake_case__					:	List[str]          =					JukeboxModel.from_pretrained(self.model_id					, min_duration=0						).eval()
          set_seed(0						)
          snake_case__					:	Optional[int]          =					torch.rand((1, 5_120, 1)						)
          with torch.no_grad():
                 snake_case__					:	Optional[Any]          =					model.vqvae.encode(__lowercase					, 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():
                 snake_case__					:	Optional[Any]          =					model.vqvae.decode(__lowercase					, 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 					_A						(						unittest.TestCase							):



   '''simple docstring'''






   _lowercase    						= (JukeboxModel,) if is_torch_available() else ()
   _lowercase    						= """openai/jukebox-5b-lyrics"""
   _lowercase    						= {
       """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
   _lowercase    						= [
       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
   ]
   _lowercase    						= [
       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
   ]

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

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

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

   _lowercase    						= [
       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
   ]
   _lowercase    						= [
       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   __lowerCAmelCase						(	self	:					List[Any]					, lowerCamelCase	:					int						)->	str:
          snake_case__					:	List[str]          =					JukeboxTokenizer.from_pretrained(__lowercase						)
          snake_case__					:	int          =					tokenizer(**self.metas						)['''input_ids''']
          return tokens




   @slow
   def   __lowerCAmelCase						(	self	:					int						)->	Tuple:
          snake_case__					:	Optional[Any]          =					JukeboxModel.from_pretrained(self.model_id					, min_duration=0						).eval()
          snake_case__					:	Tuple          =					self.prepare_inputs(self.model_id						)

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

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

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

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




   @slow
   def   __lowerCAmelCase						(	self	:					Dict						)->	Tuple:
          snake_case__					:	Optional[int]          =					'''ArthurZ/jukebox_prior_0'''
          snake_case__					:	Any          =					JukeboxPrior.from_pretrained(__lowercase					, min_duration=0						).eval().half().to("""cuda"""						)

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








import tempfile
import unittest

from diffusers import (
    DDIMScheduler,
    DDPMScheduler,
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    PNDMScheduler,
    logging,
)
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.utils.testing_utils import CaptureLogger







class 					snake_case__(				UpperCAmelCase__				):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__   :	Union[str, Any]       =				"""config.json"""




						@register_to_config
						def __init__(					self ,    __lowercase=2 ,    __lowercase=5 ,    __lowercase=(2, 5) ,    __lowercase="for diffusion" ,    __lowercase=[1, 3] ,    )       ->     int:
										pass







class 					snake_case__(				UpperCAmelCase__				):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__   :	int       =				"""config.json"""




						@register_to_config
						def __init__(					self ,    __lowercase=2 ,    __lowercase=5 ,    __lowercase=(2, 5) ,    __lowercase="for diffusion" ,    __lowercase=[1, 3] ,    )       ->     Optional[Any]:
										pass







class 					snake_case__(				UpperCAmelCase__				):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__   :	int       =				"""config.json"""




						@register_to_config
						def __init__(					self ,    __lowercase=2 ,    __lowercase=5 ,    __lowercase=(2, 5) ,    __lowercase="for diffusion" ,    __lowercase=[1, 3] ,    __lowercase=[1, 3] ,    )       ->     List[str]:
										pass







class 					snake_case__(				UpperCAmelCase__				):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__   :	Optional[Any]       =				"""config.json"""




						@register_to_config
						def __init__(					self ,    __lowercase=2 ,    __lowercase=5 ,    __lowercase=(2, 5) ,    __lowercase="for diffusion" ,    __lowercase=[1, 5] ,    __lowercase=[5, 4] ,    )       ->     List[str]:
										pass







class 					snake_case__(				unittest.TestCase				):







						'''simple docstring'''




						def     lowercase_			(					self				)       ->     str:
										with self.assertRaises(__lowercase				):
														ConfigMixin.load_config('''dummy_path'''				)




						def     lowercase_			(					self				)       ->     Optional[int]:
										lowerCAmelCase_			:   List[str]     					=					SampleObject()
										lowerCAmelCase_			:   Dict     					=					obj.config
										assert config["a"] == 2
										assert config["b"] == 5
										assert config["c"] == (2, 5)
										assert config["d"] == "for diffusion"
										assert config["e"] == [1, 3]

										# init ignore private arguments
										lowerCAmelCase_			:   List[str]     					=					SampleObject(_name_or_path='''lalala'''				)
										lowerCAmelCase_			:   Dict     					=					obj.config
										assert config["a"] == 2
										assert config["b"] == 5
										assert config["c"] == (2, 5)
										assert config["d"] == "for diffusion"
										assert config["e"] == [1, 3]

										# can override default
										lowerCAmelCase_			:   List[Any]     					=					SampleObject(c=6				)
										lowerCAmelCase_			:   Optional[int]     					=					obj.config
										assert config["a"] == 2
										assert config["b"] == 5
										assert config["c"] == 6
										assert config["d"] == "for diffusion"
										assert config["e"] == [1, 3]

										# can use positional arguments.
										lowerCAmelCase_			:   Optional[Any]     					=					SampleObject(1 ,    c=6				)
										lowerCAmelCase_			:   Dict     					=					obj.config
										assert config["a"] == 1
										assert config["b"] == 5
										assert config["c"] == 6
										assert config["d"] == "for diffusion"
										assert config["e"] == [1, 3]




						def     lowercase_			(					self				)       ->     int:
										lowerCAmelCase_			:   List[Any]     					=					SampleObject()
										lowerCAmelCase_			:   str     					=					obj.config

										assert config["a"] == 2
										assert config["b"] == 5
										assert config["c"] == (2, 5)
										assert config["d"] == "for diffusion"
										assert config["e"] == [1, 3]

										with tempfile.TemporaryDirectory() as tmpdirname:
														obj.save_config(__lowercase				)
														lowerCAmelCase_			:   int     					=					SampleObject.from_config(SampleObject.load_config(__lowercase				)				)
														lowerCAmelCase_			:   Optional[Any]     					=					new_obj.config

										# unfreeze configs
										lowerCAmelCase_			:   str     					=					dict(__lowercase				)
										lowerCAmelCase_			:   Any     					=					dict(__lowercase				)

										assert config.pop('''c'''				) == (2, 5)  # instantiated as tuple
										assert new_config.pop('''c'''				) == [2, 5]  # saved & loaded as list because of json
										config.pop('''_use_default_values'''				)
										assert config == new_config




						def     lowercase_			(					self				)       ->     List[str]:
										lowerCAmelCase_			:   Optional[Any]     					=					logging.get_logger('''diffusers.configuration_utils'''				)
										# 30 for warning
										logger.setLevel(3_0				)

										with CaptureLogger(__lowercase				) as cap_logger:
														lowerCAmelCase_			:   int     					=					DDIMScheduler.from_pretrained(
														    '''hf-internal-testing/tiny-stable-diffusion-torch''' ,    subfolder='''scheduler'''				)

										assert ddim.__class__ == DDIMScheduler
										# no warning should be thrown
										assert cap_logger.out == ""




						def     lowercase_			(					self				)       ->     Optional[Any]:
										lowerCAmelCase_			:   List[Any]     					=					logging.get_logger('''diffusers.configuration_utils'''				)
										# 30 for warning
										logger.setLevel(3_0				)

										with CaptureLogger(__lowercase				) as cap_logger:
														lowerCAmelCase_			:   Any     					=					EulerDiscreteScheduler.from_pretrained(
														    '''hf-internal-testing/tiny-stable-diffusion-torch''' ,    subfolder='''scheduler'''				)

										assert euler.__class__ == EulerDiscreteScheduler
										# no warning should be thrown
										assert cap_logger.out == ""




						def     lowercase_			(					self				)       ->     Any:
										lowerCAmelCase_			:   Tuple     					=					logging.get_logger('''diffusers.configuration_utils'''				)
										# 30 for warning
										logger.setLevel(3_0				)

										with CaptureLogger(__lowercase				) as cap_logger:
														lowerCAmelCase_			:   Optional[int]     					=					EulerAncestralDiscreteScheduler.from_pretrained(
														    '''hf-internal-testing/tiny-stable-diffusion-torch''' ,    subfolder='''scheduler'''				)

										assert euler.__class__ == EulerAncestralDiscreteScheduler
										# no warning should be thrown
										assert cap_logger.out == ""




						def     lowercase_			(					self				)       ->     List[Any]:
										lowerCAmelCase_			:   Any     					=					logging.get_logger('''diffusers.configuration_utils'''				)
										# 30 for warning
										logger.setLevel(3_0				)

										with CaptureLogger(__lowercase				) as cap_logger:
														lowerCAmelCase_			:   Dict     					=					PNDMScheduler.from_pretrained(
														    '''hf-internal-testing/tiny-stable-diffusion-torch''' ,    subfolder='''scheduler'''				)

										assert pndm.__class__ == PNDMScheduler
										# no warning should be thrown
										assert cap_logger.out == ""




						def     lowercase_			(					self				)       ->     Tuple:
										lowerCAmelCase_			:   Optional[int]     					=					logging.get_logger('''diffusers.configuration_utils'''				)
										# 30 for warning
										logger.setLevel(3_0				)

										with CaptureLogger(__lowercase				) as cap_logger:
														lowerCAmelCase_			:   int     					=					DDPMScheduler.from_pretrained(
														    '''hf-internal-testing/tiny-stable-diffusion-torch''' ,    subfolder='''scheduler''' ,    prediction_type='''sample''' ,    beta_end=8 ,    )

										with CaptureLogger(__lowercase				) as cap_logger_a:
														lowerCAmelCase_			:   Tuple     					=					DDPMScheduler.from_pretrained('''google/ddpm-celebahq-256''' ,    beta_start=8_8				)

										assert ddpm.__class__ == DDPMScheduler
										assert ddpm.config.prediction_type == "sample"
										assert ddpm.config.beta_end == 8
										assert ddpm_a.config.beta_start == 8_8

										# no warning should be thrown
										assert cap_logger.out == ""
										assert cap_logger_a.out == ""




						def     lowercase_			(					self				)       ->     int:
										lowerCAmelCase_			:   Tuple     					=					logging.get_logger('''diffusers.configuration_utils'''				)
										# 30 for warning
										logger.setLevel(3_0				)

										with CaptureLogger(__lowercase				) as cap_logger:
														lowerCAmelCase_			:   List[str]     					=					DPMSolverMultistepScheduler.from_pretrained(
														    '''hf-internal-testing/tiny-stable-diffusion-torch''' ,    subfolder='''scheduler'''				)

										assert dpm.__class__ == DPMSolverMultistepScheduler
										# no warning should be thrown
										assert cap_logger.out == ""




						def     lowercase_			(					self				)       ->     Optional[Any]:
										# let's first save a config that should be in the form
										#    a=2,
										#    b=5,
										#    c=(2, 5),
										#    d="for diffusion",
										#    e=[1, 3],

										lowerCAmelCase_			:   Dict     					=					SampleObject()

										lowerCAmelCase_			:   Optional[int]     					=					{k: v for k, v in config.config.items() if not k.startswith('''_'''				)}

										# make sure that default config has all keys in `_use_default_values`
										assert set(config_dict.keys()				) == set(config.config._use_default_values				)

										with tempfile.TemporaryDirectory() as tmpdirname:
														config.save_config(__lowercase				)

														# now loading it with SampleObject2 should put f into `_use_default_values`
														lowerCAmelCase_			:   List[Any]     					=					SampleObjecta.from_config(__lowercase				)

														assert "f" in config._use_default_values
														assert config.f == [1, 3]

										# now loading the config, should **NOT** use [1, 3] for `f`, but the default [1, 4] value
										# **BECAUSE** it is part of `config._use_default_values`
										lowerCAmelCase_			:   List[Any]     					=					SampleObjecta.from_config(config.config				)
										assert new_config.f == [5, 4]

										config.config._use_default_values.pop()
										lowerCAmelCase_			:   int     					=					SampleObjecta.from_config(config.config				)
										assert new_config_a.f == [1, 3]

										# Nevertheless "e" should still be correctly loaded to [1, 3] from SampleObject2 instead of defaulting to [1, 5]
										assert new_config_a.e == [1, 3]