


"""simple docstring"""






import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModelWithProjection

from diffusers import (
    AutoencoderKL,
    DPMSolverMultistepScheduler,
    PNDMScheduler,
    StableDiffusionImageVariationPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, load_image, load_numpy, nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()


class 	A_       (     lowerCAmelCase__      ,  lowerCAmelCase__      ,  lowerCAmelCase__      ,  unittest.TestCase   ):





  """simple docstring"""

  lowercase		:	Tuple   =						StableDiffusionImageVariationPipeline
  lowercase		:	Optional[int]   =						IMAGE_VARIATION_PARAMS
  lowercase		:	List[Any]   =						IMAGE_VARIATION_BATCH_PARAMS
  lowercase		:	Any   =						frozenset([]   )
  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
  lowercase		:	List[str]   =						frozenset([]   )
  def 							lowercase_     (						self						)				->					Optional[int]:
     torch.manual_seed(0						)
     a   :	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  ,					)
     a   :	Dict						=    PNDMScheduler(skip_prk_steps=__UpperCAmelCase						)
     torch.manual_seed(0						)
     a   :	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  ,					)
     torch.manual_seed(0						)
     a   :	Any						=    CLIPVisionConfig(
         hidden_size=32  ,					projection_dim=32  ,					intermediate_size=37  ,					layer_norm_eps=1E-0_5  ,					num_attention_heads=4  ,					num_hidden_layers=5  ,					image_size=32  ,					patch_size=4  ,					)
     a   :	str						=    CLIPVisionModelWithProjection(__UpperCAmelCase						)
     a   :	str						=    CLIPImageProcessor(crop_size=32  ,					size=32						)

     a   :	List[str]						=    {
         'unet': unet,
         'scheduler': scheduler,
         'vae': vae,
         'image_encoder': image_encoder,
         'feature_extractor': feature_extractor,
         'safety_checker': None,
     }
     return components
  def 							lowercase_     (						self  ,					__UpperCAmelCase  ,					__UpperCAmelCase=0						)				->					List[str]:
     a   :	List[str]						=    floats_tensor((1, 3, 32, 32)  ,					rng=random.Random(__UpperCAmelCase						)						)
     a   :	str						=    image.cpu().permute(0  ,					2  ,					3  ,					1						)[0]
     a   :	int						=    Image.fromarray(np.uinta(__UpperCAmelCase						)						).convert('RGB'						).resize((32, 32)						)
     if str(__UpperCAmelCase						).startswith('mps'						):
        a   :	Dict						=    torch.manual_seed(__UpperCAmelCase						)
     else:
        a   :	Tuple						=    torch.Generator(device=__UpperCAmelCase						).manual_seed(__UpperCAmelCase						)
     a   :	Optional[int]						=    {
         'image': image,
         'generator': generator,
         'num_inference_steps': 2,
         'guidance_scale': 6.0,
         'output_type': 'numpy',
     }
     return inputs
  def 							lowercase_     (						self						)				->					Any:
     a   :	Any						=    'cpu'  # ensure determinism for the device-dependent torch.Generator
     a   :	Any						=    self.get_dummy_components()
     a   :	Any						=    StableDiffusionImageVariationPipeline(**__UpperCAmelCase						)
     a   :	Tuple						=    sd_pipe.to(__UpperCAmelCase						)
     sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase						)

     a   :	int						=    self.get_dummy_inputs(__UpperCAmelCase						)
     a   :	str						=    sd_pipe(**__UpperCAmelCase						).images
     a   :	Union[str, Any]						=    image[0, -3:, -3:, -1]

     assert image.shape == (1, 64, 64, 3)
     a   :	Optional[int]						=    np.array([0.5239, 0.5723, 0.4796, 0.5049, 0.5550, 0.4685, 0.5329, 0.4891, 0.4921]						)

     assert np.abs(image_slice.flatten() - expected_slice						).max() < 1E-3
  def 							lowercase_     (						self						)				->					Dict:
     a   :	List[str]						=    'cpu'  # ensure determinism for the device-dependent torch.Generator
     a   :	List[Any]						=    self.get_dummy_components()
     a   :	Tuple						=    StableDiffusionImageVariationPipeline(**__UpperCAmelCase						)
     a   :	Union[str, Any]						=    sd_pipe.to(__UpperCAmelCase						)
     sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase						)

     a   :	int						=    self.get_dummy_inputs(__UpperCAmelCase						)
     a   :	List[Any]						=    2 * [inputs['image']]
     a   :	Dict						=    sd_pipe(**__UpperCAmelCase						)

     a   :	Tuple						=    output.images

     a   :	int						=    image[-1, -3:, -3:, -1]

     assert image.shape == (2, 64, 64, 3)
     a   :	List[Any]						=    np.array([0.6892, 0.5637, 0.5836, 0.5771, 0.6254, 0.6409, 0.5580, 0.5569, 0.5289]						)

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






  def 							lowercase_     (						self						)				->					Dict:
     super().test_inference_batch_single_identical(expected_max_diff=3E-3						)



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





  """simple docstring"""

  def 							lowercase_     (						self						)				->					Optional[int]:
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()
  def 							lowercase_     (						self  ,					__UpperCAmelCase  ,					__UpperCAmelCase="cpu"  ,					__UpperCAmelCase=torch.floataa  ,					__UpperCAmelCase=0						)				->					Any:
     a   :	Optional[Any]						=    torch.Generator(device=__UpperCAmelCase						).manual_seed(__UpperCAmelCase						)
     a   :	List[str]						=    load_image(
         'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
         '/stable_diffusion_imgvar/input_image_vermeer.png'						)
     a   :	Optional[Any]						=    np.random.RandomState(__UpperCAmelCase						).standard_normal((1, 4, 64, 64)						)
     a   :	str						=    torch.from_numpy(__UpperCAmelCase						).to(device=__UpperCAmelCase  ,					dtype=__UpperCAmelCase						)
     a   :	int						=    {
         'image': init_image,
         'latents': latents,
         'generator': generator,
         'num_inference_steps': 3,
         'guidance_scale': 7.5,
         'output_type': 'numpy',
     }
     return inputs
  def 							lowercase_     (						self						)				->					Any:
     a   :	Dict						=    StableDiffusionImageVariationPipeline.from_pretrained(
         'lambdalabs/sd-image-variations-diffusers'  ,					safety_checker=__UpperCAmelCase						)
     a   :	List[str]						=    sd_pipe.to(__UpperCAmelCase						)
     sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase						)

     a   :	str						=    self.get_inputs(__UpperCAmelCase						)
     a   :	Optional[Any]						=    sd_pipe(**__UpperCAmelCase						).images
     a   :	int						=    image[0, -3:, -3:, -1].flatten()

     assert image.shape == (1, 5_12, 5_12, 3)
     a   :	Optional[int]						=    np.array([0.8_4491, 0.9_0789, 0.7_5708, 0.7_8734, 0.8_3485, 0.7_0099, 0.6_6938, 0.6_8727, 0.6_1379]						)
     assert np.abs(image_slice - expected_slice						).max() < 6E-3
  def 							lowercase_     (						self						)				->					Any:
     a   :	Dict						=    0

     def callback_fn(__UpperCAmelCase  ,					__UpperCAmelCase  ,					__UpperCAmelCase						) -> None:
        a   :	List[Any]						=    True
        nonlocal number_of_steps
        number_of_steps += 1
        if step == 1:
           a   :	List[Any]						=    latents.detach().cpu().numpy()
           assert latents.shape == (1, 4, 64, 64)
           a   :	Tuple						=    latents[0, -3:, -3:, -1]
           a   :	Tuple						=    np.array(
               [-0.1621, 0.2837, -0.7979, -0.1221, -1.3057, 0.7681, -2.1191, 0.0464, 1.6309]						)

           assert np.abs(latents_slice.flatten() - expected_slice						).max() < 5E-2
        elif step == 2:
           a   :	List[Any]						=    latents.detach().cpu().numpy()
           assert latents.shape == (1, 4, 64, 64)
           a   :	Dict						=    latents[0, -3:, -3:, -1]
           a   :	Dict						=    np.array([0.6299, 1.7500, 1.1992, -2.1582, -1.8994, 0.7334, -0.7090, 1.0137, 1.5273]						)

           assert np.abs(latents_slice.flatten() - expected_slice						).max() < 5E-2

     a   :	List[Any]						=    False

     a   :	Any						=    StableDiffusionImageVariationPipeline.from_pretrained(
         'fusing/sd-image-variations-diffusers'  ,					safety_checker=__UpperCAmelCase  ,					torch_dtype=torch.floataa  ,					)
     pipe.to(__UpperCAmelCase						)
     pipe.set_progress_bar_config(disable=__UpperCAmelCase						)
     pipe.enable_attention_slicing()

     a   :	str						=    self.get_inputs(__UpperCAmelCase  ,					dtype=torch.floataa						)
     pipe(**__UpperCAmelCase  ,					callback=__UpperCAmelCase  ,					callback_steps=1						)
     assert callback_fn.has_been_called
     assert number_of_steps == inputs["num_inference_steps"]






  def 							lowercase_     (						self						)				->					str:
     torch.cuda.empty_cache()
     torch.cuda.reset_max_memory_allocated()
     torch.cuda.reset_peak_memory_stats()

     a   :	List[Any]						=    'fusing/sd-image-variations-diffusers'
     a   :	Tuple						=    StableDiffusionImageVariationPipeline.from_pretrained(
         __UpperCAmelCase  ,					safety_checker=__UpperCAmelCase  ,					torch_dtype=torch.floataa						)
     a   :	Dict						=    pipe.to(__UpperCAmelCase						)
     pipe.set_progress_bar_config(disable=__UpperCAmelCase						)
     pipe.enable_attention_slicing(1						)
     pipe.enable_sequential_cpu_offload()

     a   :	List[str]						=    self.get_inputs(__UpperCAmelCase  ,					dtype=torch.floataa						)
     a   :	Union[str, Any]						=    pipe(**__UpperCAmelCase						)

     a   :	Any						=    torch.cuda.max_memory_allocated()
     # make sure that less than 2.6 GB is allocated
     assert mem_bytes < 2.6 * 10**9



