





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 logging


_UpperCamelCase				=       logging.get_logger(__name__)
_UpperCamelCase				=       {'''vocab_file''': '''spiece.model'''}

_UpperCamelCase				=       {
    '''vocab_file''': {
        '''albert-base-v1''': '''https://huggingface.co/albert-base-v1/resolve/main/spiece.model''',
        '''albert-large-v1''': '''https://huggingface.co/albert-large-v1/resolve/main/spiece.model''',
        '''albert-xlarge-v1''': '''https://huggingface.co/albert-xlarge-v1/resolve/main/spiece.model''',
        '''albert-xxlarge-v1''': '''https://huggingface.co/albert-xxlarge-v1/resolve/main/spiece.model''',
        '''albert-base-v2''': '''https://huggingface.co/albert-base-v2/resolve/main/spiece.model''',
        '''albert-large-v2''': '''https://huggingface.co/albert-large-v2/resolve/main/spiece.model''',
        '''albert-xlarge-v2''': '''https://huggingface.co/albert-xlarge-v2/resolve/main/spiece.model''',
        '''albert-xxlarge-v2''': '''https://huggingface.co/albert-xxlarge-v2/resolve/main/spiece.model''',
    }
}

_UpperCamelCase				=       {
    '''albert-base-v1''': 512,
    '''albert-large-v1''': 512,
    '''albert-xlarge-v1''': 512,
    '''albert-xxlarge-v1''': 512,
    '''albert-base-v2''': 512,
    '''albert-large-v2''': 512,
    '''albert-xlarge-v2''': 512,
    '''albert-xxlarge-v2''': 512,
}

_UpperCamelCase				=       '''▁'''




class  _lowerCamelCase       (			UpperCAmelCase_			):



      """simple docstring"""






      UpperCAmelCase_  :					Tuple      			=VOCAB_FILES_NAMES
      UpperCAmelCase_  :					Union[str, Any]      			=PRETRAINED_VOCAB_FILES_MAP
      UpperCAmelCase_  :					Union[str, Any]      			=PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES







      def __init__(       self ,  UpperCAmelCase ,  UpperCAmelCase=True ,  UpperCAmelCase=True ,  UpperCAmelCase=False ,  UpperCAmelCase="[CLS]" ,  UpperCAmelCase="[SEP]" ,  UpperCAmelCase="<unk>" ,  UpperCAmelCase="[SEP]" ,  UpperCAmelCase="<pad>" ,  UpperCAmelCase="[CLS]" ,  UpperCAmelCase="[MASK]" ,  UpperCAmelCase = None ,  **UpperCAmelCase ,  )				->      None:






        '''simple docstring'''






        __snake_case      : Union[str, Any]						=				(
            AddedToken(__lowercase ,  lstrip=__lowercase ,  rstrip=__lowercase ,  normalized=__lowercase      )
            if isinstance(__lowercase ,  __lowercase      )
            else mask_token
        )

        __snake_case      : List[str]						=				{} if sp_model_kwargs is None else sp_model_kwargs

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

        __snake_case      : str						=				do_lower_case
        __snake_case      : str						=				remove_space
        __snake_case      : Union[str, Any]						=				keep_accents
        __snake_case      : Dict						=				vocab_file

        __snake_case      : str						=				spm.SentencePieceProcessor(**self.sp_model_kwargs      )
        self.sp_model.Load(__lowercase      )







      @property
      def   UpperCAmelCase						(       self      )				->      Any:






        '''simple docstring'''






        return len(self.sp_model      )







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






        '''simple docstring'''






        __snake_case      : Optional[int]						=				{self.convert_ids_to_tokens(__lowercase      ): i for i in range(self.vocab_size      )}
        vocab.update(self.added_tokens_encoder      )
        return vocab







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






        '''simple docstring'''






        __snake_case      : List[str]						=				self.__dict__.copy()
        __snake_case      : Union[str, Any]						=				None
        return state







      def __setstate__(       self ,  UpperCAmelCase      )				->      Any:






        '''simple docstring'''






        __snake_case      : int						=				d

        # for backward compatibility
        if not hasattr(self ,  "sp_model_kwargs"      ):
          __snake_case      : Dict						=				{}

        __snake_case      : Dict						=				spm.SentencePieceProcessor(**self.sp_model_kwargs      )
        self.sp_model.Load(self.vocab_file      )







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






        '''simple docstring'''






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

        if not self.keep_accents:
          __snake_case      : Optional[Any]						=				unicodedata.normalize("NFKD" ,  __lowercase      )
          __snake_case      : List[str]						=				''''''.join([c for c in outputs if not unicodedata.combining(__lowercase      )]      )
        if self.do_lower_case:
          __snake_case      : List[Any]						=				outputs.lower()

        return outputs







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






        '''simple docstring'''






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

        return new_pieces







      def   UpperCAmelCase						(       self ,  UpperCAmelCase      )				->      str:






        '''simple docstring'''






        return self.sp_model.PieceToId(__lowercase      )







      def   UpperCAmelCase						(       self ,  UpperCAmelCase      )				->      int:






        '''simple docstring'''






        return self.sp_model.IdToPiece(__lowercase      )







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






        '''simple docstring'''






        __snake_case      : int						=				[]
        __snake_case      : List[Any]						=				''''''
        __snake_case      : int						=				False
        for token in tokens:
          # make sure that special tokens are not decoded using sentencepiece model
          if token in self.all_special_tokens:
            if not prev_is_special:
              out_string += " "
            out_string += self.sp_model.decode(__lowercase      ) + token
            __snake_case      : List[str]						=				True
            __snake_case      : Optional[Any]						=				[]
          else:
            current_sub_tokens.append(__lowercase      )
            __snake_case      : Tuple						=				False
        out_string += self.sp_model.decode(__lowercase      )
        return out_string.strip()







      def   UpperCAmelCase						(       self ,  UpperCAmelCase ,  UpperCAmelCase = None      )				->      List[int]:






        '''simple docstring'''






        __snake_case      : str						=				[self.sep_token_id]
        __snake_case      : Tuple						=				[self.cls_token_id]
        if token_ids_a is None:
          return cls + token_ids_a + sep
        return cls + token_ids_a + sep + token_ids_a + sep







      def   UpperCAmelCase						(       self ,  UpperCAmelCase ,  UpperCAmelCase = None ,  UpperCAmelCase = False      )				->      List[int]:






        '''simple docstring'''






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

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







      def   UpperCAmelCase						(       self ,  UpperCAmelCase ,  UpperCAmelCase = None      )				->      List[int]:






        '''simple docstring'''






        __snake_case      : List[Any]						=				[self.sep_token_id]
        __snake_case      : Any						=				[self.cls_token_id]

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







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






        '''simple docstring'''






        if not os.path.isdir(__lowercase      ):
          logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""      )
          return
        __snake_case      : List[Any]						=				os.path.join(
            __lowercase ,  (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]      )

        if os.path.abspath(self.vocab_file      ) != os.path.abspath(__lowercase      ) and os.path.isfile(self.vocab_file      ):
          copyfile(self.vocab_file ,  __lowercase      )
        elif not os.path.isfile(self.vocab_file      ):
          with open(__lowercase ,  "wb"      ) as fi:
            __snake_case      : Optional[int]						=				self.sp_model.serialized_model_proto()
            fi.write(__lowercase      )

        return (out_vocab_file,)


