
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     __UpperCamelCase (							lowerCAmelCase__ ,						unittest.TestCase       ):






     """simple docstring"""







     lowerCAmelCase_							   =			ShapEPipeline
     lowerCAmelCase_							   =			['''prompt''']
     lowerCAmelCase_							   =			['''prompt''']
     lowerCAmelCase_							   =			[
         '''num_images_per_prompt''',
         '''num_inference_steps''',
         '''generator''',
         '''latents''',
         '''guidance_scale''',
         '''frame_size''',
         '''output_type''',
         '''return_dict''',
     ]
     lowerCAmelCase_							   =			False



     @property
     def  UpperCAmelCase__							(				self							:							Optional[int] ):







       """simple docstring"""




       return 32



     @property
     def  UpperCAmelCase__							(				self							:							Tuple ):







       """simple docstring"""




       return 32



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







       """simple docstring"""




       return self.time_input_dim * 4



     @property
     def  UpperCAmelCase__							(				self							:							int ):







       """simple docstring"""




       return 8



     @property
     def  UpperCAmelCase__							(				self							:							Any ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''' )
       return tokenizer



     @property
     def  UpperCAmelCase__							(				self							:							Any ):







       """simple docstring"""




       torch.manual_seed(0 )
       __SCREAMING_SNAKE_CASE      :	Dict					=		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(_A )



     @property
     def  UpperCAmelCase__							(				self							:							List[Any] ):







       """simple docstring"""




       torch.manual_seed(0 )

       __SCREAMING_SNAKE_CASE      :	Any					=		{
           '''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,
       }

       __SCREAMING_SNAKE_CASE      :	List[Any]					=		PriorTransformer(**_A )
       return model



     @property
     def  UpperCAmelCase__							(				self							:							str ):







       """simple docstring"""




       torch.manual_seed(0 )

       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		{
           '''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,
           ),
       }
       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		ShapERenderer(**_A )
       return model



     def  UpperCAmelCase__							(				self							:							int ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	str					=		self.dummy_prior
       __SCREAMING_SNAKE_CASE      :	int					=		self.dummy_text_encoder
       __SCREAMING_SNAKE_CASE      :	Optional[int]					=		self.dummy_tokenizer
       __SCREAMING_SNAKE_CASE      :	List[Any]					=		self.dummy_renderer

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

       return components



     def  UpperCAmelCase__							(				self							:							Optional[int]     ,						_A							:							Dict     ,						_A							:							Optional[int]=0 ):







       """simple docstring"""




       if str(_A ).startswith('''mps''' ):
         __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		torch.manual_seed(_A )
       else:
         __SCREAMING_SNAKE_CASE      :	Dict					=		torch.Generator(device=_A ).manual_seed(_A )
       __SCREAMING_SNAKE_CASE      :	int					=		{
           '''prompt''': '''horse''',
           '''generator''': generator,
           '''num_inference_steps''': 1,
           '''frame_size''': 32,
           '''output_type''': '''np''',
       }
       return inputs



     def  UpperCAmelCase__							(				self							:							List[Any] ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Any					=		'''cpu'''

       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		self.get_dummy_components()

       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		self.pipeline_class(**_A )
       __SCREAMING_SNAKE_CASE      :	str					=		pipe.to(_A )

       pipe.set_progress_bar_config(disable=_A )

       __SCREAMING_SNAKE_CASE      :	Dict					=		pipe(**self.get_dummy_inputs(_A ) )
       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		output.images[0]
       __SCREAMING_SNAKE_CASE      :	Any					=		image[0, -3:, -3:, -1]

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

       __SCREAMING_SNAKE_CASE      :	str					=		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  UpperCAmelCase__							(				self							:							str ):







       """simple docstring"""




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



     def  UpperCAmelCase__							(				self							:							Optional[Any] ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	int					=		torch_device == '''cpu'''
       __SCREAMING_SNAKE_CASE      :	List[str]					=		True

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




     def  UpperCAmelCase__							(				self							:							Any ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		self.get_dummy_components()
       __SCREAMING_SNAKE_CASE      :	Tuple					=		self.pipeline_class(**_A )
       __SCREAMING_SNAKE_CASE      :	Dict					=		pipe.to(_A )
       pipe.set_progress_bar_config(disable=_A )

       __SCREAMING_SNAKE_CASE      :	Tuple					=		1
       __SCREAMING_SNAKE_CASE      :	Any					=		2

       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		self.get_dummy_inputs(_A )

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

       __SCREAMING_SNAKE_CASE      :	str					=		pipe(**_A     ,						num_images_per_prompt=_A )[0]

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






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






     """simple docstring"""







     def  UpperCAmelCase__							(				self							:							int ):







       """simple docstring"""




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




     def  UpperCAmelCase__							(				self							:							str ):







       """simple docstring"""




       __SCREAMING_SNAKE_CASE      :	Optional[int]					=		load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/shap_e/test_shap_e_np_out.npy''' )
       __SCREAMING_SNAKE_CASE      :	str					=		ShapEPipeline.from_pretrained('''openai/shap-e''' )
       __SCREAMING_SNAKE_CASE      :	Optional[Any]					=		pipe.to(_A )
       pipe.set_progress_bar_config(disable=_A )

       __SCREAMING_SNAKE_CASE      :	Any					=		torch.Generator(device=_A ).manual_seed(0 )

       __SCREAMING_SNAKE_CASE      :	Union[str, Any]					=		pipe(
           '''a shark'''     ,						generator=_A     ,						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(_A     ,						_A )








'''simple docstring'''

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


UpperCamelCase      					={
    "configuration_clap": [
        "CLAP_PRETRAINED_MODEL_ARCHIVE_LIST",
        "ClapAudioConfig",
        "ClapConfig",
        "ClapTextConfig",
    ],
    "processing_clap": ["ClapProcessor"],
}

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      UpperCamelCase      					=[
          "CLAP_PRETRAINED_MODEL_ARCHIVE_LIST",
          "ClapModel",
          "ClapPreTrainedModel",
          "ClapTextModel",
          "ClapTextModelWithProjection",
          "ClapAudioModel",
          "ClapAudioModelWithProjection",
      ]
      UpperCamelCase      					=["ClapFeatureExtractor"]

if TYPE_CHECKING:
      from .configuration_clap import (
          CLAP_PRETRAINED_MODEL_ARCHIVE_LIST,
          ClapAudioConfig,
          ClapConfig,
          ClapTextConfig,
      )
      from .processing_clap import ClapProcessor

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .feature_extraction_clap import ClapFeatureExtractor
            from .modeling_clap import (
                CLAP_PRETRAINED_MODEL_ARCHIVE_LIST,
                ClapAudioModel,
                ClapAudioModelWithProjection,
                ClapModel,
                ClapPreTrainedModel,
                ClapTextModel,
                ClapTextModelWithProjection,
            )


else:
      import sys

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






