


import math
from typing import Optional

import numpy as np
import torch
from torch import nn

from .activations import get_activation






def   lowerCamelCase_					( _a			,     _a			,     _a = False			,     _a = 1			,     _a = 1			,     _a = 10_000			,     ):




		"""simple docstring"""






		assert len(timesteps.shape ) == 1, "Timesteps should be a 1d-array"

		lowerCAmelCase__     :    Tuple					  =      embedding_dim // 2
		lowerCAmelCase__     :    Optional[Any]					  =      -math.log(_a ) * torch.arange(
		    start=0			,     end=_a			,     dtype=torch.floataa			,     device=timesteps.device )
		lowerCAmelCase__     :    int					  =      exponent / (half_dim - downscale_freq_shift)

		lowerCAmelCase__     :    Any					  =      torch.exp(_a )
		lowerCAmelCase__     :    List[str]					  =      timesteps[:, None].float() * emb[None, :]

		# scale embeddings
		lowerCAmelCase__     :    List[Any]					  =      scale * emb

		# concat sine and cosine embeddings
		lowerCAmelCase__     :    List[str]					  =      torch.cat([torch.sin(_a ), torch.cos(_a )]			,     dim=-1 )

		# flip sine and cosine embeddings
		if flip_sin_to_cos:
				lowerCAmelCase__     :    Any					  =      torch.cat([emb[:, half_dim:], emb[:, :half_dim]]			,     dim=-1 )

		# zero pad
		if embedding_dim % 2 == 1:
				lowerCAmelCase__     :    Optional[int]					  =      torch.nn.functional.pad(_a			,     (0, 1, 0, 0) )
		return emb






def   lowerCamelCase_					( _a			,     _a			,     _a=False			,     _a=0 ):




		"""simple docstring"""






		lowerCAmelCase__     :    Union[str, Any]					  =      np.arange(_a			,     dtype=np.floataa )
		lowerCAmelCase__     :    Optional[Any]					  =      np.arange(_a			,     dtype=np.floataa )
		lowerCAmelCase__     :    Optional[int]					  =      np.meshgrid(_a			,     _a )  # here w goes first
		lowerCAmelCase__     :    List[str]					  =      np.stack(_a			,     axis=0 )

		lowerCAmelCase__     :    Dict					  =      grid.reshape([2, 1, grid_size, grid_size] )
		lowerCAmelCase__     :    int					  =      get_ad_sincos_pos_embed_from_grid(_a			,     _a )
		if cls_token and extra_tokens > 0:
				lowerCAmelCase__     :    Tuple					  =      np.concatenate([np.zeros([extra_tokens, embed_dim] ), pos_embed]			,     axis=0 )
		return pos_embed






