


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




 @property
 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       return 32




 @property
 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       return 32




 @property
 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       return self.time_input_dim * 4




 @property
 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       return 8




 @property
 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:Optional[Any]									=			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''       )
       return tokenizer




 @property
 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       torch.manual_seed(0       )
       a		:Tuple									=			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=1000   ,    )
       return CLIPTextModelWithProjection(_lowerCamelCase       )




 @property
 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       torch.manual_seed(0       )

       a		:List[str]									=			{
           '''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		:Any									=			PriorTransformer(**_lowerCamelCase       )
       return model




 @property
 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       torch.manual_seed(0       )

       a		:Dict									=			{
           '''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		:List[str]									=			ShapERenderer(**_lowerCamelCase       )
       return model




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:Union[str, Any]									=			self.dummy_prior
       a		:str									=			self.dummy_text_encoder
       a		:Optional[Any]									=			self.dummy_tokenizer
       a		:List[Any]									=			self.dummy_renderer

       a		:Any									=			HeunDiscreteScheduler(
           beta_schedule='''exp'''   ,    num_train_timesteps=1024   ,    prediction_type='''sample'''   ,    use_karras_sigmas=_lowerCamelCase   ,    clip_sample=_lowerCamelCase   ,    clip_sample_range=1.0   ,    )
       a		:int									=			{
           '''prior''': prior,
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
           '''renderer''': renderer,
           '''scheduler''': scheduler,
       }

       return components




 def 			SCREAMING_SNAKE_CASE__							(						self   ,    _lowerCamelCase   ,    _lowerCamelCase=0       ):
       if str(_lowerCamelCase       ).startswith('''mps'''       ):
             a		:List[str]									=			torch.manual_seed(_lowerCamelCase       )
       else:
             a		:Any									=			torch.Generator(device=_lowerCamelCase       ).manual_seed(_lowerCamelCase       )
       a		:List[str]									=			{
           '''prompt''': '''horse''',
           '''generator''': generator,
           '''num_inference_steps''': 1,
           '''frame_size''': 32,
           '''output_type''': '''np''',
       }
       return inputs




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:Tuple									=			'''cpu'''

       a		:Dict									=			self.get_dummy_components()

       a		:Tuple									=			self.pipeline_class(**_lowerCamelCase       )
       a		:Any									=			pipe.to(_lowerCamelCase       )

       pipe.set_progress_bar_config(disable=_lowerCamelCase       )

       a		:List[str]									=			pipe(**self.get_dummy_inputs(_lowerCamelCase       )       )
       a		:Union[str, Any]									=			output.images[0]
       a		:Tuple									=			image[0, -3:, -3:, -1]

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

       a		:Dict									=			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 			SCREAMING_SNAKE_CASE__							(						self       ):
       # NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
       self._test_inference_batch_consistent(batch_sizes=[1, 2]       )




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:Tuple									=			torch_device == '''cpu'''
       a		:List[Any]									=			True

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




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:Union[str, Any]									=			self.get_dummy_components()
       a		:str									=			self.pipeline_class(**_lowerCamelCase       )
       a		:Tuple									=			pipe.to(_lowerCamelCase       )
       pipe.set_progress_bar_config(disable=_lowerCamelCase       )

       a		:Dict									=			1
       a		:Tuple									=			2

       a		:str									=			self.get_dummy_inputs(_lowerCamelCase       )

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

       a		:List[str]									=			pipe(**_lowerCamelCase   ,    num_images_per_prompt=_lowerCamelCase       )[0]

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







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




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()




 def 			SCREAMING_SNAKE_CASE__							(						self       ):
       a		:Union[str, Any]									=			load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/shap_e/test_shap_e_np_out.npy'''       )
       a		:str									=			ShapEPipeline.from_pretrained('''openai/shap-e'''       )
       a		:Any									=			pipe.to(_lowerCamelCase       )
       pipe.set_progress_bar_config(disable=_lowerCamelCase       )

       a		:List[Any]									=			torch.Generator(device=_lowerCamelCase       ).manual_seed(0       )

       a		:Union[str, Any]									=			pipe(
           '''a shark'''   ,    generator=_lowerCamelCase   ,    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(_lowerCamelCase   ,    _lowerCamelCase       )









from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_torch_available


snake_case							:       Any										=	{
    '''configuration_gpt_neo''': ['''GPT_NEO_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''GPTNeoConfig''', '''GPTNeoOnnxConfig'''],
}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     snake_case							:       List[Any]										=	[
         '''GPT_NEO_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''GPTNeoForCausalLM''',
         '''GPTNeoForQuestionAnswering''',
         '''GPTNeoForSequenceClassification''',
         '''GPTNeoForTokenClassification''',
         '''GPTNeoModel''',
         '''GPTNeoPreTrainedModel''',
         '''load_tf_weights_in_gpt_neo''',
     ]

try:
     if not is_flax_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     snake_case							:       Any										=	[
         '''FlaxGPTNeoForCausalLM''',
         '''FlaxGPTNeoModel''',
         '''FlaxGPTNeoPreTrainedModel''',
     ]


if TYPE_CHECKING:
     from .configuration_gpt_neo import GPT_NEO_PRETRAINED_CONFIG_ARCHIVE_MAP, GPTNeoConfig, GPTNeoOnnxConfig

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_gpt_neo import (
              GPT_NEO_PRETRAINED_MODEL_ARCHIVE_LIST,
              GPTNeoForCausalLM,
              GPTNeoForQuestionAnswering,
              GPTNeoForSequenceClassification,
              GPTNeoForTokenClassification,
              GPTNeoModel,
              GPTNeoPreTrainedModel,
              load_tf_weights_in_gpt_neo,
          )

     try:
          if not is_flax_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_flax_gpt_neo import FlaxGPTNeoForCausalLM, FlaxGPTNeoModel, FlaxGPTNeoPreTrainedModel


else:
     import sys

     snake_case							:       Optional[int]										=	_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)





