

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    _lowercase				(					__lowerCamelCase,unittest.TestCase			):
   _lowercase     :       Optional[int]  					=    ShapEPipeline
   _lowercase     :       str  					=    ['prompt']
   _lowercase     :       int  					=    ['prompt']
   _lowercase     :       Optional[int]  					=    [
       'num_images_per_prompt',
       'num_inference_steps',
       'generator',
       'latents',
       'guidance_scale',
       'frame_size',
       'output_type',
       'return_dict',
   ]
   _lowercase     :       Any  					=    False

   @property
   def     UpperCamelCase					(				self	:    Any       ) ->			Dict:
         """simple docstring"""





         return 3_2

   @property
   def     UpperCamelCase					(				self	:    Dict       ) ->			Tuple:
         """simple docstring"""





         return 3_2

   @property
   def     UpperCamelCase					(				self	:    Union[str, Any]       ) ->			int:
         """simple docstring"""





         return self.time_input_dim * 4

   @property
   def     UpperCamelCase					(				self	:    Tuple       ) ->			Any:
         """simple docstring"""





         return 8

   @property
   def     UpperCamelCase					(				self	:    Optional[Any]       ) ->			str:
         """simple docstring"""





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

   @property
   def     UpperCamelCase					(				self	:    Dict       ) ->			Union[str, 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=3_7			,						layer_norm_eps=1e-05			,						num_attention_heads=4			,						num_hidden_layers=5			,						pad_token_id=1			,						vocab_size=1_0_0_0			,						)
         return CLIPTextModelWithProjection(lowerCamelCase__       )

   @property
   def     UpperCamelCase					(				self	:    List[Any]       ) ->			Union[str, Any]:
         """simple docstring"""





         torch.manual_seed(0       )

         A_							       =		{
             '''num_attention_heads''': 2,
             '''attention_head_dim''': 1_6,
             '''embedding_dim''': self.time_input_dim,
             '''num_embeddings''': 3_2,
             '''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(**lowerCamelCase__       )
         return model

   @property
   def     UpperCamelCase					(				self	:    str       ) ->			str:
         """simple docstring"""





         torch.manual_seed(0       )

         A_							       =		{
             '''param_shapes''': (
                 (self.renderer_dim, 9_3),
                 (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''': 1_2,
             '''background''': (
                 0.1,
                 0.1,
                 0.1,
             ),
         }
         A_							       =		ShapERenderer(**lowerCamelCase__       )
         return model

   def     UpperCamelCase					(				self	:    List[str]       ) ->			Optional[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_0_2_4			,						prediction_type='''sample'''			,						use_karras_sigmas=lowerCamelCase__			,						clip_sample=lowerCamelCase__			,						clip_sample_range=1.0			,						)
         A_							       =		{
             '''prior''': prior,
             '''text_encoder''': text_encoder,
             '''tokenizer''': tokenizer,
             '''renderer''': renderer,
             '''scheduler''': scheduler,
         }

         return components

   def     UpperCamelCase					(				self	:    Union[str, Any]			,						lowerCamelCase__	:    Dict			,						lowerCamelCase__	:    Optional[Any]=0       ) ->			Union[str, Any]:
         """simple docstring"""





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

   def     UpperCamelCase					(				self	:    Union[str, Any]       ) ->			List[Any]:
         """simple docstring"""





         A_							       =		'''cpu'''

         A_							       =		self.get_dummy_components()

         A_							       =		self.pipeline_class(**lowerCamelCase__       )
         A_							       =		pipe.to(lowerCamelCase__       )

         pipe.set_progress_bar_config(disable=lowerCamelCase__       )

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

         assert image.shape == (2_0, 3_2, 3_2, 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     UpperCamelCase					(				self	:    str       ) ->			List[str]:
         """simple docstring"""





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

   def     UpperCamelCase					(				self	:    List[Any]       ) ->			Any:
         """simple docstring"""





         A_							       =		torch_device == '''cpu'''
         A_							       =		True

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





   def     UpperCamelCase					(				self	:    Optional[Any]       ) ->			str:
         """simple docstring"""





         A_							       =		self.get_dummy_components()
         A_							       =		self.pipeline_class(**lowerCamelCase__       )
         A_							       =		pipe.to(lowerCamelCase__       )
         pipe.set_progress_bar_config(disable=lowerCamelCase__       )

         A_							       =		1
         A_							       =		2

         A_							       =		self.get_dummy_inputs(lowerCamelCase__       )

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

         A_							       =		pipe(**lowerCamelCase__			,						num_images_per_prompt=lowerCamelCase__       )[0]

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



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

   def     UpperCamelCase					(				self	:    int       ) ->			Dict:
         """simple docstring"""





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





   def     UpperCamelCase					(				self	:    Optional[int]       ) ->			str:
         """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(lowerCamelCase__       )
         pipe.set_progress_bar_config(disable=lowerCamelCase__       )

         A_							       =		torch.Generator(device=lowerCamelCase__       ).manual_seed(0       )

         A_							       =		pipe(
             '''a shark'''			,						generator=lowerCamelCase__			,						guidance_scale=15.0			,						num_inference_steps=6_4			,						frame_size=6_4			,						output_type='''np'''			,						).images[0]

         assert images.shape == (2_0, 6_4, 6_4, 3)

         assert_mean_pixel_difference(lowerCamelCase__			,						lowerCamelCase__       )








from __future__ import annotations







def 							_lowerCamelCase   (			SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE = None		,				SCREAMING_SNAKE_CASE = None   ):
      '''simple docstring'''







      if start is None:
            A_							       =		0

      if end is None:
            A_							       =		len(SCREAMING_SNAKE_CASE   ) - 1

      if start >= end:
            return

      A_							       =		(start + end) // 2

      slowsort(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE   )
      slowsort(SCREAMING_SNAKE_CASE		,				mid + 1		,				SCREAMING_SNAKE_CASE   )

      if sequence[end] < sequence[mid]:
            A_       ,A_							       =		sequence[mid], sequence[end]

      slowsort(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				end - 1   )


if __name__ == "__main__":
      from doctest import testmod

      testmod()





