

import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModel

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEImgaImgPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import floats_tensor, load_image, 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	(   SCREAMING_SNAKE_CASE     ,     unittest.TestCase				):
  _a		     =    ShapEImgaImgPipeline
  _a		     =    ["""image"""]
  _a		     =    ["""image"""]
  _a		     =    [
      """num_images_per_prompt""",
      """num_inference_steps""",
      """generator""",
      """latents""",
      """guidance_scale""",
      """frame_size""",
      """output_type""",
      """return_dict""",
  ]
  _a		     =    False




  @property
  def     A__	(							self   )    ->					Dict:
        '''simple docstring'''






        return 32




  @property
  def     A__	(							self   )    ->					Union[str, Any]:
        '''simple docstring'''






        return 32




  @property
  def     A__	(							self   )    ->					int:
        '''simple docstring'''






        return self.time_input_dim * 4




  @property
  def     A__	(							self   )    ->					Dict:
        '''simple docstring'''






        return 8




  @property
  def     A__	(							self   )    ->					Optional[Any]:
        '''simple docstring'''






        torch.manual_seed(0   )
        _lowercase      			=CLIPVisionConfig(
            hidden_size=self.text_embedder_hidden_size				,  image_size=64				,  projection_dim=self.text_embedder_hidden_size				,  intermediate_size=37				,  num_attention_heads=4				,  num_channels=3				,  num_hidden_layers=5				,  patch_size=1				,  )

        _lowercase      			=CLIPVisionModel(lowerCAmelCase   )
        return model




  @property
  def     A__	(							self   )    ->					List[Any]:
        '''simple docstring'''






        _lowercase      			=CLIPImageProcessor(
            crop_size=224				,  do_center_crop=lowerCAmelCase				,  do_normalize=lowerCAmelCase				,  do_resize=lowerCAmelCase				,  image_mean=[0.48145466, 0.4578275, 0.40821073]				,  image_std=[0.26862954, 0.26130258, 0.27577711]				,  resample=3				,  size=224				,  )

        return image_processor




  @property
  def     A__	(							self   )    ->					int:
        '''simple docstring'''






        torch.manual_seed(0   )

        _lowercase      			={
            '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',
            'embedding_proj_norm_type': 'layer',
            'encoder_hid_proj_type': None,
            'added_emb_type': None,
        }

        _lowercase      			=PriorTransformer(**lowerCAmelCase   )
        return model




  @property
  def     A__	(							self   )    ->					Optional[Any]:
        '''simple docstring'''






        torch.manual_seed(0   )

        _lowercase      			={
            '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,
            ),
        }
        _lowercase      			=ShapERenderer(**lowerCAmelCase   )
        return model




  def     A__	(							self   )    ->					Optional[Any]:
        '''simple docstring'''






        _lowercase      			=self.dummy_prior
        _lowercase      			=self.dummy_image_encoder
        _lowercase      			=self.dummy_image_processor
        _lowercase      			=self.dummy_renderer

        _lowercase      			=HeunDiscreteScheduler(
            beta_schedule='exp'				,  num_train_timesteps=1_024				,  prediction_type='sample'				,  use_karras_sigmas=lowerCAmelCase				,  clip_sample=lowerCAmelCase				,  clip_sample_range=1.0				,  )
        _lowercase      			={
            'prior': prior,
            'image_encoder': image_encoder,
            'image_processor': image_processor,
            'renderer': renderer,
            'scheduler': scheduler,
        }

        return components




  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase=0   )    ->					Dict:
        '''simple docstring'''






        _lowercase      			=floats_tensor((1, 3, 64, 64)				,  rng=random.Random(lowerCAmelCase   )   ).to(lowerCAmelCase   )

        if str(lowerCAmelCase   ).startswith('mps'   ):
              _lowercase      			=torch.manual_seed(lowerCAmelCase   )
        else:
              _lowercase      			=torch.Generator(device=lowerCAmelCase   ).manual_seed(lowerCAmelCase   )
        _lowercase      			={
            'image': input_image,
            'generator': generator,
            'num_inference_steps': 1,
            'frame_size': 32,
            'output_type': 'np',
        }
        return inputs




  def     A__	(							self   )    ->					List[Any]:
        '''simple docstring'''






        _lowercase      			='cpu'

        _lowercase      			=self.get_dummy_components()

        _lowercase      			=self.pipeline_class(**lowerCAmelCase   )
        _lowercase      			=pipe.to(lowerCAmelCase   )

        pipe.set_progress_bar_config(disable=lowerCAmelCase   )

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

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

        _lowercase      			=np.array(
            [
                0.00039216,
                0.00039216,
                0.00039216,
                0.00039216,
                0.00039216,
                0.00039216,
                0.00039216,
                0.00039216,
                0.00039216,
            ]   )

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




  def     A__	(							self   )    ->					str:
        '''simple docstring'''






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




  def     A__	(							self   )    ->					Dict:
        '''simple docstring'''






        _lowercase      			=torch_device == 'cpu'
        _lowercase      			=True
        self._test_inference_batch_single_identical(
            batch_size=2				,  test_max_difference=lowerCAmelCase				,  relax_max_difference=lowerCAmelCase				,  )




  def     A__	(							self   )    ->					Optional[int]:
        '''simple docstring'''






        _lowercase      			=self.get_dummy_components()
        _lowercase      			=self.pipeline_class(**lowerCAmelCase   )
        _lowercase      			=pipe.to(lowerCAmelCase   )
        pipe.set_progress_bar_config(disable=lowerCAmelCase   )

        _lowercase      			=1
        _lowercase      			=2

        _lowercase      			=self.get_dummy_inputs(lowerCAmelCase   )

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

        _lowercase      			=pipe(**lowerCAmelCase				,  num_images_per_prompt=lowerCAmelCase   )[0]

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





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




  def     A__	(							self   )    ->					int:
        '''simple docstring'''






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




  def     A__	(							self   )    ->					Tuple:
        '''simple docstring'''






        _lowercase      			=load_image(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main' '/shap_e/corgi.png'   )
        _lowercase      			=load_numpy(
            'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
            '/shap_e/test_shap_e_img2img_out.npy'   )
        _lowercase      			=ShapEImgaImgPipeline.from_pretrained('openai/shap-e-img2img'   )
        _lowercase      			=pipe.to(lowerCAmelCase   )
        pipe.set_progress_bar_config(disable=lowerCAmelCase   )

        _lowercase      			=torch.Generator(device=lowerCAmelCase   ).manual_seed(0   )

        _lowercase      			=pipe(
            lowerCAmelCase				,  generator=lowerCAmelCase				,  guidance_scale=3.0				,  num_inference_steps=64				,  frame_size=64				,  output_type='np'				,  ).images[0]

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

        assert_mean_pixel_difference(lowerCAmelCase				,  lowerCAmelCase   )





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

from ...utils import logging


lowercase_   							=  logging.get_logger(__name__)
def 	a (     A__      :  Optional[int] ,  A__      :  Union[str, Any]	)       -> int:





      """simple docstring"""



      _lowercase      			=nn.functional.normalize(A__	)
      _lowercase      			=nn.functional.normalize(A__	)
      return torch.mm(A__ ,  normalized_text_embeds.t()	)







class 					__lowerCAmelCase	(   SCREAMING_SNAKE_CASE				):
  _a		     =    CLIPConfig

  _a		     =    ["""CLIPEncoderLayer"""]




  def __init__(							self				,  lowerCAmelCase   )    ->					Union[str, Any]:
        '''simple docstring'''






        super().__init__(lowerCAmelCase   )

        _lowercase      			=CLIPVisionModel(config.vision_config   )
        _lowercase      			=nn.Linear(config.vision_config.hidden_size				,  config.projection_dim				,  bias=lowerCAmelCase   )

        _lowercase      			=nn.Parameter(torch.ones(17				,  config.projection_dim   )				,  requires_grad=lowerCAmelCase   )
        _lowercase      			=nn.Parameter(torch.ones(3				,  config.projection_dim   )				,  requires_grad=lowerCAmelCase   )

        _lowercase      			=nn.Parameter(torch.ones(17   )				,  requires_grad=lowerCAmelCase   )
        _lowercase      			=nn.Parameter(torch.ones(3   )				,  requires_grad=lowerCAmelCase   )




  @torch.no_grad()
  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase   )    ->					Optional[int]:
        '''simple docstring'''






        _lowercase      			=self.vision_model(lowerCAmelCase   )[1]  # pooled_output
        _lowercase      			=self.visual_projection(lowerCAmelCase   )

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

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

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

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

              result.append(lowerCAmelCase   )

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

        return images, has_nsfw_concepts




  @torch.no_grad()
  def     A__	(							self				,  lowerCAmelCase				,  lowerCAmelCase   )    ->					Tuple:
        '''simple docstring'''






        _lowercase      			=self.vision_model(lowerCAmelCase   )[1]  # pooled_output
        _lowercase      			=self.visual_projection(lowerCAmelCase   )

        _lowercase      			=cosine_distance(lowerCAmelCase				,  self.special_care_embeds   )
        _lowercase      			=cosine_distance(lowerCAmelCase				,  self.concept_embeds   )

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

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

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

        return images, has_nsfw_concepts


