





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   _SCREAMING_SNAKE_CASE		(      __SCREAMING_SNAKE_CASE       ,				unittest.TestCase       ):







 '''simple docstring'''
 lowercase_						=     ShapEPipeline
 lowercase_						=     ["prompt"]
 lowercase_						=     ["prompt"]
 lowercase_						=     [
     "num_images_per_prompt",
     "num_inference_steps",
     "generator",
     "latents",
     "guidance_scale",
     "frame_size",
     "output_type",
     "return_dict",
 ]
 lowercase_						=     False

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



   '''simple docstring'''



   return 32

 @property
 def 				SCREAMING_SNAKE_CASE_   (self							:  List[str])			->Optional[int]:



   '''simple docstring'''



   return 32

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



   '''simple docstring'''



   return self.time_input_dim * 4

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



   '''simple docstring'''



   return 8

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



   '''simple docstring'''



   lowerCamelCase__: Tuple         =CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip")
   return tokenizer

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



   '''simple docstring'''



   torch.manual_seed(0)
   lowerCamelCase__: Optional[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-0_5			,					num_attention_heads=4			,					num_hidden_layers=5			,					pad_token_id=1			,					vocab_size=1_000			,					)
   return CLIPTextModelWithProjection(UpperCAmelCase_)

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



   '''simple docstring'''



   torch.manual_seed(0)

   lowerCamelCase__: List[str]         ={
       "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__: Optional[Any]         =PriorTransformer(**UpperCAmelCase_)
   return model

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



   '''simple docstring'''



   torch.manual_seed(0)

   lowerCamelCase__: 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,
       ),
   }
   lowerCamelCase__: Tuple         =ShapERenderer(**UpperCAmelCase_)
   return model

 def 				SCREAMING_SNAKE_CASE_   (self							:  List[Any])			->Optional[int]:



   '''simple docstring'''



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

   lowerCamelCase__: Any         =HeunDiscreteScheduler(
       beta_schedule="exp"			,					num_train_timesteps=1_024			,					prediction_type="sample"			,					use_karras_sigmas=UpperCAmelCase_			,					clip_sample=UpperCAmelCase_			,					clip_sample_range=1.0			,					)
   lowerCamelCase__: Dict         ={
       "prior": prior,
       "text_encoder": text_encoder,
       "tokenizer": tokenizer,
       "renderer": renderer,
       "scheduler": scheduler,
   }

   return components

 def 				SCREAMING_SNAKE_CASE_   (self							:  Any			,					UpperCAmelCase_							:  Tuple			,					UpperCAmelCase_							:  List[Any]=0)			->str:



   '''simple docstring'''



   if str(UpperCAmelCase_).startswith("mps"):
     lowerCamelCase__: str         =torch.manual_seed(UpperCAmelCase_)
   else:
     lowerCamelCase__: Optional[int]         =torch.Generator(device=UpperCAmelCase_).manual_seed(UpperCAmelCase_)
   lowerCamelCase__: Tuple         ={
       "prompt": "horse",
       "generator": generator,
       "num_inference_steps": 1,
       "frame_size": 32,
       "output_type": "np",
   }
   return inputs

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



   '''simple docstring'''



   lowerCamelCase__: str         ="cpu"

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

   lowerCamelCase__: Dict         =self.pipeline_class(**UpperCAmelCase_)
   lowerCamelCase__: Any         =pipe.to(UpperCAmelCase_)

   pipe.set_progress_bar_config(disable=UpperCAmelCase_)

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

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

   lowerCamelCase__: 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 				SCREAMING_SNAKE_CASE_   (self							:  Tuple)			->Union[str, Any]:



   '''simple docstring'''



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

 def 				SCREAMING_SNAKE_CASE_   (self							:  Any)			->List[str]:



   '''simple docstring'''



   lowerCamelCase__: List[Any]         =torch_device == "cpu"
   lowerCamelCase__: int         =True

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







 def 				SCREAMING_SNAKE_CASE_   (self							:  Tuple)			->Any:



   '''simple docstring'''



   lowerCamelCase__: Tuple         =self.get_dummy_components()
   lowerCamelCase__: List[Any]         =self.pipeline_class(**UpperCAmelCase_)
   lowerCamelCase__: List[str]         =pipe.to(UpperCAmelCase_)
   pipe.set_progress_bar_config(disable=UpperCAmelCase_)

   lowerCamelCase__: Dict         =1
   lowerCamelCase__: List[str]         =2

   lowerCamelCase__: List[str]         =self.get_dummy_inputs(UpperCAmelCase_)

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

   lowerCamelCase__: Union[str, Any]         =pipe(**UpperCAmelCase_			,					num_images_per_prompt=UpperCAmelCase_)[0]

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



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







 '''simple docstring'''

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



   '''simple docstring'''



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







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



   '''simple docstring'''



   lowerCamelCase__: Union[str, Any]         =load_numpy(
       "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"
       "/shap_e/test_shap_e_np_out.npy")
   lowerCamelCase__: Any         =ShapEPipeline.from_pretrained("openai/shap-e")
   lowerCamelCase__: Union[str, Any]         =pipe.to(UpperCAmelCase_)
   pipe.set_progress_bar_config(disable=UpperCAmelCase_)

   lowerCamelCase__: Dict         =torch.Generator(device=UpperCAmelCase_).manual_seed(0)

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



"""simple docstring"""

import functools
import operator

from ...configuration_utils import PretrainedConfig
from ...utils import logging


SCREAMING_SNAKE_CASE_            =    logging.get_logger(__name__)

