


import gc
import os
import tempfile
import unittest

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNetaDConditionModel
from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin
from diffusers.models.attention_processor import (
    Attention,
    AttnProcessor,
    AttnProcessor2_0,
    LoRAAttnProcessor,
    LoRAAttnProcessor2_0,
    LoRAXFormersAttnProcessor,
    XFormersAttnProcessor,
)
from diffusers.utils import TEXT_ENCODER_ATTN_MODULE, floats_tensor, torch_device
def  __SCREAMING_SNAKE_CASE       (				lowercase__  :List[str]				)       ->  int:
						__magic_name__				:							Optional[Any]				=      {}
						for name in unet.attn_processors.keys():
												__magic_name__				:							Optional[Any]				=      None if name.endswith('attn1.processor'				) else unet.config.cross_attention_dim
												if name.startswith('mid_block'				):
																		__magic_name__				:							Optional[int]				=      unet.config.block_out_channels[-1]
												elif name.startswith('up_blocks'				):
																		__magic_name__				:							int				=      int(name[len('up_blocks.'				)]				)
																		__magic_name__				:							str				=      list(reversed(unet.config.block_out_channels				)				)[block_id]
												elif name.startswith('down_blocks'				):
																		__magic_name__				:							Tuple				=      int(name[len('down_blocks.'				)]				)
																		__magic_name__				:							List[str]				=      unet.config.block_out_channels[block_id]
												__magic_name__				:							str				=      (
												    LoRAAttnProcessor2_0 if hasattr(lowercase__      ,						'scaled_dot_product_attention'				) else LoRAAttnProcessor
												)
												__magic_name__				:							Optional[Any]				=      lora_attn_processor_class(
												    hidden_size=lowercase__      ,						cross_attention_dim=lowercase__				)
						__magic_name__				:							Optional[int]				=      AttnProcsLayers(lowercase__				)
						return lora_attn_procs, unet_lora_layers
def  __SCREAMING_SNAKE_CASE       (				lowercase__  :List[Any]				)       ->  Optional[int]:
						__magic_name__				:							Optional[Any]				=      {}
						__magic_name__				:							Any				=      (
						    LoRAAttnProcessor2_0 if hasattr(lowercase__      ,						'scaled_dot_product_attention'				) else LoRAAttnProcessor
						)
						for name, module in text_encoder.named_modules():
												if name.endswith(lowercase__				):
																		__magic_name__				:							Optional[int]				=      lora_attn_processor_class(
																		    hidden_size=module.out_proj.out_features      ,						cross_attention_dim=lowercase__				)
						return text_lora_attn_procs
def  __SCREAMING_SNAKE_CASE       (				lowercase__  :Any				)       ->  Dict:
						__magic_name__				:							str				=      create_text_encoder_lora_attn_procs(lowercase__				)
						__magic_name__				:							List[str]				=      AttnProcsLayers(lowercase__				)
						return text_encoder_lora_layers
def  __SCREAMING_SNAKE_CASE       (				lowercase__  :Optional[int]      ,						lowercase__  :str=False				)       ->  Tuple:







						for _, attn_proc in text_lora_attn_procs.items():
												# set up.weights
												for layer_name, layer_module in attn_proc.named_modules():
																		if layer_name.endswith('_lora'				):
																								__magic_name__				:							Any				=      (
																								    torch.randn_like(layer_module.up.weight				)
																								    if randn_weight
																								    else torch.zeros_like(layer_module.up.weight				)
																								)
																								__magic_name__				:							str				=      torch.nn.Parameter(lowercase__				)



