



"""simple docstring"""







import tempfile
import unittest

import torch

from diffusers import UNetaDConditionModel
from diffusers.training_utils import EMAModel
from diffusers.utils.testing_utils import enable_full_determinism, skip_mps, torch_device


enable_full_determinism()





class    UpperCamelCase     (			unittest.TestCase    ):




	"""simple docstring"""
	SCREAMING_SNAKE_CASE_				:    Optional[Any]		    =   "hf-internal-testing/tiny-stable-diffusion-pipe"
	SCREAMING_SNAKE_CASE_				:    Tuple		    =   1
	SCREAMING_SNAKE_CASE_				:    Union[str, Any]		    =   7_7
	SCREAMING_SNAKE_CASE_				:    Dict		    =   3_2
	SCREAMING_SNAKE_CASE_				:    Tuple		    =   4
	SCREAMING_SNAKE_CASE_				:    Union[str, Any]		    =   6_4
	SCREAMING_SNAKE_CASE_				:    List[str]		    =   torch.manual_seed(0    )






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_=0.9999					):
			_lowercase							:      str			=					UNetaDConditionModel.from_pretrained(self.model_id							,subfolder="""unet"""					)
			_lowercase							:      List[Any]			=					unet.to(UpperCAmelCase_					)
			_lowercase							:      List[Any]			=					EMAModel(unet.parameters()							,decay=UpperCAmelCase_							,model_cls=UpperCAmelCase_							,model_config=unet.config					)
			return unet, ema_unet






	def 					lowerCamelCase__ (			self					):
			_lowercase							:      List[str]			=					torch.randn(
			    self.batch_size							,self.num_in_channels							,self.latent_height							,self.latent_width							,generator=self.generator					).to(UpperCAmelCase_					)
			_lowercase							:      Tuple			=					torch.randint(0							,10_00							,size=(self.batch_size,)							,generator=self.generator					).to(UpperCAmelCase_					)
			_lowercase							:      Tuple			=					torch.randn(
			    self.batch_size							,self.prompt_length							,self.text_encoder_hidden_dim							,generator=self.generator					).to(UpperCAmelCase_					)
			return noisy_latents, timesteps, encoder_hidden_states






	def 					lowerCamelCase__ (			self							,UpperCAmelCase_					):
			_lowercase							:      int			=					{}
			for k, param in unet.state_dict().items():
					_lowercase							:      Any			=					torch.randn_like(UpperCAmelCase_					) + (param * torch.randn_like(UpperCAmelCase_					))
					updated_state_dict.update({k: updated_param}					)
			unet.load_state_dict(UpperCAmelCase_					)
			return unet






	def 					lowerCamelCase__ (			self					):
			_lowercase			,   _lowercase							:      List[str]			=					self.get_models()
			# Take the first (hypothetical) EMA step.
			ema_unet.step(unet.parameters()					)
			assert ema_unet.optimization_step == 1

			# Take two more.
			for _ in range(2					):
					ema_unet.step(unet.parameters()					)
			assert ema_unet.optimization_step == 3






	def 					lowerCamelCase__ (			self					):
			_lowercase			,   _lowercase							:      Tuple			=					self.get_models()
			# Since the `unet` is not being updated (i.e., backprop'd)
			# there won't be any difference between the `params` of `unet`
			# and `ema_unet` even if we call `ema_unet.step(unet.parameters())`.
			ema_unet.step(unet.parameters()					)
			_lowercase							:      Any			=					list(unet.parameters()					)
			for s_param, param in zip(ema_unet.shadow_params							,UpperCAmelCase_					):
					assert torch.allclose(UpperCAmelCase_							,UpperCAmelCase_					)

			# The above holds true even if we call `ema.step()` multiple times since
			# `unet` params are still not being updated.
			for _ in range(4					):
					ema_unet.step(unet.parameters()					)
			for s_param, param in zip(ema_unet.shadow_params							,UpperCAmelCase_					):
					assert torch.allclose(UpperCAmelCase_							,UpperCAmelCase_					)






	def 					lowerCamelCase__ (			self					):
			_lowercase			,   _lowercase							:      Tuple			=					self.get_models()
			# Here we simulate the parameter updates for `unet`. Since there might
			# be some parameters which are initialized to zero we take extra care to
			# initialize their values to something non-zero before the multiplication.
			_lowercase							:      Any			=					self.simulate_backprop(UpperCAmelCase_					)

			# Take the EMA step.
			ema_unet.step(unet_pseudo_updated_step_one.parameters()					)

			# Now the EMA'd parameters won't be equal to the original model parameters.
			_lowercase							:      Tuple			=					list(unet_pseudo_updated_step_one.parameters()					)
			for s_param, param in zip(ema_unet.shadow_params							,UpperCAmelCase_					):
					assert ~torch.allclose(UpperCAmelCase_							,UpperCAmelCase_					)

			# Ensure this is the case when we take multiple EMA steps.
			for _ in range(4					):
					ema_unet.step(unet.parameters()					)
			for s_param, param in zip(ema_unet.shadow_params							,UpperCAmelCase_					):
					assert ~torch.allclose(UpperCAmelCase_							,UpperCAmelCase_					)






	def 					lowerCamelCase__ (			self					):
			# If we call EMA step after a backpropagation consecutively for two times,
			# the shadow params from those two steps should be different.
			_lowercase			,   _lowercase							:      Any			=					self.get_models()

			# First backprop + EMA
			_lowercase							:      Tuple			=					self.simulate_backprop(UpperCAmelCase_					)
			ema_unet.step(unet_step_one.parameters()					)
			_lowercase							:      List[Any]			=					ema_unet.shadow_params

			# Second backprop + EMA
			_lowercase							:      int			=					self.simulate_backprop(UpperCAmelCase_					)
			ema_unet.step(unet_step_two.parameters()					)
			_lowercase							:      Any			=					ema_unet.shadow_params

			for step_one, step_two in zip(UpperCAmelCase_							,UpperCAmelCase_					):
					assert ~torch.allclose(UpperCAmelCase_							,UpperCAmelCase_					)






	def 					lowerCamelCase__ (			self					):
			# If there's no decay even if there are backprops, EMA steps
			# won't take any effect i.e., the shadow params would remain the
			# same.
			_lowercase			,   _lowercase							:      Dict			=					self.get_models(decay=0.0					)
			_lowercase							:      str			=					self.simulate_backprop(UpperCAmelCase_					)
			ema_unet.step(unet_step_one.parameters()					)
			_lowercase							:      int			=					ema_unet.shadow_params

			_lowercase							:      Any			=					self.simulate_backprop(UpperCAmelCase_					)
			ema_unet.step(unet_step_two.parameters()					)
			_lowercase							:      Optional[Any]			=					ema_unet.shadow_params

			for step_one, step_two in zip(UpperCAmelCase_							,UpperCAmelCase_					):
					assert torch.allclose(UpperCAmelCase_							,UpperCAmelCase_					)






	@skip_mps
	def 					lowerCamelCase__ (			self					):
			_lowercase			,   _lowercase							:      Union[str, Any]			=					self.get_models()
			_lowercase			,   _lowercase			,   _lowercase							:      Any			=					self.get_dummy_inputs()

			with tempfile.TemporaryDirectory() as tmpdir:
					ema_unet.save_pretrained(UpperCAmelCase_					)
					_lowercase							:      List[str]			=					UNetaDConditionModel.from_pretrained(UpperCAmelCase_							,model_cls=UpperCAmelCase_					)
					_lowercase							:      List[str]			=					loaded_unet.to(unet.device					)

			# Since no EMA step has been performed the outputs should match.
			_lowercase							:      Any			=					unet(UpperCAmelCase_							,UpperCAmelCase_							,UpperCAmelCase_					).sample
			_lowercase							:      List[str]			=					loaded_unet(UpperCAmelCase_							,UpperCAmelCase_							,UpperCAmelCase_					).sample

			assert torch.allclose(UpperCAmelCase_							,UpperCAmelCase_							,atol=1E-4					)








