






import os
import unicodedata
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import SPIECE_UNDERLINE, logging


A_       :Optional[int]       =						logging.get_logger(__name__)

A_       :Union[str, Any]       =						{'vocab_file': 'spiece.model'}

A_       :int       =						{
    'vocab_file': {
        'TsinghuaAI/CPM-Generate': 'https://huggingface.co/TsinghuaAI/CPM-Generate/resolve/main/spiece.model',
    }
}




class 				__A     (			lowercase_ ):





				"""simple docstring"""







				def __init__( self							,  lowerCamelCase__							,  lowerCamelCase__=False							,  lowerCamelCase__=True							,  lowerCamelCase__=False							,  lowerCamelCase__="<s>"							,  lowerCamelCase__="</s>"							,  lowerCamelCase__="<unk>"							,  lowerCamelCase__="<sep>"							,  lowerCamelCase__="<pad>"							,  lowerCamelCase__="<cls>"							,  lowerCamelCase__="<mask>"							,  lowerCamelCase__=["<eop>", "<eod>"]							,  lowerCamelCase__ = None							,  **lowerCamelCase__							,  ):




									"""simple docstring"""

									__UpperCamelCase      :					List[Any]       		=AddedToken(lowerCamelCase__							,  lstrip=lowerCamelCase__							,  rstrip=lowerCamelCase__    ) if isinstance(lowerCamelCase__							,  lowerCamelCase__    ) else mask_token

									__UpperCamelCase      :					List[Any]       		={} if sp_model_kwargs is None else sp_model_kwargs

									super().__init__(
									    do_lower_case=lowerCamelCase__							,  remove_space=lowerCamelCase__							,  keep_accents=lowerCamelCase__							,  bos_token=lowerCamelCase__							,  eos_token=lowerCamelCase__							,  unk_token=lowerCamelCase__							,  sep_token=lowerCamelCase__							,  pad_token=lowerCamelCase__							,  cls_token=lowerCamelCase__							,  mask_token=lowerCamelCase__							,  additional_special_tokens=lowerCamelCase__							,  sp_model_kwargs=self.sp_model_kwargs							,  **lowerCamelCase__							,  )

									__UpperCamelCase      :					int       		=3

									__UpperCamelCase      :					int       		=do_lower_case
									__UpperCamelCase      :					str       		=remove_space
									__UpperCamelCase      :					Optional[int]       		=keep_accents
									__UpperCamelCase      :					Optional[int]       		=vocab_file

									__UpperCamelCase      :					Optional[int]       		=spm.SentencePieceProcessor(**self.sp_model_kwargs    )
									self.sp_model.Load(lowerCamelCase__    )

									try:
														import jieba
									except ModuleNotFoundError as error:
														raise error.__class__(
														    'You need to install jieba to use CpmTokenizer or CpmTokenizerFast. '
														    'See https://pypi.org/project/jieba/ for installation.'    )
									__UpperCamelCase      :					Optional[Any]       		=jieba
									__UpperCamelCase      :					Tuple       		=str.maketrans(' \n'							,  '\u2582\u2583'    )


				@property
				# Copied from transformers.models.xlnet.tokenization_xlnet.XLNetTokenizer.vocab_size
				def       __lowercase							( self    ):




									"""simple docstring"""

									return len(self.sp_model    )


				def       __lowercase							( self    ):




									"""simple docstring"""

									__UpperCamelCase      :					int       		={self.convert_ids_to_tokens(lowerCamelCase__    ): i for i in range(self.vocab_size    )}
									vocab.update(self.added_tokens_encoder    )
									return vocab


				def __getstate__( self    ):




									"""simple docstring"""

									__UpperCamelCase      :					Optional[Any]       		=self.__dict__.copy()
									__UpperCamelCase      :					List[str]       		=None
									return state


				def __setstate__( self							,  lowerCamelCase__    ):




									"""simple docstring"""

									__UpperCamelCase      :					int       		=d

									# for backward compatibility
									if not hasattr(self							,  'sp_model_kwargs'    ):
														__UpperCamelCase      :					List[str]       		={}

									__UpperCamelCase      :					Tuple       		=spm.SentencePieceProcessor(**self.sp_model_kwargs    )
									self.sp_model.Load(self.vocab_file    )


				def       __lowercase							( self							,  lowerCamelCase__    ):




									"""simple docstring"""

									if self.remove_space:
														__UpperCamelCase      :					List[Any]       		=' '.join(inputs.strip().split()    )
									else:
														__UpperCamelCase      :					Union[str, Any]       		=inputs
									__UpperCamelCase      :					Any       		=outputs.replace('``'							,  '"'    ).replace('\'\''							,  '"'    )

									if not self.keep_accents:
														__UpperCamelCase      :					Optional[Any]       		=unicodedata.normalize('NFKD'							,  lowerCamelCase__    )
														__UpperCamelCase      :					List[Any]       		=''.join([c for c in outputs if not unicodedata.combining(lowerCamelCase__    )]    )
									if self.do_lower_case:
														__UpperCamelCase      :					int       		=outputs.lower()

									return outputs


				def       __lowercase							( self							,  lowerCamelCase__    ):




									"""simple docstring"""

									__UpperCamelCase      :					Union[str, Any]       		=self.preprocess_text(lowerCamelCase__    )
									__UpperCamelCase      :					Optional[int]       		=self.sp_model.encode(lowerCamelCase__							,  out_type=lowerCamelCase__    )
									__UpperCamelCase      :					Dict       		=[]
									for piece in pieces:
														if len(lowerCamelCase__    ) > 1 and piece[-1] == str(','    ) and piece[-2].isdigit():
																			__UpperCamelCase      :					Tuple       		=self.sp_model.EncodeAsPieces(piece[:-1].replace(lowerCamelCase__							,  ''    )    )
																			if piece[0] != SPIECE_UNDERLINE and cur_pieces[0][0] == SPIECE_UNDERLINE:
																								if len(cur_pieces[0]    ) == 1:
																													__UpperCamelCase      :					Any       		=cur_pieces[1:]
																								else:
																													__UpperCamelCase      :					Optional[Any]       		=cur_pieces[0][1:]
																			cur_pieces.append(piece[-1]    )
																			new_pieces.extend(lowerCamelCase__    )
														else:
																			new_pieces.append(lowerCamelCase__    )

									return new_pieces


				def       __lowercase							( self							,  lowerCamelCase__    ):




									"""simple docstring"""

									return self.sp_model.PieceToId(lowerCamelCase__    )


				def       __lowercase							( self							,  lowerCamelCase__    ):




									"""simple docstring"""

									return self.sp_model.IdToPiece(lowerCamelCase__    )


				def       __lowercase							( self							,  lowerCamelCase__    ):




									"""simple docstring"""

									__UpperCamelCase      :					Optional[int]       		=''.join(lowerCamelCase__    ).replace(lowerCamelCase__							,  ' '    ).strip()
									return out_string


				def       __lowercase							( self							,  lowerCamelCase__							,  lowerCamelCase__ = None    ):




									"""simple docstring"""

									__UpperCamelCase      :					Union[str, Any]       		=[self.sep_token_id]
									__UpperCamelCase      :					Optional[Any]       		=[self.cls_token_id]
									if token_ids_a is None:
														return token_ids_a + sep + cls
									return token_ids_a + sep + token_ids_a + sep + cls


				def       __lowercase							( self							,  lowerCamelCase__							,  lowerCamelCase__ = None							,  lowerCamelCase__ = False    ):




									"""simple docstring"""

									if already_has_special_tokens:
														return super().get_special_tokens_mask(
														    token_ids_a=lowerCamelCase__							,  token_ids_a=lowerCamelCase__							,  already_has_special_tokens=lowerCamelCase__    )

									if token_ids_a is not None:
														return ([0] * len(lowerCamelCase__    )) + [1] + ([0] * len(lowerCamelCase__    )) + [1, 1]
									return ([0] * len(lowerCamelCase__    )) + [1, 1]


				def       __lowercase							( self							,  lowerCamelCase__							,  lowerCamelCase__ = None    ):




									"""simple docstring"""

									__UpperCamelCase      :					Optional[Any]       		=[self.sep_token_id]
									__UpperCamelCase      :					str       		=[2]

									if token_ids_a is None:
														return len(token_ids_a + sep    ) * [0] + cls_segment_id
									return len(token_ids_a + sep    ) * [0] + len(token_ids_a + sep    ) * [1] + cls_segment_id


				def       __lowercase							( self							,  lowerCamelCase__							,  lowerCamelCase__ = None    ):




									"""simple docstring"""

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

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

									return (out_vocab_file,)







				def       __lowercase							( self							,  *lowerCamelCase__							,  **lowerCamelCase__    ):




									"""simple docstring"""

									__UpperCamelCase      :					str       		=super()._decode(*lowerCamelCase__							,  **lowerCamelCase__    )
									__UpperCamelCase      :					str       		=text.replace(' '							,  ''    ).replace('\u2582'							,  ' '    ).replace('\u2583'							,  '\n'    )
									return text





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





      '''simple docstring'''






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



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

             '''simple docstring'''







             return 32



      @property
      def       UpperCamelCase__		(self    :    str ):

             '''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'''







             torch.manual_seed(0 )
             lowercase__      			= 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   ,						)

             lowercase__      			= CLIPVisionModel(UpperCamelCase )
             return model



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

             '''simple docstring'''







             lowercase__      			= CLIPImageProcessor(
                 crop_size=224   ,						do_center_crop=UpperCamelCase   ,						do_normalize=UpperCamelCase   ,						do_resize=UpperCamelCase   ,						image_mean=[0.48_14_54_66, 0.4_57_82_75, 0.40_82_10_73]   ,						image_std=[0.26_86_29_54, 0.26_13_02_58, 0.27_57_77_11]   ,						resample=3   ,						size=224   ,						)

             return image_processor



      @property
      def       UpperCamelCase__		(self    :    str ):

             '''simple docstring'''







             torch.manual_seed(0 )

             lowercase__      			= {
                 '''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,
             }

             lowercase__      			= PriorTransformer(**UpperCamelCase )
             return model



      @property
      def       UpperCamelCase__		(self    :    int ):

             '''simple docstring'''







             torch.manual_seed(0 )

             lowercase__      			= {
                 '''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,
                 ),
             }
             lowercase__      			= ShapERenderer(**UpperCamelCase )
             return model



      def       UpperCamelCase__		(self    :    Union[str, Any] ):

             '''simple docstring'''







             lowercase__      			= self.dummy_prior
             lowercase__      			= self.dummy_image_encoder
             lowercase__      			= self.dummy_image_processor
             lowercase__      			= self.dummy_renderer

             lowercase__      			= HeunDiscreteScheduler(
                 beta_schedule='''exp'''   ,						num_train_timesteps=1024   ,						prediction_type='''sample'''   ,						use_karras_sigmas=UpperCamelCase   ,						clip_sample=UpperCamelCase   ,						clip_sample_range=1.0   ,						)
             lowercase__      			= {
                 '''prior''': prior,
                 '''image_encoder''': image_encoder,
                 '''image_processor''': image_processor,
                 '''renderer''': renderer,
                 '''scheduler''': scheduler,
             }

             return components



      def       UpperCamelCase__		(self    :    Optional[int]   ,						UpperCamelCase    :    Optional[int]   ,						UpperCamelCase    :    str=0 ):

             '''simple docstring'''







             lowercase__      			= floats_tensor((1, 3, 64, 64)   ,						rng=random.Random(UpperCamelCase ) ).to(UpperCamelCase )

             if str(UpperCamelCase ).startswith('''mps''' ):
                    lowercase__      			= torch.manual_seed(UpperCamelCase )
             else:
                    lowercase__      			= torch.Generator(device=UpperCamelCase ).manual_seed(UpperCamelCase )
             lowercase__      			= {
                 '''image''': input_image,
                 '''generator''': generator,
                 '''num_inference_steps''': 1,
                 '''frame_size''': 32,
                 '''output_type''': '''np''',
             }
             return inputs



      def       UpperCamelCase__		(self    :    Tuple ):

             '''simple docstring'''







             lowercase__      			= '''cpu'''

             lowercase__      			= self.get_dummy_components()

             lowercase__      			= self.pipeline_class(**UpperCamelCase )
             lowercase__      			= pipe.to(UpperCamelCase )

             pipe.set_progress_bar_config(disable=UpperCamelCase )

             lowercase__      			= pipe(**self.get_dummy_inputs(UpperCamelCase ) )
             lowercase__      			= output.images[0]
             lowercase__      			= image[0, -3:, -3:, -1]

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

             lowercase__      			= 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    :    Optional[Any] ):

             '''simple docstring'''







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



      def       UpperCamelCase__		(self    :    int ):

             '''simple docstring'''







             lowercase__      			= torch_device == '''cpu'''
             lowercase__      			= True
             self._test_inference_batch_single_identical(
                 batch_size=2   ,						test_max_difference=UpperCamelCase   ,						relax_max_difference=UpperCamelCase   ,						)






      def       UpperCamelCase__		(self    :    int ):

             '''simple docstring'''







             lowercase__      			= self.get_dummy_components()
             lowercase__      			= self.pipeline_class(**UpperCamelCase )
             lowercase__      			= pipe.to(UpperCamelCase )
             pipe.set_progress_bar_config(disable=UpperCamelCase )

             lowercase__      			= 1
             lowercase__      			= 2

             lowercase__      			= self.get_dummy_inputs(UpperCamelCase )

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

             lowercase__      			= pipe(**UpperCamelCase   ,						num_images_per_prompt=UpperCamelCase )[0]

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






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





      '''simple docstring'''






      def       UpperCamelCase__		(self    :    Optional[Any] ):

             '''simple docstring'''







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






      def       UpperCamelCase__		(self    :    Union[str, Any] ):

             '''simple docstring'''







             lowercase__      			= load_image(
                 '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main''' '''/shap_e/corgi.png''' )
             lowercase__      			= load_numpy(
                 '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
                 '''/shap_e/test_shap_e_img2img_out.npy''' )
             lowercase__      			= ShapEImgaImgPipeline.from_pretrained('''openai/shap-e-img2img''' )
             lowercase__      			= pipe.to(UpperCamelCase )
             pipe.set_progress_bar_config(disable=UpperCamelCase )

             lowercase__      			= torch.Generator(device=UpperCamelCase ).manual_seed(0 )

             lowercase__      			= pipe(
                 UpperCamelCase   ,						generator=UpperCamelCase   ,						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(UpperCamelCase   ,						UpperCamelCase )




