





"""simple docstring"""
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(       A   ,     unittest.TestCase  ):
       SCREAMING_SNAKE_CASE_       :				Dict     						=  ShapEPipeline
       SCREAMING_SNAKE_CASE_       :				Union[str, Any]     						=  ['''prompt''']
       SCREAMING_SNAKE_CASE_       :				Tuple     						=  ['''prompt''']
       SCREAMING_SNAKE_CASE_       :				List[Any]     						=  [
           '''num_images_per_prompt''',
           '''num_inference_steps''',
           '''generator''',
           '''latents''',
           '''guidance_scale''',
           '''frame_size''',
           '''output_type''',
           '''return_dict''',
       ]
       SCREAMING_SNAKE_CASE_       :				str     						=  False




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





             """simple docstring"""






             return 32




       @property
       def    _UpperCamelCase					(						self							)      ->				str:





             """simple docstring"""






             return 32




       @property
       def    _UpperCamelCase					(						self							)      ->				Dict:





             """simple docstring"""






             return self.time_input_dim * 4




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





             """simple docstring"""






             return 8




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





             """simple docstring"""






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




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





             """simple docstring"""






             torch.manual_seed(0							)
             __SCREAMING_SNAKE_CASE       :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=10_00		,)
             return CLIPTextModelWithProjection(SCREAMING_SNAKE_CASE__							)




       @property
       def    _UpperCamelCase					(						self							)      ->				Any:





             """simple docstring"""






             torch.manual_seed(0							)

             __SCREAMING_SNAKE_CASE       :Tuple		    =  {
                 '''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       :int		    =  PriorTransformer(**SCREAMING_SNAKE_CASE__							)
             return model




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





             """simple docstring"""






             torch.manual_seed(0							)

             __SCREAMING_SNAKE_CASE       :Tuple		    =  {
                 '''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       :Optional[Any]		    =  ShapERenderer(**SCREAMING_SNAKE_CASE__							)
             return model




       def    _UpperCamelCase					(						self							)      ->				Any:





             """simple docstring"""






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

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

             return components




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__=0							)      ->				str:





             """simple docstring"""






             if str(SCREAMING_SNAKE_CASE__							).startswith('''mps'''							):
                   __SCREAMING_SNAKE_CASE       :str		    =  torch.manual_seed(SCREAMING_SNAKE_CASE__							)
             else:
                   __SCREAMING_SNAKE_CASE       :int		    =  torch.Generator(device=SCREAMING_SNAKE_CASE__							).manual_seed(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE       :Tuple		    =  {
                 '''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       :Dict		    =  '''cpu'''

             __SCREAMING_SNAKE_CASE       :Any		    =  self.get_dummy_components()

             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  self.pipeline_class(**SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE       :Any		    =  pipe.to(SCREAMING_SNAKE_CASE__							)

             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE       :int		    =  pipe(**self.get_dummy_inputs(SCREAMING_SNAKE_CASE__							)							)
             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  output.images[0]
             __SCREAMING_SNAKE_CASE       :Tuple		    =  image[0, -3:, -3:, -1]

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

             __SCREAMING_SNAKE_CASE       :str		    =  np.array(
                 [
                     0.0_0_0_3_9_2_1_6,
                     0.0_0_0_3_9_2_1_6,
                     0.0_0_0_3_9_2_1_6,
                     0.0_0_0_3_9_2_1_6,
                     0.0_0_0_3_9_2_1_6,
                     0.0_0_0_3_9_2_1_6,
                     0.0_0_0_3_9_2_1_6,
                     0.0_0_0_3_9_2_1_6,
                     0.0_0_0_3_9_2_1_6,
                 ]							)

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




       def    _UpperCamelCase					(						self							)      ->				Any:





             """simple docstring"""






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




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :Dict		    =  torch_device == '''cpu'''
             __SCREAMING_SNAKE_CASE       :str		    =  True

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




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





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  self.get_dummy_components()
             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  self.pipeline_class(**SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE       :str		    =  pipe.to(SCREAMING_SNAKE_CASE__							)
             pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE__							)

             __SCREAMING_SNAKE_CASE       :Dict		    =  1
             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  2

             __SCREAMING_SNAKE_CASE       :Tuple		    =  self.get_dummy_inputs(SCREAMING_SNAKE_CASE__							)

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

             __SCREAMING_SNAKE_CASE       :Any		    =  pipe(**SCREAMING_SNAKE_CASE__		,num_images_per_prompt=SCREAMING_SNAKE_CASE__							)[0]

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






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




       def    _UpperCamelCase					(						self							)      ->				str:





             """simple docstring"""






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




       def    _UpperCamelCase					(						self							)      ->				Any:





             """simple docstring"""






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

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

             __SCREAMING_SNAKE_CASE       :int		    =  pipe(
                 '''a shark'''		,generator=SCREAMING_SNAKE_CASE__		,guidance_scale=1_5.0		,num_inference_steps=64		,frame_size=64		,output_type='''np'''		,).images[0]

             assert images.shape == (20, 64, 64, 3)

             assert_mean_pixel_difference(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							)






"""simple docstring"""
import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import logging


lowerCamelCase_       =		logging.get_logger(__name__)

lowerCamelCase_       =		{"vocab_file": "spiece.model"}

lowerCamelCase_       =		{
    "vocab_file": {
        "bert_for_seq_generation": (
            "https://huggingface.co/google/bert_for_seq_generation_L-24_bbc_encoder/resolve/main/spiece.model"
        ),
    }
}

lowerCamelCase_       =		{"bert_for_seq_generation": 5_1_2}






class  _SCREAMING_SNAKE_CASE(       A  ):
       SCREAMING_SNAKE_CASE_       :				str     						=  VOCAB_FILES_NAMES
       SCREAMING_SNAKE_CASE_       :				Optional[Any]     						=  PRETRAINED_VOCAB_FILES_MAP
       SCREAMING_SNAKE_CASE_       :				Optional[int]     						=  PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
       SCREAMING_SNAKE_CASE_       :				List[int]     						=  []
       SCREAMING_SNAKE_CASE_       :				List[Any]     						=  ['''input_ids''', '''attention_mask''']




       def __init__(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__="<s>"		,SCREAMING_SNAKE_CASE__="</s>"		,SCREAMING_SNAKE_CASE__="<unk>"		,SCREAMING_SNAKE_CASE__="<pad>"		,SCREAMING_SNAKE_CASE__="<::::>"		,SCREAMING_SNAKE_CASE__ = None		,**SCREAMING_SNAKE_CASE__		,)      ->				None:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  {} if sp_model_kwargs is None else sp_model_kwargs

             # Add extra_ids to the special token list
             super().__init__(
                 bos_token=SCREAMING_SNAKE_CASE__		,eos_token=SCREAMING_SNAKE_CASE__		,unk_token=SCREAMING_SNAKE_CASE__		,pad_token=SCREAMING_SNAKE_CASE__		,sep_token=SCREAMING_SNAKE_CASE__		,sp_model_kwargs=self.sp_model_kwargs		,**SCREAMING_SNAKE_CASE__		,)

             __SCREAMING_SNAKE_CASE       :Dict		    =  vocab_file

             __SCREAMING_SNAKE_CASE       :Optional[Any]		    =  spm.SentencePieceProcessor(**self.sp_model_kwargs							)
             self.sp_model.Load(SCREAMING_SNAKE_CASE__							)




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





             """simple docstring"""






             return self.sp_model.get_piece_size()




       def    _UpperCamelCase					(						self							)      ->				int:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :List[Any]		    =  {self.convert_ids_to_tokens(SCREAMING_SNAKE_CASE__							): i for i in range(self.vocab_size							)}
             vocab.update(self.added_tokens_encoder							)
             return vocab




       def __getstate__(						self							)      ->				Dict:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :int		    =  self.__dict__.copy()
             __SCREAMING_SNAKE_CASE       :Dict		    =  None
             return state




       def __setstate__(						self		,SCREAMING_SNAKE_CASE__							)      ->				Optional[Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :List[str]		    =  d

             # for backward compatibility
             if not hasattr(self		,'''sp_model_kwargs'''							):
                   __SCREAMING_SNAKE_CASE       :str		    =  {}

             __SCREAMING_SNAKE_CASE       :str		    =  spm.SentencePieceProcessor(**self.sp_model_kwargs							)
             self.sp_model.Load(self.vocab_file							)




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__							)      ->				List[str]:





             """simple docstring"""






             return self.sp_model.encode(SCREAMING_SNAKE_CASE__		,out_type=SCREAMING_SNAKE_CASE__							)




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__							)      ->				int:





             """simple docstring"""






             return self.sp_model.piece_to_id(SCREAMING_SNAKE_CASE__							)




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__							)      ->				str:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :Optional[int]		    =  self.sp_model.IdToPiece(SCREAMING_SNAKE_CASE__							)
             return token




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__							)      ->				List[Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE       :int		    =  []
             __SCREAMING_SNAKE_CASE       :List[str]		    =  ''''''
             for token in tokens:
                   # make sure that special tokens are not decoded using sentencepiece model
                   if token in self.all_special_tokens:
                         out_string += self.sp_model.decode(SCREAMING_SNAKE_CASE__							) + token
                         __SCREAMING_SNAKE_CASE       :Dict		    =  []
                   else:
                         current_sub_tokens.append(SCREAMING_SNAKE_CASE__							)
             out_string += self.sp_model.decode(SCREAMING_SNAKE_CASE__							)
             return out_string.strip()




       def    _UpperCamelCase					(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__ = None							)      ->				Tuple[str]:





             """simple docstring"""






             if not os.path.isdir(SCREAMING_SNAKE_CASE__							):
                   logger.error(f'''Vocabulary path ({save_directory}) should be a directory'''							)
                   return
             __SCREAMING_SNAKE_CASE       :Union[str, Any]		    =  os.path.join(
                 SCREAMING_SNAKE_CASE__		,(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']							)

             if os.path.abspath(self.vocab_file							) != os.path.abspath(SCREAMING_SNAKE_CASE__							) and os.path.isfile(self.vocab_file							):
                   copyfile(self.vocab_file		,SCREAMING_SNAKE_CASE__							)
             elif not os.path.isfile(self.vocab_file							):
                   with open(SCREAMING_SNAKE_CASE__		,'''wb'''							) as fi:
                         __SCREAMING_SNAKE_CASE       :Tuple		    =  self.sp_model.serialized_model_proto()
                         fi.write(SCREAMING_SNAKE_CASE__							)

             return (out_vocab_file,)