



import os
import sys

a      :  List[str]			     =			os.path.join(os.path.dirname(__file__), 'src')
sys.path.append(SRC_DIR)


from transformers import (
    AutoConfig,
    AutoModel,
    AutoModelForCausalLM,
    AutoModelForMaskedLM,
    AutoModelForQuestionAnswering,
    AutoModelForSequenceClassification,
    AutoTokenizer,
    add_start_docstrings,
)


a      :  Tuple			     =			[
    'torch',
    'numpy',
    'tokenizers',
    'filelock',
    'requests',
    'tqdm',
    'regex',
    'sentencepiece',
    'sacremoses',
    'importlib_metadata',
    'huggingface_hub',
]







@add_start_docstrings(AutoConfig.__doc__  )
def  lowerCAmelCase_	(*lowerCAmelCase__:		Optional[int]		,						**lowerCAmelCase__:		Optional[Any]  ):





			"""simple docstring"""


			return AutoConfig.from_pretrained(*lowerCAmelCase__		,						**lowerCAmelCase__  )







@add_start_docstrings(AutoTokenizer.__doc__  )
def  lowerCAmelCase_	(*lowerCAmelCase__:		Dict		,						**lowerCAmelCase__:		Tuple  ):





			"""simple docstring"""


			return AutoTokenizer.from_pretrained(*lowerCAmelCase__		,						**lowerCAmelCase__  )







@add_start_docstrings(AutoModel.__doc__  )
def  lowerCAmelCase_	(*lowerCAmelCase__:		List[Any]		,						**lowerCAmelCase__:		List[Any]  ):





			"""simple docstring"""


			return AutoModel.from_pretrained(*lowerCAmelCase__		,						**lowerCAmelCase__  )







@add_start_docstrings(AutoModelForCausalLM.__doc__  )
def  lowerCAmelCase_	(*lowerCAmelCase__:		Optional[int]		,						**lowerCAmelCase__:		Union[str, Any]  ):





			"""simple docstring"""


			return AutoModelForCausalLM.from_pretrained(*lowerCAmelCase__		,						**lowerCAmelCase__  )







@add_start_docstrings(AutoModelForMaskedLM.__doc__  )
def  lowerCAmelCase_	(*lowerCAmelCase__:		Tuple		,						**lowerCAmelCase__:		Optional[int]  ):





			"""simple docstring"""


			return AutoModelForMaskedLM.from_pretrained(*lowerCAmelCase__		,						**lowerCAmelCase__  )







@add_start_docstrings(AutoModelForSequenceClassification.__doc__  )
def  lowerCAmelCase_	(*lowerCAmelCase__:		Dict		,						**lowerCAmelCase__:		Dict  ):





			"""simple docstring"""


			return AutoModelForSequenceClassification.from_pretrained(*lowerCAmelCase__		,						**lowerCAmelCase__  )







@add_start_docstrings(AutoModelForQuestionAnswering.__doc__  )
def  lowerCAmelCase_	(*lowerCAmelCase__:		Optional[int]		,						**lowerCAmelCase__:		str  ):





			"""simple docstring"""


			return AutoModelForQuestionAnswering.from_pretrained(*lowerCAmelCase__		,						**lowerCAmelCase__  )











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

from ...utils import logging


a      :  Union[str, Any]			     =			logging.get_logger(__name__)







class 				_a				(   _lowerCAmelCase					):
			A      		=	CLIPConfig

			A      		=	['''CLIPEncoderLayer''']

			def __init__(self,   SCREAMING_SNAKE_CASE_					)       ->   Optional[Any]:
						super().__init__(SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		Tuple							      =     CLIPVisionModelWithProjection(config.vision_config					)

						UpperCAmelCase_:		Tuple							      =     nn.Linear(config.vision_config.projection_dim,   1					)
						UpperCAmelCase_:		Tuple							      =     nn.Linear(config.vision_config.projection_dim,   1					)

			@torch.no_grad()
			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_,   SCREAMING_SNAKE_CASE_=0.5,   SCREAMING_SNAKE_CASE_=0.5					)       ->   Tuple:
						UpperCAmelCase_:		Optional[Any]							      =     self.vision_model(SCREAMING_SNAKE_CASE_					)[0]

						UpperCAmelCase_:		List[str]							      =     self.p_head(SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Any							      =     nsfw_detected.flatten()
						UpperCAmelCase_:		List[Any]							      =     nsfw_detected > p_threshold
						UpperCAmelCase_:		Any							      =     nsfw_detected.tolist()

						if any(SCREAMING_SNAKE_CASE_					):
									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(SCREAMING_SNAKE_CASE_					):
									if nsfw_detected_:
												UpperCAmelCase_:		Tuple							      =     np.zeros(images[idx].shape					)

						UpperCAmelCase_:		str							      =     self.w_head(SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		List[Any]							      =     watermark_detected.flatten()
						UpperCAmelCase_:		Tuple							      =     watermark_detected > w_threshold
						UpperCAmelCase_:		int							      =     watermark_detected.tolist()

						if any(SCREAMING_SNAKE_CASE_					):
									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(SCREAMING_SNAKE_CASE_					):
									if watermark_detected_:
												UpperCAmelCase_:		Any							      =     np.zeros(images[idx].shape					)

						return images, nsfw_detected, watermark_detected