@nightly
@require_torch_gpu
class 	A_       (     unittest.TestCase   ):





  """simple docstring"""

  def 							lowercase_     (						self						)				->					Any:
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()
  def 							lowercase_     (						self  ,					__UpperCAmelCase  ,					__UpperCAmelCase="cpu"  ,					__UpperCAmelCase=torch.floataa  ,					__UpperCAmelCase=0						)				->					Tuple:
     a   :	int						=    torch.Generator(device=__UpperCAmelCase						).manual_seed(__UpperCAmelCase						)
     a   :	Optional[Any]						=    load_image(
         'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
         '/stable_diffusion_imgvar/input_image_vermeer.png'						)
     a   :	int						=    np.random.RandomState(__UpperCAmelCase						).standard_normal((1, 4, 64, 64)						)
     a   :	Tuple						=    torch.from_numpy(__UpperCAmelCase						).to(device=__UpperCAmelCase  ,					dtype=__UpperCAmelCase						)
     a   :	Any						=    {
         'image': init_image,
         'latents': latents,
         'generator': generator,
         'num_inference_steps': 50,
         'guidance_scale': 7.5,
         'output_type': 'numpy',
     }
     return inputs
  def 							lowercase_     (						self						)				->					Dict:
     a   :	List[Any]						=    StableDiffusionImageVariationPipeline.from_pretrained('fusing/sd-image-variations-diffusers'						)
     sd_pipe.to(__UpperCAmelCase						)
     sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase						)

     a   :	List[str]						=    self.get_inputs(__UpperCAmelCase						)
     a   :	Optional[Any]						=    sd_pipe(**__UpperCAmelCase						).images[0]

     a   :	str						=    load_numpy(
         'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
         '/stable_diffusion_imgvar/lambdalabs_variations_pndm.npy'						)
     a   :	Optional[Any]						=    np.abs(expected_image - image						).max()
     assert max_diff < 1E-3






  def 							lowercase_     (						self						)				->					str:
     a   :	List[Any]						=    StableDiffusionImageVariationPipeline.from_pretrained('fusing/sd-image-variations-diffusers'						)
     a   :	int						=    DPMSolverMultistepScheduler.from_config(sd_pipe.scheduler.config						)
     sd_pipe.to(__UpperCAmelCase						)
     sd_pipe.set_progress_bar_config(disable=__UpperCAmelCase						)

     a   :	Any						=    self.get_inputs(__UpperCAmelCase						)
     a   :	List[Any]						=    25
     a   :	List[str]						=    sd_pipe(**__UpperCAmelCase						).images[0]

     a   :	str						=    load_numpy(
         'https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'
         '/stable_diffusion_imgvar/lambdalabs_variations_dpm_multi.npy'						)
     a   :	Union[str, Any]						=    np.abs(expected_image - image						).max()
     assert max_diff < 1E-3










