


import gc
import random
import unittest

import numpy as np
import torch
from transformers import CLIPImageProcessor, CLIPVisionConfig, CLIPVisionModel

from diffusers import HeunDiscreteScheduler, PriorTransformer, ShapEImgaImgPipeline
from diffusers.pipelines.shap_e import ShapERenderer
from diffusers.utils import floats_tensor, load_image, 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   __lowercase				(  __snake_case ,					unittest.TestCase  ):
      UpperCamelCase								=       ShapEImgaImgPipeline
      UpperCamelCase								=       ['''image''']
      UpperCamelCase								=       ['''image''']
      UpperCamelCase								=       [
          '''num_images_per_prompt''',
          '''num_inference_steps''',
          '''generator''',
          '''latents''',
          '''guidance_scale''',
          '''frame_size''',
          '''output_type''',
          '''return_dict''',
      ]
      UpperCamelCase								=       False





      @property
      def        _lowercase				(		self						:					Dict       )    ->   Any:






         """simple docstring"""







         return 3_2





      @property
      def        _lowercase				(		self						:					Any       )    ->   int:






         """simple docstring"""







         return 3_2





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






         """simple docstring"""







         return self.time_input_dim * 4





      @property
      def        _lowercase				(		self						:					List[str]       )    ->   Optional[Any]:






         """simple docstring"""







         return 8





      @property
      def        _lowercase				(		self						:					int       )    ->   Tuple:






         """simple docstring"""







         torch.manual_seed(0       )
         UpperCAmelCase							=     CLIPVisionConfig(
             hidden_size=self.text_embedder_hidden_size					,       image_size=6_4					,       projection_dim=self.text_embedder_hidden_size					,       intermediate_size=3_7					,       num_attention_heads=4					,       num_channels=3					,       num_hidden_layers=5					,       patch_size=1					,       )

         UpperCAmelCase							=     CLIPVisionModel(__lowerCamelCase       )
         return model





      @property
      def        _lowercase				(		self						:					Tuple       )    ->   Tuple:






         """simple docstring"""







         UpperCAmelCase							=     CLIPImageProcessor(
             crop_size=2_2_4					,       do_center_crop=__lowerCamelCase					,       do_normalize=__lowerCamelCase					,       do_resize=__lowerCamelCase					,       image_mean=[0.48_145_466, 0.4_578_275, 0.40_821_073]					,       image_std=[0.26_862_954, 0.26_130_258, 0.27_577_711]					,       resample=3					,       size=2_2_4					,       )

         return image_processor





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






         """simple docstring"""







         torch.manual_seed(0       )

         UpperCAmelCase							=     {
             """num_attention_heads""": 2,
             """attention_head_dim""": 1_6,
             """embedding_dim""": self.time_input_dim,
             """num_embeddings""": 3_2,
             """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""",
             """embedding_proj_norm_type""": """layer""",
             """encoder_hid_proj_type""": None,
             """added_emb_type""": None,
         }

         UpperCAmelCase							=     PriorTransformer(**__lowerCamelCase       )
         return model





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






         """simple docstring"""







         torch.manual_seed(0       )

         UpperCAmelCase							=     {
             """param_shapes""": (
                 (self.renderer_dim, 9_3),
                 (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""": 1_2,
             """background""": (
                 0.1,
                 0.1,
                 0.1,
             ),
         }
         UpperCAmelCase							=     ShapERenderer(**__lowerCamelCase       )
         return model





      def        _lowercase				(		self						:					Dict       )    ->   int:






         """simple docstring"""







         UpperCAmelCase							=     self.dummy_prior
         UpperCAmelCase							=     self.dummy_image_encoder
         UpperCAmelCase							=     self.dummy_image_processor
         UpperCAmelCase							=     self.dummy_renderer

         UpperCAmelCase							=     HeunDiscreteScheduler(
             beta_schedule="""exp"""					,       num_train_timesteps=1_0_2_4					,       prediction_type="""sample"""					,       use_karras_sigmas=__lowerCamelCase					,       clip_sample=__lowerCamelCase					,       clip_sample_range=1.0					,       )
         UpperCAmelCase							=     {
             """prior""": prior,
             """image_encoder""": image_encoder,
             """image_processor""": image_processor,
             """renderer""": renderer,
             """scheduler""": scheduler,
         }

         return components





      def        _lowercase				(		self						:					Any					,       __lowerCamelCase						:					Dict					,       __lowerCamelCase						:					List[Any]=0       )    ->   List[Any]:






         """simple docstring"""







         UpperCAmelCase							=     floats_tensor((1, 3, 6_4, 6_4)					,       rng=random.Random(__lowerCamelCase       )       ).to(__lowerCamelCase       )

         if str(__lowerCamelCase       ).startswith("""mps"""       ):
            UpperCAmelCase							=     torch.manual_seed(__lowerCamelCase       )
         else:
            UpperCAmelCase							=     torch.Generator(device=__lowerCamelCase       ).manual_seed(__lowerCamelCase       )
         UpperCAmelCase							=     {
             """image""": input_image,
             """generator""": generator,
             """num_inference_steps""": 1,
             """frame_size""": 3_2,
             """output_type""": """np""",
         }
         return inputs





      def        _lowercase				(		self						:					str       )    ->   List[str]:






         """simple docstring"""







         UpperCAmelCase							=     """cpu"""

         UpperCAmelCase							=     self.get_dummy_components()

         UpperCAmelCase							=     self.pipeline_class(**__lowerCamelCase       )
         UpperCAmelCase							=     pipe.to(__lowerCamelCase       )

         pipe.set_progress_bar_config(disable=__lowerCamelCase       )

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

         assert image.shape == (2_0, 3_2, 3_2, 3)

         UpperCAmelCase							=     np.array(
             [
                 0.00_039_216,
                 0.00_039_216,
                 0.00_039_216,
                 0.00_039_216,
                 0.00_039_216,
                 0.00_039_216,
                 0.00_039_216,
                 0.00_039_216,
                 0.00_039_216,
             ]       )

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





      def        _lowercase				(		self						:					List[Any]       )    ->   Any:






         """simple docstring"""







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





      def        _lowercase				(		self						:					Any       )    ->   Any:






         """simple docstring"""







         UpperCAmelCase							=     torch_device == """cpu"""
         UpperCAmelCase							=     True
         self._test_inference_batch_single_identical(
             batch_size=2					,       test_max_difference=__lowerCamelCase					,       relax_max_difference=__lowerCamelCase					,       )





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






         """simple docstring"""







         UpperCAmelCase							=     self.get_dummy_components()
         UpperCAmelCase							=     self.pipeline_class(**__lowerCamelCase       )
         UpperCAmelCase							=     pipe.to(__lowerCamelCase       )
         pipe.set_progress_bar_config(disable=__lowerCamelCase       )

         UpperCAmelCase							=     1
         UpperCAmelCase							=     2

         UpperCAmelCase							=     self.get_dummy_inputs(__lowerCamelCase       )

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

         UpperCAmelCase							=     pipe(**__lowerCamelCase					,       num_images_per_prompt=__lowerCamelCase       )[0]

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






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





      def        _lowercase				(		self						:					Dict       )    ->   Optional[int]:






         """simple docstring"""







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





      def        _lowercase				(		self						:					List[str]       )    ->   Tuple:






         """simple docstring"""







         UpperCAmelCase							=     load_image(
             """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main""" """/shap_e/corgi.png"""       )
         UpperCAmelCase							=     load_numpy(
             """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
             """/shap_e/test_shap_e_img2img_out.npy"""       )
         UpperCAmelCase							=     ShapEImgaImgPipeline.from_pretrained("""openai/shap-e-img2img"""       )
         UpperCAmelCase							=     pipe.to(__lowerCamelCase       )
         pipe.set_progress_bar_config(disable=__lowerCamelCase       )

         UpperCAmelCase							=     torch.Generator(device=__lowerCamelCase       ).manual_seed(0       )

         UpperCAmelCase							=     pipe(
             __lowerCamelCase					,       generator=__lowerCamelCase					,       guidance_scale=3.0					,       num_inference_steps=6_4					,       frame_size=6_4					,       output_type="""np"""					,       ).images[0]

         assert images.shape == (2_0, 6_4, 6_4, 3)

         assert_mean_pixel_difference(__lowerCamelCase					,       __lowerCamelCase       )







from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available


__a     =			{"""tokenization_herbert""": ["""HerbertTokenizer"""]}

try:
 if not is_tokenizers_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __a     =			["""HerbertTokenizerFast"""]


if TYPE_CHECKING:
 from .tokenization_herbert import HerbertTokenizer

 try:
  if not is_tokenizers_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_herbert_fast import HerbertTokenizerFast

else:
 import sys

 __a     =			_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



