


"""simple docstring"""







from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


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

a__    :							Union[str, Any]            =					{
    '''hustvl/yolos-small''': '''https://huggingface.co/hustvl/yolos-small/resolve/main/config.json''',
    # See all YOLOS models at https://huggingface.co/models?filter=yolos
}




class 			UpperCamelCase_     (			UpperCamelCase):

	"""simple docstring"""





	snake_case__    :  int    =      "yolos"




	def __init__(			self      :    Tuple      ,      UpperCAmelCase__      :    Tuple=7_6_8      ,      UpperCAmelCase__      :    List[str]=1_2      ,      UpperCAmelCase__      :    Union[str, Any]=1_2      ,      UpperCAmelCase__      :    List[str]=3_0_7_2      ,      UpperCAmelCase__      :    Tuple="gelu"      ,      UpperCAmelCase__      :    Tuple=0.0      ,      UpperCAmelCase__      :    List[Any]=0.0      ,      UpperCAmelCase__      :    int=0.02      ,      UpperCAmelCase__      :    Any=1E-12      ,      UpperCAmelCase__      :    Tuple=[5_1_2, 8_6_4]      ,      UpperCAmelCase__      :    List[str]=1_6      ,      UpperCAmelCase__      :    int=3      ,      UpperCAmelCase__      :    Optional[Any]=True      ,      UpperCAmelCase__      :    str=1_0_0      ,      UpperCAmelCase__      :    Optional[int]=True      ,      UpperCAmelCase__      :    Optional[int]=False      ,      UpperCAmelCase__      :    Optional[Any]=1      ,      UpperCAmelCase__      :    Optional[int]=5      ,      UpperCAmelCase__      :    Dict=2      ,      UpperCAmelCase__      :    Tuple=5      ,      UpperCAmelCase__      :    Union[str, Any]=2      ,      UpperCAmelCase__      :    Dict=0.1      ,      **UpperCAmelCase__      :    int      ,      )	->     List[Any]:
						super().__init__(**UpperCAmelCase__       )

						__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_act
						__SCREAMING_SNAKE_CASE						       =    hidden_dropout_prob
						__SCREAMING_SNAKE_CASE						       =    attention_probs_dropout_prob
						__SCREAMING_SNAKE_CASE						       =    initializer_range
						__SCREAMING_SNAKE_CASE						       =    layer_norm_eps
						__SCREAMING_SNAKE_CASE						       =    image_size
						__SCREAMING_SNAKE_CASE						       =    patch_size
						__SCREAMING_SNAKE_CASE						       =    num_channels
						__SCREAMING_SNAKE_CASE						       =    qkv_bias
						__SCREAMING_SNAKE_CASE						       =    num_detection_tokens
						__SCREAMING_SNAKE_CASE						       =    use_mid_position_embeddings
						__SCREAMING_SNAKE_CASE						       =    auxiliary_loss
						# Hungarian matcher
						__SCREAMING_SNAKE_CASE						       =    class_cost
						__SCREAMING_SNAKE_CASE						       =    bbox_cost
						__SCREAMING_SNAKE_CASE						       =    giou_cost
						# Loss coefficients
						__SCREAMING_SNAKE_CASE						       =    bbox_loss_coefficient
						__SCREAMING_SNAKE_CASE						       =    giou_loss_coefficient
						__SCREAMING_SNAKE_CASE						       =    eos_coefficient




class 			UpperCamelCase_     (			UpperCamelCase):

	"""simple docstring"""





	snake_case__    :  Optional[Any]    =      version.parse("1.11")




	@property
	def   UpperCAmelCase_				(			self      :    Tuple       )	->     Mapping[str, Mapping[int, str]]:
						return OrderedDict(
						    [
						        ("pixel_values", {0: "batch", 1: "num_channels", 2: "height", 3: "width"}),
						    ]       )




	@property
	def   UpperCAmelCase_				(			self      :    str       )	->     float:
						return 1E-4




	@property
	def   UpperCAmelCase_				(			self      :    Tuple       )	->     int:
						return 1_2




"""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
