

'''simple docstring'''







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     ):
    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    snake_case_     (   self    )  -> Dict:





           '''simple docstring'''
           return 32



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





           '''simple docstring'''
           return 32



    @property
    def    snake_case_     (   self    )  -> Dict:





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



    @property
    def    snake_case_     (   self    )  -> Any:





           '''simple docstring'''
           return 8



    @property
    def    snake_case_     (   self    )  -> str:





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



    @property
    def    snake_case_     (   self    )  -> Dict:





           '''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=37						,			layer_norm_eps=1e-0_5						,			num_attention_heads=4						,			num_hidden_layers=5						,			pad_token_id=1						,			vocab_size=1000						,			)
           return CLIPTextModelWithProjection(UpperCamelCase__    )



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





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

           A_		   = {
               """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_		   = PriorTransformer(**UpperCamelCase__    )
           return model



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





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

           A_		   = {
               """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_		   = ShapERenderer(**UpperCamelCase__    )
           return model



    def    snake_case_     (   self    )  -> List[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=1024						,			prediction_type="""sample"""						,			use_karras_sigmas=UpperCamelCase__						,			clip_sample=UpperCamelCase__						,			clip_sample_range=1.0						,			)
           A_		   = {
               """prior""": prior,
               """text_encoder""": text_encoder,
               """tokenizer""": tokenizer,
               """renderer""": renderer,
               """scheduler""": scheduler,
           }

           return components



    def    snake_case_     (   self						,			UpperCamelCase__						,			UpperCamelCase__=0    )  -> List[str]:





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



    def    snake_case_     (   self    )  -> Dict:





           '''simple docstring'''
           A_		   = """cpu"""

           A_		   = self.get_dummy_components()

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

           pipe.set_progress_bar_config(disable=UpperCamelCase__    )

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

           assert image.shape == (20, 32, 32, 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    snake_case_     (   self    )  -> Union[str, 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    snake_case_     (   self    )  -> Optional[Any]:





           '''simple docstring'''
           A_		   = torch_device == """cpu"""
           A_		   = True

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







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





           '''simple docstring'''
           A_		   = self.get_dummy_components()
           A_		   = self.pipeline_class(**UpperCamelCase__    )
           A_		   = pipe.to(UpperCamelCase__    )
           pipe.set_progress_bar_config(disable=UpperCamelCase__    )

           A_		   = 1
           A_		   = 2

           A_		   = self.get_dummy_inputs(UpperCamelCase__    )

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

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

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


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



    def    snake_case_     (   self    )  -> Dict:





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







    def    snake_case_     (   self    )  -> List[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(UpperCamelCase__    )
           pipe.set_progress_bar_config(disable=UpperCamelCase__    )

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

           A_		   = pipe(
               """a shark"""						,			generator=UpperCamelCase__						,			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(UpperCamelCase__						,			UpperCamelCase__    )



'''simple docstring'''







def  UpperCAmelCase__      ( UpperCAmelCase__			)     ->     float:

       if edge <= 0 or not isinstance(UpperCAmelCase__,    UpperCAmelCase__			):
              raise ValueError("""Length must be a positive."""			)
       return 3 * ((25 + 10 * (5 ** (1 / 2))) ** (1 / 2)) * (edge**2)
def  UpperCAmelCase__      ( UpperCAmelCase__			)     ->     float:

       if edge <= 0 or not isinstance(UpperCAmelCase__,    UpperCAmelCase__			):
              raise ValueError("""Length must be a positive."""			)
       return ((15 + (7 * (5 ** (1 / 2)))) / 4) * (edge**3)


if __name__ == "__main__":
      import doctest

      doctest.testmod()
