



import json
import os
from functools import lru_cache
from typing import TYPE_CHECKING, List, Optional, Tuple, Union

import numpy as np
import regex as re

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 ...tokenization_utils import AddedToken, PreTrainedTokenizer


__lowerCamelCase      =    logging.get_logger(__name__)

__lowerCamelCase      =    {
    """vocab_file""": """vocab.json""",
    """merges_file""": """merges.txt""",
}

__lowerCamelCase      =    {
    """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""",
    },
}

__lowerCamelCase      =    {
    """Salesforce/codegen-350M-mono""": 2_0_4_8,
}




@lru_cache()
def 			UpperCAmelCase__       (							)      ->       Optional[int]:
  _A							=			(
      list(range(ord('''!'''							)  , ord('''~'''							) + 1							)							) + list(range(ord('''¡'''							)  , ord('''¬'''							) + 1							)							) + list(range(ord('''®'''							)  , ord('''ÿ'''							) + 1							)							)
  )
  _A							=			bs[:]
  _A							=			0
  for b in range(2**8							):
    if b not in bs:
      bs.append(lowercase_							)
      cs.append(2**8 + n							)
      n += 1
  _A							=			[chr(lowercase_							) for n in cs]
  return dict(zip(lowercase_  , lowercase_							)							)




def 			UpperCAmelCase__       (							__snake_case							)      ->       Any:
  _A							=			set()
  _A							=			word[0]
  for char in word[1:]:
    pairs.add((prev_char, char)							)
    _A							=			char
  return pairs



