




import gc
import random
import tempfile
import unittest

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

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.semantic_stable_diffusion import SemanticStableDiffusionPipeline as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()





class        lowerCamelCase__(   unittest.TestCase):
      def      lowerCAmelCase__		(							self:					Any  ):
             # clean up the VRAM after each test
             super().tearDown()
             gc.collect()
             torch.cuda.empty_cache()
      @property
      def      lowerCAmelCase__		(							self:					str  ):
             __lowerCamelCase  =		1
             __lowerCamelCase  =		3
             __lowerCamelCase  =		(32, 32)

             __lowerCamelCase  =		floats_tensor((batch_size, num_channels) + sizes       ,			rng=random.Random(0  )  ).to(UpperCamelCase_  )
             return image
      @property
      def      lowerCAmelCase__		(							self:					Optional[int]  ):
             torch.manual_seed(0  )
             __lowerCamelCase  =		UNetaDConditionModel(
                 block_out_channels=(32, 64)       ,			layers_per_block=2       ,			sample_size=32       ,			in_channels=4       ,			out_channels=4       ,			down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")       ,			up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")       ,			cross_attention_dim=32       ,			)
             return model
      @property
      def      lowerCAmelCase__		(							self:					Union[str, Any]  ):
             torch.manual_seed(0  )
             __lowerCamelCase  =		AutoencoderKL(
                 block_out_channels=[32, 64]       ,			in_channels=3       ,			out_channels=3       ,			down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]       ,			up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]       ,			latent_channels=4       ,			)
             return model
      @property
      def      lowerCAmelCase__		(							self:					List[str]  ):
             torch.manual_seed(0  )
             __lowerCamelCase  =		CLIPTextConfig(
                 bos_token_id=0       ,			eos_token_id=2       ,			hidden_size=32       ,			intermediate_size=37       ,			layer_norm_eps=1E-05       ,			num_attention_heads=4       ,			num_hidden_layers=5       ,			pad_token_id=1       ,			vocab_size=10_00       ,			)
             return CLIPTextModel(UpperCamelCase_  )
      @property
      def      lowerCAmelCase__		(							self:					Union[str, Any]  ):
             def extract(*UpperCamelCase_:					Dict       ,			**UpperCamelCase_:					int  ):





                    class        lowerCamelCase__:
                          def __init__(							self:					List[Any]  ):
                                 __lowerCamelCase  =		torch.ones([0]  )




                          def      lowerCAmelCase__		(							self:					Union[str, Any]       ,			UpperCamelCase_:					List[str]  ):
                                 self.pixel_values.to(UpperCamelCase_  )
                                 return self



                    return Out()

             return extract
      def      lowerCAmelCase__		(							self:					Union[str, Any]  ):
             __lowerCamelCase  =		"""cpu"""  # ensure determinism for the device-dependent torch.Generator
             __lowerCamelCase  =		self.dummy_cond_unet
             __lowerCamelCase  =		DDIMScheduler(
                 beta_start=0.0_0085       ,			beta_end=0.012       ,			beta_schedule="""scaled_linear"""       ,			clip_sample=UpperCamelCase_       ,			set_alpha_to_one=UpperCamelCase_       ,			)

             __lowerCamelCase  =		self.dummy_vae
             __lowerCamelCase  =		self.dummy_text_encoder
             __lowerCamelCase  =		CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""  )

             # make sure here that pndm scheduler skips prk
             __lowerCamelCase  =		StableDiffusionPipeline(
                 unet=UpperCamelCase_       ,			scheduler=UpperCamelCase_       ,			vae=UpperCamelCase_       ,			text_encoder=UpperCamelCase_       ,			tokenizer=UpperCamelCase_       ,			safety_checker=UpperCamelCase_       ,			feature_extractor=self.dummy_extractor       ,			)
             __lowerCamelCase  =		sd_pipe.to(UpperCamelCase_  )
             sd_pipe.set_progress_bar_config(disable=UpperCamelCase_  )

             __lowerCamelCase  =		"""A painting of a squirrel eating a burger"""

             __lowerCamelCase  =		torch.Generator(device=UpperCamelCase_  ).manual_seed(0  )
             __lowerCamelCase  =		sd_pipe([prompt]       ,			generator=UpperCamelCase_       ,			guidance_scale=6.0       ,			num_inference_steps=2       ,			output_type="""np"""  )
             __lowerCamelCase  =		output.images

             __lowerCamelCase  =		torch.Generator(device=UpperCamelCase_  ).manual_seed(0  )
             __lowerCamelCase  =		sd_pipe(
                 [prompt]       ,			generator=UpperCamelCase_       ,			guidance_scale=6.0       ,			num_inference_steps=2       ,			output_type="""np"""       ,			return_dict=UpperCamelCase_       ,			)[0]

             __lowerCamelCase  =		image[0, -3:, -3:, -1]
             __lowerCamelCase  =		image_from_tuple[0, -3:, -3:, -1]

             assert image.shape == (1, 64, 64, 3)
             __lowerCamelCase  =		np.array([0.5753, 0.6114, 0.5001, 0.5034, 0.5470, 0.4729, 0.4971, 0.4867, 0.4867]  )

             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      lowerCAmelCase__		(							self:					Any  ):
             __lowerCamelCase  =		"""cpu"""  # ensure determinism for the device-dependent torch.Generator
             __lowerCamelCase  =		self.dummy_cond_unet
             __lowerCamelCase  =		PNDMScheduler(skip_prk_steps=UpperCamelCase_  )
             __lowerCamelCase  =		self.dummy_vae
             __lowerCamelCase  =		self.dummy_text_encoder
             __lowerCamelCase  =		CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""  )

             # make sure here that pndm scheduler skips prk
             __lowerCamelCase  =		StableDiffusionPipeline(
                 unet=UpperCamelCase_       ,			scheduler=UpperCamelCase_       ,			vae=UpperCamelCase_       ,			text_encoder=UpperCamelCase_       ,			tokenizer=UpperCamelCase_       ,			safety_checker=UpperCamelCase_       ,			feature_extractor=self.dummy_extractor       ,			)
             __lowerCamelCase  =		sd_pipe.to(UpperCamelCase_  )
             sd_pipe.set_progress_bar_config(disable=UpperCamelCase_  )

             __lowerCamelCase  =		"""A painting of a squirrel eating a burger"""
             __lowerCamelCase  =		torch.Generator(device=UpperCamelCase_  ).manual_seed(0  )
             __lowerCamelCase  =		sd_pipe([prompt]       ,			generator=UpperCamelCase_       ,			guidance_scale=6.0       ,			num_inference_steps=2       ,			output_type="""np"""  )

             __lowerCamelCase  =		output.images

             __lowerCamelCase  =		torch.Generator(device=UpperCamelCase_  ).manual_seed(0  )
             __lowerCamelCase  =		sd_pipe(
                 [prompt]       ,			generator=UpperCamelCase_       ,			guidance_scale=6.0       ,			num_inference_steps=2       ,			output_type="""np"""       ,			return_dict=UpperCamelCase_       ,			)[0]

             __lowerCamelCase  =		image[0, -3:, -3:, -1]
             __lowerCamelCase  =		image_from_tuple[0, -3:, -3:, -1]

             assert image.shape == (1, 64, 64, 3)
             __lowerCamelCase  =		np.array([0.5122, 0.5712, 0.4825, 0.5053, 0.5646, 0.4769, 0.5179, 0.4894, 0.4994]  )

             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      lowerCAmelCase__		(							self:					Tuple  ):
             __lowerCamelCase  =		StableDiffusionPipeline.from_pretrained(
                 """hf-internal-testing/tiny-stable-diffusion-lms-pipe"""       ,			safety_checker=UpperCamelCase_  )
             assert isinstance(UpperCamelCase_       ,			UpperCamelCase_  )
             assert isinstance(pipe.scheduler       ,			UpperCamelCase_  )
             assert pipe.safety_checker is None

             __lowerCamelCase  =		pipe("""example prompt"""       ,			num_inference_steps=2  ).images[0]
             assert image is not None

             # check that there's no error when saving a pipeline with one of the models being None
             with tempfile.TemporaryDirectory() as tmpdirname:
                    pipe.save_pretrained(UpperCamelCase_  )
                    __lowerCamelCase  =		StableDiffusionPipeline.from_pretrained(UpperCamelCase_  )

             # sanity check that the pipeline still works
             assert pipe.safety_checker is None
             __lowerCamelCase  =		pipe("""example prompt"""       ,			num_inference_steps=2  ).images[0]
             assert image is not None




      @unittest.skipIf(torch_device != """cuda"""       ,			"""This test requires a GPU"""  )
      def      lowerCAmelCase__		(							self:					Dict  ):
             __lowerCamelCase  =		self.dummy_cond_unet
             __lowerCamelCase  =		PNDMScheduler(skip_prk_steps=UpperCamelCase_  )
             __lowerCamelCase  =		self.dummy_vae
             __lowerCamelCase  =		self.dummy_text_encoder
             __lowerCamelCase  =		CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""  )

             # put models in fp16
             __lowerCamelCase  =		unet.half()
             __lowerCamelCase  =		vae.half()
             __lowerCamelCase  =		bert.half()

             # make sure here that pndm scheduler skips prk
             __lowerCamelCase  =		StableDiffusionPipeline(
                 unet=UpperCamelCase_       ,			scheduler=UpperCamelCase_       ,			vae=UpperCamelCase_       ,			text_encoder=UpperCamelCase_       ,			tokenizer=UpperCamelCase_       ,			safety_checker=UpperCamelCase_       ,			feature_extractor=self.dummy_extractor       ,			)
             __lowerCamelCase  =		sd_pipe.to(UpperCamelCase_  )
             sd_pipe.set_progress_bar_config(disable=UpperCamelCase_  )

             __lowerCamelCase  =		"""A painting of a squirrel eating a burger"""
             __lowerCamelCase  =		sd_pipe([prompt]       ,			num_inference_steps=2       ,			output_type="""np"""  ).images

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





