

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





      @property
      def 					_lowerCamelCase  (  self		)   ->					Tuple:
         return 32





      @property
      def 					_lowerCamelCase  (  self		)   ->					Tuple:
         return 32





      @property
      def 					_lowerCamelCase  (  self		)   ->					Optional[Any]:
         return self.time_input_dim * 4





      @property
      def 					_lowerCamelCase  (  self		)   ->					Union[str, Any]:
         return 8





      @property
      def 					_lowerCamelCase  (  self		)   ->					List[str]:
         lowerCamelCase__						=			CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)
         return tokenizer





      @property
      def 					_lowerCamelCase  (  self		)   ->					Optional[int]:
         torch.manual_seed(0		)
         lowerCamelCase__						=			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(SCREAMING_SNAKE_CASE__		)





      @property
      def 					_lowerCamelCase  (  self		)   ->					Dict:
         torch.manual_seed(0		)

         lowerCamelCase__						=			{
             "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,
         }

         lowerCamelCase__						=			PriorTransformer(**SCREAMING_SNAKE_CASE__		)
         return model





      @property
      def 					_lowerCamelCase  (  self		)   ->					Optional[Any]:
         torch.manual_seed(0		)

         lowerCamelCase__						=			{
             "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,
             ),
         }
         lowerCamelCase__						=			ShapERenderer(**SCREAMING_SNAKE_CASE__		)
         return model





      def 					_lowerCamelCase  (  self		)   ->					Union[str, Any]:
         lowerCamelCase__						=			self.dummy_prior
         lowerCamelCase__						=			self.dummy_text_encoder
         lowerCamelCase__						=			self.dummy_tokenizer
         lowerCamelCase__						=			self.dummy_renderer

         lowerCamelCase__						=			HeunDiscreteScheduler(
             beta_schedule="exp"    ,						num_train_timesteps=1024    ,						prediction_type="sample"    ,						use_karras_sigmas=SCREAMING_SNAKE_CASE__    ,						clip_sample=SCREAMING_SNAKE_CASE__    ,						clip_sample_range=1.0    ,						)
         lowerCamelCase__						=			{
             "prior": prior,
             "text_encoder": text_encoder,
             "tokenizer": tokenizer,
             "renderer": renderer,
             "scheduler": scheduler,
         }

         return components





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





      def 					_lowerCamelCase  (  self		)   ->					Optional[Any]:
         lowerCamelCase__						=			"cpu"

         lowerCamelCase__						=			self.get_dummy_components()

         lowerCamelCase__						=			self.pipeline_class(**SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			pipe.to(SCREAMING_SNAKE_CASE__		)

         pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__		)

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

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

         lowerCamelCase__						=			np.array(
             [
                 0.00_03_92_16,
                 0.00_03_92_16,
                 0.00_03_92_16,
                 0.00_03_92_16,
                 0.00_03_92_16,
                 0.00_03_92_16,
                 0.00_03_92_16,
                 0.00_03_92_16,
                 0.00_03_92_16,
             ]		)

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





      def 					_lowerCamelCase  (  self		)   ->					List[Any]:
         # NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
         self._test_inference_batch_consistent(batch_sizes=[1, 2]		)





      def 					_lowerCamelCase  (  self		)   ->					str:
         lowerCamelCase__						=			torch_device == "cpu"
         lowerCamelCase__						=			True

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





      def 					_lowerCamelCase  (  self		)   ->					Optional[Any]:
         lowerCamelCase__						=			self.get_dummy_components()
         lowerCamelCase__						=			self.pipeline_class(**SCREAMING_SNAKE_CASE__		)
         lowerCamelCase__						=			pipe.to(SCREAMING_SNAKE_CASE__		)
         pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__		)

         lowerCamelCase__						=			1
         lowerCamelCase__						=			2

         lowerCamelCase__						=			self.get_dummy_inputs(SCREAMING_SNAKE_CASE__		)

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

         lowerCamelCase__						=			pipe(**SCREAMING_SNAKE_CASE__    ,						num_images_per_prompt=SCREAMING_SNAKE_CASE__		)[0]

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






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





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





      def 					_lowerCamelCase  (  self		)   ->					str:
         lowerCamelCase__						=			load_numpy(
             "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
             "/shap_e/test_shap_e_np_out.npy"		)
         lowerCamelCase__						=			ShapEPipeline.from_pretrained("openai/shap-e"		)
         lowerCamelCase__						=			pipe.to(SCREAMING_SNAKE_CASE__		)
         pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__		)

         lowerCamelCase__						=			torch.Generator(device=SCREAMING_SNAKE_CASE__		).manual_seed(0		)

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







"""simple docstring"""




# flake8: noqa
# Lint as: python3
SCREAMING_SNAKE_CASE_     :			Dict	  =			[
    '''VerificationMode''',
    '''Version''',
    '''disable_progress_bar''',
    '''enable_progress_bar''',
    '''is_progress_bar_enabled''',
    '''experimental''',
]

from .info_utils import VerificationMode
from .logging import disable_progress_bar, enable_progress_bar, is_progress_bar_enabled
from .version import Version
from .experimental import experimental