class 						_snake_case					(	a_ ):





 """simple docstring"""



 lowerCamelCase_								=	VOCAB_FILES_NAMES
 lowerCamelCase_								=	PRETRAINED_VOCAB_FILES_MAP
 lowerCamelCase_								=	PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 lowerCamelCase_								=	['input_ids', 'attention_mask']
 def __init__(  self		,			a		,			a		,			a="replace"		,			a="<|endoftext|>"		,			a="<|endoftext|>"		,			a="<|endoftext|>"		,			a=None		,			a=False		,			a=False		,			**a		,			)		->							int:




   """simple docstring"""







   _A							=			AddedToken(a		,			lstrip=a		,			rstrip=a    ) if isinstance(a		,			a    ) else bos_token
   _A							=			AddedToken(a		,			lstrip=a		,			rstrip=a    ) if isinstance(a		,			a    ) else eos_token
   _A							=			AddedToken(a		,			lstrip=a		,			rstrip=a    ) if isinstance(a		,			a    ) else unk_token
   _A							=			AddedToken(a		,			lstrip=a		,			rstrip=a    ) if isinstance(a		,			a    ) else pad_token
   super().__init__(
       errors=a		,			unk_token=a		,			bos_token=a		,			eos_token=a		,			pad_token=a		,			add_prefix_space=a		,			add_bos_token=a		,			**a		,			)
   _A							=			add_bos_token

   with open(a		,			encoding='''utf-8'''    ) as vocab_handle:
     _A							=			json.load(a    )
   _A							=			{v: k for k, v in self.encoder.items()}
   _A							=			errors  # how to handle errors in decoding
   _A							=			bytes_to_unicode()
   _A							=			{v: k for k, v in self.byte_encoder.items()}
   with open(a		,			encoding='''utf-8'''    ) as merges_handle:
     _A							=			merges_handle.read().split('''\n'''    )[1:-1]
   _A							=			[tuple(merge.split()    ) for merge in bpe_merges]
   _A							=			dict(zip(a		,			range(len(a    )    )    )    )
   _A							=			{}
   _A							=			add_prefix_space

   # Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
   _A							=			re.compile(R'''\'s|\'t|\'re|\'ve|\'m|\'ll|\'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+'''    )
 @property
 def lowercase_						(  self    )		->							Union[str, Any]:




   """simple docstring"""







   return len(self.encoder    )
 def lowercase_						(  self    )		->							Optional[Any]:




   """simple docstring"""







   return dict(self.encoder		,			**self.added_tokens_encoder    )
 def lowercase_						(  self		,			a    )		->							Union[str, Any]:




   """simple docstring"""







   if token in self.cache:
     return self.cache[token]
   _A							=			tuple(a    )
   _A							=			get_pairs(a    )

   if not pairs:
     return token

   while True:
     _A							=			min(a		,			key=lambda a    : self.bpe_ranks.get(a		,			float('''inf'''    )    )    )
     if bigram not in self.bpe_ranks:
       break
     _A							=			bigram
     _A							=			[]
     _A							=			0
     while i < len(a    ):
       try:
         _A							=			word.index(a		,			a    )
       except ValueError:
         new_word.extend(word[i:]    )
         break
       else:
         new_word.extend(word[i:j]    )
         _A							=			j

       if word[i] == first and i < len(a    ) - 1 and word[i + 1] == second:
         new_word.append(first + second    )
         i += 2
       else:
         new_word.append(word[i]    )
         i += 1
     _A							=			tuple(a    )
     _A							=			new_word
     if len(a    ) == 1:
       break
     else:
       _A							=			get_pairs(a    )
   _A							=			' '.join(a    )
   _A							=			word
   return word
 def lowercase_						(  self		,			a		,			a=None    )		->							Any:




   """simple docstring"""







   if self.add_bos_token:
     _A							=			[self.bos_token_id]
   else:
     _A							=			[]

   _A							=			bos_token_ids + token_ids_a

   if token_ids_a is None:
     return output

   return output + bos_token_ids + token_ids_a
 def lowercase_						(  self		,			a    )		->							str:




   """simple docstring"""







   _A							=			[]
   for token in re.findall(self.pat		,			a    ):
     _A							=			''.join(
         self.byte_encoder[b] for b in token.encode('''utf-8'''    )    )  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
     bpe_tokens.extend(bpe_token for bpe_token in self.bpe(a    ).split(''' '''    )    )
   return bpe_tokens
 def lowercase_						(  self		,			a    )		->							List[str]:




   """simple docstring"""







   return self.encoder.get(a		,			self.encoder.get(self.unk_token    )    )
 def lowercase_						(  self		,			a    )		->							List[str]:




   """simple docstring"""







   return self.decoder.get(a    )
 def lowercase_						(  self		,			a    )		->							List[Any]:




   """simple docstring"""







   _A							=			''.join(a    )
   _A							=			bytearray([self.byte_decoder[c] for c in text]    ).decode('''utf-8'''		,			errors=self.errors    )
   return text
 def lowercase_						(  self		,			a		,			a = None    )		->							Tuple[str]:




   """simple docstring"""







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

   with open(a		,			'''w'''		,			encoding='''utf-8'''    ) as f:
     f.write(json.dumps(self.encoder		,			indent=2		,			sort_keys=a		,			ensure_ascii=a    ) + '''\n'''    )

   _A							=			0
   with open(a		,			'''w'''		,			encoding='''utf-8'''    ) as writer:
     writer.write('''#version: 0.2\n'''    )
     for bpe_tokens, token_index in sorted(self.bpe_ranks.items()		,			key=lambda a    : kv[1]    ):
       if index != token_index:
         logger.warning(
             f'''Saving vocabulary to {merge_file}: BPE merge indices are not consecutive.'''
             ''' Please check that the tokenizer is not corrupted!'''    )
         _A							=			token_index
       writer.write(''' '''.join(a    ) + '''\n'''    )
       index += 1

   return vocab_file, merge_file
 def lowercase_						(  self		,			a		,			a=False		,			**a    )		->							Tuple:




   """simple docstring"""







   _A							=			kwargs.pop('''add_prefix_space'''		,			self.add_prefix_space    )
   if is_split_into_words or add_prefix_space:
     _A							=			' ' + text
   return (text, kwargs)
 def lowercase_						(  self		,			a		,			a = False		,			a = None		,			a = None		,			**a		,			)		->							str:




   """simple docstring"""







   _A							=			super()._decode(
       token_ids=a		,			skip_special_tokens=a		,			clean_up_tokenization_spaces=a		,			**a		,			)

   if truncate_before_pattern is not None and len(a    ) > 0:
     _A							=			self.truncate(a		,			a    )

   return decoded_text



 def lowercase_						(  self		,			a		,			a    )		->							Any:




   """simple docstring"""







   def find_re(a		,			a		,			a    ):
     _A							=			pattern.search(a		,			a    )
     return m.start() if m else -1

   _A							=			[re.compile(a		,			re.MULTILINE    ) for pattern in truncate_before_pattern]

   _A							=			list(re.finditer('''^print'''		,			a		,			re.MULTILINE    )    )

   if len(a    ) > 1:
     _A							=			completion[: prints[1].start()]

   _A							=			list(re.finditer('''^def'''		,			a		,			re.MULTILINE    )    )

   if len(a    ) > 1:
     _A							=			completion[: defs[1].start()]

   _A							=			0

   _A							=			[
       pos for pos in [find_re(a		,			a		,			a    ) for terminal in terminals] if pos != -1
   ]

   if len(a    ) > 0:
     return completion[: min(a    )]
   else:
     return completion