"""simple docstring"""







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

from ...utils import logging


UpperCAmelCase:   Dict       			=							logging.get_logger(__name__)


def       __SCREAMING_SNAKE_CASE      (   __UpperCAmelCase					,							__UpperCAmelCase	):
		_lowercase							:      List[str]			=					nn.functional.normalize(__UpperCAmelCase	)
		_lowercase							:      int			=					nn.functional.normalize(__UpperCAmelCase	)
		return torch.mm(__UpperCAmelCase					,							normalized_text_embeds.t()	)





class    UpperCamelCase     (			snake_case    ):




	"""simple docstring"""
	SCREAMING_SNAKE_CASE_				:    List[str]		    =   CLIPConfig

	SCREAMING_SNAKE_CASE_				:    List[Any]		    =   ["CLIPEncoderLayer"]






	def __init__(			self							,UpperCAmelCase_					):
			super().__init__(UpperCAmelCase_					)

			_lowercase							:      Any			=					CLIPVisionModel(config.vision_config					)
			_lowercase							:      Any			=					nn.Linear(config.vision_config.hidden_size							,config.projection_dim							,bias=UpperCAmelCase_					)

			_lowercase							:      Dict			=					nn.Parameter(torch.ones(17							,config.projection_dim					)							,requires_grad=UpperCAmelCase_					)
			_lowercase							:      Tuple			=					nn.Parameter(torch.ones(3							,config.projection_dim					)							,requires_grad=UpperCAmelCase_					)

			_lowercase							:      Optional[int]			=					nn.Parameter(torch.ones(17					)							,requires_grad=UpperCAmelCase_					)
			_lowercase							:      List[Any]			=					nn.Parameter(torch.ones(3					)							,requires_grad=UpperCAmelCase_					)






	@torch.no_grad()
	def 					lowerCamelCase__ (			self							,UpperCAmelCase_							,UpperCAmelCase_					):
			_lowercase							:      List[str]			=					self.vision_model(UpperCAmelCase_					)[1]  # pooled_output
			_lowercase							:      str			=					self.visual_projection(UpperCAmelCase_					)

			# we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
			_lowercase							:      Union[str, Any]			=					cosine_distance(UpperCAmelCase_							,self.special_care_embeds					).cpu().float().numpy()
			_lowercase							:      str			=					cosine_distance(UpperCAmelCase_							,self.concept_embeds					).cpu().float().numpy()

			_lowercase							:      List[str]			=					[]
			_lowercase							:      Optional[Any]			=					image_embeds.shape[0]
			for i in range(UpperCAmelCase_					):
					_lowercase							:      int			=					{"""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
					_lowercase							:      str			=					0.0

					for concept_idx in range(len(special_cos_dist[0]					)					):
							_lowercase							:      Optional[Any]			=					special_cos_dist[i][concept_idx]
							_lowercase							:      str			=					self.special_care_embeds_weights[concept_idx].item()
							_lowercase							:      Tuple			=					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]}					)
									_lowercase							:      Union[str, Any]			=					0.01

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

					result.append(UpperCAmelCase_					)

			_lowercase							:      List[str]			=					[len(res["""bad_concepts"""]					) > 0 for res in result]

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

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

			return images, has_nsfw_concepts






	@torch.no_grad()
	def 					lowerCamelCase__ (			self							,UpperCAmelCase_							,UpperCAmelCase_					):
			_lowercase							:      List[Any]			=					self.vision_model(UpperCAmelCase_					)[1]  # pooled_output
			_lowercase							:      Optional[int]			=					self.visual_projection(UpperCAmelCase_					)

			_lowercase							:      Union[str, Any]			=					cosine_distance(UpperCAmelCase_							,self.special_care_embeds					)
			_lowercase							:      Tuple			=					cosine_distance(UpperCAmelCase_							,self.concept_embeds					)

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

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

			_lowercase							:      Optional[Any]			=					(cos_dist - self.concept_embeds_weights) + special_adjustment
			# concept_scores = concept_scores.round(decimals=3)
			_lowercase							:      Optional[int]			=					torch.any(concept_scores > 0							,dim=1					)

			_lowercase							:      Tuple			=					0.0  # black image

			return images, has_nsfw_concepts



