code: 
  
  
  
  
  
  
  """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 						__lowerCAmelCase   (			unittest.TestCase       ):
  
  
  
  
  
  
  
    '''simple docstring'''
  
  
    __lowerCamelCase       :    Optional[Any]     	=			(JukeboxModel,) if is_torch_available() else ()
    __lowerCamelCase       :    str     	=			"openai/jukebox-1b-lyrics"
    __lowerCamelCase       :    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
    __lowerCamelCase       :    Tuple     	=			[
        18_64, 15_36, 12_13, 18_70, 13_57, 15_36, 5_19, 8_80, 13_23, 7_89, 10_82, 5_34,
        10_00, 14_45, 11_05, 11_30, 9_67, 5_15, 14_34, 16_20, 5_34, 14_95, 2_83, 14_45,
        3_33, 13_07, 5_39, 16_31, 15_28, 3_75, 14_34, 6_73, 6_27, 7_10, 7_78, 18_83,
        14_05, 12_76, 14_55, 12_28
    ]
  
    __lowerCamelCase       :    int     	=			[
        14_89, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
        6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
        6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
        6_53, 6_53, 6_53, 6_53
    ]
  
    __lowerCamelCase       :    List[str]     	=			[
        11_25, 17_51, 6_97, 17_76, 11_41, 14_76, 3_91, 6_97, 11_25, 6_84, 8_67, 4_16,
        8_44, 13_72, 12_74, 7_17, 12_74, 8_44, 12_99, 14_19, 6_97, 13_70, 3_17, 11_25,
        1_91, 14_40, 13_70, 14_40, 13_70, 2_82, 16_21, 13_70, 3_68, 3_49, 8_67, 18_72,
        12_62, 8_69, 17_28, 7_47
    ]
    __lowerCamelCase       :    Optional[int]     	=			[
        4_16, 4_16, 11_25, 11_25, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
        4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
        4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
        4_16, 4_16, 4_16, 4_16
    ]
  
    __lowerCamelCase       :    str     	=			[
        17_55, 8_42, 3_07, 18_43, 10_22, 13_95, 2_34, 15_54, 8_06, 7_39, 10_22, 4_42,
        6_16, 5_56, 2_68, 14_99, 9_33, 4_57, 14_40, 18_37, 7_55, 9_85, 3_08, 9_02,
        2_93, 14_43, 16_71, 11_41, 15_33, 5_55, 15_62, 10_61, 2_87, 4_17, 10_22, 20_08,
        11_86, 10_15, 17_77, 2_68
    ]
    __lowerCamelCase       :    Optional[int]     	=			[
        8_54, 8_42, 13_53, 1_14, 13_53, 8_42, 1_85, 8_42, 1_85, 1_14, 5_91, 8_42,
        1_85, 4_17, 1_85, 8_42, 3_07, 8_42, 5_91, 8_42, 1_85, 8_42, 3_07, 8_42,
        5_91, 8_42, 13_53, 8_42, 1_85, 8_42, 5_91, 8_42, 5_91, 1_14, 5_91, 8_42,
        1_85, 8_42, 5_91, 89
    ]
  
    __lowerCamelCase       :    str     	=			[1_05_83_04, 0, 78_64_32, 71_69, 5_07, 76, 27, 40, 30, 76]
  
    __lowerCamelCase       :    Any     	=			[
        3_90, 11_60, 10_02, 19_07, 17_88, 17_88, 17_88, 19_07, 10_02, 10_02, 18_54, 10_02,
        10_02, 10_02, 10_02, 10_02, 10_02, 11_60, 11_60, 16_06, 5_96, 5_96, 11_60, 10_02,
        15_16, 5_96, 10_02, 10_02, 10_02, 19_07, 17_88, 17_88, 17_88, 18_54, 17_88, 19_07,
        19_07, 17_88, 5_96, 16_26
    ]
    __lowerCamelCase       :    Union[str, Any]     	=			[
        12_36, 16_68, 14_84, 19_20, 18_48, 14_09, 1_39, 8_64, 18_28, 12_72, 15_99, 8_24,
        16_72, 1_39, 5_55, 14_84, 8_24, 19_20, 5_55, 5_96, 15_79, 15_99, 12_31, 15_99,
        16_37, 14_07, 2_12, 8_24, 15_99, 1_16, 14_33, 8_24, 2_58, 15_99, 14_33, 18_95,
        10_63, 14_33, 14_33, 15_99
    ]
    __lowerCamelCase       :    str     	=			[
        16_84, 18_73, 11_19, 11_89, 3_95, 6_11, 19_01, 9_72, 8_90, 13_37, 13_92, 19_27,
        96, 9_72, 6_72, 7_80, 11_19, 8_90, 1_58, 7_71, 10_73, 19_27, 3_53, 13_31,
        12_69, 14_59, 13_33, 16_45, 8_12, 15_77, 13_37, 6_06, 3_53, 9_81, 14_66, 6_19,
        1_97, 3_91, 3_02, 19_30
    ]
    __lowerCamelCase       :    Tuple     	=			[
        3_90, 11_60, 10_02, 19_07, 17_88, 17_88, 17_88, 19_07, 10_02, 10_02, 18_54, 10_02,
        10_02, 10_02, 10_02, 10_02, 10_02, 11_60, 11_60, 16_06, 5_96, 5_96, 11_60, 10_02,
        15_16, 5_96, 10_02, 10_02, 10_02, 19_07, 17_88, 17_88, 17_88, 18_54, 17_88, 19_07,
        19_07, 17_88, 5_96, 16_26
    ]
    __lowerCamelCase       :    Dict     	=			[
        -0.04_92, -0.05_24, -0.05_65, -0.06_40, -0.06_86, -0.06_84, -0.06_77, -0.06_64,
        -0.06_05, -0.04_90, -0.03_30, -0.01_68, -0.00_83, -0.00_75, -0.00_51, 0.00_25,
        0.01_36, 0.02_61, 0.03_86, 0.04_97, 0.05_80, 0.05_99, 0.05_83, 0.06_14,
        0.07_40, 0.08_89, 0.10_23, 0.11_62, 0.12_11, 0.12_12, 0.12_51, 0.13_36,
        0.15_02, 0.16_86, 0.18_83, 0.21_48, 0.23_63, 0.24_58, 0.25_07, 0.25_31
    ]
    __lowerCamelCase       :    Any     	=			[
        0.02_56, -0.05_44, 0.16_00, -0.00_32, 0.10_66, 0.08_25, -0.00_13, 0.34_40,
        0.02_10, 0.04_12, -0.17_77, -0.08_92, -0.01_64, 0.02_85, -0.06_13, -0.06_17,
        -0.01_37, -0.02_01, -0.01_75, 0.02_15, -0.06_27, 0.05_20, -0.07_30, 0.09_70,
        -0.01_00, 0.04_42, -0.05_86, 0.02_07, -0.00_15, -0.00_82
    ]
    __lowerCamelCase       :    Dict     	=			[
        0.04_15, 0.08_77, 0.00_22, -0.00_55, 0.07_51, 0.03_34, 0.03_24, -0.00_68,
        0.00_11, 0.00_17, -0.06_76, 0.06_55, -0.01_43, 0.03_99, 0.03_03, 0.07_43,
        -0.01_68, -0.03_94, -0.11_13, 0.01_24, 0.04_42, 0.02_67, -0.00_03, -0.15_36,
        -0.01_16, -0.18_37, -0.01_80, -0.10_26, -0.07_77, -0.04_56
    ]
    __lowerCamelCase       :    List[str]     	=			[
        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:		Optional[int]					):
       UpperCamelCase_     =JukeboxTokenizer.from_pretrained(self.model_id					)
       UpperCamelCase_     =tokenizer(**self.metas					)["input_ids"]
       return tokens
    @slow
    def 					UpperCamelCase__   (      self:		Dict					):
       UpperCamelCase_     =JukeboxModel.from_pretrained(self.model_id				,			min_duration=0					).eval()
       UpperCamelCase_     =self.prepare_inputs()
  
       set_seed(0					)
       UpperCamelCase_     =[torch.zeros(1				,			0				,			dtype=torch.long					).cpu() for _ in range(3					)]
       UpperCamelCase_     =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					)
       UpperCamelCase_     =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					)
       UpperCamelCase_     =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:		Dict					):
       UpperCamelCase_     =False
       UpperCamelCase_     =JukeboxModel.from_pretrained(self.model_id				,			min_duration=0					).eval()
  
       UpperCamelCase_     =self.prepare_inputs()
       set_seed(0					)
       UpperCamelCase_     =[torch.zeros(1				,			0				,			dtype=torch.long					) for _ in range(3					)]
  
       UpperCamelCase_     =model.priors[0]
       UpperCamelCase_     =0
       UpperCamelCase_     =top_prior.get_music_tokens_conds(UpperCamelCase_				,			start=UpperCamelCase_				,			end=start + top_prior.n_ctx					)
       UpperCamelCase_     =top_prior.get_metadata(labels[0].clone()				,			UpperCamelCase_				,			105_8304				,			0					)
  
       self.assertIsNone(UpperCamelCase_					)
       self.assertListEqual(metadata.numpy()[0][:10].tolist()				,			self.EXPECTED_Y_COND					)
  
       UpperCamelCase_    ,							UpperCamelCase_    ,							UpperCamelCase_     =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:		int					):
       UpperCamelCase_     =False
  
       UpperCamelCase_     =JukeboxModel.from_pretrained(self.model_id				,			min_duration=0					).eval()
       set_seed(0					)
       UpperCamelCase_     =torch.rand((1, 5120, 1)					)
       UpperCamelCase_     =list(self.prepare_inputs()					)
  
       UpperCamelCase_     =[model.vqvae.encode(UpperCamelCase_				,			start_level=2				,			bs_chunks=waveform.shape[0]					)[0], None, None]
       UpperCamelCase_     =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					)					)
  
       UpperCamelCase_     =torch.cat((zs[0], torch.zeros(1				,			2048 - zs[0].shape[-1]					))				,			dim=-1					).long()
       UpperCamelCase_     =[upper_a, model.vqvae.encode(UpperCamelCase_				,			start_level=1				,			bs_chunks=waveform.shape[0]					)[0], None]
       UpperCamelCase_     =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					)					)
  
       UpperCamelCase_     =torch.cat((zs[1], torch.zeros(1				,			2048 - zs[1].shape[-1]					))				,			dim=-1					).long()
       UpperCamelCase_     =[upper_a, upper_a, model.vqvae.encode(UpperCamelCase_				,			start_level=0				,			bs_chunks=waveform.shape[0]					)[0]]
       UpperCamelCase_     =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:		Optional[int]					):
       UpperCamelCase_     =JukeboxModel.from_pretrained(self.model_id				,			min_duration=0					).eval()
       set_seed(0					)
       UpperCamelCase_     =torch.rand((1, 5120, 1)					)
       with torch.no_grad():
          UpperCamelCase_     =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():
          UpperCamelCase_     =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'''
  
  
    __lowerCamelCase       :    Union[str, Any]     	=			(JukeboxModel,) if is_torch_available() else ()
    __lowerCamelCase       :    Optional[Any]     	=			"openai/jukebox-5b-lyrics"
    __lowerCamelCase       :    int     	=			{
        "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
    __lowerCamelCase       :    Optional[int]     	=			[
        14_89, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
        6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
        6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
        6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
        14_89, 14_89, 14_89, 14_89, 11_50, 18_53, 15_09, 11_50, 13_57, 15_09, 6, 12_72
    ]
    __lowerCamelCase       :    Optional[Any]     	=			[
        14_89, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
        6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
        6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
        6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
        6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53
    ]
  
    __lowerCamelCase       :    Dict     	=			[
        11_25, 4_16, 11_25, 11_25, 11_25, 11_25, 11_25, 4_16, 4_16, 4_16, 4_16, 4_16,
        4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
        4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
        4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
        4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16
    ]
    __lowerCamelCase       :    Optional[int]     	=			[
        4_16, 4_16, 11_25, 11_25, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
        4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
        4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
        4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
        4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16
    ]
  
    __lowerCamelCase       :    List[Any]     	=			[
        17_55, 10_61, 2_34, 17_55, 10_61, 17_55, 1_85, 2_90, 3_07, 3_07, 6_16, 6_16,
        6_16, 6_16, 6_16, 6_16, 3_07, 2_90, 4_17, 17_55, 2_34, 17_55, 1_85, 2_90,
        2_90, 2_90, 3_07, 6_16, 6_16, 6_16, 6_16, 6_16, 2_90, 2_34, 2_34, 17_55,
        2_34, 2_34, 17_55, 2_34, 1_85, 1_85, 3_07, 6_16, 6_16, 6_16, 6_16, 2_90,
        17_55, 17_55, 17_55, 2_34, 2_34, 17_55, 15_72, 2_90, 3_07, 6_16, 34, 6_16
    ]
    __lowerCamelCase       :    Optional[int]     	=			[
        8_54, 8_42, 13_53, 1_14, 13_53, 8_42, 1_85, 8_42, 1_85, 1_14, 5_91, 8_42, 1_85,
        4_17, 1_85, 8_42, 3_07, 8_42, 5_91, 8_42, 1_85, 8_42, 1_85, 8_42, 5_91, 8_42,
        13_53, 8_42, 1_85, 8_42, 5_91, 8_42, 5_91, 1_14, 5_91, 8_42, 1_85, 8_42, 5_91,
        89, 5_91, 8_42, 5_91, 8_42, 5_91, 4_17, 13_72, 8_42, 13_72, 8_42, 34, 8_42,
        1_85, 89, 5_91, 8_42, 1_85, 8_42, 5_91, 6_32
    ]
  
    __lowerCamelCase       :    str     	=			[
        14_89, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
        6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
        6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
        6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
        6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53
    ]
    __lowerCamelCase       :    Dict     	=			[
        14_89, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
        6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53,
        6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 6_53, 18_53, 11_77, 15_36, 12_28,
        7_10, 4_75, 14_89, 12_29, 12_24, 2_31, 12_24, 2_52, 14_34, 6_53, 4_75,
        11_06, 18_77, 15_99, 12_28, 16_00, 16_83, 11_82, 18_53, 4_75, 18_64,
        2_52, 12_29, 14_34, 20_01
    ]
  
    __lowerCamelCase       :    Tuple     	=			[
        11_25, 11_25, 4_16, 11_25, 11_25, 4_16, 11_25, 11_25, 4_16, 4_16, 11_25, 4_16,
        4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
        4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
        4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16,
        4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16, 4_16
    ]
    __lowerCamelCase       :    int     	=			[
        4_91, 17_55, 34, 16_13, 17_55, 4_17, 9_92, 16_13, 2_22, 8_42, 13_53, 16_13,
        8_44, 6_32, 1_85, 16_13, 8_44, 6_32, 1_85, 16_13, 1_85, 8_42, 6_77, 16_13,
        1_85, 1_14, 13_53, 16_13, 3_07, 89, 8_44, 16_13, 3_07, 13_32, 2_34, 19_79,
        3_07, 89, 13_53, 6_16, 34, 8_42, 1_85, 8_42, 34, 8_42, 1_85, 8_42,
        3_07, 1_14, 1_85, 89, 34, 12_68, 1_85, 89, 34, 8_42, 1_85, 89
    ]
    def 					UpperCamelCase__   (      self:		str				,			UpperCamelCase_:		List[str]					):
       UpperCamelCase_     =JukeboxTokenizer.from_pretrained(UpperCamelCase_					)
       UpperCamelCase_     =tokenizer(**self.metas					)["input_ids"]
       return tokens
    @slow
    def 					UpperCamelCase__   (      self:		Optional[int]					):
       UpperCamelCase_     =JukeboxModel.from_pretrained(self.model_id				,			min_duration=0					).eval()
       UpperCamelCase_     =self.prepare_inputs(self.model_id					)
  
       set_seed(0					)
       UpperCamelCase_     =[torch.zeros(1				,			0				,			dtype=torch.long					).cpu() for _ in range(3					)]
       UpperCamelCase_     =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					)
       UpperCamelCase_     =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					)
       UpperCamelCase_     =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:		Tuple					):
       UpperCamelCase_     =JukeboxModel.from_pretrained(self.model_id				,			min_duration=0					).eval()
       UpperCamelCase_     =[i.cuda() for i in self.prepare_inputs(self.model_id					)]
  
       set_seed(0					)
       model.priors[0].cuda()
       UpperCamelCase_     =[torch.zeros(1				,			0				,			dtype=torch.long					).cuda() for _ in range(3					)]
       UpperCamelCase_     =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()
       UpperCamelCase_     =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()
       UpperCamelCase_     =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]					):
       UpperCamelCase_     ="ArthurZ/jukebox_prior_0"
       UpperCamelCase_     =JukeboxPrior.from_pretrained(UpperCamelCase_				,			min_duration=0					).eval().half().to("cuda"					)
  
       UpperCamelCase_     =self.prepare_inputs(UpperCamelCase_					)[0].cuda()
       UpperCamelCase_     =model.get_metadata(UpperCamelCase_				,			0				,			7680				,			0					)
       set_seed(0					)
       UpperCamelCase_     =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]					)
  
style_context: 
  
  
  
  
  
  
  """simple docstring"""
  
  
  
  
  
  
  import unittest
  from dataclasses import dataclass
  from typing import List, Union
  
  import numpy as np
  import PIL.Image
  
  from diffusers.utils.outputs import BaseOutput
  @dataclass
  class 						__lowerCAmelCase   (			UpperCAmelCase       ):
  
  
  
  
  
  
  
    '''simple docstring'''
  
  
    __lowerCamelCase       :    Union[List[PIL.Image.Image], np.ndarray]
  
  
  
  
  class 						__lowerCAmelCase   (			unittest.TestCase       ):
  
  
  
  
  
  
  
    '''simple docstring'''
  
  
    def 					UpperCamelCase__   (      self:		Tuple					):
       UpperCamelCase_     =CustomOutput(images=np.random.rand(1				,			3				,			4				,			4					)					)
  
       # check every way of getting the attribute
       assert isinstance(outputs.images				,			np.ndarray					)
       assert outputs.images.shape == (1, 3, 4, 4)
       assert isinstance(outputs["images"]				,			np.ndarray					)
       assert outputs["images"].shape == (1, 3, 4, 4)
       assert isinstance(outputs[0]				,			np.ndarray					)
       assert outputs[0].shape == (1, 3, 4, 4)
  
       # test with a non-tensor attribute
       UpperCamelCase_     =CustomOutput(images=[PIL.Image.new("RGB"				,			(4, 4)					)]					)
  
       # check every way of getting the attribute
       assert isinstance(outputs.images				,			UpperCamelCase_					)
       assert isinstance(outputs.images[0]				,			PIL.Image.Image					)
       assert isinstance(outputs["images"]				,			UpperCamelCase_					)
       assert isinstance(outputs["images"][0]				,			PIL.Image.Image					)
       assert isinstance(outputs[0]				,			UpperCamelCase_					)
       assert isinstance(outputs[0][0]				,			PIL.Image.Image					)
  
  
  
  
  
    def 					UpperCamelCase__   (      self:		Optional[int]					):
       # test output reinitialization with a `dict` for compatibility with `accelerate`
       UpperCamelCase_     =CustomOutput({"images": np.random.rand(1				,			3				,			4				,			4					)}					)
  
       # check every way of getting the attribute
       assert isinstance(outputs.images				,			np.ndarray					)
       assert outputs.images.shape == (1, 3, 4, 4)
       assert isinstance(outputs["images"]				,			np.ndarray					)
       assert outputs["images"].shape == (1, 3, 4, 4)
       assert isinstance(outputs[0]				,			np.ndarray					)
       assert outputs[0].shape == (1, 3, 4, 4)
  
       # test with a non-tensor attribute
       UpperCamelCase_     =CustomOutput({"images": [PIL.Image.new("RGB"				,			(4, 4)					)]}					)
  
       # check every way of getting the attribute
       assert isinstance(outputs.images				,			UpperCamelCase_					)
       assert isinstance(outputs.images[0]				,			PIL.Image.Image					)
       assert isinstance(outputs["images"]				,			UpperCamelCase_					)
       assert isinstance(outputs["images"][0]				,			PIL.Image.Image					)
       assert isinstance(outputs[0]				,			UpperCamelCase_					)
       assert isinstance(outputs[0][0]				,			PIL.Image.Image					)
  