
import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import load_numpy, slow
from diffusers.utils.testing_utils import require_torch_gpu, torch_device

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference







class        lowerCAmelCase__      (   __lowercase		,	unittest.TestCase	):
						UpperCamelCase_							: Tuple	=		ShapEPipeline
						UpperCamelCase_							: Optional[int]	=		["prompt"]
						UpperCamelCase_							: Tuple	=		["prompt"]
						UpperCamelCase_							: Union[str, Any]	=		[
						    "num_images_per_prompt",
						    "num_inference_steps",
						    "generator",
						    "latents",
						    "guidance_scale",
						    "frame_size",
						    "output_type",
						    "return_dict",
						]
						UpperCamelCase_							: Union[str, Any]	=		False

						@property
						def  A_				(			self				)					->      List[Any]:


								'''simple docstring'''

								return 32

						@property
						def  A_				(			self				)					->      List[str]:


								'''simple docstring'''

								return 32

						@property
						def  A_				(			self				)					->      Tuple:


								'''simple docstring'''

								return self.time_input_dim * 4

						@property
						def  A_				(			self				)					->      int:


								'''simple docstring'''

								return 8

						@property
						def  A_				(			self				)					->      List[Any]:


								'''simple docstring'''

								_UpperCamelCase						    =	CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""				)
								return tokenizer

						@property
						def  A_				(			self				)					->      Dict:


								'''simple docstring'''

								torch.manual_seed(0				)
								_UpperCamelCase						    =	CLIPTextConfig(
								    bos_token_id=0       ,							eos_token_id=2       ,							hidden_size=self.text_embedder_hidden_size       ,							projection_dim=self.text_embedder_hidden_size       ,							intermediate_size=37       ,							layer_norm_eps=1e-05       ,							num_attention_heads=4       ,							num_hidden_layers=5       ,							pad_token_id=1       ,							vocab_size=10_00       ,							)
								return CLIPTextModelWithProjection(a				)

						@property
						def  A_				(			self				)					->      Tuple:


								'''simple docstring'''

								torch.manual_seed(0				)

								_UpperCamelCase						    =	{
								    """num_attention_heads""": 2,
								    """attention_head_dim""": 16,
								    """embedding_dim""": self.time_input_dim,
								    """num_embeddings""": 32,
								    """embedding_proj_dim""": self.text_embedder_hidden_size,
								    """time_embed_dim""": self.time_embed_dim,
								    """num_layers""": 1,
								    """clip_embed_dim""": self.time_input_dim * 2,
								    """additional_embeddings""": 0,
								    """time_embed_act_fn""": """gelu""",
								    """norm_in_type""": """layer""",
								    """encoder_hid_proj_type""": None,
								    """added_emb_type""": None,
								}

								_UpperCamelCase						    =	PriorTransformer(**a				)
								return model

						@property
						def  A_				(			self				)					->      Optional[int]:


								'''simple docstring'''

								torch.manual_seed(0				)

								_UpperCamelCase						    =	{
								    """param_shapes""": (
								        (self.renderer_dim, 93),
								        (self.renderer_dim, 8),
								        (self.renderer_dim, 8),
								        (self.renderer_dim, 8),
								    ),
								    """d_latent""": self.time_input_dim,
								    """d_hidden""": self.renderer_dim,
								    """n_output""": 12,
								    """background""": (
								        0.1,
								        0.1,
								        0.1,
								    ),
								}
								_UpperCamelCase						    =	ShapERenderer(**a				)
								return model

						def  A_				(			self				)					->      str:


								'''simple docstring'''

								_UpperCamelCase						    =	self.dummy_prior
								_UpperCamelCase						    =	self.dummy_text_encoder
								_UpperCamelCase						    =	self.dummy_tokenizer
								_UpperCamelCase						    =	self.dummy_renderer

								_UpperCamelCase						    =	HeunDiscreteScheduler(
								    beta_schedule="""exp"""       ,							num_train_timesteps=10_24       ,							prediction_type="""sample"""       ,							use_karras_sigmas=a       ,							clip_sample=a       ,							clip_sample_range=1.0       ,							)
								_UpperCamelCase						    =	{
								    """prior""": prior,
								    """text_encoder""": text_encoder,
								    """tokenizer""": tokenizer,
								    """renderer""": renderer,
								    """scheduler""": scheduler,
								}

								return components

						def  A_				(			self       ,							a       ,							a=0				)					->      Dict:


								'''simple docstring'''

								if str(a				).startswith("""mps"""				):
										_UpperCamelCase						    =	torch.manual_seed(a				)
								else:
										_UpperCamelCase						    =	torch.Generator(device=a				).manual_seed(a				)
								_UpperCamelCase						    =	{
								    """prompt""": """horse""",
								    """generator""": generator,
								    """num_inference_steps""": 1,
								    """frame_size""": 32,
								    """output_type""": """np""",
								}
								return inputs

						def  A_				(			self				)					->      str:


								'''simple docstring'''

								_UpperCamelCase						    =	"""cpu"""

								_UpperCamelCase						    =	self.get_dummy_components()

								_UpperCamelCase						    =	self.pipeline_class(**a				)
								_UpperCamelCase						    =	pipe.to(a				)

								pipe.set_progress_bar_config(disable=a				)

								_UpperCamelCase						    =	pipe(**self.get_dummy_inputs(a				)				)
								_UpperCamelCase						    =	output.images[0]
								_UpperCamelCase						    =	image[0, -3:, -3:, -1]

								assert image.shape == (20, 32, 32, 3)

								_UpperCamelCase						    =	np.array(
								    [
								        0.0003_9216,
								        0.0003_9216,
								        0.0003_9216,
								        0.0003_9216,
								        0.0003_9216,
								        0.0003_9216,
								        0.0003_9216,
								        0.0003_9216,
								        0.0003_9216,
								    ]				)

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

						def  A_				(			self				)					->      int:


								'''simple docstring'''

								self._test_inference_batch_consistent(batch_sizes=[1, 2]				)

						def  A_				(			self				)					->      Tuple:


								'''simple docstring'''

								_UpperCamelCase						    =	torch_device == """cpu"""
								_UpperCamelCase						    =	True

								self._test_inference_batch_single_identical(
								    batch_size=2       ,							test_max_difference=a       ,							relax_max_difference=a       ,							)







						def  A_				(			self				)					->      Optional[Any]:


								'''simple docstring'''

								_UpperCamelCase						    =	self.get_dummy_components()
								_UpperCamelCase						    =	self.pipeline_class(**a				)
								_UpperCamelCase						    =	pipe.to(a				)
								pipe.set_progress_bar_config(disable=a				)

								_UpperCamelCase						    =	1
								_UpperCamelCase						    =	2

								_UpperCamelCase						    =	self.get_dummy_inputs(a				)

								for key in inputs.keys():
										if key in self.batch_params:
												_UpperCamelCase						    =	batch_size * [inputs[key]]

								_UpperCamelCase						    =	pipe(**a       ,							num_images_per_prompt=a				)[0]

								assert images.shape[0] == batch_size * num_images_per_prompt







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

						def  A_				(			self				)					->      Optional[int]:


								'''simple docstring'''

								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()







						def  A_				(			self				)					->      Union[str, Any]:


								'''simple docstring'''

								_UpperCamelCase						    =	load_numpy(
								    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
								    """/shap_e/test_shap_e_np_out.npy"""				)
								_UpperCamelCase						    =	ShapEPipeline.from_pretrained("""openai/shap-e"""				)
								_UpperCamelCase						    =	pipe.to(a				)
								pipe.set_progress_bar_config(disable=a				)

								_UpperCamelCase						    =	torch.Generator(device=a				).manual_seed(0				)

								_UpperCamelCase						    =	pipe(
								    """a shark"""       ,							generator=a       ,							guidance_scale=15.0       ,							num_inference_steps=64       ,							frame_size=64       ,							output_type="""np"""       ,							).images[0]

								assert images.shape == (20, 64, 64, 3)

								assert_mean_pixel_difference(a       ,							a				)





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