class     snake_case__  (      unittest.TestCase  ):
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										torch.manual_seed(0      )
										__magic_name__				:							Dict				=      UNetaDConditionModel(
										    block_out_channels=(32, 64)     ,      layers_per_block=2     ,      sample_size=32     ,      in_channels=4     ,      out_channels=4     ,      down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')     ,      up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')     ,      cross_attention_dim=32     ,      )
										__magic_name__				:							Optional[Any]				=      DDIMScheduler(
										    beta_start=0.00085     ,      beta_end=0.012     ,      beta_schedule='scaled_linear'     ,      clip_sample=UpperCamelCase__     ,      set_alpha_to_one=UpperCamelCase__     ,      steps_offset=1     ,      )
										torch.manual_seed(0      )
										__magic_name__				:							Optional[int]				=      AutoencoderKL(
										    block_out_channels=[32, 64]     ,      in_channels=3     ,      out_channels=3     ,      down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']     ,      up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']     ,      latent_channels=4     ,      )
										__magic_name__				:							Tuple				=      CLIPTextConfig(
										    bos_token_id=0     ,      eos_token_id=2     ,      hidden_size=32     ,      intermediate_size=37     ,      layer_norm_eps=1E-0_5     ,      num_attention_heads=4     ,      num_hidden_layers=5     ,      pad_token_id=1     ,      vocab_size=1000     ,      )
										__magic_name__				:							List[str]				=      CLIPTextModel(UpperCamelCase__      )
										__magic_name__				:							int				=      CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'      )

										__magic_name__		, __magic_name__				:							Any				=      create_unet_lora_layers(UpperCamelCase__      )
										__magic_name__				:							Tuple				=      create_text_encoder_lora_layers(UpperCamelCase__      )

										__magic_name__				:							List[Any]				=      {
										    'unet': unet,
										    'scheduler': scheduler,
										    'vae': vae,
										    'text_encoder': text_encoder,
										    'tokenizer': tokenizer,
										    'safety_checker': None,
										    'feature_extractor': None,
										}
										__magic_name__				:							Tuple				=      {
										    'unet_lora_layers': unet_lora_layers,
										    'text_encoder_lora_layers': text_encoder_lora_layers,
										    'unet_lora_attn_procs': unet_lora_attn_procs,
										}
										return pipeline_components, lora_components
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										__magic_name__				:							Any				=      1
										__magic_name__				:							Optional[int]				=      10
										__magic_name__				:							List[Any]				=      4
										__magic_name__				:							Optional[Any]				=      (32, 32)

										__magic_name__				:							Union[str, Any]				=      torch.manual_seed(0      )
										__magic_name__				:							Union[str, Any]				=      floats_tensor((batch_size, num_channels) + sizes      )
										__magic_name__				:							int				=      torch.randint(1     ,      UpperCamelCase__     ,      size=(batch_size, sequence_length)     ,      generator=UpperCamelCase__      )

										__magic_name__				:							Optional[int]				=      {
										    'prompt': 'A painting of a squirrel eating a burger',
										    'generator': generator,
										    'num_inference_steps': 2,
										    'guidance_scale': 6.0,
										    'output_type': 'numpy',
										}

										return noise, input_ids, pipeline_inputs

										# copied from: https://colab.research.google.com/gist/sayakpaul/df2ef6e1ae6d8c10a49d859883b10860/scratchpad.ipynb
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										__magic_name__				:							str				=      77

										__magic_name__				:							Optional[int]				=      torch.randint(2     ,      56     ,      size=(1, max_seq_length)     ,      generator=torch.manual_seed(0      )      )

										__magic_name__				:							str				=      {}
										__magic_name__				:							int				=      inputs
										return prepared_inputs
				def  UpperCamelCase_						(  self     ,      _UpperCAmelCase      ):




										"""simple docstring"""




										__magic_name__		, __magic_name__				:							List[str]				=      self.get_dummy_components()
										LoraLoaderMixin.save_lora_weights(
										    save_directory=UpperCamelCase__     ,      unet_lora_layers=lora_components['unet_lora_layers']     ,      text_encoder_lora_layers=lora_components['text_encoder_lora_layers']     ,      )
										self.assertTrue(os.path.isfile(os.path.join(UpperCamelCase__     ,      'pytorch_lora_weights.bin'      )      )      )
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										__magic_name__		, __magic_name__				:							Optional[int]				=      self.get_dummy_components()
										__magic_name__				:							Union[str, Any]				=      StableDiffusionPipeline(**UpperCamelCase__      )
										__magic_name__				:							Tuple				=      sd_pipe.to(UpperCamelCase__      )
										sd_pipe.set_progress_bar_config(disable=UpperCamelCase__      )

										__magic_name__		, __magic_name__		, __magic_name__				:							Dict				=      self.get_dummy_inputs()

										__magic_name__				:							Any				=      sd_pipe(**UpperCamelCase__      ).images
										__magic_name__				:							Optional[int]				=      original_images[0, -3:, -3:, -1]

										with tempfile.TemporaryDirectory() as tmpdirname:
																LoraLoaderMixin.save_lora_weights(
																    save_directory=UpperCamelCase__     ,      unet_lora_layers=lora_components['unet_lora_layers']     ,      text_encoder_lora_layers=lora_components['text_encoder_lora_layers']     ,      )
																self.assertTrue(os.path.isfile(os.path.join(UpperCamelCase__     ,      'pytorch_lora_weights.bin'      )      )      )
																sd_pipe.load_lora_weights(UpperCamelCase__      )

										__magic_name__				:							Optional[int]				=      sd_pipe(**UpperCamelCase__      ).images
										__magic_name__				:							str				=      lora_images[0, -3:, -3:, -1]

										# Outputs shouldn't match.
										self.assertFalse(torch.allclose(torch.from_numpy(UpperCamelCase__      )     ,      torch.from_numpy(UpperCamelCase__      )      )      )
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										__magic_name__		, __magic_name__				:							str				=      self.get_dummy_components()
										__magic_name__				:							List[str]				=      StableDiffusionPipeline(**UpperCamelCase__      )
										__magic_name__				:							str				=      sd_pipe.to(UpperCamelCase__      )
										sd_pipe.set_progress_bar_config(disable=UpperCamelCase__      )

										__magic_name__		, __magic_name__		, __magic_name__				:							str				=      self.get_dummy_inputs()

										__magic_name__				:							int				=      sd_pipe(**UpperCamelCase__      ).images
										__magic_name__				:							Union[str, Any]				=      original_images[0, -3:, -3:, -1]

										with tempfile.TemporaryDirectory() as tmpdirname:
																LoraLoaderMixin.save_lora_weights(
																    save_directory=UpperCamelCase__     ,      unet_lora_layers=lora_components['unet_lora_layers']     ,      text_encoder_lora_layers=lora_components['text_encoder_lora_layers']     ,      safe_serialization=UpperCamelCase__     ,      )
																self.assertTrue(os.path.isfile(os.path.join(UpperCamelCase__     ,      'pytorch_lora_weights.safetensors'      )      )      )
																sd_pipe.load_lora_weights(UpperCamelCase__      )

										__magic_name__				:							Optional[Any]				=      sd_pipe(**UpperCamelCase__      ).images
										__magic_name__				:							int				=      lora_images[0, -3:, -3:, -1]

										# Outputs shouldn't match.
										self.assertFalse(torch.allclose(torch.from_numpy(UpperCamelCase__      )     ,      torch.from_numpy(UpperCamelCase__      )      )      )
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										__magic_name__		, __magic_name__				:							Dict				=      self.get_dummy_components()
										__magic_name__				:							List[str]				=      lora_components['unet_lora_attn_procs']
										__magic_name__				:							int				=      StableDiffusionPipeline(**UpperCamelCase__      )
										__magic_name__				:							str				=      sd_pipe.to(UpperCamelCase__      )
										sd_pipe.set_progress_bar_config(disable=UpperCamelCase__      )

										__magic_name__		, __magic_name__		, __magic_name__				:							List[str]				=      self.get_dummy_inputs()

										__magic_name__				:							int				=      sd_pipe(**UpperCamelCase__      ).images
										__magic_name__				:							Optional[int]				=      original_images[0, -3:, -3:, -1]

										with tempfile.TemporaryDirectory() as tmpdirname:
																__magic_name__				:							Any				=      sd_pipe.unet
																unet.set_attn_processor(UpperCamelCase__      )
																unet.save_attn_procs(UpperCamelCase__      )
																self.assertTrue(os.path.isfile(os.path.join(UpperCamelCase__     ,      'pytorch_lora_weights.bin'      )      )      )
																sd_pipe.load_lora_weights(UpperCamelCase__      )

										__magic_name__				:							int				=      sd_pipe(**UpperCamelCase__      ).images
										__magic_name__				:							List[Any]				=      lora_images[0, -3:, -3:, -1]

										# Outputs shouldn't match.
										self.assertFalse(torch.allclose(torch.from_numpy(UpperCamelCase__      )     ,      torch.from_numpy(UpperCamelCase__      )      )      )
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										__magic_name__		, __magic_name__				:							str				=      self.get_dummy_components()
										__magic_name__				:							Any				=      StableDiffusionPipeline(**UpperCamelCase__      )

										__magic_name__				:							int				=      self.get_dummy_tokens()

										# inference without lora
										__magic_name__				:							Optional[int]				=      pipe.text_encoder(**UpperCamelCase__      )[0]
										assert outputs_without_lora.shape == (1, 77, 32)

										# create lora_attn_procs with zeroed out up.weights
										__magic_name__				:							str				=      create_text_encoder_lora_attn_procs(pipe.text_encoder      )
										set_lora_up_weights(UpperCamelCase__     ,      randn_weight=UpperCamelCase__      )

										# monkey patch
										pipe._modify_text_encoder(UpperCamelCase__      )

										# verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
										del text_attn_procs
										gc.collect()

										# inference with lora
										__magic_name__				:							List[Any]				=      pipe.text_encoder(**UpperCamelCase__      )[0]
										assert outputs_with_lora.shape == (1, 77, 32)

										assert torch.allclose(
										    UpperCamelCase__     ,      UpperCamelCase__      ), "lora_up_weight are all zero, so the lora outputs should be the same to without lora outputs"

										# create lora_attn_procs with randn up.weights
										__magic_name__				:							List[str]				=      create_text_encoder_lora_attn_procs(pipe.text_encoder      )
										set_lora_up_weights(UpperCamelCase__     ,      randn_weight=UpperCamelCase__      )

										# monkey patch
										pipe._modify_text_encoder(UpperCamelCase__      )

										# verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
										del text_attn_procs
										gc.collect()

										# inference with lora
										__magic_name__				:							Optional[int]				=      pipe.text_encoder(**UpperCamelCase__      )[0]
										assert outputs_with_lora.shape == (1, 77, 32)

										assert not torch.allclose(
										    UpperCamelCase__     ,      UpperCamelCase__      ), "lora_up_weight are not zero, so the lora outputs should be different to without lora outputs"
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										__magic_name__		, __magic_name__				:							str				=      self.get_dummy_components()
										__magic_name__				:							int				=      StableDiffusionPipeline(**UpperCamelCase__      )

										__magic_name__				:							Tuple				=      self.get_dummy_tokens()

										# inference without lora
										__magic_name__				:							Optional[int]				=      pipe.text_encoder(**UpperCamelCase__      )[0]
										assert outputs_without_lora.shape == (1, 77, 32)

										# create lora_attn_procs with randn up.weights
										__magic_name__				:							Tuple				=      create_text_encoder_lora_attn_procs(pipe.text_encoder      )
										set_lora_up_weights(UpperCamelCase__     ,      randn_weight=UpperCamelCase__      )

										# monkey patch
										pipe._modify_text_encoder(UpperCamelCase__      )

										# verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
										del text_attn_procs
										gc.collect()

										# inference with lora
										__magic_name__				:							Optional[int]				=      pipe.text_encoder(**UpperCamelCase__      )[0]
										assert outputs_with_lora.shape == (1, 77, 32)

										assert not torch.allclose(
										    UpperCamelCase__     ,      UpperCamelCase__      ), "lora outputs should be different to without lora outputs"

										# remove monkey patch
										pipe._remove_text_encoder_monkey_patch()

										# inference with removed lora
										__magic_name__				:							Optional[int]				=      pipe.text_encoder(**UpperCamelCase__      )[0]
										assert outputs_without_lora_removed.shape == (1, 77, 32)

										assert torch.allclose(
										    UpperCamelCase__     ,      UpperCamelCase__      ), "remove lora monkey patch should restore the original outputs"
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										__magic_name__		, __magic_name__				:							List[Any]				=      self.get_dummy_components()
										__magic_name__				:							Union[str, Any]				=      StableDiffusionPipeline(**UpperCamelCase__      )
										__magic_name__				:							List[Any]				=      sd_pipe.to(UpperCamelCase__      )
										sd_pipe.set_progress_bar_config(disable=UpperCamelCase__      )

										__magic_name__		, __magic_name__		, __magic_name__				:							str				=      self.get_dummy_inputs()

										with tempfile.TemporaryDirectory() as tmpdirname:
																LoraLoaderMixin.save_lora_weights(
																    save_directory=UpperCamelCase__     ,      unet_lora_layers=lora_components['unet_lora_layers']     ,      text_encoder_lora_layers=lora_components['text_encoder_lora_layers']     ,      )
																self.assertTrue(os.path.isfile(os.path.join(UpperCamelCase__     ,      'pytorch_lora_weights.bin'      )      )      )
																sd_pipe.load_lora_weights(UpperCamelCase__      )

										__magic_name__				:							Union[str, Any]				=      sd_pipe(**UpperCamelCase__      ).images
										__magic_name__				:							int				=      lora_images[0, -3:, -3:, -1]

										__magic_name__				:							Dict				=      sd_pipe(**UpperCamelCase__     ,      cross_attention_kwargs={'scale': 0.5}      ).images
										__magic_name__				:							Tuple				=      lora_images_with_scale[0, -3:, -3:, -1]

										# Outputs shouldn't match.
										self.assertFalse(
										    torch.allclose(torch.from_numpy(UpperCamelCase__      )     ,      torch.from_numpy(UpperCamelCase__      )      )      )
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										with tempfile.TemporaryDirectory() as tmpdirname:
																self.create_lora_weight_file(UpperCamelCase__      )

																__magic_name__		, __magic_name__				:							int				=      self.get_dummy_components()
																__magic_name__				:							str				=      StableDiffusionPipeline(**UpperCamelCase__      )
																__magic_name__				:							Any				=      sd_pipe.to(UpperCamelCase__      )
																sd_pipe.set_progress_bar_config(disable=UpperCamelCase__      )

																# check if vanilla attention processors are used
																for _, module in sd_pipe.unet.named_modules():
																						if isinstance(UpperCamelCase__     ,      UpperCamelCase__      ):
																												self.assertIsInstance(module.processor     ,      (AttnProcessor, AttnProcessor2_0)      )

            # load LoRA weight file
																sd_pipe.load_lora_weights(UpperCamelCase__      )

																# check if lora attention processors are used
																for _, module in sd_pipe.unet.named_modules():
																						if isinstance(UpperCamelCase__     ,      UpperCamelCase__      ):
																												__magic_name__				:							Any				=      (
																												    LoRAAttnProcessor2_0 if hasattr(UpperCamelCase__     ,      'scaled_dot_product_attention'      ) else LoRAAttnProcessor
																												)
																												self.assertIsInstance(module.processor     ,      UpperCamelCase__      )
				@unittest.skipIf(torch_device != 'cuda'     ,      'This test is supposed to run on GPU'      )
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										with tempfile.TemporaryDirectory() as tmpdirname:
																self.create_lora_weight_file(UpperCamelCase__      )

																__magic_name__		, __magic_name__				:							Tuple				=      self.get_dummy_components()
																__magic_name__				:							Union[str, Any]				=      StableDiffusionPipeline(**UpperCamelCase__      )
																__magic_name__				:							List[Any]				=      sd_pipe.to(UpperCamelCase__      )
																sd_pipe.set_progress_bar_config(disable=UpperCamelCase__      )

																# enable XFormers
																sd_pipe.enable_xformers_memory_efficient_attention()

																# check if xFormers attention processors are used
																for _, module in sd_pipe.unet.named_modules():
																						if isinstance(UpperCamelCase__     ,      UpperCamelCase__      ):
																												self.assertIsInstance(module.processor     ,      UpperCamelCase__      )

            # load LoRA weight file
																sd_pipe.load_lora_weights(UpperCamelCase__      )

																# check if lora attention processors are used
																for _, module in sd_pipe.unet.named_modules():
																						if isinstance(UpperCamelCase__     ,      UpperCamelCase__      ):
																												self.assertIsInstance(module.processor     ,      UpperCamelCase__      )





				@unittest.skipIf(torch_device != 'cuda'     ,      'This test is supposed to run on GPU'      )
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										__magic_name__		, __magic_name__				:							List[Any]				=      self.get_dummy_components()
										__magic_name__				:							Optional[Any]				=      StableDiffusionPipeline(**UpperCamelCase__      )
										__magic_name__				:							Any				=      sd_pipe.to(UpperCamelCase__      )
										sd_pipe.set_progress_bar_config(disable=UpperCamelCase__      )

										__magic_name__		, __magic_name__		, __magic_name__				:							Optional[Any]				=      self.get_dummy_inputs()

										# enable XFormers
										sd_pipe.enable_xformers_memory_efficient_attention()

										__magic_name__				:							Any				=      sd_pipe(**UpperCamelCase__      ).images
										__magic_name__				:							int				=      original_images[0, -3:, -3:, -1]

										with tempfile.TemporaryDirectory() as tmpdirname:
																LoraLoaderMixin.save_lora_weights(
																    save_directory=UpperCamelCase__     ,      unet_lora_layers=lora_components['unet_lora_layers']     ,      text_encoder_lora_layers=lora_components['text_encoder_lora_layers']     ,      )
																self.assertTrue(os.path.isfile(os.path.join(UpperCamelCase__     ,      'pytorch_lora_weights.bin'      )      )      )
																sd_pipe.load_lora_weights(UpperCamelCase__      )

										__magic_name__				:							Union[str, Any]				=      sd_pipe(**UpperCamelCase__      ).images
										__magic_name__				:							List[str]				=      lora_images[0, -3:, -3:, -1]

										# Outputs shouldn't match.
										self.assertFalse(torch.allclose(torch.from_numpy(UpperCamelCase__      )     ,      torch.from_numpy(UpperCamelCase__      )      )      )






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