import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModel,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import AutoencoderKL, DDIMScheduler, DDPMScheduler, StableUnCLIPImgaImgPipeline, UNetaDConditionModel
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.pipelines.stable_diffusion.stable_unclip_image_normalizer import StableUnCLIPImageNormalizer
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import (
    enable_full_determinism,
    floats_tensor,
    load_image,
    load_numpy,
    require_torch_gpu,
    skip_mps,
    slow,
    torch_device,
)

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
    assert_mean_pixel_difference,
)


enable_full_determinism()


class 	lowerCamelCase_     (   UpperCAmelCase_						,	UpperCAmelCase_						,	UpperCAmelCase_						,	unittest.TestCase ):






    '''simple docstring'''






    a__	:						int						=      StableUnCLIPImgaImgPipeline
    a__	:						Optional[int]						=      TEXT_GUIDED_IMAGE_VARIATION_PARAMS
    a__	:						Union[str, Any]						=      TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
    a__	:						Optional[Any]						=      frozenset(
        [] )  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
    a__	:						int						=      frozenset([] )







    def 			UpperCamelCase__   (	self)       ->	Tuple:
        __UpperCamelCase     :Tuple			  =		32
        __UpperCamelCase     :Optional[int]			  =		embedder_hidden_size

        # image encoding components

        __UpperCamelCase     :Union[str, Any]			  =		CLIPImageProcessor(crop_size=32 , size=32)

        torch.manual_seed(0)
        __UpperCamelCase     :Union[str, Any]			  =		CLIPVisionModelWithProjection(
            CLIPVisionConfig(
                hidden_size=__lowercase , projection_dim=__lowercase , num_hidden_layers=5 , num_attention_heads=4 , image_size=32 , intermediate_size=37 , patch_size=1 , ))

        # regular denoising components

        torch.manual_seed(0)
        __UpperCamelCase     :str			  =		StableUnCLIPImageNormalizer(embedding_dim=__lowercase)
        __UpperCamelCase     :Optional[int]			  =		DDPMScheduler(beta_schedule='''squaredcos_cap_v2''')

        torch.manual_seed(0)
        __UpperCamelCase     :Union[str, Any]			  =		CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''')

        torch.manual_seed(0)
        __UpperCamelCase     :Dict			  =		CLIPTextModel(
            CLIPTextConfig(
                bos_token_id=0 , eos_token_id=2 , hidden_size=__lowercase , projection_dim=32 , 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 , ))

        torch.manual_seed(0)
        __UpperCamelCase     :List[Any]			  =		UNetaDConditionModel(
            sample_size=32 , in_channels=4 , out_channels=4 , down_block_types=('''CrossAttnDownBlock2D''', '''DownBlock2D''') , up_block_types=('''UpBlock2D''', '''CrossAttnUpBlock2D''') , block_out_channels=(32, 64) , attention_head_dim=(2, 4) , class_embed_type='''projection''' , projection_class_embeddings_input_dim=embedder_projection_dim * 2 , cross_attention_dim=__lowercase , layers_per_block=1 , upcast_attention=__lowercase , use_linear_projection=__lowercase , )

        torch.manual_seed(0)
        __UpperCamelCase     :Tuple			  =		DDIMScheduler(
            beta_schedule='''scaled_linear''' , beta_start=0.0_00_85 , beta_end=0.0_12 , prediction_type='''v_prediction''' , set_alpha_to_one=__lowercase , steps_offset=1 , )

        torch.manual_seed(0)
        __UpperCamelCase     :List[str]			  =		AutoencoderKL()

        __UpperCamelCase     :Tuple			  =		{
            # image encoding components
            '''feature_extractor''': feature_extractor,
            '''image_encoder''': image_encoder.eval(),
            # image noising components
            '''image_normalizer''': image_normalizer.eval(),
            '''image_noising_scheduler''': image_noising_scheduler,
            # regular denoising components
            '''tokenizer''': tokenizer,
            '''text_encoder''': text_encoder.eval(),
            '''unet''': unet.eval(),
            '''scheduler''': scheduler,
            '''vae''': vae.eval(),
        }

        return components







    def 			UpperCamelCase__   (	self , __lowercase , __lowercase=0 , __lowercase=True)       ->	str:
        if str(__lowercase).startswith('''mps'''):
            __UpperCamelCase     :Union[str, Any]			  =		torch.manual_seed(__lowercase)
        else:
            __UpperCamelCase     :int			  =		torch.Generator(device=__lowercase).manual_seed(__lowercase)

        __UpperCamelCase     :int			  =		floats_tensor((1, 3, 32, 32) , rng=random.Random(__lowercase)).to(__lowercase)

        if pil_image:
            __UpperCamelCase     :List[Any]			  =		input_image * 0.5 + 0.5
            __UpperCamelCase     :Optional[Any]			  =		input_image.clamp(0 , 1)
            __UpperCamelCase     :int			  =		input_image.cpu().permute(0 , 2 , 3 , 1).float().numpy()
            __UpperCamelCase     :Optional[Any]			  =		DiffusionPipeline.numpy_to_pil(__lowercase)[0]

        return {
            "prompt": "An anime racoon running a marathon",
            "image": input_image,
            "generator": generator,
            "num_inference_steps": 2,
            "output_type": "np",
        }







    @skip_mps
    def 			UpperCamelCase__   (	self)       ->	Union[str, Any]:
        __UpperCamelCase     :Dict			  =		'''cpu'''  # ensure determinism for the device-dependent torch.Generator
        __UpperCamelCase     :Tuple			  =		self.get_dummy_components()
        __UpperCamelCase     :Any			  =		StableUnCLIPImgaImgPipeline(**__lowercase)
        __UpperCamelCase     :Optional[Any]			  =		sd_pipe.to(__lowercase)
        sd_pipe.set_progress_bar_config(disable=__lowercase)

        __UpperCamelCase     :List[Any]			  =		self.get_dummy_inputs(__lowercase)
        inputs.update({'''image_embeds''': None})
        __UpperCamelCase     :Any			  =		sd_pipe(**__lowercase).images
        __UpperCamelCase     :List[str]			  =		image[0, -3:, -3:, -1]

        assert image.shape == (1, 32, 32, 3)
        __UpperCamelCase     :List[Any]			  =		np.array([0.38_72, 0.72_24, 0.56_01, 0.47_41, 0.68_72, 0.58_14, 0.46_36, 0.38_67, 0.50_78])

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







    def 			UpperCamelCase__   (	self)       ->	str:
        __UpperCamelCase     :Optional[Any]			  =		torch_device in ['''cpu''', '''mps''']

        self._test_attention_slicing_forward_pass(test_max_difference=__lowercase)







    def 			UpperCamelCase__   (	self)       ->	List[Any]:
        __UpperCamelCase     :Optional[Any]			  =		torch_device in ['''cpu''', '''mps''']

        self._test_inference_batch_single_identical(test_max_difference=__lowercase)







    @unittest.skipIf(
        torch_device != '''cuda''' or not is_xformers_available() , reason='''XFormers attention is only available with CUDA and `xformers` installed''' , )
    def 			UpperCamelCase__   (	self)       ->	Union[str, Any]:
        self._test_xformers_attention_forwardGenerator_pass(test_max_difference=__lowercase)


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






    '''simple docstring'''






    def 			UpperCamelCase__   (	self)       ->	Union[str, Any]:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()







    def 			UpperCamelCase__   (	self)       ->	Union[str, Any]:
        __UpperCamelCase     :int			  =		load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png''')

        __UpperCamelCase     :Any			  =		load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_l_img2img_anime_turtle_fp16.npy''')

        __UpperCamelCase     :List[Any]			  =		StableUnCLIPImgaImgPipeline.from_pretrained(
            '''fusing/stable-unclip-2-1-l-img2img''' , torch_dtype=torch.floataa)
        pipe.to(__lowercase)
        pipe.set_progress_bar_config(disable=__lowercase)
        # stable unclip will oom when integration tests are run on a V100,
        # so turn on memory savings
        pipe.enable_attention_slicing()
        pipe.enable_sequential_cpu_offload()

        __UpperCamelCase     :int			  =		torch.Generator(device='''cpu''').manual_seed(0)
        __UpperCamelCase     :Dict			  =		pipe(__lowercase , '''anime turle''' , generator=__lowercase , output_type='''np''')

        __UpperCamelCase     :Dict			  =		output.images[0]

        assert image.shape == (768, 768, 3)

        assert_mean_pixel_difference(__lowercase , __lowercase)







    def 			UpperCamelCase__   (	self)       ->	List[str]:
        __UpperCamelCase     :Optional[Any]			  =		load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png''')

        __UpperCamelCase     :Any			  =		load_numpy(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/stable_unclip_2_1_h_img2img_anime_turtle_fp16.npy''')

        __UpperCamelCase     :Any			  =		StableUnCLIPImgaImgPipeline.from_pretrained(
            '''fusing/stable-unclip-2-1-h-img2img''' , torch_dtype=torch.floataa)
        pipe.to(__lowercase)
        pipe.set_progress_bar_config(disable=__lowercase)
        # stable unclip will oom when integration tests are run on a V100,
        # so turn on memory savings
        pipe.enable_attention_slicing()
        pipe.enable_sequential_cpu_offload()

        __UpperCamelCase     :int			  =		torch.Generator(device='''cpu''').manual_seed(0)
        __UpperCamelCase     :Optional[int]			  =		pipe(__lowercase , '''anime turle''' , generator=__lowercase , output_type='''np''')

        __UpperCamelCase     :List[Any]			  =		output.images[0]

        assert image.shape == (768, 768, 3)

        assert_mean_pixel_difference(__lowercase , __lowercase)







    def 			UpperCamelCase__   (	self)       ->	List[str]:
        __UpperCamelCase     :Dict			  =		load_image(
            '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/stable_unclip/turtle.png''')

        torch.cuda.empty_cache()
        torch.cuda.reset_max_memory_allocated()
        torch.cuda.reset_peak_memory_stats()

        __UpperCamelCase     :List[Any]			  =		StableUnCLIPImgaImgPipeline.from_pretrained(
            '''fusing/stable-unclip-2-1-h-img2img''' , torch_dtype=torch.floataa)
        __UpperCamelCase     :Union[str, Any]			  =		pipe.to(__lowercase)
        pipe.set_progress_bar_config(disable=__lowercase)
        pipe.enable_attention_slicing()
        pipe.enable_sequential_cpu_offload()

        __UpperCamelCase     :Optional[Any]			  =		pipe(
            __lowercase , '''anime turtle''' , num_inference_steps=2 , output_type='''np''' , )

        __UpperCamelCase     :int			  =		torch.cuda.max_memory_allocated()
        # make sure that less than 7 GB is allocated
        assert mem_bytes < 7 * 10**9