import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMParallelScheduler,
    DDPMParallelScheduler,
    StableDiffusionParadigmsPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    require_torch_gpu,
)

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




class       __lowerCAmelCase     (				a_,		a_,		unittest.TestCase  ):



 '''simple docstring'''
 _lowerCamelCase  :		Any 			=       StableDiffusionParadigmsPipeline
 _lowerCamelCase  :		Any 			=       TEXT_TO_IMAGE_PARAMS
 _lowerCamelCase  :		str 			=       TEXT_TO_IMAGE_BATCH_PARAMS
 _lowerCamelCase  :		int 			=       TEXT_TO_IMAGE_IMAGE_PARAMS
 _lowerCamelCase  :		List[Any] 			=       TEXT_TO_IMAGE_IMAGE_PARAMS

 def     __snake_case    (	self				:						int     )       ->							List[str]:
        torch.manual_seed(0     )
        _UpperCamelCase   :				Union[str, Any]       =     UNetaDConditionModel(
            block_out_channels=(3_2, 6_4)				,	layers_per_block=2				,	sample_size=3_2				,	in_channels=4				,	out_channels=4				,	down_block_types=('DownBlock2D', 'CrossAttnDownBlock2D')				,	up_block_types=('CrossAttnUpBlock2D', 'UpBlock2D')				,	cross_attention_dim=3_2				,	attention_head_dim=(2, 4)				,	use_linear_projection=lowerCamelCase				,	)
        _UpperCamelCase   :				str       =     DDIMParallelScheduler(
            beta_start=0.0_0_0_8_5				,	beta_end=0.0_1_2				,	beta_schedule='scaled_linear'				,	clip_sample=lowerCamelCase				,	set_alpha_to_one=lowerCamelCase				,	)
        torch.manual_seed(0     )
        _UpperCamelCase   :				str       =     AutoencoderKL(
            block_out_channels=[3_2, 6_4]				,	in_channels=3				,	out_channels=3				,	down_block_types=['DownEncoderBlock2D', 'DownEncoderBlock2D']				,	up_block_types=['UpDecoderBlock2D', 'UpDecoderBlock2D']				,	latent_channels=4				,	sample_size=1_2_8				,	)
        torch.manual_seed(0     )
        _UpperCamelCase   :				int       =     CLIPTextConfig(
            bos_token_id=0				,	eos_token_id=2				,	hidden_size=3_2				,	intermediate_size=3_7				,	layer_norm_eps=1e-05				,	num_attention_heads=4				,	num_hidden_layers=5				,	pad_token_id=1				,	vocab_size=1_0_0_0				,	hidden_act='gelu'				,	projection_dim=5_1_2				,	)
        _UpperCamelCase   :				Any       =     CLIPTextModel(lowerCamelCase     )
        _UpperCamelCase   :				Optional[Any]       =     CLIPTokenizer.from_pretrained('hf-internal-testing/tiny-random-clip'     )

        _UpperCamelCase   :				Tuple       =     {
            'unet': unet,
            'scheduler': scheduler,
            'vae': vae,
            'text_encoder': text_encoder,
            'tokenizer': tokenizer,
            'safety_checker': None,
            'feature_extractor': None,
        }
        return components

 def     __snake_case    (	self				:						List[Any]				,	lowerCamelCase				:						Any				,	lowerCamelCase				:						str=0     )       ->							Dict:
        if str(lowerCamelCase     ).startswith('mps'     ):
               _UpperCamelCase   :				List[Any]       =     torch.manual_seed(lowerCamelCase     )
        else:
               _UpperCamelCase   :				Union[str, Any]       =     torch.Generator(device=lowerCamelCase     ).manual_seed(lowerCamelCase     )
        _UpperCamelCase   :				Optional[Any]       =     {
            'prompt': 'a photograph of an astronaut riding a horse',
            'generator': generator,
            'num_inference_steps': 1_0,
            'guidance_scale': 6.0,
            'output_type': 'numpy',
            'parallel': 3,
            'debug': True,
        }
        return inputs

 def     __snake_case    (	self				:						List[Any]     )       ->							str:
        _UpperCamelCase   :				List[str]       =     'cpu'  # ensure determinism for the device-dependent torch.Generator
        _UpperCamelCase   :				List[Any]       =     self.get_dummy_components()
        _UpperCamelCase   :				Tuple       =     StableDiffusionParadigmsPipeline(**lowerCamelCase     )
        _UpperCamelCase   :				str       =     sd_pipe.to(lowerCamelCase     )
        sd_pipe.set_progress_bar_config(disable=lowerCamelCase     )

        _UpperCamelCase   :				Any       =     self.get_dummy_inputs(lowerCamelCase     )
        _UpperCamelCase   :				str       =     sd_pipe(**lowerCamelCase     ).images
        _UpperCamelCase   :				str       =     image[0, -3:, -3:, -1]
        assert image.shape == (1, 6_4, 6_4, 3)

        _UpperCamelCase   :				Any       =     np.array([0.4_7_7_3, 0.5_4_1_7, 0.4_7_2_3, 0.4_9_2_5, 0.5_6_3_1, 0.4_7_5_2, 0.5_2_4_0, 0.4_9_3_5, 0.5_0_2_3]     )

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

 def     __snake_case    (	self				:						List[Any]     )       ->							Union[str, Any]:
        _UpperCamelCase   :				Dict       =     'cpu'  # ensure determinism for the device-dependent torch.Generator
        _UpperCamelCase   :				str       =     self.get_dummy_components()
        torch.manual_seed(0     )
        _UpperCamelCase   :				Union[str, Any]       =     DDPMParallelScheduler()
        torch.manual_seed(0     )
        _UpperCamelCase   :				List[Any]       =     StableDiffusionParadigmsPipeline(**lowerCamelCase     )
        _UpperCamelCase   :				str       =     sd_pipe.to(lowerCamelCase     )
        sd_pipe.set_progress_bar_config(disable=lowerCamelCase     )

        _UpperCamelCase   :				int       =     self.get_dummy_inputs(lowerCamelCase     )
        _UpperCamelCase   :				Optional[int]       =     sd_pipe(**lowerCamelCase     ).images
        _UpperCamelCase   :				Dict       =     image[0, -3:, -3:, -1]
        assert image.shape == (1, 6_4, 6_4, 3)

        _UpperCamelCase   :				List[Any]       =     np.array([0.3_5_7_3, 0.4_4_2_0, 0.4_9_6_0, 0.4_7_9_9, 0.3_7_9_6, 0.3_8_7_9, 0.4_8_1_9, 0.4_3_6_5, 0.4_4_6_8]     )

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

 def     __snake_case    (	self				:						Any     )       ->							List[Any]:
        super().test_inference_batch_consistent(batch_sizes=[1, 2]     )

 def     __snake_case    (	self				:						Union[str, Any]     )       ->							Tuple:
        super().test_inference_batch_single_identical(batch_size=2				,	expected_max_diff=3e-3     )

 def     __snake_case    (	self				:						List[str]     )       ->							List[str]:
        _UpperCamelCase   :				List[str]       =     'cpu'  # ensure determinism for the device-dependent torch.Generator
        _UpperCamelCase   :				Tuple       =     self.get_dummy_components()
        _UpperCamelCase   :				Any       =     StableDiffusionParadigmsPipeline(**lowerCamelCase     )
        _UpperCamelCase   :				List[Any]       =     sd_pipe.to(lowerCamelCase     )
        sd_pipe.set_progress_bar_config(disable=lowerCamelCase     )

        _UpperCamelCase   :				Tuple       =     self.get_dummy_inputs(lowerCamelCase     )
        _UpperCamelCase   :				int       =     'french fries'
        _UpperCamelCase   :				Optional[Any]       =     sd_pipe(**lowerCamelCase				,	negative_prompt=lowerCamelCase     )
        _UpperCamelCase   :				int       =     output.images
        _UpperCamelCase   :				str       =     image[0, -3:, -3:, -1]

        assert image.shape == (1, 6_4, 6_4, 3)

        _UpperCamelCase   :				Dict       =     np.array([0.4_7_7_1, 0.5_4_2_0, 0.4_6_8_3, 0.4_9_1_8, 0.5_6_3_6, 0.4_7_2_5, 0.5_2_3_0, 0.4_9_2_3, 0.5_0_1_5]     )

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







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



 '''simple docstring'''

 def     __snake_case    (	self				:						Optional[Any]     )       ->							Optional[int]:
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()

 def     __snake_case    (	self				:						Dict				,	lowerCamelCase				:						Any=0     )       ->							Dict:
        _UpperCamelCase   :				List[str]       =     torch.Generator(device=lowerCamelCase     ).manual_seed(lowerCamelCase     )
        _UpperCamelCase   :				Tuple       =     {
            'prompt': 'a photograph of an astronaut riding a horse',
            'generator': generator,
            'num_inference_steps': 1_0,
            'guidance_scale': 7.5,
            'output_type': 'numpy',
            'parallel': 3,
            'debug': True,
        }
        return inputs

 def     __snake_case    (	self				:						List[Any]     )       ->							str:
        _UpperCamelCase   :				List[str]       =     'stabilityai/stable-diffusion-2-base'
        _UpperCamelCase   :				int       =     DDIMParallelScheduler.from_pretrained(lowerCamelCase				,	subfolder='scheduler'     )
        _UpperCamelCase   :				List[Any]       =     StableDiffusionParadigmsPipeline.from_pretrained(lowerCamelCase				,	scheduler=lowerCamelCase				,	safety_checker=lowerCamelCase     )
        pipe.to(lowerCamelCase     )
        pipe.set_progress_bar_config(disable=lowerCamelCase     )
        pipe.enable_attention_slicing()

        _UpperCamelCase   :				Optional[Any]       =     self.get_inputs()
        _UpperCamelCase   :				Optional[int]       =     pipe(**lowerCamelCase     ).images
        _UpperCamelCase   :				Optional[int]       =     image[0, -3:, -3:, -1].flatten()

        assert image.shape == (1, 5_1_2, 5_1_2, 3)

        _UpperCamelCase   :				int       =     np.array([0.9_6_2_2, 0.9_6_0_2, 0.9_7_4_8, 0.9_5_9_1, 0.9_6_3_0, 0.9_6_9_1, 0.9_6_6_1, 0.9_6_3_1, 0.9_7_4_1]     )

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



