

'''simple docstring'''

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        A_ (				__lowerCAmelCase       ,   unittest.TestCase							):






     '''simple docstring'''



     UpperCAmelCase_						:	Dict  		=  ShapEImgaImgPipeline
     UpperCAmelCase_						:	Optional[Any]  		=  ["""image"""]
     UpperCAmelCase_						:	str  		=  ["""image"""]
     UpperCAmelCase_						:	Dict  		=  [
         """num_images_per_prompt""",
         """num_inference_steps""",
         """generator""",
         """latents""",
         """guidance_scale""",
         """frame_size""",
         """output_type""",
         """return_dict""",
     ]
     UpperCAmelCase_						:	str  		=  False


     @property
     def    UpperCAmelCase_       (		self :     List[str]       )						->	str:
      return 32


     @property
     def    UpperCAmelCase_       (		self :     Optional[int]       )						->	List[Any]:
      return 32


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


     @property
     def    UpperCAmelCase_       (		self :     List[Any]       )						->	List[Any]:
      return 8


     @property
     def    UpperCAmelCase_       (		self :     Optional[int]       )						->	Optional[int]:
      torch.manual_seed(0       )
      UpperCAmelCase       :	str											= CLIPVisionConfig(
          hidden_size=self.text_embedder_hidden_size     ,			image_size=64     ,			projection_dim=self.text_embedder_hidden_size     ,			intermediate_size=37     ,			num_attention_heads=4     ,			num_channels=3     ,			num_hidden_layers=5     ,			patch_size=1     ,			)

      UpperCAmelCase       :	List[str]											= CLIPVisionModel(lowercase_       )
      return model


     @property
     def    UpperCAmelCase_       (		self :     Any       )						->	Union[str, Any]:
      UpperCAmelCase       :	Optional[Any]											= CLIPImageProcessor(
          crop_size=224     ,			do_center_crop=lowercase_     ,			do_normalize=lowercase_     ,			do_resize=lowercase_     ,			image_mean=[0.4814_5466, 0.457_8275, 0.4082_1073]     ,			image_std=[0.2686_2954, 0.2613_0258, 0.2757_7711]     ,			resample=3     ,			size=224     ,			)

      return image_processor


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

      UpperCAmelCase       :	Optional[int]											= {
          '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',
          'embedding_proj_norm_type': 'layer',
          'encoder_hid_proj_type': None,
          'added_emb_type': None,
      }

      UpperCAmelCase       :	Optional[Any]											= PriorTransformer(**lowercase_       )
      return model


     @property
     def    UpperCAmelCase_       (		self :     Union[str, Any]       )						->	Optional[Any]:
      torch.manual_seed(0       )

      UpperCAmelCase       :	List[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,
          ),
      }
      UpperCAmelCase       :	Optional[Any]											= ShapERenderer(**lowercase_       )
      return model


     def    UpperCAmelCase_       (		self :     Tuple       )						->	List[str]:
      UpperCAmelCase       :	List[Any]											= self.dummy_prior
      UpperCAmelCase       :	Optional[Any]											= self.dummy_image_encoder
      UpperCAmelCase       :	str											= self.dummy_image_processor
      UpperCAmelCase       :	Tuple											= self.dummy_renderer

      UpperCAmelCase       :	int											= HeunDiscreteScheduler(
          beta_schedule='exp'     ,			num_train_timesteps=1_024     ,			prediction_type='sample'     ,			use_karras_sigmas=lowercase_     ,			clip_sample=lowercase_     ,			clip_sample_range=1.0     ,			)
      UpperCAmelCase       :	str											= {
          'prior': prior,
          'image_encoder': image_encoder,
          'image_processor': image_processor,
          'renderer': renderer,
          'scheduler': scheduler,
      }

      return components


     def    UpperCAmelCase_       (		self :     Any     ,			lowercase_ :     List[str]     ,			lowercase_ :     Optional[Any]=0       )						->	int:
      UpperCAmelCase       :	Optional[Any]											= floats_tensor((1, 3, 64, 64)     ,			rng=random.Random(lowercase_       )       ).to(lowercase_       )

      if str(lowercase_       ).startswith('mps'       ):
       UpperCAmelCase       :	Any											= torch.manual_seed(lowercase_       )
      else:
       UpperCAmelCase       :	Optional[Any]											= torch.Generator(device=lowercase_       ).manual_seed(lowercase_       )
      UpperCAmelCase       :	Dict											= {
          'image': input_image,
          'generator': generator,
          'num_inference_steps': 1,
          'frame_size': 32,
          'output_type': 'np',
      }
      return inputs


     def    UpperCAmelCase_       (		self :     Optional[Any]       )						->	str:
      UpperCAmelCase       :	Dict											= 'cpu'

      UpperCAmelCase       :	Dict											= self.get_dummy_components()

      UpperCAmelCase       :	str											= self.pipeline_class(**lowercase_       )
      UpperCAmelCase       :	Tuple											= pipe.to(lowercase_       )

      pipe.set_progress_bar_config(disable=lowercase_       )

      UpperCAmelCase       :	Any											= pipe(**self.get_dummy_inputs(lowercase_       )       )
      UpperCAmelCase       :	List[str]											= output.images[0]
      UpperCAmelCase       :	Tuple											= image[0, -3:, -3:, -1]

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

      UpperCAmelCase       :	int											= np.array(
          [
              0.0003_9216,
              0.0003_9216,
              0.0003_9216,
              0.0003_9216,
              0.0003_9216,
              0.0003_9216,
              0.0003_9216,
              0.0003_9216,
              0.0003_9216,
          ]       )

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


     def    UpperCAmelCase_       (		self :     List[str]       )						->	List[str]:
      self._test_inference_batch_consistent(batch_sizes=[1, 2]       )


     def    UpperCAmelCase_       (		self :     Dict       )						->	List[str]:
      UpperCAmelCase       :	str											= torch_device == 'cpu'
      UpperCAmelCase       :	Dict											= True
      self._test_inference_batch_single_identical(
          batch_size=2     ,			test_max_difference=lowercase_     ,			relax_max_difference=lowercase_     ,			)




     def    UpperCAmelCase_       (		self :     int       )						->	Tuple:
      UpperCAmelCase       :	Tuple											= self.get_dummy_components()
      UpperCAmelCase       :	List[Any]											= self.pipeline_class(**lowercase_       )
      UpperCAmelCase       :	List[str]											= pipe.to(lowercase_       )
      pipe.set_progress_bar_config(disable=lowercase_       )

      UpperCAmelCase       :	Dict											= 1
      UpperCAmelCase       :	Tuple											= 2

      UpperCAmelCase       :	Tuple											= self.get_dummy_inputs(lowercase_       )

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

      UpperCAmelCase       :	List[str]											= 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    UpperCAmelCase_       (		self :     List[Any]       )						->	Dict:
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()




     def    UpperCAmelCase_       (		self :     Tuple       )						->	str:
      UpperCAmelCase       :	Optional[Any]											= load_image(
          'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main' '/shap_e/corgi.png'       )
      UpperCAmelCase       :	Any											= load_numpy(
          'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
          '/shap_e/test_shap_e_img2img_out.npy'       )
      UpperCAmelCase       :	List[str]											= ShapEImgaImgPipeline.from_pretrained('openai/shap-e-img2img'       )
      UpperCAmelCase       :	List[str]											= pipe.to(lowercase_       )
      pipe.set_progress_bar_config(disable=lowercase_       )

      UpperCAmelCase       :	Tuple											= torch.Generator(device=lowercase_       ).manual_seed(0       )

      UpperCAmelCase       :	Optional[Any]											= pipe(
          lowercase_     ,			generator=lowercase_     ,			guidance_scale=3.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_       )





