






import PIL.Image
import PIL.ImageOps
from packaging import version
from PIL import Image


if version.parse(version.parse(PIL.__version__).base_version) >= version.parse("""9.1.0"""):
   snake_case    				=		{
       """linear""": PIL.Image.Resampling.BILINEAR,
       """bilinear""": PIL.Image.Resampling.BILINEAR,
       """bicubic""": PIL.Image.Resampling.BICUBIC,
       """lanczos""": PIL.Image.Resampling.LANCZOS,
       """nearest""": PIL.Image.Resampling.NEAREST,
   }
else:
   snake_case    				=		{
       """linear""": PIL.Image.LINEAR,
       """bilinear""": PIL.Image.BILINEAR,
       """bicubic""": PIL.Image.BICUBIC,
       """lanczos""": PIL.Image.LANCZOS,
       """nearest""": PIL.Image.NEAREST,
   }


def   lowerCamelCase__			(     lowercase				):

  """simple docstring"""






  SCREAMING_SNAKE_CASE       :      Tuple							=    (images / 2 + 0.5).clamp(0						,     1				)
  SCREAMING_SNAKE_CASE       :      List[str]							=    images.cpu().permute(0						,     2						,     3						,     1				).float().numpy()
  SCREAMING_SNAKE_CASE       :      int							=    numpy_to_pil(lowercase				)
  return images


def   lowerCamelCase__			(     lowercase				):

  """simple docstring"""






  if images.ndim == 3:
    SCREAMING_SNAKE_CASE       :      Optional[int]							=    images[None, ...]
  SCREAMING_SNAKE_CASE       :      str							=    (images * 255).round().astype("uint8"				)
  if images.shape[-1] == 1:
    # special case for grayscale (single channel) images
    SCREAMING_SNAKE_CASE       :      Any							=    [Image.fromarray(image.squeeze()						,     mode="L"				) for image in images]
  else:
    SCREAMING_SNAKE_CASE       :      List[Any]							=    [Image.fromarray(lowercase				) for image in images]

  return pil_images














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






