



import unittest

from transformers import JukeboxTokenizer
from transformers.testing_utils import require_torch

class   _A   ( unittest.TestCase	):


						'''simple docstring'''





						_snake_case			:     List[Any]      							=      JukeboxTokenizer
						_snake_case			:     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        """,
						}



						@require_torch
						def     _snake_case	(	self	:							Tuple	):






							'''simple docstring'''



							import torch

							__lowercase		       =			JukeboxTokenizer.from_pretrained("openai/jukebox-1b-lyrics"	)
							__lowercase		       =			tokenizer(**self.metas	)["input_ids"]
							# fmt: off
							__lowercase		       =			[
							    torch.tensor([[
							        0, 0, 0, 7_169, 507, 9, 76, 39, 31, 46, 76, 27,
							        76, 46, 44, 27, 48, 31, 38, 38, 31, 44, 76, 32,
							        44, 41, 39, 76, 27, 40, 76, 27, 40, 46, 35, 43,
							        47, 31, 76, 38, 27, 40, 30, 64, 78, 76, 76, 76,
							        76, 76, 76, 76, 76, 23, 34, 41, 76, 45, 27, 35,
							        30, 76, 71, 20, 49, 41, 76, 48, 27, 45, 46, 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, 76, 76, 76, 76, 76, 76, 76,
							        19, 46, 27, 40, 30, 76, 35, 40, 76, 46, 34, 31,
							        76, 30, 31, 45, 31, 44, 46, 63, 76, 63, 76, 63,
							        76, 63, 76, 14, 31, 27, 44, 76, 46, 34, 31, 39,
							        64, 76, 41, 40, 76, 46, 34, 31, 76, 45, 27, 40,
							        30, 64, 78, 76, 76, 76, 76, 76, 76, 76, 76, 8,
							        27, 38, 32, 76, 45, 47, 40, 37, 76, 27, 76, 45,
							        34, 27, 46, 46, 31, 44, 31, 30, 76, 48, 35, 45,
							        27, 33, 31, 76, 38, 35, 31, 45, 64, 76, 49, 34,
							        41, 45, 31, 76, 32, 44, 41, 49, 40, 64, 78, 76,
							        76, 76, 76, 76, 76, 76, 76, 1, 40, 30, 76, 49,
							        44, 35, 40, 37, 38, 31, 30, 76, 38, 35, 42, 64,
							        76, 27, 40, 30, 76, 45, 40, 31, 31, 44, 76, 41,
							        32, 76, 29, 41, 38, 30, 76, 29, 41, 39, 39, 27,
							        40, 30, 64, 78, 76, 76, 76, 76, 76, 76, 76, 76,
							        20, 31, 38, 38, 76, 46, 34, 27, 46, 76, 35, 46,
							        45, 76, 45, 29, 47, 38, 42, 46, 41, 44, 76, 49,
							        31, 38, 38, 76, 46, 34, 41, 45, 31, 76, 42, 27,
							        45, 45, 35, 41, 40, 45, 76, 44, 31, 27, 30, 78,
							        76, 76, 76, 76, 76, 76, 76, 76, 23, 34, 35, 29,
							        34, 76, 51, 31, 46, 76, 45, 47, 44, 48, 35, 48,
							        31, 64, 76, 45, 46, 27, 39, 42, 31, 30, 76, 41,
							        40, 76, 46, 34, 31, 45, 31, 76, 38, 35, 32, 31,
							        38, 31, 45, 45, 76, 46, 34, 35, 40, 33, 45, 64,
							        78, 76, 76, 76, 76, 76, 76, 76, 76, 20, 34, 31,
							        76, 34, 27, 40, 30, 76, 46, 34, 27, 46, 76, 39,
							        41, 29, 37, 31, 30, 76, 46, 34, 31, 39, 64, 76,
							        27, 40, 30, 76, 46, 34, 31, 76, 34, 31, 27, 44,
							        46, 76, 46, 34, 27, 46, 76, 32, 31, 30, 66, 78,
							        76, 76, 76, 76, 76, 76, 76, 76, 1, 40, 30, 76,
							        41, 40, 76, 46, 34, 31, 76, 42, 31, 30, 31, 45,
							        46, 27, 38, 64, 76, 46, 34, 31, 45, 31, 76, 49,
							        41, 44, 30, 45, 76, 27, 42, 42, 31, 27, 44, 65,
							        78, 76, 76, 76, 76, 76, 76, 76, 76, 13, 51, 76,
							        40, 27, 39, 31, 76, 35, 45, 76, 15, 52, 51, 39,
							        27, 40, 30, 35, 27, 45, 64, 76, 11, 35, 40, 33,
							        76, 41, 32, 76, 11, 35, 40, 33, 45, 66, 78, 76,
							        76, 76, 76, 76, 76, 76, 76, 12, 41, 41, 37, 76,
							        41, 40, 76, 39, 51, 76, 23, 41, 44, 37, 45, 64,
							        76, 51, 31, 76, 13, 35, 33, 34, 46, 51, 64, 76,
							        27, 40, 30, 76, 30, 31, 45, 42, 27, 35, 44, 67,
							        78, 76, 76, 76, 76, 76, 76, 76, 76, 14, 41, 46,
							        34, 35, 40, 33, 76, 28, 31, 45, 35, 30, 31, 76,
							        44, 31, 39, 27, 35, 40, 45, 63, 76, 18, 41, 47,
							        40, 30, 76, 46, 34, 31, 76, 30, 31, 29, 27, 51,
							        78, 76, 76, 76, 76, 76, 76, 76, 76, 15, 32, 76,
							        46, 34, 27, 46, 76, 29, 41, 38, 41, 45, 45, 27,
							        38, 76, 23, 44, 31, 29, 37, 64, 76, 28, 41, 47,
							        40, 30, 38, 31, 45, 45, 76, 27, 40, 30, 76, 28,
							        27, 44, 31, 78, 76, 76, 76, 76, 76, 76, 76, 76,
							        20, 34, 31, 76, 38, 41, 40, 31, 76, 27, 40, 30,
							        76, 38, 31, 48, 31, 38, 76, 45, 27, 40, 30, 45,
							        76, 45, 46, 44, 31, 46, 29, 34, 76, 32, 27, 44,
							        76, 27, 49, 27, 51, 78, 76, 76, 76, 76, 76, 76,
							        76, 76]]	),
							    torch.tensor([[0, 0, 0, 1_069, 11]]	),
							    torch.tensor([[0, 0, 0, 1_069, 11]]	),
							]
							# fmt: on
							self.assertTrue(torch.allclose(tokens[0]							,      EXPECTED_OUTPUT[0]	)	)
							self.assertTrue(torch.allclose(tokens[1]							,      EXPECTED_OUTPUT[1]	)	)
							self.assertTrue(torch.allclose(tokens[2]							,      EXPECTED_OUTPUT[2]	)	)



						@require_torch
						def     _snake_case	(	self	:							Dict	):






							'''simple docstring'''



							import torch

							__lowercase		       =			JukeboxTokenizer.from_pretrained("openai/jukebox-5b-lyrics"	)
							__lowercase		       =			tokenizer(**self.metas	)["input_ids"]
							# fmt: off
							__lowercase		       =			[
							    torch.tensor([[
							        0, 0, 0, 1_069, 11, -1, -1, -1, -1, 9, 77, 39,
							        31, 46, 77, 27, 77, 46, 44, 27, 48, 31, 38, 38,
							        31, 44, 77, 32, 44, 41, 39, 77, 27, 40, 77, 27,
							        40, 46, 35, 43, 47, 31, 77, 38, 27, 40, 30, 64,
							        79, 77, 77, 77, 77, 77, 77, 77, 77, 23, 34, 41,
							        77, 45, 27, 35, 30, 77, 72, 20, 49, 41, 77, 48,
							        27, 45, 46, 77, 27, 40, 30, 77, 46, 44, 47, 40,
							        37, 38, 31, 45, 45, 77, 38, 31, 33, 45, 77, 41,
							        32, 77, 45, 46, 41, 40, 31, 79, 77, 77, 77, 77,
							        77, 77, 77, 77, 19, 46, 27, 40, 30, 77, 35, 40,
							        77, 46, 34, 31, 77, 30, 31, 45, 31, 44, 46, 63,
							        77, 63, 77, 63, 77, 63, 77, 14, 31, 27, 44, 77,
							        46, 34, 31, 39, 64, 77, 41, 40, 77, 46, 34, 31,
							        77, 45, 27, 40, 30, 64, 79, 77, 77, 77, 77, 77,
							        77, 77, 77, 8, 27, 38, 32, 77, 45, 47, 40, 37,
							        77, 27, 77, 45, 34, 27, 46, 46, 31, 44, 31, 30,
							        77, 48, 35, 45, 27, 33, 31, 77, 38, 35, 31, 45,
							        64, 77, 49, 34, 41, 45, 31, 77, 32, 44, 41, 49,
							        40, 64, 79, 77, 77, 77, 77, 77, 77, 77, 77, 1,
							        40, 30, 77, 49, 44, 35, 40, 37, 38, 31, 30, 77,
							        38, 35, 42, 64, 77, 27, 40, 30, 77, 45, 40, 31,
							        31, 44, 77, 41, 32, 77, 29, 41, 38, 30, 77, 29,
							        41, 39, 39, 27, 40, 30, 64, 79, 77, 77, 77, 77,
							        77, 77, 77, 77, 20, 31, 38, 38, 77, 46, 34, 27,
							        46, 77, 35, 46, 45, 77, 45, 29, 47, 38, 42, 46,
							        41, 44, 77, 49, 31, 38, 38, 77, 46, 34, 41, 45,
							        31, 77, 42, 27, 45, 45, 35, 41, 40, 45, 77, 44,
							        31, 27, 30, 79, 77, 77, 77, 77, 77, 77, 77, 77,
							        23, 34, 35, 29, 34, 77, 51, 31, 46, 77, 45, 47,
							        44, 48, 35, 48, 31, 64, 77, 45, 46, 27, 39, 42,
							        31, 30, 77, 41, 40, 77, 46, 34, 31, 45, 31, 77,
							        38, 35, 32, 31, 38, 31, 45, 45, 77, 46, 34, 35,
							        40, 33, 45, 64, 79, 77, 77, 77, 77, 77, 77, 77,
							        77, 20, 34, 31, 77, 34, 27, 40, 30, 77, 46, 34,
							        27, 46, 77, 39, 41, 29, 37, 31, 30, 77, 46, 34,
							        31, 39, 64, 77, 27, 40, 30, 77, 46, 34, 31, 77,
							        34, 31, 27, 44, 46, 77, 46, 34, 27, 46, 77, 32,
							        31, 30, 66, 79, 77, 77, 77, 77, 77, 77, 77, 77,
							        1, 40, 30, 77, 41, 40, 77, 46, 34, 31, 77, 42,
							        31, 30, 31, 45, 46, 27, 38, 64, 77, 46, 34, 31,
							        45, 31, 77, 49, 41, 44, 30, 45, 77, 27, 42, 42,
							        31, 27, 44, 65, 79, 77, 77, 77, 77, 77, 77, 77,
							        77, 13, 51, 77, 40, 27, 39, 31, 77, 35, 45, 77,
							        15, 52, 51, 39, 27, 40, 30, 35, 27, 45, 64, 77,
							        11, 35, 40, 33, 77, 41, 32, 77, 11, 35, 40, 33,
							        45, 66, 79, 77, 77, 77, 77, 77, 77, 77, 77, 12,
							        41, 41, 37, 77, 41, 40, 77, 39, 51, 77, 23, 41,
							        44, 37, 45, 64, 77, 51, 31, 77, 13, 35, 33, 34,
							        46, 51, 64, 77, 27, 40, 30, 77, 30, 31, 45, 42,
							        27, 35, 44, 67, 79, 77, 77, 77, 77, 77, 77, 77,
							        77, 14, 41, 46, 34, 35, 40, 33, 77, 28, 31, 45,
							        35, 30, 31, 77, 44, 31, 39, 27, 35, 40, 45, 63,
							        77, 18, 41, 47, 40, 30, 77, 46, 34, 31, 77, 30,
							        31, 29, 27, 51, 79, 77, 77, 77, 77, 77, 77, 77,
							        77, 15, 32, 77, 46, 34, 27, 46, 77, 29, 41, 38,
							        41, 45, 45, 27, 38, 77, 23, 44, 31, 29, 37, 64,
							        77, 28, 41, 47, 40, 30, 38, 31, 45, 45, 77, 27,
							        40, 30, 77, 28, 27, 44, 31, 79, 77, 77, 77, 77,
							        77, 77, 77, 77, 20, 34, 31, 77, 38, 41, 40, 31,
							        77, 27, 40, 30, 77, 38, 31, 48, 31, 38, 77, 45,
							        27, 40, 30, 45, 77, 45, 46, 44, 31, 46, 29, 34,
							        77, 32, 27, 44, 77, 27, 49, 27, 51, 79, 77, 77,
							        77, 77, 77, 77, 77, 77]]	),
							    torch.tensor([[0, 0, 0, 1_069, 11, -1, -1, -1, -1]]	),
							    torch.tensor([[0, 0, 0, 1_069, 11, -1, -1, -1, -1]]	),
							]
							# fmt: on
							self.assertTrue(torch.allclose(tokens[0]							,      EXPECTED_OUTPUT[0]	)	)
							self.assertTrue(torch.allclose(tokens[1]							,      EXPECTED_OUTPUT[1]	)	)
							self.assertTrue(torch.allclose(tokens[2]							,      EXPECTED_OUTPUT[2]	)	)




import torch
from torch import nn

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin

class   _A   ( _lowercase  ,						_lowercase	):


						'''simple docstring'''





						@register_to_config
						def __init__(	self	:							Optional[Any]							,      *,
						    lowerCamelCase	:							int = 4							,      lowerCamelCase	:							int = 768							,      lowerCamelCase	:							int							,      lowerCamelCase	:							Optional[int]							,      ):






							'''simple docstring'''



							super().__init__()

							__lowercase		       =			nn.Parameter(torch.zeros(lowerCamelCase	)	)

							# parameters for additional clip time embeddings
							__lowercase		       =			nn.Linear(lowerCamelCase							,      lowerCamelCase	)
							__lowercase		       =			nn.Linear(lowerCamelCase							,      lowerCamelCase	)

							# parameters for encoder hidden states
							__lowercase		       =			clip_extra_context_tokens
							__lowercase		       =			nn.Linear(
							    lowerCamelCase							,      self.clip_extra_context_tokens * cross_attention_dim	)
							__lowercase		       =			nn.Linear(lowerCamelCase							,      lowerCamelCase	)
							__lowercase		       =			nn.LayerNorm(lowerCamelCase	)



						def     _snake_case	(	self	:							Union[str, Any]							,      *, lowerCamelCase	:							Any							,      lowerCamelCase	:							Tuple							,      lowerCamelCase	:							Optional[int]							,      lowerCamelCase	:							Tuple	):






							'''simple docstring'''



							if do_classifier_free_guidance:
								# Add the classifier free guidance embeddings to the image embeddings
								__lowercase		       =			image_embeddings.shape[0]
								__lowercase		       =			self.learned_classifier_free_guidance_embeddings.unsqueeze(0	)
								__lowercase		       =			classifier_free_guidance_embeddings.expand(
								    lowerCamelCase							,      -1	)
								__lowercase		       =			torch.cat([classifier_free_guidance_embeddings, image_embeddings]							,      dim=0	)

							# The image embeddings batch size and the text embeddings batch size are equal
							assert image_embeddings.shape[0] == prompt_embeds.shape[0]

							__lowercase		       =			prompt_embeds.shape[0]

							# "Specifically, we modify the architecture described in Nichol et al. (2021) by projecting and
							# adding CLIP embeddings to the existing timestep embedding, ...
							__lowercase		       =			self.embedding_proj(lowerCamelCase	)
							__lowercase		       =			self.clip_image_embeddings_project_to_time_embeddings(lowerCamelCase	)
							__lowercase		       =			time_projected_image_embeddings + time_projected_prompt_embeds

							# ... and by projecting CLIP embeddings into four
							# extra tokens of context that are concatenated to the sequence of outputs from the GLIDE text encoder"
							__lowercase		       =			self.clip_extra_context_tokens_proj(lowerCamelCase	)
							__lowercase		       =			clip_extra_context_tokens.reshape(lowerCamelCase							,      -1							,      self.clip_extra_context_tokens	)
							__lowercase		       =			clip_extra_context_tokens.permute(0							,      2							,      1	)

							__lowercase		       =			self.encoder_hidden_states_proj(lowerCamelCase	)
							__lowercase		       =			self.text_encoder_hidden_states_norm(lowerCamelCase	)
							__lowercase		       =			torch.cat([clip_extra_context_tokens, text_encoder_hidden_states]							,      dim=1	)

							return text_encoder_hidden_states, additive_clip_time_embeddings



