


import torch
from torch import nn

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




class 				UpperCamelCase_			(				UpperCamelCase    ,						UpperCamelCase):

	"""simple docstring"""





	@register_to_config
	def __init__(				self    :       Optional[int]    ,    *,
	    UpperCAmelCase__    :       int = 4    ,    UpperCAmelCase__    :       int = 7_6_8    ,    UpperCAmelCase__    :       int    ,    UpperCAmelCase__    :       Optional[Any]    ,    ) ->			str:
						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 						UpperCAmelCase_   (				self    :       int    ,    *, UpperCAmelCase__    :       str    ,    UpperCAmelCase__    :       Tuple    ,    UpperCAmelCase__    :       List[Any]    ,    UpperCAmelCase__    :       List[Any]     ) ->			str:
						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"""







from math import asin, atan, cos, radians, sin, sqrt, tan

a__       :							List[str]        =  6_37_81_37.0
a__       :							Tuple        =  6_35_67_52.31_42_45
a__       :							str        =  6_3_7_8_1_3_7




def      UpperCAmelCase__ (lowerCAmelCase_						,		lowerCAmelCase_						,		lowerCAmelCase_						,		lowerCAmelCase_     ):




					'''simple docstring'''


					__SCREAMING_SNAKE_CASE		     =       (AXIS_A - AXIS_B) / AXIS_A
					__SCREAMING_SNAKE_CASE		     =       atan((1 - flattening) * tan(radians(lowerCAmelCase_     )     )     )
					__SCREAMING_SNAKE_CASE		     =       atan((1 - flattening) * tan(radians(lowerCAmelCase_     )     )     )
					__SCREAMING_SNAKE_CASE		     =       radians(lowerCAmelCase_     )
					__SCREAMING_SNAKE_CASE		     =       radians(lowerCAmelCase_     )
					# Equation
					__SCREAMING_SNAKE_CASE		     =       sin((phi_a - phi_a) / 2     )
					__SCREAMING_SNAKE_CASE		     =       sin((lambda_a - lambda_a) / 2     )
					# Square both values
					sin_sq_phi *= sin_sq_phi
					sin_sq_lambda *= sin_sq_lambda
					__SCREAMING_SNAKE_CASE		     =       sqrt(sin_sq_phi + (cos(lowerCAmelCase_     ) * cos(lowerCAmelCase_     ) * sin_sq_lambda)     )
					return 2 * RADIUS * asin(lowerCAmelCase_     )


if __name__ == "__main__":
		import doctest

		doctest.testmod()
