






import operator as op


snake_case    				=		"""scaler.pt"""
snake_case    				=		"""pytorch_model"""
snake_case    				=		"""random_states"""
snake_case    				=		"""optimizer"""
snake_case    				=		"""scheduler"""
snake_case    				=		"""pytorch_model.bin"""
snake_case    				=		"""pytorch_model.bin.index.json"""
snake_case    				=		"""model.safetensors"""
snake_case    				=		"""model.safetensors.index.json"""
snake_case    				=		"""1.10.2"""
snake_case    				=		"""py38"""
snake_case    				=		"""4.17.0"""
snake_case    				=		["""ml.p3.16xlarge""", """ml.p3dn.24xlarge""", """ml.p4dn.24xlarge"""]
snake_case    				=		["""FULL_SHARD""", """SHARD_GRAD_OP""", """NO_SHARD""", """HYBRID_SHARD""", """HYBRID_SHARD_ZERO2"""]
snake_case    				=		["""TRANSFORMER_BASED_WRAP""", """SIZE_BASED_WRAP""", """NO_WRAP"""]
snake_case    				=		["""BACKWARD_PRE""", """BACKWARD_POST""", """NO_PREFETCH"""]
snake_case    				=		["""FULL_STATE_DICT""", """LOCAL_STATE_DICT""", """SHARDED_STATE_DICT"""]
snake_case    				=		"""2.0.1"""
snake_case    				=		["""pdsh""", """standard""", """openmpi""", """mvapich"""]
snake_case    				=		["""default""", """reduce-overhead""", """max-autotune"""]

snake_case    				=		{""">""": op.gt, """>=""": op.ge, """==""": op.eq, """!=""": op.ne, """<=""": op.le, """<""": op.lt}

# These are the args for `torch.distributed.launch` for pytorch < 1.9
snake_case    				=		[
    """nnodes""",
    """nproc_per_node""",
    """rdzv_backend""",
    """rdzv_endpoint""",
    """rdzv_id""",
    """rdzv_conf""",
    """standalone""",
    """max_restarts""",
    """monitor_interval""",
    """start_method""",
    """role""",
    """module""",
    """m""",
    """no_python""",
    """run_path""",
    """log_dir""",
    """r""",
    """redirects""",
    """t""",
    """tee""",
    """node_rank""",
    """master_addr""",
    """master_port""",
]

snake_case    				=		["""DEEPSPEED""", """MULTI_GPU""", """FSDP""", """MEGATRON_LM"""]
snake_case    				=		["""DEEPSPEED""", """MULTI_XPU""", """FSDP"""]














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

from ...utils import logging


snake_case    				=		logging.get_logger(__name__)


class  SCREAMING_SNAKE_CASE			(			lowerCAmelCase					):
     '''simple docstring'''


     UpperCamelCase_							:		Union[str, Any]			=	CLIPConfig

     UpperCamelCase_							:		Optional[int]			=	['''CLIPEncoderLayer''']


     def __init__(				self       :	List[Any]					, UpperCAmelCase_       :	CLIPConfig      ):
       super().__init__(UpperCAmelCase_      )

       SCREAMING_SNAKE_CASE       :      Any							=    CLIPVisionModelWithProjection(config.vision_config      )

       SCREAMING_SNAKE_CASE       :      List[Any]							=    nn.Linear(config.vision_config.projection_dim					, 1      )
       SCREAMING_SNAKE_CASE       :      List[str]							=    nn.Linear(config.vision_config.projection_dim					, 1      )






     @torch.no_grad()
     def 				_A       (				self       :	Union[str, Any]					, UpperCAmelCase_       :	int					, UpperCAmelCase_       :	Dict					, UpperCAmelCase_       :	List[Any]=0.5					, UpperCAmelCase_       :	Optional[Any]=0.5      ):
       SCREAMING_SNAKE_CASE       :      str							=    self.vision_model(UpperCAmelCase_      )[0]

       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    self.p_head(UpperCAmelCase_      )
       SCREAMING_SNAKE_CASE       :      List[str]							=    nsfw_detected.flatten()
       SCREAMING_SNAKE_CASE       :      List[str]							=    nsfw_detected > p_threshold
       SCREAMING_SNAKE_CASE       :      List[Any]							=    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       :      Dict							=    np.zeros(images[idx].shape      )

       SCREAMING_SNAKE_CASE       :      Any							=    self.w_head(UpperCAmelCase_      )
       SCREAMING_SNAKE_CASE       :      List[Any]							=    watermark_detected.flatten()
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    watermark_detected > w_threshold
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    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       :      Optional[int]							=    np.zeros(images[idx].shape      )

       return images, nsfw_detected, watermark_detected