"""simple docstring"""
from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowercase__					       =	logging.get_logger(__name__)

lowercase__					       =	{
    'MIT/ast-finetuned-audioset-10-10-0.4593': (
        'https://huggingface.co/MIT/ast-finetuned-audioset-10-10-0.4593/resolve/main/config.json'
    ),
}




class    __snake_case       (     __lowerCAmelCase				):
      a__        =		"""audio-spectrogram-transformer"""



      def __init__(						self    ,			lowercase=7_68    ,			lowercase=12    ,			lowercase=12    ,			lowercase=30_72    ,			lowercase="gelu"    ,			lowercase=0.0    ,			lowercase=0.0    ,			lowercase=0.02    ,			lowercase=1e-12    ,			lowercase=16    ,			lowercase=True    ,			lowercase=10    ,			lowercase=10    ,			lowercase=10_24    ,			lowercase=1_28    ,			**lowercase    ,			)	->			str:





        '''simple docstring'''





        super().__init__(**lowercase)

        a__:     Any										=	hidden_size
        a__:     int										=	num_hidden_layers
        a__:     Union[str, Any]										=	num_attention_heads
        a__:     Any										=	intermediate_size
        a__:     Union[str, Any]										=	hidden_act
        a__:     int										=	hidden_dropout_prob
        a__:     str										=	attention_probs_dropout_prob
        a__:     str										=	initializer_range
        a__:     Tuple										=	layer_norm_eps
        a__:     Any										=	patch_size
        a__:     int										=	qkv_bias
        a__:     Optional[Any]										=	frequency_stride
        a__:     int										=	time_stride
        a__:     List[str]										=	max_length
        a__:     Tuple										=	num_mel_bins


