





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 					lowercase__ (			__lowerCamelCase    ,       unittest.TestCase       ):

   '''simple docstring'''


   a		:							Tuple     	=	ShapEPipeline
   a		:							str     	=	["prompt"]
   a		:							Optional[int]     	=	["prompt"]
   a		:							Optional[Any]     	=	[
       "num_images_per_prompt",
       "num_inference_steps",
       "generator",
       "latents",
       "guidance_scale",
       "frame_size",
       "output_type",
       "return_dict",
   ]
   a		:							Union[str, Any]     	=	False




   @property
   def    UpperCamelCase__						(					self							)				->   str:






       """simple docstring"""
       return 32




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






       """simple docstring"""
       return 32




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






       """simple docstring"""
       return self.time_input_dim * 4




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






       """simple docstring"""
       return 8




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






       """simple docstring"""
       UpperCamelCase__							:       List[Any]      =   CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''							)
       return tokenizer




   @property
   def    UpperCamelCase__						(					self							)				->   int:






       """simple docstring"""
       torch.manual_seed(0							)
       UpperCamelCase__							:       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=1000,	)
       return CLIPTextModelWithProjection(__magic_name__							)




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






       """simple docstring"""
       torch.manual_seed(0							)

       UpperCamelCase__							:       List[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,
       }

       UpperCamelCase__							:       int      =   PriorTransformer(**__magic_name__							)
       return model




   @property
   def    UpperCamelCase__						(					self							)				->   str:






       """simple docstring"""
       torch.manual_seed(0							)

       UpperCamelCase__							:       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,
           ),
       }
       UpperCamelCase__							:       Union[str, Any]      =   ShapERenderer(**__magic_name__							)
       return model




   def    UpperCamelCase__						(					self							)				->   List[str]:






       """simple docstring"""
       UpperCamelCase__							:       int      =   self.dummy_prior
       UpperCamelCase__							:       List[Any]      =   self.dummy_text_encoder
       UpperCamelCase__							:       Optional[int]      =   self.dummy_tokenizer
       UpperCamelCase__							:       Tuple      =   self.dummy_renderer

       UpperCamelCase__							:       Optional[Any]      =   HeunDiscreteScheduler(
           beta_schedule='''exp''',	num_train_timesteps=1024,	prediction_type='''sample''',	use_karras_sigmas=__magic_name__,	clip_sample=__magic_name__,	clip_sample_range=1.0,	)
       UpperCamelCase__							:       List[str]      =   {
           '''prior''': prior,
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
           '''renderer''': renderer,
           '''scheduler''': scheduler,
       }

       return components




   def    UpperCamelCase__						(					self,	__magic_name__,	__magic_name__=0							)				->   Optional[Any]:






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




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






       """simple docstring"""
       UpperCamelCase__							:       str      =   '''cpu'''

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

       UpperCamelCase__							:       Tuple      =   self.pipeline_class(**__magic_name__							)
       UpperCamelCase__							:       int      =   pipe.to(__magic_name__							)

       pipe.set_progress_bar_config(disable=__magic_name__							)

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

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

       UpperCamelCase__							:       Dict      =   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							)				->   str:






       """simple docstring"""
       # NOTE: Larger batch sizes cause this test to timeout, only test on smaller batches
       self._test_inference_batch_consistent(batch_sizes=[1, 2]							)




   def    UpperCamelCase__						(					self							)				->   Optional[Any]:






       """simple docstring"""
       UpperCamelCase__							:       int      =   torch_device == '''cpu'''
       UpperCamelCase__							:       Optional[int]      =   True

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




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






       """simple docstring"""
       UpperCamelCase__							:       Any      =   self.get_dummy_components()
       UpperCamelCase__							:       int      =   self.pipeline_class(**__magic_name__							)
       UpperCamelCase__							:       Tuple      =   pipe.to(__magic_name__							)
       pipe.set_progress_bar_config(disable=__magic_name__							)

       UpperCamelCase__							:       List[Any]      =   1
       UpperCamelCase__							:       Optional[Any]      =   2

       UpperCamelCase__							:       List[str]      =   self.get_dummy_inputs(__magic_name__							)

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

       UpperCamelCase__							:       Optional[Any]      =   pipe(**__magic_name__,	num_images_per_prompt=__magic_name__							)[0]

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






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

   '''simple docstring'''


   def    UpperCamelCase__						(					self							)				->   Tuple:






       """simple docstring"""
       # clean up the VRAM after each test
       super().tearDown()
       gc.collect()
       torch.cuda.empty_cache()




   def    UpperCamelCase__						(					self							)				->   Tuple:






       """simple docstring"""
       UpperCamelCase__							:       Optional[Any]      =   load_numpy(
           '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
           '''/shap_e/test_shap_e_np_out.npy'''							)
       UpperCamelCase__							:       List[Any]      =   ShapEPipeline.from_pretrained('''openai/shap-e'''							)
       UpperCamelCase__							:       Tuple      =   pipe.to(__magic_name__							)
       pipe.set_progress_bar_config(disable=__magic_name__							)

       UpperCamelCase__							:       Union[str, Any]      =   torch.Generator(device=__magic_name__							).manual_seed(0							)

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








