
import gc
import unittest

import numpy as np
import torch

from diffusers import (
    AudioDiffusionPipeline,
    AutoencoderKL,
    DDIMScheduler,
    DDPMScheduler,
    DiffusionPipeline,
    Mel,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()

class __snake_case		(				unittest.TestCase					):
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()
    @property
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           torch.manual_seed(0				)
           lowercase							:					int	     =				UNetaDModel(
               sample_size=(32, 64)		,in_channels=1		,out_channels=1		,layers_per_block=2		,block_out_channels=(128, 128)		,down_block_types=("""AttnDownBlock2D""", """DownBlock2D""")		,up_block_types=("""UpBlock2D""", """AttnUpBlock2D""")		,)
           return model
    @property
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           torch.manual_seed(0				)
           lowercase							:					List[str]	     =				UNetaDConditionModel(
               sample_size=(64, 32)		,in_channels=1		,out_channels=1		,layers_per_block=2		,block_out_channels=(128, 128)		,down_block_types=("""CrossAttnDownBlock2D""", """DownBlock2D""")		,up_block_types=("""UpBlock2D""", """CrossAttnUpBlock2D""")		,cross_attention_dim=10		,)
           return model
    @property
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           torch.manual_seed(0				)
           lowercase							:					Any	     =				AutoencoderKL(
               sample_size=(128, 64)		,in_channels=1		,out_channels=1		,latent_channels=1		,layers_per_block=2		,block_out_channels=(128, 128)		,down_block_types=("""DownEncoderBlock2D""", """DownEncoderBlock2D""")		,up_block_types=("""UpDecoderBlock2D""", """UpDecoderBlock2D""")		,)
           lowercase							:					str	     =				UNetaDModel(
               sample_size=(64, 32)		,in_channels=1		,out_channels=1		,layers_per_block=2		,block_out_channels=(128, 128)		,down_block_types=("""AttnDownBlock2D""", """DownBlock2D""")		,up_block_types=("""UpBlock2D""", """AttnUpBlock2D""")		,)
           return vqvae, unet
    @slow
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					Dict	     =				"""cpu"""  # ensure determinism for the device-dependent torch.Generator
           lowercase							:					str	     =				Mel(
               x_res=self.dummy_unet.config.sample_size[1]		,y_res=self.dummy_unet.config.sample_size[0]		,)

           lowercase							:					List[str]	     =				DDPMScheduler()
           lowercase							:					Dict	     =				AudioDiffusionPipeline(vqvae=snake_case		,unet=self.dummy_unet		,mel=snake_case		,scheduler=snake_case				)
           lowercase							:					Optional[Any]	     =				pipe.to(snake_case				)
           pipe.set_progress_bar_config(disable=snake_case				)

           lowercase							:					Optional[Any]	     =				torch.Generator(device=snake_case				).manual_seed(42				)
           lowercase							:					str	     =				pipe(generator=snake_case		,steps=4				)
           lowercase							:					List[Any]	     =				output.audios[0]
           lowercase							:					Dict	     =				output.images[0]

           lowercase							:					Optional[Any]	     =				torch.Generator(device=snake_case				).manual_seed(42				)
           lowercase							:					Optional[int]	     =				pipe(generator=snake_case		,steps=4		,return_dict=snake_case				)
           lowercase							:					Dict	     =				output[0][0]

           assert audio.shape == (1, (self.dummy_unet.config.sample_size[1] - 1) * mel.hop_length)
           assert (
               image.height == self.dummy_unet.config.sample_size[0]
               and image.width == self.dummy_unet.config.sample_size[1]
           )
           lowercase							:					List[str]	     =				np.frombuffer(image.tobytes()		,dtype="""uint8"""				)[:10]
           lowercase							:					List[Any]	     =				np.frombuffer(image_from_tuple.tobytes()		,dtype="""uint8"""				)[:10]
           lowercase							:					int	     =				np.array([69, 255, 255, 255, 0, 0, 77, 181, 12, 127]				)

           assert np.abs(image_slice.flatten() - expected_slice				).max() == 0
           assert np.abs(image_from_tuple_slice.flatten() - expected_slice				).max() == 0

           lowercase							:					List[Any]	     =				Mel(
               x_res=self.dummy_vqvae_and_unet[0].config.sample_size[1]		,y_res=self.dummy_vqvae_and_unet[0].config.sample_size[0]		,)

           lowercase							:					int	     =				DDIMScheduler()
           lowercase							:					Union[str, Any]	     =				self.dummy_vqvae_and_unet
           lowercase							:					int	     =				AudioDiffusionPipeline(
               vqvae=self.dummy_vqvae_and_unet[0]		,unet=dummy_vqvae_and_unet[1]		,mel=snake_case		,scheduler=snake_case				)
           lowercase							:					str	     =				pipe.to(snake_case				)
           pipe.set_progress_bar_config(disable=snake_case				)

           np.random.seed(0				)
           lowercase							:					Union[str, Any]	     =				np.random.uniform(-1		,1		,((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)				)
           lowercase							:					Union[str, Any]	     =				torch.Generator(device=snake_case				).manual_seed(42				)
           lowercase							:					Union[str, Any]	     =				pipe(raw_audio=snake_case		,generator=snake_case		,start_step=5		,steps=10				)
           lowercase							:					Optional[int]	     =				output.images[0]

           assert (
               image.height == self.dummy_vqvae_and_unet[0].config.sample_size[0]
               and image.width == self.dummy_vqvae_and_unet[0].config.sample_size[1]
           )
           lowercase							:					Dict	     =				np.frombuffer(image.tobytes()		,dtype="""uint8"""				)[:10]
           lowercase							:					List[Any]	     =				np.array([120, 117, 110, 109, 138, 167, 138, 148, 132, 121]				)

           assert np.abs(image_slice.flatten() - expected_slice				).max() == 0

           lowercase							:					Optional[Any]	     =				self.dummy_unet_condition
           lowercase							:					Optional[int]	     =				AudioDiffusionPipeline(
               vqvae=self.dummy_vqvae_and_unet[0]		,unet=snake_case		,mel=snake_case		,scheduler=snake_case				)
           lowercase							:					List[Any]	     =				pipe.to(snake_case				)
           pipe.set_progress_bar_config(disable=snake_case				)

           np.random.seed(0				)
           lowercase							:					int	     =				torch.rand((1, 1, 10)				)
           lowercase							:					List[Any]	     =				pipe(generator=snake_case		,encoding=snake_case				)
           lowercase							:					Tuple	     =				output.images[0]
           lowercase							:					Any	     =				np.frombuffer(image.tobytes()		,dtype="""uint8"""				)[:10]
           lowercase							:					Dict	     =				np.array([107, 103, 120, 127, 142, 122, 113, 122, 97, 111]				)

           assert np.abs(image_slice.flatten() - expected_slice				).max() == 0






@slow
@require_torch_gpu
class __snake_case		(				unittest.TestCase					):
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           super().tearDown()
           gc.collect()
           torch.cuda.empty_cache()
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					Tuple	     =				torch_device

           lowercase							:					Dict	     =				DiffusionPipeline.from_pretrained("""teticio/audio-diffusion-ddim-256"""				)
           lowercase							:					Optional[int]	     =				pipe.to(snake_case				)
           pipe.set_progress_bar_config(disable=snake_case				)

           lowercase							:					int	     =				torch.Generator(device=snake_case				).manual_seed(42				)
           lowercase							:					Optional[Any]	     =				pipe(generator=snake_case				)
           lowercase							:					Union[str, Any]	     =				output.audios[0]
           lowercase							:					Optional[Any]	     =				output.images[0]

           assert audio.shape == (1, (pipe.unet.config.sample_size[1] - 1) * pipe.mel.hop_length)
           assert image.height == pipe.unet.config.sample_size[0] and image.width == pipe.unet.config.sample_size[1]
           lowercase							:					Optional[Any]	     =				np.frombuffer(image.tobytes()		,dtype="""uint8"""				)[:10]
           lowercase							:					Dict	     =				np.array([151, 167, 154, 144, 122, 134, 121, 105, 70, 26]				)

           assert np.abs(image_slice.flatten() - expected_slice				).max() == 0


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


lowercase  : Union[str, Any]						   =							logging.get_logger(__name__)

lowercase  : Union[str, Any]						   =							{"""vocab_file""": """spiece.model"""}

lowercase  : Any						   =							{
    """vocab_file""": {
        """xlnet-base-cased""": """https://huggingface.co/xlnet-base-cased/resolve/main/spiece.model""",
        """xlnet-large-cased""": """https://huggingface.co/xlnet-large-cased/resolve/main/spiece.model""",
    }
}

lowercase  : List[Any]						   =							{
    """xlnet-base-cased""": None,
    """xlnet-large-cased""": None,
}

# Segments (not really needed)
lowercase  : Optional[int]						   =							0
lowercase  : Any						   =							1
lowercase  : Any						   =							2
lowercase  : str						   =							3
lowercase  : Optional[Any]						   =							4

class __snake_case		(				lowerCAmelCase					):
    _a		:						Optional[Any]=    VOCAB_FILES_NAMES
    _a		:						Optional[Any]=    PRETRAINED_VOCAB_FILES_MAP
    _a		:						List[str]=    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
    _a		:						Union[str, Any]=    "left"
    def __init__(  self		,snake_case		,snake_case=False		,snake_case=True		,snake_case=False		,snake_case="<s>"		,snake_case="</s>"		,snake_case="<unk>"		,snake_case="<sep>"		,snake_case="<pad>"		,snake_case="<cls>"		,snake_case="<mask>"		,snake_case=["<eop>", "<eod>"]		,snake_case = None		,**snake_case		,):



           '''simple docstring'''







           lowercase							:					str	     =				AddedToken(snake_case		,lstrip=snake_case		,rstrip=snake_case				) if isinstance(snake_case		,snake_case				) else mask_token

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

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

           lowercase							:					str	     =				3

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

           lowercase							:					Any	     =				spm.SentencePieceProcessor(**self.sp_model_kwargs				)
           self.sp_model.Load(snake_case				)
    @property
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           return len(self.sp_model				)
    def 				_SCREAMING_SNAKE_CASE	(  self				):



           '''simple docstring'''







           lowercase							:					Any	     =				{self.convert_ids_to_tokens(snake_case				): i for i in range(self.vocab_size				)}
           vocab.update(self.added_tokens_encoder				)
           return vocab
    def __getstate__(  self				):



           '''simple docstring'''







           lowercase							:					Any	     =				self.__dict__.copy()
           lowercase							:					List[str]	     =				None
           return state
    def __setstate__(  self		,snake_case				):



           '''simple docstring'''







           lowercase							:					Dict	     =				d

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

           lowercase							:					List[str]	     =				spm.SentencePieceProcessor(**self.sp_model_kwargs				)
           self.sp_model.Load(self.vocab_file				)
    def 				_SCREAMING_SNAKE_CASE	(  self		,snake_case				):



           '''simple docstring'''







           if self.remove_space:
                  lowercase							:					Optional[int]	     =				""" """.join(inputs.strip().split()				)
           else:
                  lowercase							:					Optional[int]	     =				inputs
           lowercase							:					Dict	     =				outputs.replace("""``"""		,"""\""""				).replace("""''"""		,"""\""""				)

           if not self.keep_accents:
                  lowercase							:					int	     =				unicodedata.normalize("""NFKD"""		,snake_case				)
                  lowercase							:					Any	     =				"""""".join([c for c in outputs if not unicodedata.combining(snake_case				)]				)
           if self.do_lower_case:
                  lowercase							:					Any	     =				outputs.lower()

           return outputs
    def 				_SCREAMING_SNAKE_CASE	(  self		,snake_case				):



           '''simple docstring'''







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

           return new_pieces
    def 				_SCREAMING_SNAKE_CASE	(  self		,snake_case				):



           '''simple docstring'''







           return self.sp_model.PieceToId(snake_case				)
    def 				_SCREAMING_SNAKE_CASE	(  self		,snake_case				):



           '''simple docstring'''







           return self.sp_model.IdToPiece(snake_case				)
    def 				_SCREAMING_SNAKE_CASE	(  self		,snake_case				):



           '''simple docstring'''







           lowercase							:					Optional[int]	     =				"""""".join(snake_case				).replace(snake_case		,""" """				).strip()
           return out_string
    def 				_SCREAMING_SNAKE_CASE	(  self		,snake_case		,snake_case = False		,snake_case = None		,snake_case = True		,**snake_case		,):



           '''simple docstring'''







           lowercase							:					List[str]	     =				kwargs.pop("""use_source_tokenizer"""		,snake_case				)

           lowercase							:					List[str]	     =				self.convert_ids_to_tokens(snake_case		,skip_special_tokens=snake_case				)

           # To avoid mixing byte-level and unicode for byte-level BPT
           # we need to build string separately for added tokens and byte-level tokens
           # cf. https://github.com/huggingface/transformers/issues/1133
           lowercase							:					Optional[int]	     =				[]
           lowercase							:					Optional[int]	     =				[]
           for token in filtered_tokens:
                  if skip_special_tokens and token in self.all_special_ids:
                         continue
                  if token in self.added_tokens_encoder:
                         if current_sub_text:
                                sub_texts.append(self.convert_tokens_to_string(snake_case				)				)
                                lowercase							:					Union[str, Any]	     =				[]
                         sub_texts.append(snake_case				)
                  else:
                         current_sub_text.append(snake_case				)
           if current_sub_text:
                  sub_texts.append(self.convert_tokens_to_string(snake_case				)				)

           # Mimic the behavior of the Rust tokenizer:
           # By default, there are no spaces between special tokens
           lowercase							:					Any	     =				"""""".join(snake_case				)

           lowercase							:					str	     =				(
               clean_up_tokenization_spaces
               if clean_up_tokenization_spaces is not None
               else self.clean_up_tokenization_spaces
           )
           if clean_up_tokenization_spaces:
                  lowercase							:					int	     =				self.clean_up_tokenization(snake_case				)
                  return clean_text
           else:
                  return text
    def 				_SCREAMING_SNAKE_CASE	(  self		,snake_case		,snake_case = None				):



           '''simple docstring'''







           lowercase							:					int	     =				[self.sep_token_id]
           lowercase							:					str	     =				[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 				_SCREAMING_SNAKE_CASE	(  self		,snake_case		,snake_case = None		,snake_case = False				):



           '''simple docstring'''







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

           if token_ids_a is not None:
                  return ([0] * len(snake_case				)) + [1] + ([0] * len(snake_case				)) + [1, 1]
           return ([0] * len(snake_case				)) + [1, 1]
    def 				_SCREAMING_SNAKE_CASE	(  self		,snake_case		,snake_case = None				):



           '''simple docstring'''







           lowercase							:					List[str]	     =				[self.sep_token_id]
           lowercase							:					Any	     =				[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 				_SCREAMING_SNAKE_CASE	(  self		,snake_case		,snake_case = None				):



           '''simple docstring'''







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

           if os.path.abspath(self.vocab_file				) != os.path.abspath(snake_case				) and os.path.isfile(self.vocab_file				):
                  copyfile(self.vocab_file		,snake_case				)
           elif not os.path.isfile(self.vocab_file				):
                  with open(snake_case		,"""wb"""				) as fi:
                         lowercase							:					str	     =				self.sp_model.serialized_model_proto()
                         fi.write(snake_case				)

           return (out_vocab_file,)