@nightly
@require_torch_gpu
class        lowerCamelCase__(   unittest.TestCase):
      def      lowerCAmelCase__		(							self:					Optional[int]  ):
             # clean up the VRAM after each test
             super().tearDown()
             gc.collect()
             torch.cuda.empty_cache()
      def      lowerCAmelCase__		(							self:					Optional[Any]  ):
             __lowerCamelCase  =		"""cuda"""
             __lowerCamelCase  =		StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""  )
             __lowerCamelCase  =		pipe.to(UpperCamelCase_  )
             pipe.set_progress_bar_config(disable=UpperCamelCase_  )

             __lowerCamelCase  =		"""a photo of a cat"""
             __lowerCamelCase  =		{
                 """editing_prompt""": ["""sunglasses"""],
                 """reverse_editing_direction""": [False],
                 """edit_warmup_steps""": 10,
                 """edit_guidance_scale""": 6,
                 """edit_threshold""": 0.95,
                 """edit_momentum_scale""": 0.5,
                 """edit_mom_beta""": 0.6,
             }

             __lowerCamelCase  =		3
             __lowerCamelCase  =		7

             # no sega enabled
             __lowerCamelCase  =		torch.Generator(UpperCamelCase_  )
             generator.manual_seed(UpperCamelCase_  )
             __lowerCamelCase  =		pipe(
                 [prompt]       ,			generator=UpperCamelCase_       ,			guidance_scale=UpperCamelCase_       ,			num_inference_steps=50       ,			output_type="""np"""       ,			width=5_12       ,			height=5_12       ,			)

             __lowerCamelCase  =		output.images
             __lowerCamelCase  =		image[0, -3:, -3:, -1]
             __lowerCamelCase  =		[
                 0.3467_3113,
                 0.3849_2733,
                 0.3759_7352,
                 0.3408_6335,
                 0.3565_0748,
                 0.3557_9205,
                 0.338_4763,
                 0.3434_0236,
                 0.357_3271,
             ]

             assert image.shape == (1, 5_12, 5_12, 3)

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

             # with sega enabled
             # generator = torch.manual_seed(seed)
             generator.manual_seed(UpperCamelCase_  )
             __lowerCamelCase  =		pipe(
                 [prompt]       ,			generator=UpperCamelCase_       ,			guidance_scale=UpperCamelCase_       ,			num_inference_steps=50       ,			output_type="""np"""       ,			width=5_12       ,			height=5_12       ,			**UpperCamelCase_       ,			)

             __lowerCamelCase  =		output.images
             __lowerCamelCase  =		image[0, -3:, -3:, -1]
             __lowerCamelCase  =		[
                 0.4188_7826,
                 0.3772_8766,
                 0.3013_8272,
                 0.4141_6335,
                 0.4166_4985,
                 0.3628_3392,
                 0.3619_1246,
                 0.4336_4465,
                 0.4300_1732,
             ]

             assert image.shape == (1, 5_12, 5_12, 3)

             assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1E-2
      def      lowerCAmelCase__		(							self:					int  ):
             __lowerCamelCase  =		"""cuda"""
             __lowerCamelCase  =		StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""  )
             __lowerCamelCase  =		pipe.to(UpperCamelCase_  )
             pipe.set_progress_bar_config(disable=UpperCamelCase_  )

             __lowerCamelCase  =		"""an image of a crowded boulevard, realistic, 4k"""
             __lowerCamelCase  =		{
                 """editing_prompt""": """crowd, crowded, people""",
                 """reverse_editing_direction""": True,
                 """edit_warmup_steps""": 10,
                 """edit_guidance_scale""": 8.3,
                 """edit_threshold""": 0.9,
                 """edit_momentum_scale""": 0.5,
                 """edit_mom_beta""": 0.6,
             }

             __lowerCamelCase  =		9
             __lowerCamelCase  =		7

             # no sega enabled
             __lowerCamelCase  =		torch.Generator(UpperCamelCase_  )
             generator.manual_seed(UpperCamelCase_  )
             __lowerCamelCase  =		pipe(
                 [prompt]       ,			generator=UpperCamelCase_       ,			guidance_scale=UpperCamelCase_       ,			num_inference_steps=50       ,			output_type="""np"""       ,			width=5_12       ,			height=5_12       ,			)

             __lowerCamelCase  =		output.images
             __lowerCamelCase  =		image[0, -3:, -3:, -1]
             __lowerCamelCase  =		[
                 0.4349_7998,
                 0.9181_4065,
                 0.754_0739,
                 0.5558_0205,
                 0.846_7265,
                 0.538_9691,
                 0.6257_4506,
                 0.5889_7763,
                 0.5092_6757,
             ]

             assert image.shape == (1, 5_12, 5_12, 3)

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

             # with sega enabled
             # generator = torch.manual_seed(seed)
             generator.manual_seed(UpperCamelCase_  )
             __lowerCamelCase  =		pipe(
                 [prompt]       ,			generator=UpperCamelCase_       ,			guidance_scale=UpperCamelCase_       ,			num_inference_steps=50       ,			output_type="""np"""       ,			width=5_12       ,			height=5_12       ,			**UpperCamelCase_       ,			)

             __lowerCamelCase  =		output.images
             __lowerCamelCase  =		image[0, -3:, -3:, -1]
             __lowerCamelCase  =		[
                 0.308_9719,
                 0.3050_0144,
                 0.2901_6042,
                 0.3063_0964,
                 0.32_5687,
                 0.2941_9225,
                 0.290_8091,
                 0.2872_3598,
                 0.2769_6294,
             ]

             assert image.shape == (1, 5_12, 5_12, 3)

             assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1E-2
      def      lowerCAmelCase__		(							self:					int  ):
             __lowerCamelCase  =		"""cuda"""
             __lowerCamelCase  =		StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""  )
             __lowerCamelCase  =		pipe.to(UpperCamelCase_  )
             pipe.set_progress_bar_config(disable=UpperCamelCase_  )

             __lowerCamelCase  =		"""a castle next to a river"""
             __lowerCamelCase  =		{
                 """editing_prompt""": ["""boat on a river, boat""", """monet, impression, sunrise"""],
                 """reverse_editing_direction""": False,
                 """edit_warmup_steps""": [15, 18],
                 """edit_guidance_scale""": 6,
                 """edit_threshold""": [0.9, 0.8],
                 """edit_momentum_scale""": 0.5,
                 """edit_mom_beta""": 0.6,
             }

             __lowerCamelCase  =		48
             __lowerCamelCase  =		7

             # no sega enabled
             __lowerCamelCase  =		torch.Generator(UpperCamelCase_  )
             generator.manual_seed(UpperCamelCase_  )
             __lowerCamelCase  =		pipe(
                 [prompt]       ,			generator=UpperCamelCase_       ,			guidance_scale=UpperCamelCase_       ,			num_inference_steps=50       ,			output_type="""np"""       ,			width=5_12       ,			height=5_12       ,			)

             __lowerCamelCase  =		output.images
             __lowerCamelCase  =		image[0, -3:, -3:, -1]
             __lowerCamelCase  =		[
                 0.7516_3555,
                 0.7603_7145,
                 0.6_1785,
                 0.918_9673,
                 0.862_7701,
                 0.8518_9694,
                 0.851_2813,
                 0.8701_2076,
                 0.831_2857,
             ]

             assert image.shape == (1, 5_12, 5_12, 3)

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

             # with sega enabled
             # generator = torch.manual_seed(seed)
             generator.manual_seed(UpperCamelCase_  )
             __lowerCamelCase  =		pipe(
                 [prompt]       ,			generator=UpperCamelCase_       ,			guidance_scale=UpperCamelCase_       ,			num_inference_steps=50       ,			output_type="""np"""       ,			width=5_12       ,			height=5_12       ,			**UpperCamelCase_       ,			)

             __lowerCamelCase  =		output.images
             __lowerCamelCase  =		image[0, -3:, -3:, -1]
             __lowerCamelCase  =		[
                 0.7355_3365,
                 0.753_7271,
                 0.7434_1905,
                 0.6648_0356,
                 0.647_2925,
                 0.6303_9416,
                 0.6481_2905,
                 0.674_9717,
                 0.651_7102,
             ]

             assert image.shape == (1, 5_12, 5_12, 3)

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




      def      lowerCAmelCase__		(							self:					Tuple  ):
             __lowerCamelCase  =		"""cuda"""
             __lowerCamelCase  =		StableDiffusionPipeline.from_pretrained("""runwayml/stable-diffusion-v1-5"""       ,			torch_dtype=torch.floataa  )
             __lowerCamelCase  =		pipe.to(UpperCamelCase_  )
             pipe.set_progress_bar_config(disable=UpperCamelCase_  )

             __lowerCamelCase  =		"""a photo of a cat"""
             __lowerCamelCase  =		{
                 """editing_prompt""": ["""sunglasses"""],
                 """reverse_editing_direction""": [False],
                 """edit_warmup_steps""": 10,
                 """edit_guidance_scale""": 6,
                 """edit_threshold""": 0.95,
                 """edit_momentum_scale""": 0.5,
                 """edit_mom_beta""": 0.6,
             }

             __lowerCamelCase  =		3
             __lowerCamelCase  =		7

             # no sega enabled
             __lowerCamelCase  =		torch.Generator(UpperCamelCase_  )
             generator.manual_seed(UpperCamelCase_  )
             __lowerCamelCase  =		pipe(
                 [prompt]       ,			generator=UpperCamelCase_       ,			guidance_scale=UpperCamelCase_       ,			num_inference_steps=50       ,			output_type="""np"""       ,			width=5_12       ,			height=5_12       ,			)

             __lowerCamelCase  =		output.images
             __lowerCamelCase  =		image[0, -3:, -3:, -1]
             __lowerCamelCase  =		[
                 0.3488_7695,
                 0.387_6953,
                 0.375,
                 0.3442_3828,
                 0.358_1543,
                 0.3571_7773,
                 0.338_3789,
                 0.3457_0312,
                 0.35_9375,
             ]

             assert image.shape == (1, 5_12, 5_12, 3)

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

             # with sega enabled
             # generator = torch.manual_seed(seed)
             generator.manual_seed(UpperCamelCase_  )
             __lowerCamelCase  =		pipe(
                 [prompt]       ,			generator=UpperCamelCase_       ,			guidance_scale=UpperCamelCase_       ,			num_inference_steps=50       ,			output_type="""np"""       ,			width=5_12       ,			height=5_12       ,			**UpperCamelCase_       ,			)

             __lowerCamelCase  =		output.images
             __lowerCamelCase  =		image[0, -3:, -3:, -1]
             __lowerCamelCase  =		[
                 0.4228_5156,
                 0.3691_4062,
                 0.2907_7148,
                 0.4204_1016,
                 0.4191_8945,
                 0.3549_8047,
                 0.361_8164,
                 0.442_3828,
                 0.4311_5234,
             ]

             assert image.shape == (1, 5_12, 5_12, 3)

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







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):
      UpperCAmelCase__ :      Optional[Any]       						=  (JukeboxModel,) if is_torch_available() else ()
      UpperCAmelCase__ :      Union[str, Any]       						=  'openai/jukebox-1b-lyrics'
      UpperCAmelCase__ :      Optional[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
      UpperCAmelCase__ :      Dict       						=  [
          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
      ]

      UpperCAmelCase__ :      List[Any]       						=  [
          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
      ]

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

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

      UpperCAmelCase__ :      List[Any]       						=  [105_8304, 0, 78_6432, 7169, 507, 76, 27, 40, 30, 76]

      UpperCAmelCase__ :      Union[str, Any]       						=  [
          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
      ]
      UpperCAmelCase__ :      List[str]       						=  [
          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
      ]
      UpperCAmelCase__ :      str       						=  [
          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
      ]
      UpperCAmelCase__ :      Union[str, Any]       						=  [
          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
      ]
      UpperCAmelCase__ :      Optional[int]       						=  [
          -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
      ]
      UpperCAmelCase__ :      Union[str, 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
      ]
      UpperCAmelCase__ :      str       						=  [
          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
      ]
      UpperCAmelCase__ :      int       						=  [
          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:					Dict  ):
             __lowerCamelCase  =		JukeboxTokenizer.from_pretrained(self.model_id  )
             __lowerCamelCase  =		tokenizer(**self.metas  )["""input_ids"""]
             return tokens
      @slow
      def      lowerCAmelCase__		(							self:					Union[str, Any]  ):
             __lowerCamelCase  =		JukeboxModel.from_pretrained(self.model_id       ,			min_duration=0  ).eval()
             __lowerCamelCase  =		self.prepare_inputs()

             set_seed(0  )
             __lowerCamelCase  =		[torch.zeros(1       ,			0       ,			dtype=torch.long  ).cpu() for _ in range(3  )]
             __lowerCamelCase  =		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  )
             __lowerCamelCase  =		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  )
             __lowerCamelCase  =		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      lowerCAmelCase__		(							self:					Union[str, Any]  ):
             __lowerCamelCase  =		False
             __lowerCamelCase  =		JukeboxModel.from_pretrained(self.model_id       ,			min_duration=0  ).eval()

             __lowerCamelCase  =		self.prepare_inputs()
             set_seed(0  )
             __lowerCamelCase  =		[torch.zeros(1       ,			0       ,			dtype=torch.long  ) for _ in range(3  )]

             __lowerCamelCase  =		model.priors[0]
             __lowerCamelCase  =		0
             __lowerCamelCase  =		top_prior.get_music_tokens_conds(UpperCamelCase_       ,			start=UpperCamelCase_       ,			end=start + top_prior.n_ctx  )
             __lowerCamelCase  =		top_prior.get_metadata(labels[0].clone()       ,			UpperCamelCase_       ,			1_05_83_04       ,			0  )

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

             __lowerCamelCase,  __lowerCamelCase,  __lowerCamelCase  =		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      lowerCAmelCase__		(							self:					List[Any]  ):
             __lowerCamelCase  =		False

             __lowerCamelCase  =		JukeboxModel.from_pretrained(self.model_id       ,			min_duration=0  ).eval()
             set_seed(0  )
             __lowerCamelCase  =		torch.rand((1, 51_20, 1)  )
             __lowerCamelCase  =		list(self.prepare_inputs()  )

             __lowerCamelCase  =		[model.vqvae.encode(UpperCamelCase_       ,			start_level=2       ,			bs_chunks=waveform.shape[0]  )[0], None, None]
             __lowerCamelCase  =		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  )  )

             __lowerCamelCase  =		torch.cat((zs[0], torch.zeros(1       ,			20_48 - zs[0].shape[-1]  ))       ,			dim=-1  ).long()
             __lowerCamelCase  =		[upper_a, model.vqvae.encode(UpperCamelCase_       ,			start_level=1       ,			bs_chunks=waveform.shape[0]  )[0], None]
             __lowerCamelCase  =		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  )  )

             __lowerCamelCase  =		torch.cat((zs[1], torch.zeros(1       ,			20_48 - zs[1].shape[-1]  ))       ,			dim=-1  ).long()
             __lowerCamelCase  =		[upper_a, upper_a, model.vqvae.encode(UpperCamelCase_       ,			start_level=0       ,			bs_chunks=waveform.shape[0]  )[0]]
             __lowerCamelCase  =		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      lowerCAmelCase__		(							self:					List[Any]  ):
             __lowerCamelCase  =		JukeboxModel.from_pretrained(self.model_id       ,			min_duration=0  ).eval()
             set_seed(0  )
             __lowerCamelCase  =		torch.rand((1, 51_20, 1)  )
             with torch.no_grad():
                    __lowerCamelCase  =		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():
                    __lowerCamelCase  =		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):
      UpperCAmelCase__ :      List[Any]       						=  (JukeboxModel,) if is_torch_available() else ()
      UpperCAmelCase__ :      Optional[int]       						=  'openai/jukebox-5b-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]       						=  [
          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
      ]
      UpperCAmelCase__ :      Union[str, Any]       						=  [
          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
      ]

      UpperCAmelCase__ :      Optional[int]       						=  [
          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
      ]
      UpperCAmelCase__ :      List[str]       						=  [
          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
      ]

      UpperCAmelCase__ :      Optional[Any]       						=  [
          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
      ]
      UpperCAmelCase__ :      str       						=  [
          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
      ]

      UpperCAmelCase__ :      List[Any]       						=  [
          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
      ]
      UpperCAmelCase__ :      Dict       						=  [
          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
      ]

      UpperCAmelCase__ :      Optional[int]       						=  [
          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
      ]
      UpperCAmelCase__ :      str       						=  [
          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:					Optional[Any]       ,			UpperCamelCase_:					Optional[int]  ):
             __lowerCamelCase  =		JukeboxTokenizer.from_pretrained(UpperCamelCase_  )
             __lowerCamelCase  =		tokenizer(**self.metas  )["""input_ids"""]
             return tokens
      @slow
      def      lowerCAmelCase__		(							self:					Tuple  ):
             __lowerCamelCase  =		JukeboxModel.from_pretrained(self.model_id       ,			min_duration=0  ).eval()
             __lowerCamelCase  =		self.prepare_inputs(self.model_id  )

             set_seed(0  )
             __lowerCamelCase  =		[torch.zeros(1       ,			0       ,			dtype=torch.long  ).cpu() for _ in range(3  )]
             __lowerCamelCase  =		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  )
             __lowerCamelCase  =		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  )
             __lowerCamelCase  =		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      lowerCAmelCase__		(							self:					List[Any]  ):
             __lowerCamelCase  =		JukeboxModel.from_pretrained(self.model_id       ,			min_duration=0  ).eval()
             __lowerCamelCase  =		[i.cuda() for i in self.prepare_inputs(self.model_id  )]

             set_seed(0  )
             model.priors[0].cuda()
             __lowerCamelCase  =		[torch.zeros(1       ,			0       ,			dtype=torch.long  ).cuda() for _ in range(3  )]
             __lowerCamelCase  =		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()
             __lowerCamelCase  =		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()
             __lowerCamelCase  =		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      lowerCAmelCase__		(							self:					Dict  ):
             __lowerCamelCase  =		"""ArthurZ/jukebox_prior_0"""
             __lowerCamelCase  =		JukeboxPrior.from_pretrained(UpperCamelCase_       ,			min_duration=0  ).eval().half().to("""cuda"""  )

             __lowerCamelCase  =		self.prepare_inputs(UpperCamelCase_  )[0].cuda()
             __lowerCamelCase  =		model.get_metadata(UpperCamelCase_       ,			0       ,			76_80       ,			0  )
             set_seed(0  )
             __lowerCamelCase  =		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]  )

