"""simple docstring"""







from dataclasses import asdict, dataclass
from typing import Optional

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


a__    :							Tuple            =					logging.get_logger(__name__)

# TODO Update this
a__    :							Any            =					{
    '''facebook/esm-1b''': '''https://huggingface.co/facebook/esm-1b/resolve/main/config.json''',
    # See all ESM models at https://huggingface.co/models?filter=esm
}




class 			UpperCamelCase_     (			UpperCamelCase):

	"""simple docstring"""





	snake_case__    :  List[str]    =      "esm"




	def __init__(			self      :    str      ,      UpperCAmelCase__      :    List[str]=None      ,      UpperCAmelCase__      :    Optional[int]=None      ,      UpperCAmelCase__      :    List[Any]=None      ,      UpperCAmelCase__      :    str=7_6_8      ,      UpperCAmelCase__      :    List[str]=1_2      ,      UpperCAmelCase__      :    str=1_2      ,      UpperCAmelCase__      :    List[str]=3_0_7_2      ,      UpperCAmelCase__      :    Tuple=0.1      ,      UpperCAmelCase__      :    Optional[int]=0.1      ,      UpperCAmelCase__      :    Optional[int]=1_0_2_6      ,      UpperCAmelCase__      :    Union[str, Any]=0.02      ,      UpperCAmelCase__      :    Tuple=1E-12      ,      UpperCAmelCase__      :    Optional[Any]="absolute"      ,      UpperCAmelCase__      :    Any=True      ,      UpperCAmelCase__      :    List[Any]=None      ,      UpperCAmelCase__      :    Optional[Any]=False      ,      UpperCAmelCase__      :    Optional[int]=False      ,      UpperCAmelCase__      :    Dict=None      ,      UpperCAmelCase__      :    Any=None      ,      **UpperCAmelCase__      :    str      ,      )	->     int:
						super().__init__(pad_token_id=UpperCAmelCase__      ,      mask_token_id=UpperCAmelCase__      ,      **UpperCAmelCase__       )

						__SCREAMING_SNAKE_CASE						       =    vocab_size
						__SCREAMING_SNAKE_CASE						       =    hidden_size
						__SCREAMING_SNAKE_CASE						       =    num_hidden_layers
						__SCREAMING_SNAKE_CASE						       =    num_attention_heads
						__SCREAMING_SNAKE_CASE						       =    intermediate_size
						__SCREAMING_SNAKE_CASE						       =    hidden_dropout_prob
						__SCREAMING_SNAKE_CASE						       =    attention_probs_dropout_prob
						__SCREAMING_SNAKE_CASE						       =    max_position_embeddings
						__SCREAMING_SNAKE_CASE						       =    initializer_range
						__SCREAMING_SNAKE_CASE						       =    layer_norm_eps
						__SCREAMING_SNAKE_CASE						       =    position_embedding_type
						__SCREAMING_SNAKE_CASE						       =    use_cache
						__SCREAMING_SNAKE_CASE						       =    emb_layer_norm_before
						__SCREAMING_SNAKE_CASE						       =    token_dropout
						__SCREAMING_SNAKE_CASE						       =    is_folding_model
						if is_folding_model:
											if esmfold_config is None:
																logger.info("No esmfold_config supplied for folding model, using default values."       )
																__SCREAMING_SNAKE_CASE						       =    EsmFoldConfig()
											elif isinstance(UpperCAmelCase__      ,      UpperCAmelCase__       ):
																__SCREAMING_SNAKE_CASE						       =    EsmFoldConfig(**UpperCAmelCase__       )
											__SCREAMING_SNAKE_CASE						       =    esmfold_config
											if vocab_list is None:
																logger.warning("No vocab_list supplied for folding model, assuming the ESM-2 vocabulary!"       )
																__SCREAMING_SNAKE_CASE						       =    get_default_vocab_list()
											else:
																__SCREAMING_SNAKE_CASE						       =    vocab_list
						else:
											__SCREAMING_SNAKE_CASE						       =    None
											__SCREAMING_SNAKE_CASE						       =    None
						if self.esmfold_config is not None and getattr(self.esmfold_config      ,      "use_esm_attn_map"      ,      UpperCAmelCase__       ):
											raise ValueError("The HuggingFace port of ESMFold does not support use_esm_attn_map at this time!"       )




	def   UpperCAmelCase_				(			self      :    Optional[Any]       )	->     Dict:
						__SCREAMING_SNAKE_CASE						       =    super().to_dict()
						if isinstance(self.esmfold_config      ,      UpperCAmelCase__       ):
											__SCREAMING_SNAKE_CASE						       =    self.esmfold_config.to_dict()
						return output