'''simple docstring'''







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

from ...utils import logging


SCREAMING_SNAKE_CASE__     : Optional[int]				=   logging.get_logger(__name__)







def  SCREAMING_SNAKE_CASE__ (				SCREAMING_SNAKE_CASE ,		SCREAMING_SNAKE_CASE     ):
		SCREAMING_SNAKE_CASE_						:List[Any]										=      nn.functional.normalize(SCREAMING_SNAKE_CASE     )
		SCREAMING_SNAKE_CASE_						:str										=      nn.functional.normalize(SCREAMING_SNAKE_CASE     )
		return torch.mm(SCREAMING_SNAKE_CASE ,		normalized_text_embeds.t()     )





class 				__lowerCAmelCase(							lowerCAmelCase__      ):
			__snake_case  :		int		=      CLIPConfig

			__snake_case  :		Optional[Any]		=      ['CLIPEncoderLayer']






			def __init__(   self  :							int      ,       SCREAMING_SNAKE_CASE  :							CLIPConfig	):




					"""simple docstring"""







					super().__init__(SCREAMING_SNAKE_CASE	)

					SCREAMING_SNAKE_CASE_						:List[Any]										=      CLIPVisionModel(config.vision_config	)
					SCREAMING_SNAKE_CASE_						:Dict										=      nn.Linear(config.vision_config.hidden_size      ,       config.projection_dim      ,       bias=SCREAMING_SNAKE_CASE	)

					SCREAMING_SNAKE_CASE_						:List[Any]										=      nn.Parameter(torch.ones(17      ,       config.projection_dim	)      ,       requires_grad=SCREAMING_SNAKE_CASE	)
					SCREAMING_SNAKE_CASE_						:Optional[int]										=      nn.Parameter(torch.ones(3      ,       config.projection_dim	)      ,       requires_grad=SCREAMING_SNAKE_CASE	)

					SCREAMING_SNAKE_CASE_						:Union[str, Any]										=      nn.Parameter(torch.ones(17	)      ,       requires_grad=SCREAMING_SNAKE_CASE	)
					SCREAMING_SNAKE_CASE_						:str										=      nn.Parameter(torch.ones(3	)      ,       requires_grad=SCREAMING_SNAKE_CASE	)






			@torch.no_grad()
			def 				_lowercase	(   self  :							Tuple      ,       SCREAMING_SNAKE_CASE  :							Tuple      ,       SCREAMING_SNAKE_CASE  :							Tuple	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:Optional[int]										=      self.vision_model(SCREAMING_SNAKE_CASE	)[1]  # pooled_output
					SCREAMING_SNAKE_CASE_						:Dict										=      self.visual_projection(SCREAMING_SNAKE_CASE	)

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

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

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

							for concept_idx in range(len(cos_dist[0]	)	):
									SCREAMING_SNAKE_CASE_						:int										=      cos_dist[i][concept_idx]
									SCREAMING_SNAKE_CASE_						:Tuple										=      self.concept_embeds_weights[concept_idx].item()
									SCREAMING_SNAKE_CASE_						:Union[str, Any]										=      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	)

					SCREAMING_SNAKE_CASE_						:Any										=      [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]	):
											SCREAMING_SNAKE_CASE_						:Tuple										=      torch.zeros_like(images[idx]	)  # black image
									else:
											SCREAMING_SNAKE_CASE_						:str										=      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 				_lowercase	(   self  :							str      ,       SCREAMING_SNAKE_CASE  :							torch.FloatTensor      ,       SCREAMING_SNAKE_CASE  :							torch.FloatTensor	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:Tuple										=      self.vision_model(SCREAMING_SNAKE_CASE	)[1]  # pooled_output
					SCREAMING_SNAKE_CASE_						:Tuple										=      self.visual_projection(SCREAMING_SNAKE_CASE	)

					SCREAMING_SNAKE_CASE_						:List[str]										=      cosine_distance(SCREAMING_SNAKE_CASE      ,       self.special_care_embeds	)
					SCREAMING_SNAKE_CASE_						:int										=      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
					SCREAMING_SNAKE_CASE_						:Dict										=      0.0

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

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

					SCREAMING_SNAKE_CASE_						:Optional[int]										=      0.0  # black image

					return images, has_nsfw_concepts
