





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__   (							__UpperCAmelCase			,  unittest.TestCase     ):




   """simple docstring"""



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

   @property
   def 							__lowercase		( self)      ->						Dict:







          '''simple docstring'''

          return 32

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







          '''simple docstring'''

          return 32

   @property
   def 							__lowercase		( self)      ->						Union[str, Any]:







          '''simple docstring'''

          return self.time_input_dim * 4

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







          '''simple docstring'''

          return 8

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







          '''simple docstring'''

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

   @property
   def 							__lowercase		( self)      ->						int:







          '''simple docstring'''

          torch.manual_seed(0)
          a__						:     Union[str, Any]			=    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(lowercase)

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







          '''simple docstring'''

          torch.manual_seed(0)

          a__						:     Dict			=    {
              '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__						:     int			=    PriorTransformer(**lowercase)
          return model

   @property
   def 							__lowercase		( self)      ->						Any:







          '''simple docstring'''

          torch.manual_seed(0)

          a__						:     Optional[int]			=    {
              '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__						:     Any			=    ShapERenderer(**lowercase)
          return model

   def 							__lowercase		( self)      ->						int:







          '''simple docstring'''

          a__						:     List[str]			=    self.dummy_prior
          a__						:     List[Any]			=    self.dummy_text_encoder
          a__						:     str			=    self.dummy_tokenizer
          a__						:     Union[str, Any]			=    self.dummy_renderer

          a__						:     Optional[Any]			=    HeunDiscreteScheduler(
              beta_schedule='exp' ,  num_train_timesteps=1024 ,  prediction_type='sample' ,  use_karras_sigmas=lowercase ,  clip_sample=lowercase ,  clip_sample_range=1.0 ,  )
          a__						:     Union[str, Any]			=    {
              'prior': prior,
              'text_encoder': text_encoder,
              'tokenizer': tokenizer,
              'renderer': renderer,
              'scheduler': scheduler,
          }

          return components

   def 							__lowercase		( self ,  lowercase ,  lowercase=0)      ->						List[Any]:







          '''simple docstring'''

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

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







          '''simple docstring'''

          a__						:     int			=    'cpu'

          a__						:     Optional[int]			=    self.get_dummy_components()

          a__						:     int			=    self.pipeline_class(**lowercase)
          a__						:     int			=    pipe.to(lowercase)

          pipe.set_progress_bar_config(disable=lowercase)

          a__						:     Optional[Any]			=    pipe(**self.get_dummy_inputs(lowercase))
          a__						:     Union[str, Any]			=    output.images[0]
          a__						:     Optional[int]			=    image[0, -3:, -3:, -1]

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

          a__						:     int			=    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 							__lowercase		( self)      ->						int:







          '''simple docstring'''

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

   def 							__lowercase		( self)      ->						Optional[int]:







          '''simple docstring'''

          a__						:     List[Any]			=    torch_device == 'cpu'
          a__						:     Dict			=    True

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

   def 							__lowercase		( self)      ->						int:







          '''simple docstring'''

          a__						:     Any			=    self.get_dummy_components()
          a__						:     Optional[int]			=    self.pipeline_class(**lowercase)
          a__						:     Dict			=    pipe.to(lowercase)
          pipe.set_progress_bar_config(disable=lowercase)

          a__						:     Dict			=    1
          a__						:     int			=    2

          a__						:     Union[str, Any]			=    self.get_dummy_inputs(lowercase)

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

          a__						:     Optional[int]			=    pipe(**lowercase ,  num_images_per_prompt=lowercase)[0]

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





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




   """simple docstring"""



   def 							__lowercase		( self)      ->						Tuple:







          '''simple docstring'''

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

   def 							__lowercase		( self)      ->						Optional[int]:







          '''simple docstring'''

          a__						:     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__						:     str			=    pipe.to(lowercase)
          pipe.set_progress_bar_config(disable=lowercase)

          a__						:     int			=    torch.Generator(device=lowercase).manual_seed(0)

          a__						:     List[str]			=    pipe(
              'a shark' ,  generator=lowercase ,  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(lowercase ,  lowercase)









from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
    is_vision_available,
)


lowercase				:	List[str]      						=  {
    """configuration_clip""": [
        """CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP""",
        """CLIPConfig""",
        """CLIPOnnxConfig""",
        """CLIPTextConfig""",
        """CLIPVisionConfig""",
    ],
    """processing_clip""": ["""CLIPProcessor"""],
    """tokenization_clip""": ["""CLIPTokenizer"""],
}

try:
    if not is_tokenizers_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    lowercase				:	Tuple      						=  ["""CLIPTokenizerFast"""]

try:
    if not is_vision_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    lowercase				:	Optional[int]      						=  ["""CLIPFeatureExtractor"""]
    lowercase				:	Union[str, Any]      						=  ["""CLIPImageProcessor"""]

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    lowercase				:	Dict      						=  [
        """CLIP_PRETRAINED_MODEL_ARCHIVE_LIST""",
        """CLIPModel""",
        """CLIPPreTrainedModel""",
        """CLIPTextModel""",
        """CLIPTextModelWithProjection""",
        """CLIPVisionModel""",
        """CLIPVisionModelWithProjection""",
    ]

try:
    if not is_tf_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    lowercase				:	List[str]      						=  [
        """TF_CLIP_PRETRAINED_MODEL_ARCHIVE_LIST""",
        """TFCLIPModel""",
        """TFCLIPPreTrainedModel""",
        """TFCLIPTextModel""",
        """TFCLIPVisionModel""",
    ]

try:
    if not is_flax_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    lowercase				:	Any      						=  [
        """FlaxCLIPModel""",
        """FlaxCLIPPreTrainedModel""",
        """FlaxCLIPTextModel""",
        """FlaxCLIPTextPreTrainedModel""",
        """FlaxCLIPVisionModel""",
        """FlaxCLIPVisionPreTrainedModel""",
    ]


if TYPE_CHECKING:
    from .configuration_clip import (
        CLIP_PRETRAINED_CONFIG_ARCHIVE_MAP,
        CLIPConfig,
        CLIPOnnxConfig,
        CLIPTextConfig,
        CLIPVisionConfig,
    )
    from .processing_clip import CLIPProcessor
    from .tokenization_clip import CLIPTokenizer

    try:
        if not is_tokenizers_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .tokenization_clip_fast import CLIPTokenizerFast

    try:
        if not is_vision_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .feature_extraction_clip import CLIPFeatureExtractor
        from .image_processing_clip import CLIPImageProcessor

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_clip import (
            CLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
            CLIPModel,
            CLIPPreTrainedModel,
            CLIPTextModel,
            CLIPTextModelWithProjection,
            CLIPVisionModel,
            CLIPVisionModelWithProjection,
        )

    try:
        if not is_tf_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_tf_clip import (
            TF_CLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
            TFCLIPModel,
            TFCLIPPreTrainedModel,
            TFCLIPTextModel,
            TFCLIPVisionModel,
        )

    try:
        if not is_flax_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_flax_clip import (
            FlaxCLIPModel,
            FlaxCLIPPreTrainedModel,
            FlaxCLIPTextModel,
            FlaxCLIPTextPreTrainedModel,
            FlaxCLIPVisionModel,
            FlaxCLIPVisionPreTrainedModel,
        )


else:
    import sys

    lowercase				:	List[Any]      						=  _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)


