






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 a_    (    snake_case__		,      unittest.TestCase							):
   _snake_case        =							ShapEPipeline
   _snake_case        =							['''prompt''']
   _snake_case        =							['''prompt''']
   _snake_case        =							[
       '''num_images_per_prompt''',
       '''num_inference_steps''',
       '''generator''',
       '''latents''',
       '''guidance_scale''',
       '''frame_size''',
       '''output_type''',
       '''return_dict''',
   ]
   _snake_case        =							False







   @property
   def     SCREAMING_SNAKE_CASE__					(self)			->      List[Any]:
      """simple docstring"""






      return 3_2







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






      return 3_2







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






      return self.time_input_dim * 4







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






      return 8







   @property
   def     SCREAMING_SNAKE_CASE__					(self)			->      List[Any]:
      """simple docstring"""






      __snake_case				:      List[Any]     							= CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip')
      return tokenizer







   @property
   def     SCREAMING_SNAKE_CASE__					(self)			->      Optional[int]:
      """simple docstring"""






      torch.manual_seed(0)
      __snake_case				:      Optional[int]     							= 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(_A)







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






      torch.manual_seed(0)

      __snake_case				:      int     							= {
          '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,
      }

      __snake_case				:      Optional[Any]     							= PriorTransformer(**_A)
      return model







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






      torch.manual_seed(0)

      __snake_case				:      List[str]     							= {
          '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,
          ),
      }
      __snake_case				:      List[Any]     							= ShapERenderer(**_A)
      return model







   def     SCREAMING_SNAKE_CASE__					(self)			->      Dict:
      """simple docstring"""






      __snake_case				:      List[str]     							= self.dummy_prior
      __snake_case				:      Optional[int]     							= self.dummy_text_encoder
      __snake_case				:      List[Any]     							= self.dummy_tokenizer
      __snake_case				:      str     							= self.dummy_renderer

      __snake_case				:      List[Any]     							= HeunDiscreteScheduler(
          beta_schedule='exp'					,   num_train_timesteps=1_0_2_4					,   prediction_type='sample'					,   use_karras_sigmas=_A					,   clip_sample=_A					,   clip_sample_range=1.0					,   )
      __snake_case				:      Any     							= {
          'prior': prior,
          'text_encoder': text_encoder,
          'tokenizer': tokenizer,
          'renderer': renderer,
          'scheduler': scheduler,
      }

      return components







   def     SCREAMING_SNAKE_CASE__					(self					,   __a					,   __a=0)			->      int:
      """simple docstring"""






      if str(_A).startswith('mps'):
         __snake_case				:      List[Any]     							= torch.manual_seed(_A)
      else:
         __snake_case				:      Dict     							= torch.Generator(device=_A).manual_seed(_A)
      __snake_case				:      int     							= {
          'prompt': 'horse',
          'generator': generator,
          'num_inference_steps': 1,
          'frame_size': 3_2,
          'output_type': 'np',
      }
      return inputs







   def     SCREAMING_SNAKE_CASE__					(self)			->      int:
      """simple docstring"""






      __snake_case				:      Tuple     							= 'cpu'

      __snake_case				:      Any     							= self.get_dummy_components()

      __snake_case				:      Tuple     							= self.pipeline_class(**_A)
      __snake_case				:      List[str]     							= pipe.to(_A)

      pipe.set_progress_bar_config(disable=_A)

      __snake_case				:      Tuple     							= pipe(**self.get_dummy_inputs(_A))
      __snake_case				:      int     							= output.images[0]
      __snake_case				:      str     							= image[0, -3:, -3:, -1]

      assert image.shape == (2_0, 3_2, 3_2, 3)

      __snake_case				:      Any     							= np.array(
          [
              0.00_039_216,
              0.00_039_216,
              0.00_039_216,
              0.00_039_216,
              0.00_039_216,
              0.00_039_216,
              0.00_039_216,
              0.00_039_216,
              0.00_039_216,
          ])

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







   def     SCREAMING_SNAKE_CASE__					(self)			->      str:
      """simple docstring"""






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







   def     SCREAMING_SNAKE_CASE__					(self)			->      str:
      """simple docstring"""






      __snake_case				:      List[str]     							= torch_device == 'cpu'
      __snake_case				:      Any     							= True

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







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






      __snake_case				:      Any     							= self.get_dummy_components()
      __snake_case				:      Any     							= self.pipeline_class(**_A)
      __snake_case				:      Dict     							= pipe.to(_A)
      pipe.set_progress_bar_config(disable=_A)

      __snake_case				:      Any     							= 1
      __snake_case				:      Dict     							= 2

      __snake_case				:      Tuple     							= self.get_dummy_inputs(_A)

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

      __snake_case				:      Optional[int]     							= pipe(**_A					,   num_images_per_prompt=_A)[0]

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




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







   def     SCREAMING_SNAKE_CASE__					(self)			->      int:
      """simple docstring"""






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







   def     SCREAMING_SNAKE_CASE__					(self)			->      List[str]:
      """simple docstring"""






      __snake_case				:      List[str]     							= load_numpy(
          'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
          '/shap_e/test_shap_e_np_out.npy')
      __snake_case				:      Dict     							= ShapEPipeline.from_pretrained('openai/shap-e')
      __snake_case				:      int     							= pipe.to(_A)
      pipe.set_progress_bar_config(disable=_A)

      __snake_case				:      str     							= torch.Generator(device=_A).manual_seed(0)

      __snake_case				:      Tuple     							= pipe(
          'a shark'					,   generator=_A					,   guidance_scale=1_5.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(_A					,   _A)







'''simple docstring'''






from functools import lru_cache




@lru_cache
def        _SCREAMING_SNAKE_CASE  (				A  :       int  )    -> int:







   """simple docstring"""
   if num < 0:
      raise ValueError('Number should not be negative.'  )

   return 1 if num in (0, 1) else num * factorial(num - 1  )


if __name__ == "__main__":
      import doctest

      doctest.testmod()