'''simple docstring'''




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

from ...utils import logging


__snake_case  : Dict		       =    logging.get_logger(__name__)





def 	_UpperCAmelCase       (				_UpperCamelCase       :      Optional[int],  _UpperCamelCase       :      str  )      -> Any:
    A_      =  nn.functional.normalize(_UpperCamelCase  )
    A_      =  nn.functional.normalize(_UpperCamelCase  )
    return torch.mm(_UpperCamelCase,  normalized_text_embeds.t()  )




class       __UpperCAmelCase							(							_UpperCamelCase	):





       '''simple docstring'''



       __lowercase   :   Optional[Any]  					=							CLIPConfig

       __lowercase   :   str  					=							['CLIPEncoderLayer']

       def __init__( self       ,	_SCREAMING_SNAKE_CASE					)						->	Any:
           super().__init__(_SCREAMING_SNAKE_CASE					)

           A_      =  CLIPVisionModel(config.vision_config					)
           A_      =  nn.Linear(config.vision_config.hidden_size       ,	config.projection_dim       ,	bias=_SCREAMING_SNAKE_CASE					)

           A_      =  nn.Parameter(torch.ones(17       ,	config.projection_dim					)       ,	requires_grad=_SCREAMING_SNAKE_CASE					)
           A_      =  nn.Parameter(torch.ones(3       ,	config.projection_dim					)       ,	requires_grad=_SCREAMING_SNAKE_CASE					)

           A_      =  nn.Parameter(torch.ones(17					)       ,	requires_grad=_SCREAMING_SNAKE_CASE					)
           A_      =  nn.Parameter(torch.ones(3					)       ,	requires_grad=_SCREAMING_SNAKE_CASE					)

       @torch.no_grad()
       def    __A			( self       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE					)						->	int:
           A_      =  self.vision_model(_SCREAMING_SNAKE_CASE					)[1]  # pooled_output
           A_      =  self.visual_projection(_SCREAMING_SNAKE_CASE					)

           # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
           A_      =  cosine_distance(_SCREAMING_SNAKE_CASE       ,	self.special_care_embeds					).cpu().float().numpy()
           A_      =  cosine_distance(_SCREAMING_SNAKE_CASE       ,	self.concept_embeds					).cpu().float().numpy()

           A_      =  []
           A_      =  image_embeds.shape[0]
           for i in range(_SCREAMING_SNAKE_CASE					):
               A_      =  {'''special_scores''': {}, '''special_care''': [], '''concept_scores''': {}, '''bad_concepts''': []}

               # increase this value to create a stronger `nfsw` filter
               # at the cost of increasing the possibility of filtering benign images
               A_      =  0.0

               for concept_idx in range(len(special_cos_dist[0]					)					):
                   A_      =  special_cos_dist[i][concept_idx]
                   A_      =  self.special_care_embeds_weights[concept_idx].item()
                   A_      =  round(concept_cos - concept_threshold + adjustment       ,	3					)
                   if result_img["special_scores"][concept_idx] > 0:
                       result_img["special_care"].append({concept_idx, result_img['''special_scores'''][concept_idx]}					)
                       A_      =  0.01

               for concept_idx in range(len(cos_dist[0]					)					):
                   A_      =  cos_dist[i][concept_idx]
                   A_      =  self.concept_embeds_weights[concept_idx].item()
                   A_      =  round(concept_cos - concept_threshold + adjustment       ,	3					)
                   if result_img["concept_scores"][concept_idx] > 0:
                       result_img["bad_concepts"].append(_SCREAMING_SNAKE_CASE					)

               result.append(_SCREAMING_SNAKE_CASE					)

           A_      =  [len(res['''bad_concepts''']					) > 0 for res in result]

           for idx, has_nsfw_concept in enumerate(_SCREAMING_SNAKE_CASE					):
               if has_nsfw_concept:
                   if torch.is_tensor(_SCREAMING_SNAKE_CASE					) or torch.is_tensor(images[0]					):
                       A_      =  torch.zeros_like(images[idx]					)  # black image
                   else:
                       A_      =  np.zeros(images[idx].shape					)  # black image

           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.'''					)

           return images, has_nsfw_concepts

       @torch.no_grad()
       def    __A			( self       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE					)						->	List[Any]:
           A_      =  self.vision_model(_SCREAMING_SNAKE_CASE					)[1]  # pooled_output
           A_      =  self.visual_projection(_SCREAMING_SNAKE_CASE					)

           A_      =  cosine_distance(_SCREAMING_SNAKE_CASE       ,	self.special_care_embeds					)
           A_      =  cosine_distance(_SCREAMING_SNAKE_CASE       ,	self.concept_embeds					)

           # increase this value to create a stronger `nsfw` filter
           # at the cost of increasing the possibility of filtering benign images
           A_      =  0.0

           A_      =  special_cos_dist - self.special_care_embeds_weights + adjustment
           # special_scores = special_scores.round(decimals=3)
           A_      =  torch.any(special_scores > 0       ,	dim=1					)
           A_      =  special_care * 0.01
           A_      =  special_adjustment.unsqueeze(1					).expand(-1       ,	cos_dist.shape[1]					)

           A_      =  (cos_dist - self.concept_embeds_weights) + special_adjustment
           # concept_scores = concept_scores.round(decimals=3)
           A_      =  torch.any(concept_scores > 0       ,	dim=1					)

           A_      =  0.0  # black image

           return images, has_nsfw_concepts







'''simple docstring'''




import tempfile
import unittest

import numpy as np
import torch

from diffusers import ScoreSdeVeScheduler




class       __UpperCAmelCase							(							unittest.TestCase	):





       '''simple docstring'''



       __lowercase   :   str  					=							(ScoreSdeVeScheduler,)
       __lowercase   :   List[str]  					=							()

       @property
       def    __A			( self					)						->	List[str]:
           A_      =  4
           A_      =  3
           A_      =  8
           A_      =  8

           A_      =  torch.rand((batch_size, num_channels, height, width)					)

           return sample

       @property
       def    __A			( self					)						->	Optional[Any]:
           A_      =  4
           A_      =  3
           A_      =  8
           A_      =  8

           A_      =  batch_size * num_channels * height * width
           A_      =  torch.arange(_SCREAMING_SNAKE_CASE					)
           A_      =  sample.reshape(_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE					)
           A_      =  sample / num_elems
           A_      =  sample.permute(3       ,	0       ,	1       ,	2					)

           return sample

       def    __A			( self					)						->	Union[str, Any]:
           def model(_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	*_SCREAMING_SNAKE_CASE					):
               return sample * t / (t + 1)

           return model

       def    __A			( self       ,	**_SCREAMING_SNAKE_CASE					)						->	Tuple:
           A_      =  {
               '''num_train_timesteps''': 2000,
               '''snr''': 0.15,
               '''sigma_min''': 0.01,
               '''sigma_max''': 1348,
               '''sampling_eps''': 1E-5,
           }

           config.update(**_SCREAMING_SNAKE_CASE					)
           return config

       def    __A			( self       ,	_SCREAMING_SNAKE_CASE=0       ,	**_SCREAMING_SNAKE_CASE					)						->	Any:
           A_      =  dict(self.forward_default_kwargs					)

           for scheduler_class in self.scheduler_classes:
               A_      =  self.dummy_sample
               A_      =  0.1 * sample

               A_      =  self.get_scheduler_config(**_SCREAMING_SNAKE_CASE					)
               A_      =  scheduler_class(**_SCREAMING_SNAKE_CASE					)

               with tempfile.TemporaryDirectory() as tmpdirname:
                   scheduler.save_config(_SCREAMING_SNAKE_CASE					)
                   A_      =  scheduler_class.from_pretrained(_SCREAMING_SNAKE_CASE					)

               A_      =  scheduler.step_pred(
                   _SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	generator=torch.manual_seed(0					)       ,	**_SCREAMING_SNAKE_CASE					).prev_sample
               A_      =  new_scheduler.step_pred(
                   _SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	generator=torch.manual_seed(0					)       ,	**_SCREAMING_SNAKE_CASE					).prev_sample

               assert torch.sum(torch.abs(output - new_output					)					) < 1E-5, "Scheduler outputs are not identical"

               A_      =  scheduler.step_correct(_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	generator=torch.manual_seed(0					)       ,	**_SCREAMING_SNAKE_CASE					).prev_sample
               A_      =  new_scheduler.step_correct(
                   _SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	generator=torch.manual_seed(0					)       ,	**_SCREAMING_SNAKE_CASE					).prev_sample

               assert torch.sum(torch.abs(output - new_output					)					) < 1E-5, "Scheduler correction are not identical"

       def    __A			( self       ,	_SCREAMING_SNAKE_CASE=0       ,	**_SCREAMING_SNAKE_CASE					)						->	Any:
           A_      =  dict(self.forward_default_kwargs					)
           kwargs.update(_SCREAMING_SNAKE_CASE					)

           for scheduler_class in self.scheduler_classes:
               A_      =  self.dummy_sample
               A_      =  0.1 * sample

               A_      =  self.get_scheduler_config()
               A_      =  scheduler_class(**_SCREAMING_SNAKE_CASE					)

               with tempfile.TemporaryDirectory() as tmpdirname:
                   scheduler.save_config(_SCREAMING_SNAKE_CASE					)
                   A_      =  scheduler_class.from_pretrained(_SCREAMING_SNAKE_CASE					)

               A_      =  scheduler.step_pred(
                   _SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	generator=torch.manual_seed(0					)       ,	**_SCREAMING_SNAKE_CASE					).prev_sample
               A_      =  new_scheduler.step_pred(
                   _SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	generator=torch.manual_seed(0					)       ,	**_SCREAMING_SNAKE_CASE					).prev_sample

               assert torch.sum(torch.abs(output - new_output					)					) < 1E-5, "Scheduler outputs are not identical"

               A_      =  scheduler.step_correct(_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	generator=torch.manual_seed(0					)       ,	**_SCREAMING_SNAKE_CASE					).prev_sample
               A_      =  new_scheduler.step_correct(
                   _SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	generator=torch.manual_seed(0					)       ,	**_SCREAMING_SNAKE_CASE					).prev_sample

               assert torch.sum(torch.abs(output - new_output					)					) < 1E-5, "Scheduler correction are not identical"

       def    __A			( self					)						->	Any:
           for timesteps in [10, 100, 1000]:
               self.check_over_configs(num_train_timesteps=_SCREAMING_SNAKE_CASE					)

       def    __A			( self					)						->	Dict:
           for sigma_min, sigma_max in zip([0.0_001, 0.001, 0.01]       ,	[1, 100, 1000]					):
               self.check_over_configs(sigma_min=_SCREAMING_SNAKE_CASE       ,	sigma_max=_SCREAMING_SNAKE_CASE					)

       def    __A			( self					)						->	Optional[int]:
           for t in [0.1, 0.5, 0.75]:
               self.check_over_forward(time_step=_SCREAMING_SNAKE_CASE					)

       def    __A			( self					)						->	int:
           A_      =  dict(self.forward_default_kwargs					)

           A_      =  self.scheduler_classes[0]
           A_      =  self.get_scheduler_config()
           A_      =  scheduler_class(**_SCREAMING_SNAKE_CASE					)

           A_      =  3

           A_      =  self.dummy_model()
           A_      =  self.dummy_sample_deter

           scheduler.set_sigmas(_SCREAMING_SNAKE_CASE					)
           scheduler.set_timesteps(_SCREAMING_SNAKE_CASE					)
           A_      =  torch.manual_seed(0					)

           for i, t in enumerate(scheduler.timesteps					):
               A_      =  scheduler.sigmas[i]

               for _ in range(scheduler.config.correct_steps					):
                   with torch.no_grad():
                       A_      =  model(_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE					)
                   A_      =  scheduler.step_correct(_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	generator=_SCREAMING_SNAKE_CASE       ,	**_SCREAMING_SNAKE_CASE					).prev_sample

               with torch.no_grad():
                   A_      =  model(_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE					)

               A_      =  scheduler.step_pred(_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	_SCREAMING_SNAKE_CASE       ,	generator=_SCREAMING_SNAKE_CASE       ,	**_SCREAMING_SNAKE_CASE					)
               A_  ,A_      =  output.prev_sample, output.prev_sample_mean

           A_      =  torch.sum(torch.abs(_SCREAMING_SNAKE_CASE					)					)
           A_      =  torch.mean(torch.abs(_SCREAMING_SNAKE_CASE					)					)

           assert np.isclose(result_sum.item()       ,	14_372_758_528.0					)
           assert np.isclose(result_mean.item()       ,	18_714_530.0					)

       def    __A			( self					)						->	List[Any]:
           A_      =  dict(self.forward_default_kwargs					)

           A_      =  kwargs.pop('''num_inference_steps'''       ,	_SCREAMING_SNAKE_CASE					)

           for scheduler_class in self.scheduler_classes:
               A_      =  self.get_scheduler_config()
               A_      =  scheduler_class(**_SCREAMING_SNAKE_CASE					)

               A_      =  self.dummy_sample
               A_      =  0.1 * sample

               if num_inference_steps is not None and hasattr(_SCREAMING_SNAKE_CASE       ,	'''set_timesteps'''					):
                   scheduler.set_timesteps(_SCREAMING_SNAKE_CASE					)
               elif num_inference_steps is not None and not hasattr(_SCREAMING_SNAKE_CASE       ,	'''set_timesteps'''					):
                   A_      =  num_inference_steps

               A_      =  scheduler.step_pred(_SCREAMING_SNAKE_CASE       ,	0       ,	_SCREAMING_SNAKE_CASE       ,	generator=torch.manual_seed(0					)       ,	**_SCREAMING_SNAKE_CASE					).prev_sample
               A_      =  scheduler.step_pred(_SCREAMING_SNAKE_CASE       ,	1       ,	_SCREAMING_SNAKE_CASE       ,	generator=torch.manual_seed(0					)       ,	**_SCREAMING_SNAKE_CASE					).prev_sample

               self.assertEqual(output_a.shape       ,	sample.shape					)
               self.assertEqual(output_a.shape       ,	output_a.shape					)