SCREAMING_SNAKE_CASE_            =    {
    '''asapp/sew-d-tiny-100k''': '''https://huggingface.co/asapp/sew-d-tiny-100k/resolve/main/config.json''',
    # See all SEW-D models at https://huggingface.co/models?filter=sew-d
}
class   _UpperCAmelCase				(		SCREAMING_SNAKE_CASE_				):
      __SCREAMING_SNAKE_CASE :				List[str]				   =						"sew-d"




      def __init__(				self ,							lowercase_=3_2 ,							lowercase_=7_6_8 ,							lowercase_=1_2 ,							lowercase_=1_2 ,							lowercase_=3_0_7_2 ,							lowercase_=2 ,							lowercase_=5_1_2 ,							lowercase_=2_5_6 ,							lowercase_=True ,							lowercase_=True ,							lowercase_=("p2c", "c2p") ,							lowercase_="layer_norm" ,							lowercase_="gelu_python" ,							lowercase_=0.1 ,							lowercase_=0.1 ,							lowercase_=0.1 ,							lowercase_=0.0 ,							lowercase_=0.1 ,							lowercase_=0.0_2 ,							lowercase_=1E-7 ,							lowercase_=1E-5 ,							lowercase_="group" ,							lowercase_="gelu" ,							lowercase_=(6_4, 1_2_8, 1_2_8, 1_2_8, 1_2_8, 2_5_6, 2_5_6, 2_5_6, 2_5_6, 5_1_2, 5_1_2, 5_1_2, 5_1_2) ,							lowercase_=(5, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1, 2, 1) ,							lowercase_=(1_0, 3, 1, 3, 1, 3, 1, 3, 1, 2, 1, 2, 1) ,							lowercase_=False ,							lowercase_=1_2_8 ,							lowercase_=1_6 ,							lowercase_=True ,							lowercase_=0.0_5 ,							lowercase_=1_0 ,							lowercase_=2 ,							lowercase_=0.0 ,							lowercase_=1_0 ,							lowercase_=0 ,							lowercase_="mean" ,							lowercase_=False ,							lowercase_=False ,							lowercase_=2_5_6 ,							lowercase_=0 ,							lowercase_=1 ,							lowercase_=2 ,							**lowercase_ ,							)    ->       Dict:
          super().__init__(**lowercase_ ,							pad_token_id=lowercase_ ,							bos_token_id=lowercase_ ,							eos_token_id=lowercase_				)
          UpperCAmelCase          =       hidden_size
          UpperCAmelCase          =       feat_extract_norm
          UpperCAmelCase          =       feat_extract_activation
          UpperCAmelCase          =       list(lowercase_				)
          UpperCAmelCase          =       list(lowercase_				)
          UpperCAmelCase          =       list(lowercase_				)
          UpperCAmelCase          =       conv_bias
          UpperCAmelCase          =       num_conv_pos_embeddings
          UpperCAmelCase          =       num_conv_pos_embedding_groups
          UpperCAmelCase          =       len(self.conv_dim				)
          UpperCAmelCase          =       num_hidden_layers
          UpperCAmelCase          =       intermediate_size
          UpperCAmelCase          =       squeeze_factor
          UpperCAmelCase          =       max_position_embeddings
          UpperCAmelCase          =       position_buckets
          UpperCAmelCase          =       share_att_key
          UpperCAmelCase          =       relative_attention
          UpperCAmelCase          =       norm_rel_ebd
          UpperCAmelCase          =       list(lowercase_				)
          UpperCAmelCase          =       hidden_act
          UpperCAmelCase          =       num_attention_heads
          UpperCAmelCase          =       hidden_dropout
          UpperCAmelCase          =       attention_dropout
          UpperCAmelCase          =       activation_dropout
          UpperCAmelCase          =       feat_proj_dropout
          UpperCAmelCase          =       final_dropout
          UpperCAmelCase          =       layer_norm_eps
          UpperCAmelCase          =       feature_layer_norm_eps
          UpperCAmelCase          =       initializer_range
          UpperCAmelCase          =       vocab_size

          if (
              (len(self.conv_stride				) != self.num_feat_extract_layers)
              or (len(self.conv_kernel				) != self.num_feat_extract_layers)
              or (len(self.conv_dim				) != self.num_feat_extract_layers)
          ):
              raise ValueError(
                  'Configuration for convolutional layers is incorrect.'
                  'It is required that `len(config.conv_dim)` == `len(config.conv_stride)` == `len(config.conv_kernel)`,'
                  F"but is `len(config.conv_dim) = {len(self.conv_dim				)}`, `len(config.conv_stride)"
                  F"= {len(self.conv_stride				)}`, `len(config.conv_kernel) = {len(self.conv_kernel				)}`."				)

          # fine-tuning config parameters for SpecAugment: https://arxiv.org/abs/1904.08779
          UpperCAmelCase          =       apply_spec_augment
          UpperCAmelCase          =       mask_time_prob
          UpperCAmelCase          =       mask_time_length
          UpperCAmelCase          =       mask_time_min_masks
          UpperCAmelCase          =       mask_feature_prob
          UpperCAmelCase          =       mask_feature_length
          UpperCAmelCase          =       mask_feature_min_masks

          # ctc loss
          UpperCAmelCase          =       ctc_loss_reduction
          UpperCAmelCase          =       ctc_zero_infinity

          # sequence classification
          UpperCAmelCase          =       use_weighted_layer_sum
          UpperCAmelCase          =       classifier_proj_size




      @property
      def      a_   (				self				)    ->       Tuple:
          return functools.reduce(operator.mul ,							self.conv_stride ,							1				)