@dataclass
class 			UpperCamelCase_     :

	"""simple docstring"""





	snake_case__    :  str    =      None
	snake_case__    :  bool    =      True
	snake_case__    :  bool    =      False
	snake_case__    :  bool    =      False
	snake_case__    :  bool    =      False
	snake_case__    :  float    =      0

	snake_case__    :  bool    =      True
	snake_case__    :  bool    =      False

	snake_case__    :  int    =      128
	snake_case__    :  "TrunkConfig"    =      None




	def   UpperCAmelCase_				(			self      :    List[Any]       )	->     List[str]:
						if self.trunk is None:
											__SCREAMING_SNAKE_CASE						       =    TrunkConfig()
						elif isinstance(self.trunk      ,      UpperCAmelCase__       ):
											__SCREAMING_SNAKE_CASE						       =    TrunkConfig(**self.trunk       )




	def   UpperCAmelCase_				(			self      :    Optional[int]       )	->     Any:
						__SCREAMING_SNAKE_CASE						       =    asdict(self       )
						__SCREAMING_SNAKE_CASE						       =    self.trunk.to_dict()
						return output




@dataclass
class 			UpperCamelCase_     :

	"""simple docstring"""





	snake_case__    :  int    =      48
	snake_case__    :  int    =      1024
	snake_case__    :  int    =      128
	snake_case__    :  int    =      32
	snake_case__    :  int    =      32
	snake_case__    :  int    =      32
	snake_case__    :  float    =      0
	snake_case__    :  float    =      0
	snake_case__    :  bool    =      False
	snake_case__    :  int    =      4
	snake_case__    :  Optional[int]    =      128
	snake_case__    :  "StructureModuleConfig"    =      None




	def   UpperCAmelCase_				(			self      :    Dict       )	->     str:
						if self.structure_module is None:
											__SCREAMING_SNAKE_CASE						       =    StructureModuleConfig()
						elif isinstance(self.structure_module      ,      UpperCAmelCase__       ):
											__SCREAMING_SNAKE_CASE						       =    StructureModuleConfig(**self.structure_module       )

						if self.max_recycles <= 0:
											raise ValueError(F"""`max_recycles` should be positive, got {self.max_recycles}."""       )
						if self.sequence_state_dim % self.sequence_state_dim != 0:
											raise ValueError(
											    "`sequence_state_dim` should be a round multiple of `sequence_state_dim`, got"
											    F""" {self.sequence_state_dim} and {self.sequence_state_dim}."""       )
						if self.pairwise_state_dim % self.pairwise_state_dim != 0:
											raise ValueError(
											    "`pairwise_state_dim` should be a round multiple of `pairwise_state_dim`, got"
											    F""" {self.pairwise_state_dim} and {self.pairwise_state_dim}."""       )

						__SCREAMING_SNAKE_CASE						       =    self.sequence_state_dim // self.sequence_head_width
						__SCREAMING_SNAKE_CASE						       =    self.pairwise_state_dim // self.pairwise_head_width

						if self.sequence_state_dim != sequence_num_heads * self.sequence_head_width:
											raise ValueError(
											    "`sequence_state_dim` should be equal to `sequence_num_heads * sequence_head_width, got"
											    F""" {self.sequence_state_dim} != {sequence_num_heads} * {self.sequence_head_width}."""       )
						if self.pairwise_state_dim != pairwise_num_heads * self.pairwise_head_width:
											raise ValueError(
											    "`pairwise_state_dim` should be equal to `pairwise_num_heads * pairwise_head_width, got"
											    F""" {self.pairwise_state_dim} != {pairwise_num_heads} * {self.pairwise_head_width}."""       )
						if self.pairwise_state_dim % 2 != 0:
											raise ValueError(F"""`pairwise_state_dim` should be even, got {self.pairwise_state_dim}."""       )

						if self.dropout >= 0.4:
											raise ValueError(F"""`dropout` should not be greater than 0.4, got {self.dropout}."""       )




	def   UpperCAmelCase_				(			self      :    int       )	->     Optional[int]:
						__SCREAMING_SNAKE_CASE						       =    asdict(self       )
						__SCREAMING_SNAKE_CASE						       =    self.structure_module.to_dict()
						return output




