





import unittest
from unittest import skip

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


if is_torch_available():
     import torch

     from transformers import JukeboxModel, JukeboxPrior, JukeboxTokenizer




@require_torch
class 			_lowerCamelCase	(      unittest.TestCase      ):



      """simple docstring"""






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

      UpperCAmelCase_  : str				      =[
          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
      ]

      UpperCAmelCase_  : str				      =[
          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
      ]
      UpperCAmelCase_  : int				      =[
          416, 416, 1_125, 1_125, 416, 416, 416, 416, 416, 416, 416, 416,
          416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
          416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416, 416,
          416, 416, 416, 416
      ]

      UpperCAmelCase_  : Tuple				      =[
          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
      ]
      UpperCAmelCase_  : Dict				      =[
          854, 842, 1_353, 114, 1_353, 842, 185, 842, 185, 114, 591, 842,
          185, 417, 185, 842, 307, 842, 591, 842, 185, 842, 307, 842,
          591, 842, 1_353, 842, 185, 842, 591, 842, 591, 114, 591, 842,
          185, 842, 591, 89
      ]

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

      UpperCAmelCase_  : List[str]				      =[
          390, 1_160, 1_002, 1_907, 1_788, 1_788, 1_788, 1_907, 1_002, 1_002, 1_854, 1_002,
          1_002, 1_002, 1_002, 1_002, 1_002, 1_160, 1_160, 1_606, 596, 596, 1_160, 1_002,
          1_516, 596, 1_002, 1_002, 1_002, 1_907, 1_788, 1_788, 1_788, 1_854, 1_788, 1_907,
          1_907, 1_788, 596, 1_626
      ]
      UpperCAmelCase_  : Union[str, Any]				      =[
          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
      ]
      UpperCAmelCase_  : List[Any]				      =[
          1_684, 1_873, 1_119, 1_189, 395, 611, 1_901, 972, 890, 1_337, 1_392, 1_927,
          96, 972, 672, 780, 1_119, 890, 158, 771, 1_073, 1_927, 353, 1_331,
          1_269, 1_459, 1_333, 1_645, 812, 1_577, 1_337, 606, 353, 981, 1_466, 619,
          197, 391, 302, 1_930
      ]
      UpperCAmelCase_  : Optional[Any]				      =[
          390, 1_160, 1_002, 1_907, 1_788, 1_788, 1_788, 1_907, 1_002, 1_002, 1_854, 1_002,
          1_002, 1_002, 1_002, 1_002, 1_002, 1_160, 1_160, 1_606, 596, 596, 1_160, 1_002,
          1_516, 596, 1_002, 1_002, 1_002, 1_907, 1_788, 1_788, 1_788, 1_854, 1_788, 1_907,
          1_907, 1_788, 596, 1_626
      ]
      UpperCAmelCase_  : Union[str, Any]				      =[
          -0.0492, -0.0524, -0.0565, -0.0640, -0.0686, -0.0684, -0.0677, -0.0664,
          -0.0605, -0.0490, -0.0330, -0.0168, -0.0083, -0.0075, -0.0051, 0.0025,
          0.0136, 0.0261, 0.0386, 0.0497, 0.0580, 0.0599, 0.0583, 0.0614,
          0.0740, 0.0889, 0.1023, 0.1162, 0.1211, 0.1212, 0.1251, 0.1336,
          0.1502, 0.1686, 0.1883, 0.2148, 0.2363, 0.2458, 0.2507, 0.2531
      ]
      UpperCAmelCase_  : List[Any]				      =[
          0.0256, -0.0544, 0.1600, -0.0032, 0.1066, 0.0825, -0.0013, 0.3440,
          0.0210, 0.0412, -0.1777, -0.0892, -0.0164, 0.0285, -0.0613, -0.0617,
          -0.0137, -0.0201, -0.0175, 0.0215, -0.0627, 0.0520, -0.0730, 0.0970,
          -0.0100, 0.0442, -0.0586, 0.0207, -0.0015, -0.0082
      ]
      UpperCAmelCase_  : Tuple				      =[
          0.0415, 0.0877, 0.0022, -0.0055, 0.0751, 0.0334, 0.0324, -0.0068,
          0.0011, 0.0017, -0.0676, 0.0655, -0.0143, 0.0399, 0.0303, 0.0743,
          -0.0168, -0.0394, -0.1113, 0.0124, 0.0442, 0.0267, -0.0003, -0.1536,
          -0.0116, -0.1837, -0.0180, -0.1026, -0.0777, -0.0456
      ]
      UpperCAmelCase_  : Optional[Any]				      =[
          76, 27, 40, 30, 76, 46, 44, 47, 40, 37, 38, 31, 45, 45, 76, 38, 31, 33,
          45, 76, 41, 32, 76, 45, 46, 41, 40, 31, 78, 76
      ]







      def   UpperCAmelCase						(	self				)    ->				str:






        '''simple docstring'''






        __snake_case				:			Optional[int]            =    JukeboxTokenizer.from_pretrained(self.model_id				)
        __snake_case				:			Optional[int]            =    tokenizer(**self.metas				)["input_ids"]
        return tokens







      @slow
      def   UpperCAmelCase						(	self				)    ->				Union[str, Any]:






        '''simple docstring'''






        __snake_case				:			Any            =    JukeboxModel.from_pretrained(self.model_id			,  min_duration=0				).eval()
        __snake_case				:			List[Any]            =    self.prepare_inputs()

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

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

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







      @slow
      def   UpperCAmelCase						(	self				)    ->				int:






        '''simple docstring'''






        __snake_case				:			List[str]            =    False
        __snake_case				:			str            =    JukeboxModel.from_pretrained(self.model_id			,  min_duration=0				).eval()

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

        __snake_case				:			Tuple            =    model.priors[0]
        __snake_case				:			Dict            =    0
        __snake_case				:			Any            =    top_prior.get_music_tokens_conds(UpperCAmelCase			,  start=UpperCAmelCase			,  end=start + top_prior.n_ctx				)
        __snake_case				:			int            =    top_prior.get_metadata(labels[0].clone()			,  UpperCAmelCase			,  1058304			,  0				)

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

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







      @slow
      def   UpperCAmelCase						(	self				)    ->				Dict:






        '''simple docstring'''






        __snake_case				:			List[str]            =    False

        __snake_case				:			Dict            =    JukeboxModel.from_pretrained(self.model_id			,  min_duration=0				).eval()
        set_seed(0				)
        __snake_case				:			List[Any]            =    torch.rand((1, 5120, 1)				)
        __snake_case				:			Any            =    list(self.prepare_inputs()				)

        __snake_case				:			int            =    [model.vqvae.encode(UpperCAmelCase			,  start_level=2			,  bs_chunks=waveform.shape[0]				)[0], None, None]
        __snake_case				:			Optional[int]            =    model._sample(
            UpperCAmelCase			,  UpperCAmelCase			,  sample_levels=[0]			,  save_results=UpperCAmelCase			,  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			,  2048 - zs[0].shape[-1]				))			,  dim=-1				).long()
        __snake_case				:			Tuple            =    [upper_a, model.vqvae.encode(UpperCAmelCase			,  start_level=1			,  bs_chunks=waveform.shape[0]				)[0], None]
        __snake_case				:			Dict            =    model._sample(
            UpperCAmelCase			,  UpperCAmelCase			,  sample_levels=[1]			,  save_results=UpperCAmelCase			,  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				:			Dict            =    torch.cat((zs[1], torch.zeros(1			,  2048 - zs[1].shape[-1]				))			,  dim=-1				).long()
        __snake_case				:			int            =    [upper_a, upper_a, model.vqvae.encode(UpperCAmelCase			,  start_level=0			,  bs_chunks=waveform.shape[0]				)[0]]
        __snake_case				:			List[Any]            =    model._sample(
            UpperCAmelCase			,  UpperCAmelCase			,  sample_levels=[2]			,  save_results=UpperCAmelCase			,  sample_length=40 * model.priors[2].raw_to_tokens				)
        torch.testing.assert_allclose(zs[2][0][:40].cpu()			,  torch.tensor(self.EXPECTED_PRIMED_2				)				)







      @slow
      def   UpperCAmelCase						(	self				)    ->				str:






        '''simple docstring'''






        __snake_case				:			List[str]            =    JukeboxModel.from_pretrained(self.model_id			,  min_duration=0				).eval()
        set_seed(0				)
        __snake_case				:			Tuple            =    torch.rand((1, 5120, 1)				)
        with torch.no_grad():
          __snake_case				:			Dict            =    model.vqvae.encode(UpperCAmelCase			,  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(UpperCAmelCase			,  start_level=2			,  bs_chunks=x.shape[0]				)
        torch.testing.assert_allclose(x[0, :40, 0]			,  torch.tensor(self.EXPECTED_VQVAE_DECODE				)			,  atol=1E-4			,  rtol=1E-4				)




@require_torch
class 			_lowerCamelCase	(      unittest.TestCase      ):



      """simple docstring"""






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

      # fmt: off
      UpperCAmelCase_  : int				      =[
          1_489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          1_489, 1_489, 1_489, 1_489, 1_150, 1_853, 1_509, 1_150, 1_357, 1_509, 6, 1_272
      ]
      UpperCAmelCase_  : Any				      =[
          1_489, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653,
          653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653, 653
      ]

      UpperCAmelCase_  : int				      =[
          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
      ]
      UpperCAmelCase_  : Optional[Any]				      =[
          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
      ]

      UpperCAmelCase_  : Dict				      =[
          1_755, 1_061, 234, 1_755, 1_061, 1_755, 185, 290, 307, 307, 616, 616,
          616, 616, 616, 616, 307, 290, 417, 1_755, 234, 1_755, 185, 290,
          290, 290, 307, 616, 616, 616, 616, 616, 290, 234, 234, 1_755,
          234, 234, 1_755, 234, 185, 185, 307, 616, 616, 616, 616, 290,
          1_755, 1_755, 1_755, 234, 234, 1_755, 1_572, 290, 307, 616, 34, 616
      ]
      UpperCAmelCase_  : Optional[int]				      =[
          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
      ]

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

      UpperCAmelCase_  : str				      =[
          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
      ]
      UpperCAmelCase_  : List[str]				      =[
          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			,  UpperCAmelCase				)    ->				Union[str, Any]:






        '''simple docstring'''






        __snake_case				:			str            =    JukeboxTokenizer.from_pretrained(UpperCAmelCase				)
        __snake_case				:			List[str]            =    tokenizer(**self.metas				)["input_ids"]
        return tokens







      @slow
      def   UpperCAmelCase						(	self				)    ->				List[Any]:






        '''simple docstring'''






        __snake_case				:			Optional[Any]            =    JukeboxModel.from_pretrained(self.model_id			,  min_duration=0				).eval()
        __snake_case				:			Optional[Any]            =    self.prepare_inputs(self.model_id				)

        set_seed(0				)
        __snake_case				:			List[str]            =    [torch.zeros(1			,  0			,  dtype=torch.long				).cpu() for _ in range(3				)]
        __snake_case				:			Any            =    model._sample(UpperCAmelCase			,  UpperCAmelCase			,  [0]			,  sample_length=60 * model.priors[0].raw_to_tokens			,  save_results=UpperCAmelCase				)
        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(UpperCAmelCase			,  UpperCAmelCase			,  [1]			,  sample_length=60 * model.priors[1].raw_to_tokens			,  save_results=UpperCAmelCase				)
        self.assertIn(zs[1][0].detach().cpu().tolist()			,  [self.EXPECTED_OUTPUT_1, self.EXPECTED_OUTPUT_1_PT_2]				)

        set_seed(0				)
        __snake_case				:			int            =    model._sample(UpperCAmelCase			,  UpperCAmelCase			,  [2]			,  sample_length=60 * model.priors[2].raw_to_tokens			,  save_results=UpperCAmelCase				)
        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				)    ->				Union[str, Any]:






        '''simple docstring'''






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







      @slow
      def   UpperCAmelCase						(	self				)    ->				List[Any]:






        '''simple docstring'''






        __snake_case				:			str            =    "ArthurZ/jukebox_prior_0"
        __snake_case				:			Optional[Any]            =    JukeboxPrior.from_pretrained(UpperCAmelCase			,  min_duration=0				).eval().half().to("cuda"				)

        __snake_case				:			List[str]            =    self.prepare_inputs(UpperCAmelCase				)[0].cuda()
        __snake_case				:			Union[str, Any]            =    model.get_metadata(UpperCAmelCase			,  0			,  7680			,  0				)
        set_seed(0				)
        __snake_case				:			List[Any]            =    model.sample(1			,  metadata=UpperCAmelCase			,  sample_tokens=60				)
        self.assertIn(outputs[0].cpu().tolist()			,  [self.EXPECTED_GPU_OUTPUTS_2, self.EXPECTED_GPU_OUTPUTS_2_PT_2]				)








import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import PriorTransformer, UnCLIPPipeline, UnCLIPScheduler, UNetaDConditionModel, UNetaDModel
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()




class 			_lowerCamelCase	(      a	,				unittest.TestCase      ):



      """simple docstring"""






      UpperCAmelCase_  : Union[str, Any]				      =UnCLIPPipeline
      UpperCAmelCase_  : Optional[int]				      =TEXT_TO_IMAGE_PARAMS - {
          "negative_prompt",
          "height",
          "width",
          "negative_prompt_embeds",
          "guidance_scale",
          "prompt_embeds",
          "cross_attention_kwargs",
      }
      UpperCAmelCase_  : Optional[Any]				      =TEXT_TO_IMAGE_BATCH_PARAMS
      UpperCAmelCase_  : Any				      =[
          "generator",
          "return_dict",
          "prior_num_inference_steps",
          "decoder_num_inference_steps",
          "super_res_num_inference_steps",
      ]
      UpperCAmelCase_  : str				      =False







      @property
      def   UpperCAmelCase						(	self				)    ->				Optional[int]:






        '''simple docstring'''






        return 32







      @property
      def   UpperCAmelCase						(	self				)    ->				str:






        '''simple docstring'''






        return 32







      @property
      def   UpperCAmelCase						(	self				)    ->				Tuple:






        '''simple docstring'''






        return self.time_input_dim







      @property
      def   UpperCAmelCase						(	self				)    ->				Optional[int]:






        '''simple docstring'''






        return self.time_input_dim * 4







      @property
      def   UpperCAmelCase						(	self				)    ->				Union[str, Any]:






        '''simple docstring'''






        return 100







      @property
      def   UpperCAmelCase						(	self				)    ->				Union[str, Any]:






        '''simple docstring'''






        __snake_case				:			Optional[int]            =    CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"				)
        return tokenizer







      @property
      def   UpperCAmelCase						(	self				)    ->				Optional[int]:






        '''simple docstring'''






        torch.manual_seed(0				)
        __snake_case				:			Optional[int]            =    CLIPTextConfig(
            bos_token_id=0			,  eos_token_id=2			,  hidden_size=self.text_embedder_hidden_size			,  projection_dim=self.text_embedder_hidden_size			,  intermediate_size=37			,  layer_norm_eps=1E-05			,  num_attention_heads=4			,  num_hidden_layers=5			,  pad_token_id=1			,  vocab_size=1000			,  )
        return CLIPTextModelWithProjection(UpperCAmelCase				)







      @property
      def   UpperCAmelCase						(	self				)    ->				Dict:






        '''simple docstring'''






        torch.manual_seed(0				)

        __snake_case				:			List[Any]            =    {
            "num_attention_heads": 2,
            "attention_head_dim": 12,
            "embedding_dim": self.text_embedder_hidden_size,
            "num_layers": 1,
        }

        __snake_case				:			str            =    PriorTransformer(**UpperCAmelCase				)
        return model







      @property
      def   UpperCAmelCase						(	self				)    ->				Any:






        '''simple docstring'''






        torch.manual_seed(0				)

        __snake_case				:			List[str]            =    {
            "clip_embeddings_dim": self.text_embedder_hidden_size,
            "time_embed_dim": self.time_embed_dim,
            "cross_attention_dim": self.cross_attention_dim,
        }

        __snake_case				:			Union[str, Any]            =    UnCLIPTextProjModel(**UpperCAmelCase				)
        return model







      @property
      def   UpperCAmelCase						(	self				)    ->				Dict:






        '''simple docstring'''






        torch.manual_seed(0				)

        __snake_case				:			Tuple            =    {
            "sample_size": 32,
            # RGB in channels
            "in_channels": 3,
            # Out channels is double in channels because predicts mean and variance
            "out_channels": 6,
            "down_block_types": ("ResnetDownsampleBlock2D", "SimpleCrossAttnDownBlock2D"),
            "up_block_types": ("SimpleCrossAttnUpBlock2D", "ResnetUpsampleBlock2D"),
            "mid_block_type": "UNetMidBlock2DSimpleCrossAttn",
            "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
            "layers_per_block": 1,
            "cross_attention_dim": self.cross_attention_dim,
            "attention_head_dim": 4,
            "resnet_time_scale_shift": "scale_shift",
            "class_embed_type": "identity",
        }

        __snake_case				:			Optional[int]            =    UNetaDConditionModel(**UpperCAmelCase				)
        return model







      @property
      def   UpperCAmelCase						(	self				)    ->				Dict:






        '''simple docstring'''






        return {
            "sample_size": 64,
            "layers_per_block": 1,
            "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
            "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
            "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
            "in_channels": 6,
            "out_channels": 3,
        }







      @property
      def   UpperCAmelCase						(	self				)    ->				Optional[int]:






        '''simple docstring'''






        torch.manual_seed(0				)

        __snake_case				:			Optional[int]            =    UNetaDModel(**self.dummy_super_res_kwargs				)
        return model







      @property
      def   UpperCAmelCase						(	self				)    ->				Tuple:






        '''simple docstring'''






        torch.manual_seed(1				)

        __snake_case				:			List[str]            =    UNetaDModel(**self.dummy_super_res_kwargs				)
        return model







      def   UpperCAmelCase						(	self				)    ->				Union[str, Any]:






        '''simple docstring'''






        __snake_case				:			List[str]            =    self.dummy_prior
        __snake_case				:			str            =    self.dummy_decoder
        __snake_case				:			Optional[Any]            =    self.dummy_text_proj
        __snake_case				:			List[Any]            =    self.dummy_text_encoder
        __snake_case				:			List[Any]            =    self.dummy_tokenizer
        __snake_case				:			Optional[int]            =    self.dummy_super_res_first
        __snake_case				:			Tuple            =    self.dummy_super_res_last

        __snake_case				:			int            =    UnCLIPScheduler(
            variance_type="fixed_small_log"			,  prediction_type="sample"			,  num_train_timesteps=1000			,  clip_sample_range=5.0			,  )

        __snake_case				:			Union[str, Any]            =    UnCLIPScheduler(
            variance_type="learned_range"			,  prediction_type="epsilon"			,  num_train_timesteps=1000			,  )

        __snake_case				:			Any            =    UnCLIPScheduler(
            variance_type="fixed_small_log"			,  prediction_type="epsilon"			,  num_train_timesteps=1000			,  )

        __snake_case				:			Dict            =    {
            "prior": prior,
            "decoder": decoder,
            "text_proj": text_proj,
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
            "super_res_first": super_res_first,
            "super_res_last": super_res_last,
            "prior_scheduler": prior_scheduler,
            "decoder_scheduler": decoder_scheduler,
            "super_res_scheduler": super_res_scheduler,
        }

        return components







      def   UpperCAmelCase						(	self			,  UpperCAmelCase			,  UpperCAmelCase=0				)    ->				int:






        '''simple docstring'''






        if str(UpperCAmelCase				).startswith("mps"				):
          __snake_case				:			int            =    torch.manual_seed(UpperCAmelCase				)
        else:
          __snake_case				:			Optional[int]            =    torch.Generator(device=UpperCAmelCase				).manual_seed(UpperCAmelCase				)
        __snake_case				:			List[str]            =    {
            "prompt": "horse",
            "generator": generator,
            "prior_num_inference_steps": 2,
            "decoder_num_inference_steps": 2,
            "super_res_num_inference_steps": 2,
            "output_type": "numpy",
        }
        return inputs







      def   UpperCAmelCase						(	self				)    ->				Tuple:






        '''simple docstring'''






        __snake_case				:			Dict            =    "cpu"

        __snake_case				:			int            =    self.get_dummy_components()

        __snake_case				:			Optional[int]            =    self.pipeline_class(**UpperCAmelCase				)
        __snake_case				:			Optional[int]            =    pipe.to(UpperCAmelCase				)

        pipe.set_progress_bar_config(disable=UpperCAmelCase				)

        __snake_case				:			Any            =    pipe(**self.get_dummy_inputs(UpperCAmelCase				)				)
        __snake_case				:			List[Any]            =    output.images

        __snake_case				:			int            =    pipe(
            **self.get_dummy_inputs(UpperCAmelCase				)			,  return_dict=UpperCAmelCase			,  )[0]

        __snake_case				:			Tuple            =    image[0, -3:, -3:, -1]
        __snake_case				:			Dict            =    image_from_tuple[0, -3:, -3:, -1]

        assert image.shape == (1, 64, 64, 3)

        __snake_case				:			int            =    np.array(
            [
                0.9_997,
                0.9_988,
                0.0_028,
                0.9_997,
                0.9_984,
                0.9_965,
                0.0_029,
                0.9_986,
                0.0_025,
            ]				)

        assert np.abs(image_slice.flatten() - expected_slice				).max() < 1E-2
        assert np.abs(image_from_tuple_slice.flatten() - expected_slice				).max() < 1E-2







      def   UpperCAmelCase						(	self				)    ->				Union[str, Any]:






        '''simple docstring'''






        __snake_case				:			List[str]            =    torch.device("cpu"				)




        class 			_lowerCamelCase	:



              """simple docstring"""






              UpperCAmelCase_  : str				      =1




        __snake_case				:			Any            =    self.get_dummy_components()

        __snake_case				:			List[Any]            =    self.pipeline_class(**UpperCAmelCase				)
        __snake_case				:			List[Any]            =    pipe.to(UpperCAmelCase				)

        __snake_case				:			Any            =    components["prior"]
        __snake_case				:			List[str]            =    components["decoder"]
        __snake_case				:			List[str]            =    components["super_res_first"]
        __snake_case				:			Union[str, Any]            =    components["tokenizer"]
        __snake_case				:			Optional[Any]            =    components["text_encoder"]

        __snake_case				:			List[str]            =    torch.Generator(device=UpperCAmelCase				).manual_seed(0				)
        __snake_case				:			List[str]            =    prior.dtype
        __snake_case				:			Dict            =    1

        __snake_case				:			str            =    (batch_size, prior.config.embedding_dim)
        __snake_case				:			Any            =    pipe.prepare_latents(
            UpperCAmelCase			,  dtype=UpperCAmelCase			,  device=UpperCAmelCase			,  generator=UpperCAmelCase			,  latents=UpperCAmelCase			,  scheduler=DummyScheduler()				)
        __snake_case				:			List[Any]            =    (batch_size, decoder.config.in_channels, decoder.config.sample_size, decoder.config.sample_size)
        __snake_case				:			List[Any]            =    pipe.prepare_latents(
            UpperCAmelCase			,  dtype=UpperCAmelCase			,  device=UpperCAmelCase			,  generator=UpperCAmelCase			,  latents=UpperCAmelCase			,  scheduler=DummyScheduler()				)

        __snake_case				:			Optional[int]            =    (
            batch_size,
            super_res_first.config.in_channels // 2,
            super_res_first.config.sample_size,
            super_res_first.config.sample_size,
        )
        __snake_case				:			Union[str, Any]            =    pipe.prepare_latents(
            UpperCAmelCase			,  dtype=UpperCAmelCase			,  device=UpperCAmelCase			,  generator=UpperCAmelCase			,  latents=UpperCAmelCase			,  scheduler=DummyScheduler()				)

        pipe.set_progress_bar_config(disable=UpperCAmelCase				)

        __snake_case				:			List[Any]            =    "this is a prompt example"

        __snake_case				:			Dict            =    torch.Generator(device=UpperCAmelCase				).manual_seed(0				)
        __snake_case				:			Optional[Any]            =    pipe(
            [prompt]			,  generator=UpperCAmelCase			,  prior_num_inference_steps=2			,  decoder_num_inference_steps=2			,  super_res_num_inference_steps=2			,  prior_latents=UpperCAmelCase			,  decoder_latents=UpperCAmelCase			,  super_res_latents=UpperCAmelCase			,  output_type="np"			,  )
        __snake_case				:			Tuple            =    output.images

        __snake_case				:			Dict            =    tokenizer(
            UpperCAmelCase			,  padding="max_length"			,  max_length=tokenizer.model_max_length			,  return_tensors="pt"			,  )
        __snake_case				:			Union[str, Any]            =    text_encoder(text_inputs.input_ids				)
        __snake_case				:			Dict            =    text_inputs.attention_mask

        __snake_case				:			int            =    torch.Generator(device=UpperCAmelCase				).manual_seed(0				)
        __snake_case				:			List[Any]            =    pipe(
            generator=UpperCAmelCase			,  prior_num_inference_steps=2			,  decoder_num_inference_steps=2			,  super_res_num_inference_steps=2			,  prior_latents=UpperCAmelCase			,  decoder_latents=UpperCAmelCase			,  super_res_latents=UpperCAmelCase			,  text_model_output=UpperCAmelCase			,  text_attention_mask=UpperCAmelCase			,  output_type="np"			,  )[0]

        # make sure passing text embeddings manually is identical
        assert np.abs(image - image_from_text				).max() < 1E-4







      @skip_mps
      def   UpperCAmelCase						(	self				)    ->				str:






        '''simple docstring'''






        __snake_case				:			Union[str, Any]            =    torch_device == "cpu"

        self._test_attention_slicing_forward_pass(test_max_difference=UpperCAmelCase			,  expected_max_diff=0.01				)







      @skip_mps
      def   UpperCAmelCase						(	self				)    ->				Dict:






        '''simple docstring'''






        __snake_case				:			int            =    torch_device == "cpu"
        __snake_case				:			str            =    True
        __snake_case				:			List[Any]            =    [
            "prior_num_inference_steps",
            "decoder_num_inference_steps",
            "super_res_num_inference_steps",
        ]

        self._test_inference_batch_single_identical(
            test_max_difference=UpperCAmelCase			,  relax_max_difference=UpperCAmelCase			,  additional_params_copy_to_batched_inputs=UpperCAmelCase			,  )







      def   UpperCAmelCase						(	self				)    ->				List[str]:






        '''simple docstring'''






        __snake_case				:			Dict            =    [
            "prior_num_inference_steps",
            "decoder_num_inference_steps",
            "super_res_num_inference_steps",
        ]

        if torch_device == "mps":
          # TODO: MPS errors with larger batch sizes
          __snake_case				:			List[str]            =    [2, 3]
          self._test_inference_batch_consistent(
              batch_sizes=UpperCAmelCase			,  additional_params_copy_to_batched_inputs=UpperCAmelCase			,  )
        else:
          self._test_inference_batch_consistent(
              additional_params_copy_to_batched_inputs=UpperCAmelCase				)







      @skip_mps
      def   UpperCAmelCase						(	self				)    ->				Tuple:






        '''simple docstring'''






        return super().test_dict_tuple_outputs_equivalent()







      @skip_mps
      def   UpperCAmelCase						(	self				)    ->				Optional[Any]:






        '''simple docstring'''






        return super().test_save_load_local()







      @skip_mps
      def   UpperCAmelCase						(	self				)    ->				Optional[Any]:






        '''simple docstring'''






        return super().test_save_load_optional_components()




@nightly
class 			_lowerCamelCase	(      unittest.TestCase      ):



      """simple docstring"""






      def   UpperCAmelCase						(	self				)    ->				Dict:






        '''simple docstring'''






        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()







      def   UpperCAmelCase						(	self				)    ->				Any:






        '''simple docstring'''






        __snake_case				:			List[str]            =    load_numpy(
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
            "/unclip/karlo_v1_alpha_horse_cpu.npy"				)

        __snake_case				:			Optional[int]            =    UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha"				)
        pipeline.set_progress_bar_config(disable=UpperCAmelCase				)

        __snake_case				:			Optional[int]            =    torch.manual_seed(0				)
        __snake_case				:			int            =    pipeline(
            "horse"			,  num_images_per_prompt=1			,  generator=UpperCAmelCase			,  output_type="np"			,  )

        __snake_case				:			Any            =    output.images[0]

        assert image.shape == (256, 256, 3)
        assert np.abs(expected_image - image				).max() < 1E-1




@slow
@require_torch_gpu
class 			_lowerCamelCase	(      unittest.TestCase      ):



      """simple docstring"""






      def   UpperCAmelCase						(	self				)    ->				Optional[int]:






        '''simple docstring'''






        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()







      def   UpperCAmelCase						(	self				)    ->				int:






        '''simple docstring'''






        __snake_case				:			Optional[int]            =    load_numpy(
            "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
            "/unclip/karlo_v1_alpha_horse_fp16.npy"				)

        __snake_case				:			List[Any]            =    UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha"			,  torch_dtype=torch.floataa				)
        __snake_case				:			int            =    pipeline.to(UpperCAmelCase				)
        pipeline.set_progress_bar_config(disable=UpperCAmelCase				)

        __snake_case				:			List[str]            =    torch.Generator(device="cpu"				).manual_seed(0				)
        __snake_case				:			Dict            =    pipeline(
            "horse"			,  generator=UpperCAmelCase			,  output_type="np"			,  )

        __snake_case				:			Union[str, Any]            =    output.images[0]

        assert image.shape == (256, 256, 3)

        assert_mean_pixel_difference(UpperCAmelCase			,  UpperCAmelCase				)







      def   UpperCAmelCase						(	self				)    ->				Any:






        '''simple docstring'''






        torch.cuda.empty_cache()
        torch.cuda.reset_max_memory_allocated()
        torch.cuda.reset_peak_memory_stats()

        __snake_case				:			List[str]            =    UnCLIPPipeline.from_pretrained("kakaobrain/karlo-v1-alpha"			,  torch_dtype=torch.floataa				)
        __snake_case				:			Optional[int]            =    pipe.to(UpperCAmelCase				)
        pipe.set_progress_bar_config(disable=UpperCAmelCase				)
        pipe.enable_attention_slicing()
        pipe.enable_sequential_cpu_offload()

        __snake_case				:			str            =    pipe(
            "horse"			,  num_images_per_prompt=1			,  prior_num_inference_steps=2			,  decoder_num_inference_steps=2			,  super_res_num_inference_steps=2			,  output_type="np"			,  )

        __snake_case				:			Optional[Any]            =    torch.cuda.max_memory_allocated()
        # make sure that less than 7 GB is allocated
        assert mem_bytes < 7 * 10**9

