





import torch
from torch import nn

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin
class 			UpperCamelCase_		(  _lowerCamelCase							,  _lowerCamelCase					):




					@register_to_config
					def __init__(		self				,						*,
					    lowerCAmelCase_ = 4				,						lowerCAmelCase_ = 768				,						lowerCAmelCase_				,						lowerCAmelCase_				,						)		->							List[Any]:
										super().__init__()

										_snake_case									=				nn.Parameter(torch.zeros(lowerCAmelCase_  )  )

										# parameters for additional clip time embeddings
										_snake_case									=				nn.Linear(lowerCAmelCase_				,						lowerCAmelCase_  )
										_snake_case									=				nn.Linear(lowerCAmelCase_				,						lowerCAmelCase_  )

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






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

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

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

										return text_encoder_hidden_states, additive_clip_time_embeddings










from typing import TYPE_CHECKING

from ...utils import _LazyModule


UpperCAmelCase_      							=     {"""tokenization_byt5""": ["""ByT5Tokenizer"""]}


if TYPE_CHECKING:
					from .tokenization_byta import ByTaTokenizer
else:
					import sys

					UpperCAmelCase_      							=     _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



