






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 					SCREAMING_SNAKE_CASE_  (       _a  ,						unittest.TestCase ):
       """simple docstring"""


       __lowerCAmelCase    :		Tuple					=ShapEPipeline
       __lowerCAmelCase    :		Union[str, Any]					=['''prompt''']
       __lowerCAmelCase    :		List[str]					=['''prompt''']
       __lowerCAmelCase    :		Optional[int]					=[
           '''num_images_per_prompt''',
           '''num_inference_steps''',
           '''generator''',
           '''latents''',
           '''guidance_scale''',
           '''frame_size''',
           '''output_type''',
           '''return_dict''',
       ]
       __lowerCAmelCase    :		Dict					=False





       @property
       def 				UpperCamelCase__						(				self	:Dict):



          """simple docstring"""

          return 32





       @property
       def 				UpperCamelCase__						(				self	:Tuple):



          """simple docstring"""

          return 32





       @property
       def 				UpperCamelCase__						(				self	:Dict):



          """simple docstring"""

          return self.time_input_dim * 4





       @property
       def 				UpperCamelCase__						(				self	:Union[str, Any]):



          """simple docstring"""

          return 8





       @property
       def 				UpperCamelCase__						(				self	:Optional[Any]):



          """simple docstring"""

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





       @property
       def 				UpperCamelCase__						(				self	:Union[str, Any]):



          """simple docstring"""

          torch.manual_seed(0)
          _lowercase              =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(snake_case)





       @property
       def 				UpperCamelCase__						(				self	:str):



          """simple docstring"""

          torch.manual_seed(0)

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

          _lowercase              =PriorTransformer(**snake_case)
          return model





       @property
       def 				UpperCamelCase__						(				self	:Any):



          """simple docstring"""

          torch.manual_seed(0)

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





       def 				UpperCamelCase__						(				self	:Tuple):



          """simple docstring"""

          _lowercase              =self.dummy_prior
          _lowercase              =self.dummy_text_encoder
          _lowercase              =self.dummy_tokenizer
          _lowercase              =self.dummy_renderer

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

          return components





       def 				UpperCamelCase__						(				self	:List[str],      snake_case	:str,      snake_case	:Tuple=0):



          """simple docstring"""

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





       def 				UpperCamelCase__						(				self	:Optional[Any]):



          """simple docstring"""

          _lowercase              ='cpu'

          _lowercase              =self.get_dummy_components()

          _lowercase              =self.pipeline_class(**snake_case)
          _lowercase              =pipe.to(snake_case)

          pipe.set_progress_bar_config(disable=snake_case)

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

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

          _lowercase              =np.array(
              [
                  0.0_0_0_3_9_2_1_6,
                  0.0_0_0_3_9_2_1_6,
                  0.0_0_0_3_9_2_1_6,
                  0.0_0_0_3_9_2_1_6,
                  0.0_0_0_3_9_2_1_6,
                  0.0_0_0_3_9_2_1_6,
                  0.0_0_0_3_9_2_1_6,
                  0.0_0_0_3_9_2_1_6,
                  0.0_0_0_3_9_2_1_6,
              ])

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





       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

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





       def 				UpperCamelCase__						(				self	:Tuple):



          """simple docstring"""

          _lowercase              =torch_device == 'cpu'
          _lowercase              =True

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





       def 				UpperCamelCase__						(				self	:Union[str, Any]):



          """simple docstring"""

          _lowercase              =self.get_dummy_components()
          _lowercase              =self.pipeline_class(**snake_case)
          _lowercase              =pipe.to(snake_case)
          pipe.set_progress_bar_config(disable=snake_case)

          _lowercase              =1
          _lowercase              =2

          _lowercase              =self.get_dummy_inputs(snake_case)

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

          _lowercase              =pipe(**snake_case,      num_images_per_prompt=snake_case)[0]

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






@slow
@require_torch_gpu
class 					SCREAMING_SNAKE_CASE_  (       unittest.TestCase ):
       """simple docstring"""


       def 				UpperCamelCase__						(				self	:Union[str, Any]):



          """simple docstring"""

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





       def 				UpperCamelCase__						(				self	:int):



          """simple docstring"""

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

          _lowercase              =torch.Generator(device=snake_case).manual_seed(0)

          _lowercase              =pipe(
              'a shark',      generator=snake_case,      guidance_scale=1_5.0,      num_inference_steps=64,      frame_size=64,      output_type='np',      ).images[0]

          assert images.shape == (20, 64, 64, 3)

          assert_mean_pixel_difference(snake_case,      snake_case)










from ...configuration_utils import PretrainedConfig
from ...utils import logging


_SCREAMING_SNAKE_CASE   =  logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE   =  {
    "uw-madison/mra-base-512-4": "https://huggingface.co/uw-madison/mra-base-512-4/resolve/main/config.json",
}






class 					SCREAMING_SNAKE_CASE_  (       _a ):
       """simple docstring"""


       __lowerCAmelCase    :		str					='''mra'''





       def __init__(				self	:Any,      snake_case	:List[str]=5_0265,      snake_case	:List[Any]=768,      snake_case	:Optional[Any]=12,      snake_case	:Optional[Any]=12,      snake_case	:str=3072,      snake_case	:Tuple="gelu",      snake_case	:Optional[int]=0.1,      snake_case	:int=0.1,      snake_case	:Any=512,      snake_case	:Union[str, Any]=1,      snake_case	:Union[str, Any]=0.0_2,      snake_case	:List[Any]=1e-5,      snake_case	:Optional[int]="absolute",      snake_case	:Optional[int]=4,      snake_case	:str="full",      snake_case	:Optional[int]=0,      snake_case	:List[Any]=0,      snake_case	:int=1,      snake_case	:List[Any]=0,      snake_case	:Dict=2,      **snake_case	:Dict,      ):



          """simple docstring"""

          super().__init__(pad_token_id=snake_case,      bos_token_id=snake_case,      eos_token_id=snake_case,      **snake_case)

          _lowercase              =vocab_size
          _lowercase              =max_position_embeddings
          _lowercase              =hidden_size
          _lowercase              =num_hidden_layers
          _lowercase              =num_attention_heads
          _lowercase              =intermediate_size
          _lowercase              =hidden_act
          _lowercase              =hidden_dropout_prob
          _lowercase              =attention_probs_dropout_prob
          _lowercase              =initializer_range
          _lowercase              =type_vocab_size
          _lowercase              =layer_norm_eps
          _lowercase              =position_embedding_type
          _lowercase              =block_per_row
          _lowercase              =approx_mode
          _lowercase              =initial_prior_first_n_blocks
          _lowercase              =initial_prior_diagonal_n_blocks


