




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 							UpperCamelCase_       (					UpperCAmelCase__						,   unittest.TestCase	):

    '''simple docstring'''






    UpperCAmelCase__				    = ShapEPipeline
    UpperCAmelCase__				    = ['''prompt''']
    UpperCAmelCase__				    = ['''prompt''']
    UpperCAmelCase__				    = [
        '''num_images_per_prompt''',
        '''num_inference_steps''',
        '''generator''',
        '''latents''',
        '''guidance_scale''',
        '''frame_size''',
        '''output_type''',
        '''return_dict''',
    ]
    UpperCAmelCase__				    = False


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





           '''simple docstring'''

           return 32


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





           '''simple docstring'''

           return 32


    @property
    def 			SCREAMING_SNAKE_CASE				(     self :      Tuple)		->Optional[Any]:





           '''simple docstring'''

           return self.time_input_dim * 4


    @property
    def 			SCREAMING_SNAKE_CASE				(     self :      Optional[Any])		->str:





           '''simple docstring'''

           return 8


    @property
    def 			SCREAMING_SNAKE_CASE				(     self :      List[str])		->Optional[int]:





           '''simple docstring'''

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


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





           '''simple docstring'''

           torch.manual_seed(0)
           A__ 					=       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=1_000   ,					)
           return CLIPTextModelWithProjection(UpperCAmelCase__)


    @property
    def 			SCREAMING_SNAKE_CASE				(     self :      Tuple)		->Union[str, Any]:





           '''simple docstring'''

           torch.manual_seed(0)

           A__ 					=       {
               '''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,
           }

           A__ 					=       PriorTransformer(**UpperCAmelCase__)
           return model


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





           '''simple docstring'''

           torch.manual_seed(0)

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


    def 			SCREAMING_SNAKE_CASE				(     self :      List[Any])		->List[Any]:





           '''simple docstring'''

           A__ 					=       self.dummy_prior
           A__ 					=       self.dummy_text_encoder
           A__ 					=       self.dummy_tokenizer
           A__ 					=       self.dummy_renderer

           A__ 					=       HeunDiscreteScheduler(
               beta_schedule='''exp'''   ,					num_train_timesteps=1_024   ,					prediction_type='''sample'''   ,					use_karras_sigmas=UpperCAmelCase__   ,					clip_sample=UpperCAmelCase__   ,					clip_sample_range=1.0   ,					)
           A__ 					=       {
               '''prior''': prior,
               '''text_encoder''': text_encoder,
               '''tokenizer''': tokenizer,
               '''renderer''': renderer,
               '''scheduler''': scheduler,
           }

           return components


    def 			SCREAMING_SNAKE_CASE				(     self :      Optional[int]   ,					UpperCAmelCase__ :      int   ,					UpperCAmelCase__ :      List[Any]=0)		->Union[str, Any]:





           '''simple docstring'''

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


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





           '''simple docstring'''

           A__ 					=       '''cpu'''

           A__ 					=       self.get_dummy_components()

           A__ 					=       self.pipeline_class(**UpperCAmelCase__)
           A__ 					=       pipe.to(UpperCAmelCase__)

           pipe.set_progress_bar_config(disable=UpperCAmelCase__)

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

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

           A__ 					=       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 			SCREAMING_SNAKE_CASE				(     self :      Any)		->int:





           '''simple docstring'''

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


    def 			SCREAMING_SNAKE_CASE				(     self :      List[Any])		->List[str]:





           '''simple docstring'''

           A__ 					=       torch_device == '''cpu'''
           A__ 					=       True

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





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





           '''simple docstring'''

           A__ 					=       self.get_dummy_components()
           A__ 					=       self.pipeline_class(**UpperCAmelCase__)
           A__ 					=       pipe.to(UpperCAmelCase__)
           pipe.set_progress_bar_config(disable=UpperCAmelCase__)

           A__ 					=       1
           A__ 					=       2

           A__ 					=       self.get_dummy_inputs(UpperCAmelCase__)

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

           A__ 					=       pipe(**UpperCAmelCase__   ,					num_images_per_prompt=UpperCAmelCase__)[0]

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


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

    '''simple docstring'''






    def 			SCREAMING_SNAKE_CASE				(     self :      str)		->List[Any]:





           '''simple docstring'''

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





    def 			SCREAMING_SNAKE_CASE				(     self :      Dict)		->Optional[int]:





           '''simple docstring'''

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

           A__ 					=       torch.Generator(device=UpperCAmelCase__).manual_seed(0)

           A__ 					=       pipe(
               '''a shark'''   ,					generator=UpperCAmelCase__   ,					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(UpperCAmelCase__   ,					UpperCAmelCase__)




"""simple docstring"""


import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
      import torch

if is_vision_available():
      from PIL import Image

      from transformers import PoolFormerImageProcessor



class 	__a	(   unittest.TestCase						):






       def __init__(	self			:				Optional[Any]							,   UpperCAmelCase_			:				Any							,   UpperCAmelCase_			:				Optional[Any]=7							,   UpperCAmelCase_			:				List[Any]=3							,   UpperCAmelCase_			:				Optional[int]=30							,   UpperCAmelCase_			:				Any=400							,   UpperCAmelCase_			:				Any=True							,   UpperCAmelCase_			:				Dict=None							,   UpperCAmelCase_			:				Dict=0.9							,   UpperCAmelCase_			:				Tuple=None							,   UpperCAmelCase_			:				Tuple=True							,   UpperCAmelCase_			:				Optional[int]=[0.5, 0.5, 0.5]							,   UpperCAmelCase_			:				Union[str, Any]=[0.5, 0.5, 0.5]							,   )->	List[str]:







              """simple docstring"""
              UpperCamelCase       			=		size if size is not None else {"shortest_edge": 30}
              UpperCamelCase       			=		crop_size if crop_size is not None else {"height": 30, "width": 30}
              UpperCamelCase       			=		parent
              UpperCamelCase       			=		batch_size
              UpperCamelCase       			=		num_channels
              UpperCamelCase       			=		min_resolution
              UpperCamelCase       			=		max_resolution
              UpperCamelCase       			=		do_resize_and_center_crop
              UpperCamelCase       			=		size
              UpperCamelCase       			=		crop_pct
              UpperCamelCase       			=		crop_size
              UpperCamelCase       			=		do_normalize
              UpperCamelCase       			=		image_mean
              UpperCamelCase       			=		image_std






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[Any]      )->	str:







              """simple docstring"""
              return {
                  "size": self.size,
                  "do_resize_and_center_crop": self.do_resize_and_center_crop,
                  "crop_pct": self.crop_pct,
                  "crop_size": self.crop_size,
                  "do_normalize": self.do_normalize,
                  "image_mean": self.image_mean,
                  "image_std": self.image_std,
              }





@require_torch
@require_vision
class 	__a	(   _lowerCAmelCase     ,  unittest.TestCase						):
       UpperCamelCase_		:       List[Any]				=      PoolFormerImageProcessor if is_vision_available() else None






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[Any]      )->	str:







              """simple docstring"""
              UpperCamelCase       			=		PoolFormerImageProcessingTester(self      )






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







              """simple docstring"""
              return self.image_processor_tester.prepare_image_processor_dict()






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Dict      )->	Optional[int]:







              """simple docstring"""
              UpperCamelCase       			=		self.image_processing_class(**self.image_processor_dict      )
              self.assertTrue(hasattr(UpperCAmelCase_							,   "do_resize_and_center_crop"      )      )
              self.assertTrue(hasattr(UpperCAmelCase_							,   "size"      )      )
              self.assertTrue(hasattr(UpperCAmelCase_							,   "crop_pct"      )      )
              self.assertTrue(hasattr(UpperCAmelCase_							,   "do_normalize"      )      )
              self.assertTrue(hasattr(UpperCAmelCase_							,   "image_mean"      )      )
              self.assertTrue(hasattr(UpperCAmelCase_							,   "image_std"      )      )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				int      )->	List[Any]:







              """simple docstring"""
              UpperCamelCase       			=		self.image_processing_class.from_dict(self.image_processor_dict      )
              self.assertEqual(image_processor.size							,   {"shortest_edge": 30}      )
              self.assertEqual(image_processor.crop_size							,   {"height": 30, "width": 30}      )

              UpperCamelCase       			=		self.image_processing_class.from_dict(self.image_processor_dict							,   size=42							,   crop_size=84      )
              self.assertEqual(image_processor.size							,   {"shortest_edge": 42}      )
              self.assertEqual(image_processor.crop_size							,   {"height": 84, "width": 84}      )






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







              """simple docstring"""
              pass






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







              """simple docstring"""
              # Initialize image_processing
              UpperCamelCase       			=		self.image_processing_class(**self.image_processor_dict      )
              # create random PIL images
              UpperCamelCase       			=		prepare_image_inputs(self.image_processor_tester							,   equal_resolution=UpperCAmelCase_      )
              for image in image_inputs:
                     self.assertIsInstance(UpperCAmelCase_							,   Image.Image      )

              # Test not batched input
              UpperCamelCase       			=		image_processing(image_inputs[0]							,   return_tensors="pt"      ).pixel_values

              self.assertEqual(
                  encoded_images.shape							,   (
                      1,
                      self.image_processor_tester.num_channels,
                      self.image_processor_tester.crop_size["height"],
                      self.image_processor_tester.crop_size["width"],
                  )							,   )

              # Test batched
              UpperCamelCase       			=		image_processing(UpperCAmelCase_							,   return_tensors="pt"      ).pixel_values
              self.assertEqual(
                  encoded_images.shape							,   (
                      self.image_processor_tester.batch_size,
                      self.image_processor_tester.num_channels,
                      self.image_processor_tester.crop_size["height"],
                      self.image_processor_tester.crop_size["width"],
                  )							,   )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				Optional[int]      )->	List[str]:







              """simple docstring"""
              # Initialize image_processing
              UpperCamelCase       			=		self.image_processing_class(**self.image_processor_dict      )
              # create random numpy tensors
              UpperCamelCase       			=		prepare_image_inputs(self.image_processor_tester							,   equal_resolution=UpperCAmelCase_							,   numpify=UpperCAmelCase_      )
              for image in image_inputs:
                     self.assertIsInstance(UpperCAmelCase_							,   np.ndarray      )

              # Test not batched input
              UpperCamelCase       			=		image_processing(image_inputs[0]							,   return_tensors="pt"      ).pixel_values
              self.assertEqual(
                  encoded_images.shape							,   (
                      1,
                      self.image_processor_tester.num_channels,
                      self.image_processor_tester.crop_size["height"],
                      self.image_processor_tester.crop_size["width"],
                  )							,   )

              # Test batched
              UpperCamelCase       			=		image_processing(UpperCAmelCase_							,   return_tensors="pt"      ).pixel_values
              self.assertEqual(
                  encoded_images.shape							,   (
                      self.image_processor_tester.batch_size,
                      self.image_processor_tester.num_channels,
                      self.image_processor_tester.crop_size["height"],
                      self.image_processor_tester.crop_size["width"],
                  )							,   )






       def 						_SCREAMING_SNAKE_CASE							(	self			:				str      )->	str:







              """simple docstring"""
              # Initialize image_processing
              UpperCamelCase       			=		self.image_processing_class(**self.image_processor_dict      )
              # create random PyTorch tensors
              UpperCamelCase       			=		prepare_image_inputs(self.image_processor_tester							,   equal_resolution=UpperCAmelCase_							,   torchify=UpperCAmelCase_      )
              for image in image_inputs:
                     self.assertIsInstance(UpperCAmelCase_							,   torch.Tensor      )

              # Test not batched input
              UpperCamelCase       			=		image_processing(image_inputs[0]							,   return_tensors="pt"      ).pixel_values
              self.assertEqual(
                  encoded_images.shape							,   (
                      1,
                      self.image_processor_tester.num_channels,
                      self.image_processor_tester.crop_size["height"],
                      self.image_processor_tester.crop_size["width"],
                  )							,   )

              # Test batched
              UpperCamelCase       			=		image_processing(UpperCAmelCase_							,   return_tensors="pt"      ).pixel_values
              self.assertEqual(
                  encoded_images.shape							,   (
                      self.image_processor_tester.batch_size,
                      self.image_processor_tester.num_channels,
                      self.image_processor_tester.crop_size["height"],
                      self.image_processor_tester.crop_size["width"],
                  )							,   )





