


import torch
from torch import nn

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


class 			lowerCamelCase_		(			lowercase       ,  lowercase   ):





				"""simple docstring"""







				@register_to_config
				def __init__(						self       ,		*,
				    UpperCAmelCase__ = 4       ,		UpperCAmelCase__ = 768       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		):
								super().__init__()

								SCREAMING_SNAKE_CASE__     							=				nn.Parameter(torch.zeros(UpperCAmelCase__     )     )

								# parameters for additional clip time embeddings
								SCREAMING_SNAKE_CASE__     							=				nn.Linear(UpperCAmelCase__       ,		UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				nn.Linear(UpperCAmelCase__       ,		UpperCAmelCase__     )

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






				def   lowerCAmelCase__		(						self       ,		*, UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__     ):
								if do_classifier_free_guidance:
												# Add the classifier free guidance embeddings to the image embeddings
												SCREAMING_SNAKE_CASE__     							=				image_embeddings.shape[0]
												SCREAMING_SNAKE_CASE__     							=				self.learned_classifier_free_guidance_embeddings.unsqueeze(0     )
												SCREAMING_SNAKE_CASE__     							=				classifier_free_guidance_embeddings.expand(
												    UpperCAmelCase__       ,		-1     )
												SCREAMING_SNAKE_CASE__     							=				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]

								SCREAMING_SNAKE_CASE__     							=				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, ...
								SCREAMING_SNAKE_CASE__     							=				self.embedding_proj(UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				self.clip_image_embeddings_project_to_time_embeddings(UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				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"
								SCREAMING_SNAKE_CASE__     							=				self.clip_extra_context_tokens_proj(UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				clip_extra_context_tokens.reshape(UpperCAmelCase__       ,		-1       ,		self.clip_extra_context_tokens     )
								SCREAMING_SNAKE_CASE__     							=				clip_extra_context_tokens.permute(0       ,		2       ,		1     )

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

								return text_encoder_hidden_states, additive_clip_time_embeddings




"""simple docstring"""






import json
from typing import Iterator, List, Union

from tokenizers import AddedToken, Regex, Tokenizer, decoders, normalizers, pre_tokenizers, trainers
from tokenizers.implementations.base_tokenizer import BaseTokenizer
from tokenizers.models import Unigram
from tokenizers.processors import TemplateProcessing


class 			lowerCamelCase_		(			lowercase   ):





				"""simple docstring"""







				def __init__(						self       ,		UpperCAmelCase__ = "▁"       ,		UpperCAmelCase__ = True       ,		UpperCAmelCase__ = "<unk>"       ,		UpperCAmelCase__ = "</s>"       ,		UpperCAmelCase__ = "<pad>"       ,		):
								SCREAMING_SNAKE_CASE__     							=				{
								    "pad": {"id": 0, "token": pad_token},
								    "eos": {"id": 1, "token": eos_token},
								    "unk": {"id": 2, "token": unk_token},
								}

								SCREAMING_SNAKE_CASE__     							=				[None] * len(self.special_tokens     )
								for token_dict in self.special_tokens.values():
												SCREAMING_SNAKE_CASE__     							=				token_dict["token"]

								SCREAMING_SNAKE_CASE__     							=				Tokenizer(Unigram()     )

								SCREAMING_SNAKE_CASE__     							=				normalizers.Sequence(
								    [
								        normalizers.Nmt(),
								        normalizers.NFKC(),
								        normalizers.Replace(Regex(" {2,}"     )       ,		" "     ),
								        normalizers.Lowercase(),
								    ]     )
								SCREAMING_SNAKE_CASE__     							=				pre_tokenizers.Sequence(
								    [
								        pre_tokenizers.Metaspace(replacement=UpperCAmelCase__       ,		add_prefix_space=UpperCAmelCase__     ),
								        pre_tokenizers.Digits(individual_digits=UpperCAmelCase__     ),
								        pre_tokenizers.Punctuation(),
								    ]     )
								SCREAMING_SNAKE_CASE__     							=				decoders.Metaspace(replacement=UpperCAmelCase__       ,		add_prefix_space=UpperCAmelCase__     )

								SCREAMING_SNAKE_CASE__     							=				TemplateProcessing(
								    single=f'''$A {self.special_tokens["eos"]["token"]}'''       ,		special_tokens=[(self.special_tokens["eos"]["token"], self.special_tokens["eos"]["id"])]       ,		)

								SCREAMING_SNAKE_CASE__     							=				{
								    "model": "SentencePieceUnigram",
								    "replacement": replacement,
								    "add_prefix_space": add_prefix_space,
								}

								super().__init__(UpperCAmelCase__       ,		UpperCAmelCase__     )



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__ = 8000       ,		UpperCAmelCase__ = True       ,		):
								SCREAMING_SNAKE_CASE__     							=				trainers.UnigramTrainer(
								    vocab_size=UpperCAmelCase__       ,		special_tokens=self.special_tokens_list       ,		show_progress=UpperCAmelCase__       ,		)

								if isinstance(UpperCAmelCase__       ,		UpperCAmelCase__     ):
												SCREAMING_SNAKE_CASE__     							=				[files]
								self._tokenizer.train(UpperCAmelCase__       ,		trainer=UpperCAmelCase__     )

								self.add_unk_id()



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__ = 8000       ,		UpperCAmelCase__ = True       ,		):
								SCREAMING_SNAKE_CASE__     							=				trainers.UnigramTrainer(
								    vocab_size=UpperCAmelCase__       ,		special_tokens=self.special_tokens_list       ,		show_progress=UpperCAmelCase__       ,		)

								self._tokenizer.train_from_iterator(UpperCAmelCase__       ,		trainer=UpperCAmelCase__     )

								self.add_unk_id()






				def   lowerCAmelCase__		(						self     ):
								SCREAMING_SNAKE_CASE__     							=				json.loads(self._tokenizer.to_str()     )

								SCREAMING_SNAKE_CASE__     							=				self.special_tokens["unk"]["id"]

								SCREAMING_SNAKE_CASE__     							=				Tokenizer.from_str(json.dumps(UpperCAmelCase__     )     )