from ...utils import logging


snake_case__     							= logging.get_logger(__name__)
def  __SCREAMING_SNAKE_CASE       (				lowercase__  :Optional[Any]      ,						lowercase__  :List[str]				)       ->  Tuple:
						__magic_name__				:							List[str]				=      nn.functional.normalize(lowercase__				)
						__magic_name__				:							Union[str, Any]				=      nn.functional.normalize(lowercase__				)
						return torch.mm(lowercase__      ,						normalized_text_embeds.t()				)



class     snake_case__  (      _A  ):
				__lowerCAmelCase		=      CLIPConfig

				__lowerCAmelCase		=      ['''CLIPEncoderLayer''']
				def __init__(  self     ,      _UpperCAmelCase      ):




										"""simple docstring"""




										super().__init__(_UpperCAmelCase      )

										__magic_name__				:							List[Any]				=      CLIPVisionModel(config.vision_config      )
										__magic_name__				:							Tuple				=      nn.Linear(config.vision_config.hidden_size     ,      config.projection_dim     ,      bias=_UpperCAmelCase      )

										__magic_name__				:							str				=      nn.Parameter(torch.ones(17     ,      config.projection_dim      )     ,      requires_grad=_UpperCAmelCase      )
										__magic_name__				:							Optional[Any]				=      nn.Parameter(torch.ones(3     ,      config.projection_dim      )     ,      requires_grad=_UpperCAmelCase      )

										__magic_name__				:							Optional[Any]				=      nn.Parameter(torch.ones(17      )     ,      requires_grad=_UpperCAmelCase      )
										__magic_name__				:							str				=      nn.Parameter(torch.ones(3      )     ,      requires_grad=_UpperCAmelCase      )
				@torch.no_grad()
				def  UpperCamelCase_						(  self     ,      _UpperCAmelCase     ,      _UpperCAmelCase      ):




										"""simple docstring"""




										__magic_name__				:							Union[str, Any]				=      self.vision_model(_UpperCAmelCase      )[1]  # pooled_output
										__magic_name__				:							Any				=      self.visual_projection(_UpperCAmelCase      )

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

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

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

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

																result.append(_UpperCAmelCase      )

										__magic_name__				:							Union[str, Any]				=      [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]      ):
																												__magic_name__				:							Union[str, Any]				=      torch.zeros_like(images[idx]      )  # black image
																						else:
																												__magic_name__				:							Optional[int]				=      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  UpperCamelCase_						(  self     ,      _UpperCAmelCase     ,      _UpperCAmelCase      ):




										"""simple docstring"""




										__magic_name__				:							List[str]				=      self.vision_model(_UpperCAmelCase      )[1]  # pooled_output
										__magic_name__				:							List[str]				=      self.visual_projection(_UpperCAmelCase      )

										__magic_name__				:							Union[str, Any]				=      cosine_distance(_UpperCAmelCase     ,      self.special_care_embeds      )
										__magic_name__				:							Optional[int]				=      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
										__magic_name__				:							int				=      0.0

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

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

										__magic_name__				:							Dict				=      0.0  # black image

										return images, has_nsfw_concepts