import json
import re
from typing import TYPE_CHECKING, List, Optional, Tuple, Union

import numpy as np

from ...utils import is_tf_available, is_torch_available, logging


if TYPE_CHECKING:
  if is_torch_available():
    import torch
  if is_tf_available():
    import tensorflow as tf

from tokenizers import pre_tokenizers

from ...tokenization_utils_base import BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from .tokenization_codegen import CodeGenTokenizer


UpperCAmelCase_  							=				logging.get_logger(__name__)

UpperCAmelCase_  							=				{'vocab_file': 'vocab.json', 'merges_file': 'merges.txt', 'tokenizer_file': 'tokenizer.json'}

UpperCAmelCase_  							=				{
    'vocab_file': {
        'Salesforce/codegen-350M-mono': 'https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/vocab.json',
    },
    'merges_file': {
        'Salesforce/codegen-350M-mono': 'https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/merges.txt',
    },
    'tokenizer_file': {
        'Salesforce/codegen-350M-mono': (
            'https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/tokenizer.json'
        ),
    },
}

UpperCAmelCase_  							=				{
    'Salesforce/codegen-350M-mono': 2048,
}






class 					lowercase__ (			__lowerCamelCase       ):

   '''simple docstring'''


   a		:							Union[str, Any]     	=	VOCAB_FILES_NAMES
   a		:							List[str]     	=	PRETRAINED_VOCAB_FILES_MAP
   a		:							Optional[int]     	=	PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
   a		:							List[Any]     	=	["input_ids", "attention_mask"]
   a		:							Any     	=	CodeGenTokenizer




   def __init__(					self,	__magic_name__=None,	__magic_name__=None,	__magic_name__=None,	__magic_name__="<|endoftext|>",	__magic_name__="<|endoftext|>",	__magic_name__="<|endoftext|>",	__magic_name__=False,	**__magic_name__,	)				->   Optional[int]:






       """simple docstring"""
       super().__init__(
           __magic_name__,	__magic_name__,	tokenizer_file=__magic_name__,	unk_token=__magic_name__,	bos_token=__magic_name__,	eos_token=__magic_name__,	add_prefix_space=__magic_name__,	**__magic_name__,	)

       if kwargs.pop('''add_bos_token''',	__magic_name__							):
           UpperCamelCase__							:       List[Any]      =   kwargs.pop('''name_or_path''',	''''''							)
           raise ValueError(
               '''Currenty GPT2\'s fast tokenizer does NOT support adding a BOS token.'''
               '''Instead you should use GPT2\'s slow tokenizer class `CodeGenTokenizer` as follows: \n'''
               f"`CodeGenTokenizer.from_pretrained('{model_id}')`\nor\n"
               f"`AutoTokenizer.from_pretrained('{model_id}', use_fast=False)`\n"
               '''This issue will be fixed soon, see: https://github.com/huggingface/tokenizers/pull/1005.'''
               ''' so that the fast tokenizer works correctly.'''							)

       UpperCamelCase__							:       Dict      =   json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()							)
       if pre_tok_state.get('''add_prefix_space''',	__magic_name__							) != add_prefix_space:
           UpperCamelCase__							:       List[str]      =   getattr(__magic_name__,	pre_tok_state.pop('''type'''							)							)
           UpperCamelCase__							:       int      =   add_prefix_space
           UpperCamelCase__							:       List[str]      =   pre_tok_class(**__magic_name__							)

       UpperCamelCase__							:       List[Any]      =   add_prefix_space




   def    UpperCamelCase__						(					self,	*__magic_name__,	**__magic_name__							)				->   BatchEncoding:






       """simple docstring"""
       UpperCamelCase__							:       List[Any]      =   kwargs.get('''is_split_into_words''',	__magic_name__							)
       assert self.add_prefix_space or not is_split_into_words, (
           f"You need to instantiate {self.__class__.__name__} with add_prefix_space=True "
           "to use it with pretokenized inputs."
       )

       return super()._batch_encode_plus(*__magic_name__,	**__magic_name__							)




   def    UpperCamelCase__						(					self,	*__magic_name__,	**__magic_name__							)				->   BatchEncoding:






       """simple docstring"""
       UpperCamelCase__							:       Dict      =   kwargs.get('''is_split_into_words''',	__magic_name__							)

       assert self.add_prefix_space or not is_split_into_words, (
           f"You need to instantiate {self.__class__.__name__} with add_prefix_space=True "
           "to use it with pretokenized inputs."
       )

       return super()._encode_plus(*__magic_name__,	**__magic_name__							)




   def    UpperCamelCase__						(					self,	__magic_name__,	__magic_name__ = None							)				->   Tuple[str]:






       """simple docstring"""
       UpperCamelCase__							:       List[Any]      =   self._tokenizer.model.save(__magic_name__,	name=__magic_name__							)
       return tuple(__magic_name__							)




   def    UpperCamelCase__						(					self,	__magic_name__,	__magic_name__ = False,	__magic_name__ = None,	__magic_name__ = None,	**__magic_name__,	)				->   str:






       """simple docstring"""
       UpperCamelCase__							:       Dict      =   super().decode(
           token_ids=__magic_name__,	skip_special_tokens=__magic_name__,	clean_up_tokenization_spaces=__magic_name__,	**__magic_name__,	)

       if truncate_before_pattern is not None and len(__magic_name__							) > 0:
           UpperCamelCase__							:       str      =   self.truncate(__magic_name__,	__magic_name__							)

       return decoded_text




   def    UpperCamelCase__						(					self,	__magic_name__,	__magic_name__							)				->   List[str]:






       """simple docstring"""
       def find_re(__magic_name__,	__magic_name__,	__magic_name__							):
           UpperCamelCase__							:       Dict      =   pattern.search(__magic_name__,	__magic_name__							)
           return m.start() if m else -1

       UpperCamelCase__							:       Union[str, Any]      =   [re.compile(__magic_name__,	re.MULTILINE							) for pattern in truncate_before_pattern]

       UpperCamelCase__							:       List[Any]      =   list(re.finditer('''^print''',	__magic_name__,	re.MULTILINE							)							)

       if len(__magic_name__							) > 1:
           UpperCamelCase__							:       Tuple      =   completion[: prints[1].start()]

       UpperCamelCase__							:       Optional[int]      =   list(re.finditer('''^def''',	__magic_name__,	re.MULTILINE							)							)

       if len(__magic_name__							) > 1:
           UpperCamelCase__							:       Optional[Any]      =   completion[: defs[1].start()]

       UpperCamelCase__							:       str      =   0

       UpperCamelCase__							:       Tuple      =   [
           pos for pos in [find_re(__magic_name__,	__magic_name__,	__magic_name__							) for terminal in terminals] if pos != -1
       ]

       if len(__magic_name__							) > 0:
           return completion[: min(__magic_name__							)]
       else:
           return completion