from ...utils import logging


lowerCamelCase__	=							logging.get_logger(__name__)


def 							__A(lowerCAmelCase						,  lowerCAmelCase )		->				List[str]:

		"""simple docstring"""






		_UpperCamelCase						    =	nn.functional.normalize(lowerCAmelCase )
		_UpperCamelCase						    =	nn.functional.normalize(lowerCAmelCase )
		return torch.mm(lowerCAmelCase						,  normalized_text_embeds.t() )







class        lowerCAmelCase__      (   __lowercase	):
						UpperCamelCase_							: Tuple	=		CLIPConfig

						UpperCamelCase_							: str	=		["CLIPEncoderLayer"]

						def __init__(			self       ,							a				)					->      Optional[Any]:


								'''simple docstring'''

								super().__init__(a				)

								_UpperCamelCase						    =	CLIPVisionModel(config.vision_config				)
								_UpperCamelCase						    =	nn.Linear(config.vision_config.hidden_size       ,							config.projection_dim       ,							bias=a				)

								_UpperCamelCase						    =	nn.Parameter(torch.ones(17       ,							config.projection_dim				)       ,							requires_grad=a				)
								_UpperCamelCase						    =	nn.Parameter(torch.ones(3       ,							config.projection_dim				)       ,							requires_grad=a				)

								_UpperCamelCase						    =	nn.Parameter(torch.ones(17				)       ,							requires_grad=a				)
								_UpperCamelCase						    =	nn.Parameter(torch.ones(3				)       ,							requires_grad=a				)

						@torch.no_grad()
						def  A_				(			self       ,							a       ,							a				)					->      str:


								'''simple docstring'''

								_UpperCamelCase						    =	self.vision_model(a				)[1]  # pooled_output
								_UpperCamelCase						    =	self.visual_projection(a				)

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

								_UpperCamelCase						    =	[]
								_UpperCamelCase						    =	image_embeds.shape[0]
								for i in range(a				):
										_UpperCamelCase						    =	{"""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
										_UpperCamelCase						    =	0.0

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

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

										result.append(a				)

								_UpperCamelCase						    =	[len(res["""bad_concepts"""]				) > 0 for res in result]

								return images, has_nsfw_concepts







						@torch.no_grad()
						def  A_				(			self       ,							a       ,							a				)					->      Any:


								'''simple docstring'''

								_UpperCamelCase						    =	self.vision_model(a				)[1]  # pooled_output
								_UpperCamelCase						    =	self.visual_projection(a				)

								_UpperCamelCase						    =	cosine_distance(a       ,							self.special_care_embeds				)
								_UpperCamelCase						    =	cosine_distance(a       ,							self.concept_embeds				)

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

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

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

								return images, has_nsfw_concepts



