



"""simple docstring"""






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

from ...utils import logging


_lowercase			:			int										=						logging.get_logger(__name__)

class 	__SCREAMING_SNAKE_CASE     (				lowerCAmelCase_       ):


					'''simple docstring'''





					_a             =							CLIPConfig

					_a             =							['CLIPEncoderLayer']



					def __init__(	self							:  int,  lowerCamelCase							:  CLIPConfig					)-> List[Any]:
								super().__init__(lowerCamelCase					)

								lowerCamelCase__		:		Dict		     =CLIPVisionModelWithProjection(config.vision_config					)

								lowerCamelCase__		:		Dict		     =nn.Linear(config.vision_config.projection_dim,  1					)
								lowerCamelCase__		:		Union[str, Any]		     =nn.Linear(config.vision_config.projection_dim,  1					)




					@torch.no_grad()
					def 		snake_case	(	self							:  Dict,  lowerCamelCase							:  Optional[Any],  lowerCamelCase							:  List[str],  lowerCamelCase							:  List[Any]=0.5,  lowerCamelCase							:  Optional[Any]=0.5					)-> Optional[int]:
								lowerCamelCase__		:		Dict		     =self.vision_model(lowerCamelCase					)[0]

								lowerCamelCase__		:		List[Any]		     =self.p_head(lowerCamelCase					)
								lowerCamelCase__		:		Optional[Any]		     =nsfw_detected.flatten()
								lowerCamelCase__		:		Any		     =nsfw_detected > p_threshold
								lowerCamelCase__		:		Optional[Any]		     =nsfw_detected.tolist()

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

								lowerCamelCase__		:		str		     =self.w_head(lowerCamelCase					)
								lowerCamelCase__		:		List[str]		     =watermark_detected.flatten()
								lowerCamelCase__		:		Optional[Any]		     =watermark_detected > w_threshold
								lowerCamelCase__		:		int		     =watermark_detected.tolist()

								if any(lowerCamelCase					):
											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(lowerCamelCase					):
											if watermark_detected_:
														lowerCamelCase__		:		Optional[int]		     =np.zeros(images[idx].shape					)

								return images, nsfw_detected, watermark_detected









"""simple docstring"""






import unittest

import numpy as np
import torch

from diffusers import VersatileDiffusionImageVariationPipeline
from diffusers.utils.testing_utils import load_image, require_torch_gpu, slow, torch_device


_lowercase			:			Tuple										=						False

class 	__SCREAMING_SNAKE_CASE     (				unittest.TestCase       ):


					'''simple docstring'''





					pass





@slow
@require_torch_gpu
class 	__SCREAMING_SNAKE_CASE     (				unittest.TestCase       ):


					'''simple docstring'''





					def 		snake_case	(	self							:  List[Any]					)-> Dict:
								lowerCamelCase__		:		str		     =VersatileDiffusionImageVariationPipeline.from_pretrained('''shi-labs/versatile-diffusion'''					)
								pipe.to(lowerCamelCase					)
								pipe.set_progress_bar_config(disable=lowerCamelCase					)

								lowerCamelCase__		:		int		     =load_image(
								    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/versatile_diffusion/benz.jpg'''					)
								lowerCamelCase__		:		Dict		     =torch.manual_seed(0					)
								lowerCamelCase__		:		str		     =pipe(
								    image=lowerCamelCase,  generator=lowerCamelCase,  guidance_scale=7.5,  num_inference_steps=50,  output_type='''numpy''',  ).images

								lowerCamelCase__		:		Dict		     =image[0, 253:256, 253:256, -1]

								assert image.shape == (1, 512, 512, 3)
								lowerCamelCase__		:		List[Any]		     =np.array([0.0_441, 0.0_469, 0.0_507, 0.0_575, 0.0_632, 0.0_650, 0.0_865, 0.0_909, 0.0_945]					)

								assert np.abs(image_slice.flatten() - expected_slice					).max() < 1E-2