@dataclass
class 			UpperCamelCase_     :

	"""simple docstring"""





	snake_case__    :  int    =      384
	snake_case__    :  int    =      128
	snake_case__    :  int    =      16
	snake_case__    :  int    =      128
	snake_case__    :  int    =      12
	snake_case__    :  int    =      4
	snake_case__    :  int    =      8
	snake_case__    :  float    =      0.1
	snake_case__    :  int    =      8
	snake_case__    :  int    =      1
	snake_case__    :  int    =      2
	snake_case__    :  int    =      7
	snake_case__    :  int    =      10
	snake_case__    :  float    =      1E-8
	snake_case__    :  float    =      1E5




	def   UpperCAmelCase_				(			self      :    List[Any]       )	->     List[str]:
						return asdict(self       )




def        UpperCAmelCase__	():




					'''simple docstring'''


					return (
					    "<cls>",
					    "<pad>",
					    "<eos>",
					    "<unk>",
					    "L",
					    "A",
					    "G",
					    "V",
					    "S",
					    "E",
					    "R",
					    "T",
					    "I",
					    "D",
					    "P",
					    "K",
					    "Q",
					    "N",
					    "F",
					    "Y",
					    "M",
					    "H",
					    "W",
					    "C",
					    "X",
					    "B",
					    "U",
					    "Z",
					    "O",
					    ".",
					    "-",
					    "<null_1>",
					    "<mask>",
					)




"""simple docstring"""







import numpy as np
import torch
import torch.nn as nn
from transformers import CLIPConfig, CLIPVisionModelWithProjection, PreTrainedModel

from ...utils import logging


a__    :							Optional[Any]            =					logging.get_logger(__name__)




class 			UpperCamelCase_     (			UpperCamelCase):

	"""simple docstring"""





	snake_case__    :  int    =      CLIPConfig

	snake_case__    :  str    =      ["CLIPEncoderLayer"]




	def __init__(			self      :    Optional[int]      ,      UpperCAmelCase__      :    CLIPConfig       )	->     Dict:
						super().__init__(UpperCAmelCase__       )

						__SCREAMING_SNAKE_CASE						       =    CLIPVisionModelWithProjection(config.vision_config       )

						__SCREAMING_SNAKE_CASE						       =    nn.Linear(config.vision_config.projection_dim      ,      1       )
						__SCREAMING_SNAKE_CASE						       =    nn.Linear(config.vision_config.projection_dim      ,      1       )




	@torch.no_grad()
	def   UpperCAmelCase_				(			self      :    str      ,      UpperCAmelCase__      :    List[str]      ,      UpperCAmelCase__      :    str      ,      UpperCAmelCase__      :    int=0.5      ,      UpperCAmelCase__      :    Optional[int]=0.5       )	->     Union[str, Any]:
						__SCREAMING_SNAKE_CASE						       =    self.vision_model(UpperCAmelCase__       )[0]

						__SCREAMING_SNAKE_CASE						       =    self.p_head(UpperCAmelCase__       )
						__SCREAMING_SNAKE_CASE						       =    nsfw_detected.flatten()
						__SCREAMING_SNAKE_CASE						       =    nsfw_detected > p_threshold
						__SCREAMING_SNAKE_CASE						       =    nsfw_detected.tolist()

						if any(UpperCAmelCase__       ):
											logger.warning(
											    "Potential NSFW content was detected in one or more images. A black image will be returned instead."
											    " Try again with a different prompt and/or seed."       )

						for idx, nsfw_detected_ in enumerate(UpperCAmelCase__       ):
											if nsfw_detected_:
																__SCREAMING_SNAKE_CASE						       =    np.zeros(images[idx].shape       )

						__SCREAMING_SNAKE_CASE						       =    self.w_head(UpperCAmelCase__       )
						__SCREAMING_SNAKE_CASE						       =    watermark_detected.flatten()
						__SCREAMING_SNAKE_CASE						       =    watermark_detected > w_threshold
						__SCREAMING_SNAKE_CASE						       =    watermark_detected.tolist()

						if any(UpperCAmelCase__       ):
											logger.warning(
											    "Potential watermarked content was detected in one or more images. A black image will be returned instead."
											    " Try again with a different prompt and/or seed."       )

						for idx, watermark_detected_ in enumerate(UpperCAmelCase__       ):
											if watermark_detected_:
																__SCREAMING_SNAKE_CASE						       =    np.zeros(images[idx].shape       )

						return images, nsfw_detected, watermark_detected