def   lowerCamelCase_					( _a			,     _a ):




		"""simple docstring"""






		if embed_dim % 2 != 0:
				raise ValueError('''embed_dim must be divisible by 2''' )

		# use half of dimensions to encode grid_h
		lowerCAmelCase__     :    Optional[int]					  =      get_ad_sincos_pos_embed_from_grid(embed_dim // 2			,     grid[0] )  # (H*W, D/2)
		lowerCAmelCase__     :    int					  =      get_ad_sincos_pos_embed_from_grid(embed_dim // 2			,     grid[1] )  # (H*W, D/2)

		lowerCAmelCase__     :    Tuple					  =      np.concatenate([emb_h, emb_w]			,     axis=1 )  # (H*W, D)
		return emb






def   lowerCamelCase_					( _a			,     _a ):




		"""simple docstring"""






		if embed_dim % 2 != 0:
				raise ValueError('''embed_dim must be divisible by 2''' )

		lowerCAmelCase__     :    Dict					  =      np.arange(embed_dim // 2			,     dtype=np.floataa )
		omega /= embed_dim / 2.0
		lowerCAmelCase__     :    Optional[Any]					  =      1.0 / 10_000**omega  # (D/2,)

		lowerCAmelCase__     :    Optional[Any]					  =      pos.reshape(-1 )  # (M,)
		lowerCAmelCase__     :    Union[str, Any]					  =      np.einsum('''m,d->md'''			,     _a			,     _a )  # (M, D/2), outer product

		lowerCAmelCase__     :    Tuple					  =      np.sin(_a )  # (M, D/2)
		lowerCAmelCase__     :    int					  =      np.cos(_a )  # (M, D/2)

		lowerCAmelCase__     :    Any					  =      np.concatenate([emb_sin, emb_cos]			,     axis=1 )  # (M, D)
		return emb





class       _a							(       nn.Module):




	def __init__(      self		:	Any	,      _SCREAMING_SNAKE_CASE		:	int=224	,      _SCREAMING_SNAKE_CASE		:	List[Any]=224	,      _SCREAMING_SNAKE_CASE		:	Union[str, Any]=16	,      _SCREAMING_SNAKE_CASE		:	List[str]=3	,      _SCREAMING_SNAKE_CASE		:	int=768	,      _SCREAMING_SNAKE_CASE		:	List[str]=False	,      _SCREAMING_SNAKE_CASE		:	Union[str, Any]=True	,      _SCREAMING_SNAKE_CASE		:	List[str]=True	,      )-> List[Any]:
			super().__init__()

			lowerCAmelCase__     :    Tuple					  =      (height // patch_size) * (width // patch_size)
			lowerCAmelCase__     :    List[Any]					  =      flatten
			lowerCAmelCase__     :    Any					  =      layer_norm

			lowerCAmelCase__     :    Tuple					  =      nn.Convad(
			    _SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	,      kernel_size=(patch_size, patch_size)	,      stride=_SCREAMING_SNAKE_CASE	,      bias=_SCREAMING_SNAKE_CASE	)
			if layer_norm:
					lowerCAmelCase__     :    Union[str, Any]					  =      nn.LayerNorm(_SCREAMING_SNAKE_CASE	,      elementwise_affine=_SCREAMING_SNAKE_CASE	,      eps=1E-6	)
			else:
					lowerCAmelCase__     :    Any					  =      None

			lowerCAmelCase__     :    Optional[int]					  =      get_ad_sincos_pos_embed(_SCREAMING_SNAKE_CASE	,      int(num_patches**0.5	)	)
			self.register_buffer('''pos_embed'''	,      torch.from_numpy(_SCREAMING_SNAKE_CASE	).float().unsqueeze(0	)	,      persistent=_SCREAMING_SNAKE_CASE	)




	def 							UpperCAmelCase__(      self		:	Optional[Any]	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]	)-> Dict:
			lowerCAmelCase__     :    List[Any]					  =      self.proj(_SCREAMING_SNAKE_CASE	)
			if self.flatten:
					lowerCAmelCase__     :    Optional[Any]					  =      latent.flatten(2	).transpose(1	,      2	)  # BCHW -> BNC
			if self.layer_norm:
					lowerCAmelCase__     :    Optional[int]					  =      self.norm(_SCREAMING_SNAKE_CASE	)
			return latent + self.pos_embed





class       _a							(       nn.Module):




	def __init__(      self		:	Tuple	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	str = "silu"	,      _SCREAMING_SNAKE_CASE		:	int = None	,      _SCREAMING_SNAKE_CASE		:	Optional[str] = None	,      _SCREAMING_SNAKE_CASE		:	int=None	,      )-> Tuple:
			super().__init__()

			lowerCAmelCase__     :    Optional[Any]					  =      nn.Linear(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)

			if cond_proj_dim is not None:
					lowerCAmelCase__     :    int					  =      nn.Linear(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	,      bias=_SCREAMING_SNAKE_CASE	)
			else:
					lowerCAmelCase__     :    Dict					  =      None

			lowerCAmelCase__     :    Dict					  =      get_activation(_SCREAMING_SNAKE_CASE	)

			if out_dim is not None:
					lowerCAmelCase__     :    Union[str, Any]					  =      out_dim
			else:
					lowerCAmelCase__     :    Tuple					  =      time_embed_dim
			lowerCAmelCase__     :    Tuple					  =      nn.Linear(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)

			if post_act_fn is None:
					lowerCAmelCase__     :    str					  =      None
			else:
					lowerCAmelCase__     :    Union[str, Any]					  =      get_activation(_SCREAMING_SNAKE_CASE	)




	def 							UpperCAmelCase__(      self		:	Any	,      _SCREAMING_SNAKE_CASE		:	Optional[int]	,      _SCREAMING_SNAKE_CASE		:	Tuple=None	)-> Optional[int]:
			if condition is not None:
					lowerCAmelCase__     :    List[Any]					  =      sample + self.cond_proj(_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Optional[Any]					  =      self.linear_a(_SCREAMING_SNAKE_CASE	)

			if self.act is not None:
					lowerCAmelCase__     :    List[Any]					  =      self.act(_SCREAMING_SNAKE_CASE	)

			lowerCAmelCase__     :    Tuple					  =      self.linear_a(_SCREAMING_SNAKE_CASE	)

			if self.post_act is not None:
					lowerCAmelCase__     :    Optional[int]					  =      self.post_act(_SCREAMING_SNAKE_CASE	)
			return sample





class       _a							(       nn.Module):




	def __init__(      self		:	Any	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	bool	,      _SCREAMING_SNAKE_CASE		:	float	)-> int:
			super().__init__()
			lowerCAmelCase__     :    Any					  =      num_channels
			lowerCAmelCase__     :    List[Any]					  =      flip_sin_to_cos
			lowerCAmelCase__     :    Optional[int]					  =      downscale_freq_shift




	def 							UpperCAmelCase__(      self		:	str	,      _SCREAMING_SNAKE_CASE		:	List[Any]	)-> Optional[int]:
			lowerCAmelCase__     :    str					  =      get_timestep_embedding(
			    _SCREAMING_SNAKE_CASE	,      self.num_channels	,      flip_sin_to_cos=self.flip_sin_to_cos	,      downscale_freq_shift=self.downscale_freq_shift	,      )
			return t_emb





class       _a							(       nn.Module):




	def __init__(      self		:	List[str]	,      _SCREAMING_SNAKE_CASE		:	int = 256	,      _SCREAMING_SNAKE_CASE		:	float = 1.0	,      _SCREAMING_SNAKE_CASE		:	Union[str, Any]=True	,      _SCREAMING_SNAKE_CASE		:	Union[str, Any]=True	,      _SCREAMING_SNAKE_CASE		:	Tuple=False	)-> List[Any]:
			super().__init__()
			lowerCAmelCase__     :    Optional[int]					  =      nn.Parameter(torch.randn(_SCREAMING_SNAKE_CASE	) * scale	,      requires_grad=_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    str					  =      log
			lowerCAmelCase__     :    Optional[int]					  =      flip_sin_to_cos

			if set_W_to_weight:
					# to delete later
					lowerCAmelCase__     :    Any					  =      nn.Parameter(torch.randn(_SCREAMING_SNAKE_CASE	) * scale	,      requires_grad=_SCREAMING_SNAKE_CASE	)

					lowerCAmelCase__     :    Optional[int]					  =      self.W




	def 							UpperCAmelCase__(      self		:	int	,      _SCREAMING_SNAKE_CASE		:	Optional[int]	)-> str:
			if self.log:
					lowerCAmelCase__     :    Union[str, Any]					  =      torch.log(_SCREAMING_SNAKE_CASE	)

			lowerCAmelCase__     :    str					  =      x[:, None] * self.weight[None, :] * 2 * np.pi

			if self.flip_sin_to_cos:
					lowerCAmelCase__     :    Optional[int]					  =      torch.cat([torch.cos(_SCREAMING_SNAKE_CASE	), torch.sin(_SCREAMING_SNAKE_CASE	)]	,      dim=-1	)
			else:
					lowerCAmelCase__     :    Dict					  =      torch.cat([torch.sin(_SCREAMING_SNAKE_CASE	), torch.cos(_SCREAMING_SNAKE_CASE	)]	,      dim=-1	)
			return out





class       _a							(       nn.Module):




	def __init__(      self		:	Dict	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	int	,      )-> Dict:
			super().__init__()

			lowerCAmelCase__     :    List[str]					  =      height
			lowerCAmelCase__     :    List[str]					  =      width
			lowerCAmelCase__     :    int					  =      num_embed
			lowerCAmelCase__     :    int					  =      embed_dim

			lowerCAmelCase__     :    Dict					  =      nn.Embedding(self.num_embed	,      _SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Dict					  =      nn.Embedding(self.height	,      _SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Any					  =      nn.Embedding(self.width	,      _SCREAMING_SNAKE_CASE	)




	def 							UpperCAmelCase__(      self		:	Dict	,      _SCREAMING_SNAKE_CASE		:	Union[str, Any]	)-> Optional[int]:
			lowerCAmelCase__     :    Union[str, Any]					  =      self.emb(_SCREAMING_SNAKE_CASE	)

			lowerCAmelCase__     :    Optional[Any]					  =      self.height_emb(torch.arange(self.height	,      device=index.device	).view(1	,      self.height	)	)

			# 1 x H x D -> 1 x H x 1 x D
			lowerCAmelCase__     :    Dict					  =      height_emb.unsqueeze(2	)

			lowerCAmelCase__     :    Any					  =      self.width_emb(torch.arange(self.width	,      device=index.device	).view(1	,      self.width	)	)

			# 1 x W x D -> 1 x 1 x W x D
			lowerCAmelCase__     :    Optional[int]					  =      width_emb.unsqueeze(1	)

			lowerCAmelCase__     :    Any					  =      height_emb + width_emb

			# 1 x H x W x D -> 1 x L xD
			lowerCAmelCase__     :    List[Any]					  =      pos_emb.view(1	,      self.height * self.width	,      -1	)

			lowerCAmelCase__     :    int					  =      emb + pos_emb[:, : emb.shape[1], :]

			return emb





class       _a							(       nn.Module):




	def __init__(      self		:	str	,      _SCREAMING_SNAKE_CASE		:	str	,      _SCREAMING_SNAKE_CASE		:	List[Any]	,      _SCREAMING_SNAKE_CASE		:	List[Any]	)-> List[str]:
			super().__init__()
			lowerCAmelCase__     :    int					  =      dropout_prob > 0
			lowerCAmelCase__     :    str					  =      nn.Embedding(num_classes + use_cfg_embedding	,      _SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    str					  =      num_classes
			lowerCAmelCase__     :    Tuple					  =      dropout_prob




	def 							UpperCAmelCase__(      self		:	Optional[int]	,      _SCREAMING_SNAKE_CASE		:	Optional[int]	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]=None	)-> int:
			if force_drop_ids is None:
					lowerCAmelCase__     :    Dict					  =      torch.rand(labels.shape[0]	,      device=labels.device	) < self.dropout_prob
			else:
					lowerCAmelCase__     :    List[str]					  =      torch.tensor(force_drop_ids == 1	)
			lowerCAmelCase__     :    Tuple					  =      torch.where(_SCREAMING_SNAKE_CASE	,      self.num_classes	,      _SCREAMING_SNAKE_CASE	)
			return labels




	def 							UpperCAmelCase__(      self		:	int	,      _SCREAMING_SNAKE_CASE		:	torch.LongTensor	,      _SCREAMING_SNAKE_CASE		:	List[str]=None	)-> List[str]:
			lowerCAmelCase__     :    Union[str, Any]					  =      self.dropout_prob > 0
			if (self.training and use_dropout) or (force_drop_ids is not None):
					lowerCAmelCase__     :    List[str]					  =      self.token_drop(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Optional[Any]					  =      self.embedding_table(_SCREAMING_SNAKE_CASE	)
			return embeddings





class       _a							(       nn.Module):




	def __init__(      self		:	int	,      _SCREAMING_SNAKE_CASE		:	int = 1024	,      _SCREAMING_SNAKE_CASE		:	int = 768	,      _SCREAMING_SNAKE_CASE		:	int = 768	,      _SCREAMING_SNAKE_CASE		:	int = 10	,      )-> Dict:
			super().__init__()

			lowerCAmelCase__     :    List[Any]					  =      num_image_text_embeds
			lowerCAmelCase__     :    Optional[Any]					  =      nn.Linear(_SCREAMING_SNAKE_CASE	,      self.num_image_text_embeds * cross_attention_dim	)
			lowerCAmelCase__     :    Optional[Any]					  =      nn.Linear(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)




	def 							UpperCAmelCase__(      self		:	List[str]	,      _SCREAMING_SNAKE_CASE		:	torch.FloatTensor	,      _SCREAMING_SNAKE_CASE		:	torch.FloatTensor	)-> List[str]:
			lowerCAmelCase__     :    int					  =      text_embeds.shape[0]

			# image
			lowerCAmelCase__     :    List[str]					  =      self.image_embeds(_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    List[str]					  =      image_text_embeds.reshape(_SCREAMING_SNAKE_CASE	,      self.num_image_text_embeds	,      -1	)

			# text
			lowerCAmelCase__     :    Dict					  =      self.text_proj(_SCREAMING_SNAKE_CASE	)

			return torch.cat([image_text_embeds, text_embeds]	,      dim=1	)





class       _a							(       nn.Module):




	def __init__(      self		:	Optional[int]	,      _SCREAMING_SNAKE_CASE		:	int = 768	,      _SCREAMING_SNAKE_CASE		:	int = 768	,      _SCREAMING_SNAKE_CASE		:	int = 32	,      )-> Optional[Any]:
			super().__init__()

			lowerCAmelCase__     :    int					  =      num_image_text_embeds
			lowerCAmelCase__     :    Any					  =      nn.Linear(_SCREAMING_SNAKE_CASE	,      self.num_image_text_embeds * cross_attention_dim	)
			lowerCAmelCase__     :    List[str]					  =      nn.LayerNorm(_SCREAMING_SNAKE_CASE	)




	def 							UpperCAmelCase__(      self		:	List[Any]	,      _SCREAMING_SNAKE_CASE		:	torch.FloatTensor	)-> Optional[Any]:
			lowerCAmelCase__     :    Any					  =      image_embeds.shape[0]

			# image
			lowerCAmelCase__     :    int					  =      self.image_embeds(_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Optional[Any]					  =      image_embeds.reshape(_SCREAMING_SNAKE_CASE	,      self.num_image_text_embeds	,      -1	)
			lowerCAmelCase__     :    Optional[Any]					  =      self.norm(_SCREAMING_SNAKE_CASE	)
			return image_embeds





class       _a							(       nn.Module):




	def __init__(      self		:	List[Any]	,      _SCREAMING_SNAKE_CASE		:	Any	,      _SCREAMING_SNAKE_CASE		:	Dict	,      _SCREAMING_SNAKE_CASE		:	str=0.1	)-> Optional[Any]:
			super().__init__()

			lowerCAmelCase__     :    Dict					  =      Timesteps(num_channels=256	,      flip_sin_to_cos=_SCREAMING_SNAKE_CASE	,      downscale_freq_shift=1	)
			lowerCAmelCase__     :    Dict					  =      TimestepEmbedding(in_channels=256	,      time_embed_dim=_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Dict					  =      LabelEmbedding(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)




	def 							UpperCAmelCase__(      self		:	List[Any]	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	List[Any]	,      _SCREAMING_SNAKE_CASE		:	Any=None	)-> Dict:
			lowerCAmelCase__     :    Optional[int]					  =      self.time_proj(_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    str					  =      self.timestep_embedder(timesteps_proj.to(dtype=_SCREAMING_SNAKE_CASE	)	)  # (N, D)

			lowerCAmelCase__     :    Optional[Any]					  =      self.class_embedder(_SCREAMING_SNAKE_CASE	)  # (N, D)

			lowerCAmelCase__     :    List[str]					  =      timesteps_emb + class_labels  # (N, D)

			return conditioning





class       _a							(       nn.Module):




	def __init__(      self		:	List[Any]	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	int	,      _SCREAMING_SNAKE_CASE		:	int = 64	)-> Union[str, Any]:
			super().__init__()
			lowerCAmelCase__     :    List[Any]					  =      nn.LayerNorm(_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Tuple					  =      AttentionPooling(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Any					  =      nn.Linear(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    int					  =      nn.LayerNorm(_SCREAMING_SNAKE_CASE	)




	def 							UpperCAmelCase__(      self		:	Optional[int]	,      _SCREAMING_SNAKE_CASE		:	List[Any]	)-> Any:
			lowerCAmelCase__     :    List[Any]					  =      self.norma(_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Optional[Any]					  =      self.pool(_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    int					  =      self.proj(_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    List[str]					  =      self.norma(_SCREAMING_SNAKE_CASE	)
			return hidden_states





class       _a							(       nn.Module):




	def __init__(      self		:	List[str]	,      _SCREAMING_SNAKE_CASE		:	int = 768	,      _SCREAMING_SNAKE_CASE		:	int = 768	,      _SCREAMING_SNAKE_CASE		:	int = 1536	)-> Dict:
			super().__init__()
			lowerCAmelCase__     :    Union[str, Any]					  =      nn.Linear(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Any					  =      nn.LayerNorm(_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Tuple					  =      nn.Linear(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)




	def 							UpperCAmelCase__(      self		:	List[Any]	,      _SCREAMING_SNAKE_CASE		:	torch.FloatTensor	,      _SCREAMING_SNAKE_CASE		:	torch.FloatTensor	)-> Optional[int]:
			# text
			lowerCAmelCase__     :    List[Any]					  =      self.text_proj(_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Union[str, Any]					  =      self.text_norm(_SCREAMING_SNAKE_CASE	)

			# image
			lowerCAmelCase__     :    Optional[int]					  =      self.image_proj(_SCREAMING_SNAKE_CASE	)

			return time_image_embeds + time_text_embeds





class       _a							(       nn.Module):




	def __init__(      self		:	Optional[Any]	,      _SCREAMING_SNAKE_CASE		:	int = 768	,      _SCREAMING_SNAKE_CASE		:	int = 1536	)-> Any:
			super().__init__()
			lowerCAmelCase__     :    str					  =      nn.Linear(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Optional[int]					  =      nn.LayerNorm(_SCREAMING_SNAKE_CASE	)




	def 							UpperCAmelCase__(      self		:	Any	,      _SCREAMING_SNAKE_CASE		:	torch.FloatTensor	)-> Any:
			# image
			lowerCAmelCase__     :    Optional[int]					  =      self.image_proj(_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Dict					  =      self.image_norm(_SCREAMING_SNAKE_CASE	)
			return time_image_embeds





class       _a							(       nn.Module):




	def __init__(      self		:	Any	,      _SCREAMING_SNAKE_CASE		:	int = 768	,      _SCREAMING_SNAKE_CASE		:	int = 1536	)-> int:
			super().__init__()
			lowerCAmelCase__     :    Union[str, Any]					  =      nn.Linear(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Dict					  =      nn.LayerNorm(_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Optional[Any]					  =      nn.Sequential(
			    nn.Convad(3	,      16	,      3	,      padding=1	)	,      nn.SiLU()	,      nn.Convad(16	,      16	,      3	,      padding=1	)	,      nn.SiLU()	,      nn.Convad(16	,      32	,      3	,      padding=1	,      stride=2	)	,      nn.SiLU()	,      nn.Convad(32	,      32	,      3	,      padding=1	)	,      nn.SiLU()	,      nn.Convad(32	,      96	,      3	,      padding=1	,      stride=2	)	,      nn.SiLU()	,      nn.Convad(96	,      96	,      3	,      padding=1	)	,      nn.SiLU()	,      nn.Convad(96	,      256	,      3	,      padding=1	,      stride=2	)	,      nn.SiLU()	,      nn.Convad(256	,      4	,      3	,      padding=1	)	,      )




	def 							UpperCAmelCase__(      self		:	int	,      _SCREAMING_SNAKE_CASE		:	torch.FloatTensor	,      _SCREAMING_SNAKE_CASE		:	torch.FloatTensor	)-> Dict:
			# image
			lowerCAmelCase__     :    Dict					  =      self.image_proj(_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Tuple					  =      self.image_norm(_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    List[str]					  =      self.input_hint_block(_SCREAMING_SNAKE_CASE	)
			return time_image_embeds, hint





class       _a							(       nn.Module):




	def __init__(      self		:	Tuple	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]	,      _SCREAMING_SNAKE_CASE		:	str	,      _SCREAMING_SNAKE_CASE		:	Dict=None	)-> Optional[int]:
			super().__init__()
			lowerCAmelCase__     :    Optional[int]					  =      dtype
			lowerCAmelCase__     :    Optional[Any]					  =      nn.Parameter(torch.randn(1	,      _SCREAMING_SNAKE_CASE	) / embed_dim**0.5	)
			lowerCAmelCase__     :    Any					  =      nn.Linear(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	,      dtype=self.dtype	)
			lowerCAmelCase__     :    int					  =      nn.Linear(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	,      dtype=self.dtype	)
			lowerCAmelCase__     :    str					  =      nn.Linear(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	,      dtype=self.dtype	)
			lowerCAmelCase__     :    List[str]					  =      num_heads
			lowerCAmelCase__     :    int					  =      embed_dim // self.num_heads




	def 							UpperCAmelCase__(      self		:	Dict	,      _SCREAMING_SNAKE_CASE		:	Union[str, Any]	)-> Union[str, Any]:
			lowerCAmelCase__  ,			lowerCAmelCase__  ,			lowerCAmelCase__     :    Union[str, Any]					  =      x.size()

			def shape(_SCREAMING_SNAKE_CASE		:	List[Any]	):
					# (bs, length, width) --> (bs, length, n_heads, dim_per_head)
					lowerCAmelCase__     :    List[str]					  =      x.view(_SCREAMING_SNAKE_CASE	,      -1	,      self.num_heads	,      self.dim_per_head	)
					# (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head)
					lowerCAmelCase__     :    Dict					  =      x.transpose(1	,      2	)
					# (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head)
					lowerCAmelCase__     :    Dict					  =      x.reshape(bs * self.num_heads	,      -1	,      self.dim_per_head	)
					# (bs*n_heads, length, dim_per_head) --> (bs*n_heads, dim_per_head, length)
					lowerCAmelCase__     :    List[Any]					  =      x.transpose(1	,      2	)
					return x

			lowerCAmelCase__     :    Dict					  =      x.mean(dim=1	,      keepdim=_SCREAMING_SNAKE_CASE	) + self.positional_embedding.to(x.dtype	)
			lowerCAmelCase__     :    Dict					  =      torch.cat([class_token, x]	,      dim=1	)  # (bs, length+1, width)

			# (bs*n_heads, class_token_length, dim_per_head)
			lowerCAmelCase__     :    Optional[int]					  =      shape(self.q_proj(_SCREAMING_SNAKE_CASE	)	)
			# (bs*n_heads, length+class_token_length, dim_per_head)
			lowerCAmelCase__     :    List[str]					  =      shape(self.k_proj(_SCREAMING_SNAKE_CASE	)	)
			lowerCAmelCase__     :    Optional[Any]					  =      shape(self.v_proj(_SCREAMING_SNAKE_CASE	)	)

			# (bs*n_heads, class_token_length, length+class_token_length):
			lowerCAmelCase__     :    List[Any]					  =      1 / math.sqrt(math.sqrt(self.dim_per_head	)	)
			lowerCAmelCase__     :    str					  =      torch.einsum('''bct,bcs->bts'''	,      q * scale	,      k * scale	)  # More stable with f16 than dividing afterwards
			lowerCAmelCase__     :    Any					  =      torch.softmax(weight.float()	,      dim=-1	).type(weight.dtype	)

			# (bs*n_heads, dim_per_head, class_token_length)
			lowerCAmelCase__     :    Tuple					  =      torch.einsum('''bts,bcs->bct'''	,      _SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	)

			# (bs, length+1, width)
			lowerCAmelCase__     :    List[Any]					  =      a.reshape(_SCREAMING_SNAKE_CASE	,      -1	,      1	).transpose(1	,      2	)

			return a[:, 0, :]  # cls_token








import copy

from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCamelCase		=	logging.get_logger(__name__)

lowerCamelCase		=	{
    '''facebook/sam-vit-huge''': '''https://huggingface.co/facebook/sam-vit-huge/resolve/main/config.json''',
    '''facebook/sam-vit-large''': '''https://huggingface.co/facebook/sam-vit-large/resolve/main/config.json''',
    '''facebook/sam-vit-base''': '''https://huggingface.co/facebook/sam-vit-base/resolve/main/config.json''',
}





class       _a							(       _lowercase):




	def __init__(      self		:	Union[str, Any]	,      _SCREAMING_SNAKE_CASE		:	Tuple=256	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]=1024	,      _SCREAMING_SNAKE_CASE		:	List[Any]=16	,      _SCREAMING_SNAKE_CASE		:	Optional[int]=16	,      _SCREAMING_SNAKE_CASE		:	Union[str, Any]=4	,      _SCREAMING_SNAKE_CASE		:	Any="gelu"	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]=1E-6	,      **_SCREAMING_SNAKE_CASE		:	List[Any]	,      )-> Tuple:
			super().__init__(**_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Any					  =      hidden_size
			lowerCAmelCase__     :    Tuple					  =      image_size
			lowerCAmelCase__     :    List[Any]					  =      patch_size
			lowerCAmelCase__     :    Any					  =      image_size // patch_size
			lowerCAmelCase__     :    Dict					  =      mask_input_channels
			lowerCAmelCase__     :    Dict					  =      num_point_embeddings
			lowerCAmelCase__     :    str					  =      hidden_act
			lowerCAmelCase__     :    List[Any]					  =      layer_norm_eps





class       _a							(       _lowercase):




	def __init__(      self		:	Tuple	,      _SCREAMING_SNAKE_CASE		:	Union[str, Any]=256	,      _SCREAMING_SNAKE_CASE		:	Dict="relu"	,      _SCREAMING_SNAKE_CASE		:	str=2048	,      _SCREAMING_SNAKE_CASE		:	List[Any]=2	,      _SCREAMING_SNAKE_CASE		:	Dict=8	,      _SCREAMING_SNAKE_CASE		:	str=2	,      _SCREAMING_SNAKE_CASE		:	Tuple=3	,      _SCREAMING_SNAKE_CASE		:	Optional[int]=3	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]=256	,      _SCREAMING_SNAKE_CASE		:	Dict=1E-6	,      **_SCREAMING_SNAKE_CASE		:	Any	,      )-> List[Any]:
			super().__init__(**_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    List[str]					  =      hidden_size
			lowerCAmelCase__     :    Union[str, Any]					  =      hidden_act
			lowerCAmelCase__     :    str					  =      mlp_dim
			lowerCAmelCase__     :    Optional[int]					  =      num_hidden_layers
			lowerCAmelCase__     :    int					  =      num_attention_heads
			lowerCAmelCase__     :    List[str]					  =      attention_downsample_rate
			lowerCAmelCase__     :    str					  =      num_multimask_outputs
			lowerCAmelCase__     :    Optional[int]					  =      iou_head_depth
			lowerCAmelCase__     :    Union[str, Any]					  =      iou_head_hidden_dim
			lowerCAmelCase__     :    Tuple					  =      layer_norm_eps





class       _a							(       _lowercase):




	def __init__(      self		:	List[str]	,      _SCREAMING_SNAKE_CASE		:	int=768	,      _SCREAMING_SNAKE_CASE		:	str=256	,      _SCREAMING_SNAKE_CASE		:	int=12	,      _SCREAMING_SNAKE_CASE		:	Tuple=12	,      _SCREAMING_SNAKE_CASE		:	Dict=3	,      _SCREAMING_SNAKE_CASE		:	Optional[int]=1024	,      _SCREAMING_SNAKE_CASE		:	Tuple=16	,      _SCREAMING_SNAKE_CASE		:	List[Any]="gelu"	,      _SCREAMING_SNAKE_CASE		:	Optional[int]=1E-06	,      _SCREAMING_SNAKE_CASE		:	int=0.0	,      _SCREAMING_SNAKE_CASE		:	int=1E-10	,      _SCREAMING_SNAKE_CASE		:	List[Any]=True	,      _SCREAMING_SNAKE_CASE		:	Any=4.0	,      _SCREAMING_SNAKE_CASE		:	List[str]=True	,      _SCREAMING_SNAKE_CASE		:	int=True	,      _SCREAMING_SNAKE_CASE		:	List[str]=14	,      _SCREAMING_SNAKE_CASE		:	Dict=[2, 5, 8, 11]	,      _SCREAMING_SNAKE_CASE		:	List[Any]=128	,      _SCREAMING_SNAKE_CASE		:	Dict=None	,      **_SCREAMING_SNAKE_CASE		:	str	,      )-> Any:
			super().__init__(**_SCREAMING_SNAKE_CASE	)

			lowerCAmelCase__     :    Optional[Any]					  =      hidden_size
			lowerCAmelCase__     :    Any					  =      output_channels
			lowerCAmelCase__     :    Optional[int]					  =      num_hidden_layers
			lowerCAmelCase__     :    Dict					  =      num_attention_heads
			lowerCAmelCase__     :    List[Any]					  =      num_channels
			lowerCAmelCase__     :    int					  =      image_size
			lowerCAmelCase__     :    Tuple					  =      patch_size
			lowerCAmelCase__     :    Union[str, Any]					  =      hidden_act
			lowerCAmelCase__     :    List[str]					  =      layer_norm_eps
			lowerCAmelCase__     :    Any					  =      attention_dropout
			lowerCAmelCase__     :    int					  =      initializer_range
			lowerCAmelCase__     :    List[Any]					  =      qkv_bias
			lowerCAmelCase__     :    List[Any]					  =      mlp_ratio
			lowerCAmelCase__     :    List[Any]					  =      use_abs_pos
			lowerCAmelCase__     :    Union[str, Any]					  =      use_rel_pos
			lowerCAmelCase__     :    Optional[Any]					  =      window_size
			lowerCAmelCase__     :    Union[str, Any]					  =      global_attn_indexes
			lowerCAmelCase__     :    Union[str, Any]					  =      num_pos_feats
			lowerCAmelCase__     :    str					  =      int(hidden_size * mlp_ratio	) if mlp_dim is None else mlp_dim





class       _a							(       _lowercase):
	_a						:						Tuple   					=							'''sam'''
	_a						:						Optional[Any]   					=							True




	def __init__(      self		:	Optional[Any]	,      _SCREAMING_SNAKE_CASE		:	Any=None	,      _SCREAMING_SNAKE_CASE		:	Optional[Any]=None	,      _SCREAMING_SNAKE_CASE		:	int=None	,      _SCREAMING_SNAKE_CASE		:	Tuple=0.02	,      **_SCREAMING_SNAKE_CASE		:	int	,      )-> Optional[int]:
			super().__init__(**_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Tuple					  =      vision_config if vision_config is not None else {}
			lowerCAmelCase__     :    List[str]					  =      prompt_encoder_config if prompt_encoder_config is not None else {}
			lowerCAmelCase__     :    List[str]					  =      mask_decoder_config if mask_decoder_config is not None else {}

			if isinstance(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	):
					lowerCAmelCase__     :    Any					  =      vision_config.to_dict()
			if isinstance(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	):
					lowerCAmelCase__     :    Any					  =      prompt_encoder_config.to_dict()
			if isinstance(_SCREAMING_SNAKE_CASE	,      _SCREAMING_SNAKE_CASE	):
					lowerCAmelCase__     :    List[Any]					  =      mask_decoder_config.to_dict()

			lowerCAmelCase__     :    Any					  =      SamVisionConfig(**_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    List[Any]					  =      SamPromptEncoderConfig(**_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    Dict					  =      SamMaskDecoderConfig(**_SCREAMING_SNAKE_CASE	)
			lowerCAmelCase__     :    List[Any]					  =      initializer_range




	def 							UpperCAmelCase__(      self		:	Any	)-> int:
			lowerCAmelCase__     :    Optional[int]					  =      copy.deepcopy(self.__dict__	)
			lowerCAmelCase__     :    Any					  =      self.vision_config.to_dict()
			lowerCAmelCase__     :    Optional[Any]					  =      self.prompt_encoder_config.to_dict()
			lowerCAmelCase__     :    str					  =      self.mask_decoder_config.to_dict()
			lowerCAmelCase__     :    Any					  =      self.__class__.model_type
			return output




