






from typing import Dict, Iterable, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_DEFAULT_MEAN,
    IMAGENET_DEFAULT_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging


snake_case    				=		logging.get_logger(__name__)


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


     UpperCamelCase_							:		int			=	['''pixel_values''']


     def __init__(				self       :	Union[str, Any]					, UpperCAmelCase_       :	bool = True					, UpperCAmelCase_       :	Dict[str, int] = None					, UpperCAmelCase_       :	PILImageResampling = PILImageResampling.BICUBIC					, UpperCAmelCase_       :	bool = True					, UpperCAmelCase_       :	Dict[str, int] = None					, UpperCAmelCase_       :	bool = True					, UpperCAmelCase_       :	Union[int, float] = 1 / 255					, UpperCAmelCase_       :	bool = True					, UpperCAmelCase_       :	Optional[Union[float, Iterable[float]]] = IMAGENET_DEFAULT_MEAN					, UpperCAmelCase_       :	Optional[Union[float, Iterable[float]]] = IMAGENET_DEFAULT_STD					, **UpperCAmelCase_       :	Optional[int]					, ):
       super().__init__(**UpperCAmelCase_      )
       SCREAMING_SNAKE_CASE       :      str							=    size if size is not None else {"shortest_edge": 224}
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    get_size_dict(UpperCAmelCase_					, default_to_square=UpperCAmelCase_      )
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    crop_size if crop_size is not None else {"height": 224, "width": 224}
       SCREAMING_SNAKE_CASE       :      Tuple							=    get_size_dict(UpperCAmelCase_					, param_name="crop_size"      )

       SCREAMING_SNAKE_CASE       :      Dict							=    do_resize
       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    size
       SCREAMING_SNAKE_CASE       :      Any							=    resample
       SCREAMING_SNAKE_CASE       :      Dict							=    do_center_crop
       SCREAMING_SNAKE_CASE       :      Any							=    crop_size
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    do_rescale
       SCREAMING_SNAKE_CASE       :      int							=    rescale_factor
       SCREAMING_SNAKE_CASE       :      Any							=    do_normalize
       SCREAMING_SNAKE_CASE       :      Dict							=    image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN
       SCREAMING_SNAKE_CASE       :      List[Any]							=    image_std if image_std is not None else IMAGENET_DEFAULT_STD


     def 				_A       (				self       :	Dict					, UpperCAmelCase_       :	np.ndarray					, UpperCAmelCase_       :	Dict[str, int]					, UpperCAmelCase_       :	PILImageResampling = PILImageResampling.BICUBIC					, UpperCAmelCase_       :	Optional[Union[str, ChannelDimension]] = None					, **UpperCAmelCase_       :	Optional[Any]					, ):
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    get_size_dict(UpperCAmelCase_					, default_to_square=UpperCAmelCase_      )
       # size_dict is a dict with either keys "height" and "width" or "shortest_edge"
       if "shortest_edge" in size:
         SCREAMING_SNAKE_CASE       :      List[Any]							=    int((256 / 224) * size["shortest_edge"]      )
         SCREAMING_SNAKE_CASE       :      int							=    get_resize_output_image_size(UpperCAmelCase_					, size=UpperCAmelCase_					, default_to_square=UpperCAmelCase_      )
         SCREAMING_SNAKE_CASE       :      List[str]							=    {"height": output_size[0], "width": output_size[1]}
       if "height" not in size_dict or "width" not in size_dict:
         raise ValueError(
             f'''Size dict must have keys \'height\' and \'width\' or \'shortest_edge\'. Got {size_dict.keys()}'''      )
       return resize(
           UpperCAmelCase_					, size=(size_dict["height"], size_dict["width"])					, resample=UpperCAmelCase_					, data_format=UpperCAmelCase_					, **UpperCAmelCase_      )


     def 				_A       (				self       :	int					, UpperCAmelCase_       :	np.ndarray					, UpperCAmelCase_       :	Dict[str, int]					, UpperCAmelCase_       :	Optional[Union[str, ChannelDimension]] = None					, **UpperCAmelCase_       :	Optional[int]					, ):
       SCREAMING_SNAKE_CASE       :      Dict							=    get_size_dict(UpperCAmelCase_      )
       if "height" not in size or "width" not in size:
         raise ValueError(f'''Size dict must have keys \'height\' and \'width\'. Got {size.keys()}'''      )
       return center_crop(UpperCAmelCase_					, size=(size["height"], size["width"])					, data_format=UpperCAmelCase_					, **UpperCAmelCase_      )


     def 				_A       (				self       :	Union[str, Any]					, UpperCAmelCase_       :	np.ndarray					, UpperCAmelCase_       :	Union[int, float]					, UpperCAmelCase_       :	Optional[Union[str, ChannelDimension]] = None					, **UpperCAmelCase_       :	Any					, ):
       return rescale(UpperCAmelCase_					, scale=UpperCAmelCase_					, data_format=UpperCAmelCase_					, **UpperCAmelCase_      )


     def 				_A       (				self       :	Union[str, Any]					, UpperCAmelCase_       :	np.ndarray					, UpperCAmelCase_       :	Union[float, List[float]]					, UpperCAmelCase_       :	Union[float, List[float]]					, UpperCAmelCase_       :	Optional[Union[str, ChannelDimension]] = None					, **UpperCAmelCase_       :	Any					, ):
       return normalize(UpperCAmelCase_					, mean=UpperCAmelCase_					, std=UpperCAmelCase_					, data_format=UpperCAmelCase_					, **UpperCAmelCase_      )






     def 				_A       (				self       :	str					, UpperCAmelCase_       :	ImageInput					, UpperCAmelCase_       :	Optional[bool] = None					, UpperCAmelCase_       :	Optional[Dict[str, int]] = None					, UpperCAmelCase_       :	PILImageResampling = None					, UpperCAmelCase_       :	Optional[bool] = None					, UpperCAmelCase_       :	Optional[Dict[str, int]] = None					, UpperCAmelCase_       :	Optional[bool] = None					, UpperCAmelCase_       :	Optional[float] = None					, UpperCAmelCase_       :	Optional[bool] = None					, UpperCAmelCase_       :	Optional[Union[float, Iterable[float]]] = None					, UpperCAmelCase_       :	Optional[Union[float, Iterable[float]]] = None					, UpperCAmelCase_       :	Optional[TensorType] = None					, UpperCAmelCase_       :	ChannelDimension = ChannelDimension.FIRST					, **UpperCAmelCase_       :	Dict					, ):
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    do_resize if do_resize is not None else self.do_resize
       SCREAMING_SNAKE_CASE       :      int							=    resample if resample is not None else self.resample
       SCREAMING_SNAKE_CASE       :      Tuple							=    do_center_crop if do_center_crop is not None else self.do_center_crop
       SCREAMING_SNAKE_CASE       :      List[str]							=    do_rescale if do_rescale is not None else self.do_rescale
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    rescale_factor if rescale_factor is not None else self.rescale_factor
       SCREAMING_SNAKE_CASE       :      Optional[int]							=    do_normalize if do_normalize is not None else self.do_normalize
       SCREAMING_SNAKE_CASE       :      Union[str, Any]							=    image_mean if image_mean is not None else self.image_mean
       SCREAMING_SNAKE_CASE       :      Optional[Any]							=    image_std if image_std is not None else self.image_std

       SCREAMING_SNAKE_CASE       :      Optional[int]							=    size if size is not None else self.size
       SCREAMING_SNAKE_CASE       :      int							=    get_size_dict(UpperCAmelCase_					, default_to_square=UpperCAmelCase_      )
       SCREAMING_SNAKE_CASE       :      Tuple							=    crop_size if crop_size is not None else self.crop_size
       SCREAMING_SNAKE_CASE       :      Any							=    get_size_dict(UpperCAmelCase_					, param_name="crop_size"      )

       SCREAMING_SNAKE_CASE       :      Dict							=    make_list_of_images(UpperCAmelCase_      )

       if not valid_images(UpperCAmelCase_      ):
         raise ValueError(
             "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
             "torch.Tensor, tf.Tensor or jax.ndarray."      )

       if do_resize and size is None:
         raise ValueError("Size must be specified if do_resize is True."      )

       if do_center_crop and crop_size is None:
         raise ValueError("Crop size must be specified if do_center_crop is True."      )

       if do_rescale and rescale_factor is None:
         raise ValueError("Rescale factor must be specified if do_rescale is True."      )

       if do_normalize and (image_mean is None or image_std is None):
         raise ValueError("Image mean and std must be specified if do_normalize is True."      )

       # All transformations expect numpy arrays.
       SCREAMING_SNAKE_CASE       :      List[Any]							=    [to_numpy_array(UpperCAmelCase_      ) for image in images]

       if do_resize:
         SCREAMING_SNAKE_CASE       :      Tuple							=    [self.resize(UpperCAmelCase_					, UpperCAmelCase_					, UpperCAmelCase_      ) for image in images]

       if do_center_crop:
         SCREAMING_SNAKE_CASE       :      List[str]							=    [self.center_crop(UpperCAmelCase_					, UpperCAmelCase_      ) for image in images]

       if do_rescale:
         SCREAMING_SNAKE_CASE       :      int							=    [self.rescale(UpperCAmelCase_					, UpperCAmelCase_      ) for image in images]

       if do_normalize:
         SCREAMING_SNAKE_CASE       :      Dict							=    [self.normalize(UpperCAmelCase_					, UpperCAmelCase_					, UpperCAmelCase_      ) for image in images]

       SCREAMING_SNAKE_CASE       :      List[str]							=    [to_channel_dimension_format(UpperCAmelCase_					, UpperCAmelCase_      ) for image in images]

       SCREAMING_SNAKE_CASE       :      Optional[int]							=    {"pixel_values": images}
       return BatchFeature(data=UpperCAmelCase_					, tensor_type=UpperCAmelCase_      )














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






