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




  @property
  def 	__lowercase(      self     :	int				)->       Dict:




       """simple docstring"""
       return 32




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




       """simple docstring"""
       return 32




  @property
  def 	__lowercase(      self     :	Dict				)->       Dict:




       """simple docstring"""
       return self.time_input_dim * 4




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




       """simple docstring"""
       return 8




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




       """simple docstring"""
       SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'				)
       return tokenizer




  @property
  def 	__lowercase(      self     :	str				)->       Dict:




       """simple docstring"""
       torch.manual_seed(0				)
       SCREAMING_SNAKE_CASE__				:				Dict			=		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-0_5     , num_attention_heads=4     , num_hidden_layers=5     , pad_token_id=1     , vocab_size=1000     , )
       return CLIPTextModelWithProjection(a_				)




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




       """simple docstring"""
       torch.manual_seed(0				)

       SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		{
           '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,
       }

       SCREAMING_SNAKE_CASE__				:				List[Any]			=		PriorTransformer(**a_				)
       return model




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




       """simple docstring"""
       torch.manual_seed(0				)

       SCREAMING_SNAKE_CASE__				:				Tuple			=		{
           '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,
           ),
       }
       SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		ShapERenderer(**a_				)
       return model




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




       """simple docstring"""
       SCREAMING_SNAKE_CASE__				:				int			=		self.dummy_prior
       SCREAMING_SNAKE_CASE__				:				Optional[int]			=		self.dummy_text_encoder
       SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		self.dummy_tokenizer
       SCREAMING_SNAKE_CASE__				:				Any			=		self.dummy_renderer

       SCREAMING_SNAKE_CASE__				:				Tuple			=		HeunDiscreteScheduler(
           beta_schedule='exp'     , num_train_timesteps=1024     , prediction_type='sample'     , use_karras_sigmas=a_     , clip_sample=a_     , clip_sample_range=1.0     , )
       SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		{
           'prior': prior,
           'text_encoder': text_encoder,
           'tokenizer': tokenizer,
           'renderer': renderer,
           'scheduler': scheduler,
       }

       return components




  def 	__lowercase(      self     :	Optional[Any]     , a_     :	Union[str, Any]     , a_     :	List[Any]=0				)->       List[str]:




       """simple docstring"""
       if str(a_				).startswith('mps'				):
            SCREAMING_SNAKE_CASE__				:				Optional[int]			=		torch.manual_seed(a_				)
       else:
            SCREAMING_SNAKE_CASE__				:				List[Any]			=		torch.Generator(device=a_				).manual_seed(a_				)
       SCREAMING_SNAKE_CASE__				:				Tuple			=		{
           'prompt': 'horse',
           'generator': generator,
           'num_inference_steps': 1,
           'frame_size': 32,
           'output_type': 'np',
       }
       return inputs




  def 	__lowercase(      self     :	Any				)->       Dict:




       """simple docstring"""
       SCREAMING_SNAKE_CASE__				:				List[str]			=		'cpu'

       SCREAMING_SNAKE_CASE__				:				str			=		self.get_dummy_components()

       SCREAMING_SNAKE_CASE__				:				Any			=		self.pipeline_class(**a_				)
       SCREAMING_SNAKE_CASE__				:				List[str]			=		pipe.to(a_				)

       pipe.set_progress_bar_config(disable=a_				)

       SCREAMING_SNAKE_CASE__				:				Any			=		pipe(**self.get_dummy_inputs(a_				)				)
       SCREAMING_SNAKE_CASE__				:				Dict			=		output.images[0]
       SCREAMING_SNAKE_CASE__				:				Any			=		image[0, -3:, -3:, -1]

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

       SCREAMING_SNAKE_CASE__				:				Tuple			=		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 	__lowercase(      self     :	Union[str, Any]				)->       Any:




       """simple docstring"""
       # NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
       self._test_inference_batch_consistent(batch_sizes=[1, 2]				)




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




       """simple docstring"""
       SCREAMING_SNAKE_CASE__				:				Any			=		torch_device == 'cpu'
       SCREAMING_SNAKE_CASE__				:				Optional[int]			=		True

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




  def 	__lowercase(      self     :	Dict				)->       List[Any]:




       """simple docstring"""
       SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		self.get_dummy_components()
       SCREAMING_SNAKE_CASE__				:				List[Any]			=		self.pipeline_class(**a_				)
       SCREAMING_SNAKE_CASE__				:				Any			=		pipe.to(a_				)
       pipe.set_progress_bar_config(disable=a_				)

       SCREAMING_SNAKE_CASE__				:				Optional[int]			=		1
       SCREAMING_SNAKE_CASE__				:				str			=		2

       SCREAMING_SNAKE_CASE__				:				Union[str, Any]			=		self.get_dummy_inputs(a_				)

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

       SCREAMING_SNAKE_CASE__				:				List[str]			=		pipe(**a_     , num_images_per_prompt=a_				)[0]

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




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




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




       """simple docstring"""
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()




  def 	__lowercase(      self     :	Any				)->       Tuple:




       """simple docstring"""
       SCREAMING_SNAKE_CASE__				:				Dict			=		load_numpy(
           'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
           '/shap_e/test_shap_e_np_out.npy'				)
       SCREAMING_SNAKE_CASE__				:				Dict			=		ShapEPipeline.from_pretrained('openai/shap-e'				)
       SCREAMING_SNAKE_CASE__				:				Optional[Any]			=		pipe.to(a_				)
       pipe.set_progress_bar_config(disable=a_				)

       SCREAMING_SNAKE_CASE__				:				Dict			=		torch.Generator(device=a_				).manual_seed(0				)

       SCREAMING_SNAKE_CASE__				:				int			=		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 math







def     _a				(      lowercase__  :			int ):







     '''simple docstring'''


     assert isinstance(lowercase__     ,  lowercase__ ) and (
         number >= 0
     ), "'number' must been an int and positive"

     if 1 < number < 4:
          # 2 and 3 are primes
          return True
     elif number < 2 or not number % 2:
          # Negatives, 0, 1 and all even numbers are not primes
          return False

     SCREAMING_SNAKE_CASE__				:				Tuple			=		range(3     ,  int(math.sqrt(lowercase__ ) + 1 )     ,  2 )
     return not any(not number % i for i in odd_numbers )







def     _a				(      lowercase__  :			List[str]     ,  lowercase__  :			Any=1     ,  **lowercase__  :			Any ):







     '''simple docstring'''


     SCREAMING_SNAKE_CASE__				:				int			=		factor * value
     SCREAMING_SNAKE_CASE__				:				Dict			=		value

     while not is_prime(lowercase__ ):
          value += 1 if not ("desc" in kwargs and kwargs["desc"] is True) else -1

     if value == first_value_val:
          return next_prime(value + 1     ,  **lowercase__ )
     return value



