


"""simple docstring"""






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

     '''simple docstring'''


     def 							_UpperCAmelCase      (      self							)       ->							int:
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()


     @property
     def 							_UpperCAmelCase      (      self							)       ->							str:
          torch.manual_seed(0							)
          _a    	=				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 							_UpperCAmelCase      (      self							)       ->							Any:
          torch.manual_seed(0							)
          _a    	=				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 							_UpperCAmelCase      (      self							)       ->							str:
          torch.manual_seed(0							)
          _a    	=				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''')    ,				)
          _a    	=				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 							_UpperCAmelCase      (      self							)       ->							int:
          _a    	=				'''cpu'''  # ensure determinism for the device-dependent torch.Generator
          _a    	=				Mel(
              x_res=self.dummy_unet.config.sample_size[1]    ,				y_res=self.dummy_unet.config.sample_size[0]    ,				)

          _a    	=				DDPMScheduler()
          _a    	=				AudioDiffusionPipeline(vqvae=__UpperCAmelCase    ,				unet=self.dummy_unet    ,				mel=__UpperCAmelCase    ,				scheduler=__UpperCAmelCase							)
          _a    	=				pipe.to(__UpperCAmelCase							)
          pipe.set_progress_bar_config(disable=__UpperCAmelCase							)

          _a    	=				torch.Generator(device=__UpperCAmelCase							).manual_seed(42							)
          _a    	=				pipe(generator=__UpperCAmelCase    ,				steps=4							)
          _a    	=				output.audios[0]
          _a    	=				output.images[0]

          _a    	=				torch.Generator(device=__UpperCAmelCase							).manual_seed(42							)
          _a    	=				pipe(generator=__UpperCAmelCase    ,				steps=4    ,				return_dict=__UpperCAmelCase							)
          _a    	=				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]
          )
          _a    	=				np.frombuffer(image.tobytes()    ,				dtype='''uint8'''							)[:10]
          _a    	=				np.frombuffer(image_from_tuple.tobytes()    ,				dtype='''uint8'''							)[:10]
          _a    	=				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

          _a    	=				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]    ,				)

          _a    	=				DDIMScheduler()
          _a    	=				self.dummy_vqvae_and_unet
          _a    	=				AudioDiffusionPipeline(
              vqvae=self.dummy_vqvae_and_unet[0]    ,				unet=dummy_vqvae_and_unet[1]    ,				mel=__UpperCAmelCase    ,				scheduler=__UpperCAmelCase							)
          _a    	=				pipe.to(__UpperCAmelCase							)
          pipe.set_progress_bar_config(disable=__UpperCAmelCase							)

          np.random.seed(0							)
          _a    	=				np.random.uniform(-1    ,				1    ,				((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)							)
          _a    	=				torch.Generator(device=__UpperCAmelCase							).manual_seed(42							)
          _a    	=				pipe(raw_audio=__UpperCAmelCase    ,				generator=__UpperCAmelCase    ,				start_step=5    ,				steps=10							)
          _a    	=				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]
          )
          _a    	=				np.frombuffer(image.tobytes()    ,				dtype='''uint8'''							)[:10]
          _a    	=				np.array([120, 117, 110, 109, 138, 167, 138, 148, 132, 121]							)

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

          _a    	=				self.dummy_unet_condition
          _a    	=				AudioDiffusionPipeline(
              vqvae=self.dummy_vqvae_and_unet[0]    ,				unet=__UpperCAmelCase    ,				mel=__UpperCAmelCase    ,				scheduler=__UpperCAmelCase							)
          _a    	=				pipe.to(__UpperCAmelCase							)
          pipe.set_progress_bar_config(disable=__UpperCAmelCase							)

          np.random.seed(0							)
          _a    	=				torch.rand((1, 1, 10)							)
          _a    	=				pipe(generator=__UpperCAmelCase    ,				encoding=__UpperCAmelCase							)
          _a    	=				output.images[0]
          _a    	=				np.frombuffer(image.tobytes()    ,				dtype='''uint8'''							)[:10]
          _a    	=				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     __lowerCamelCase  (					unittest.TestCase				):

     '''simple docstring'''


     def 							_UpperCAmelCase      (      self							)       ->							str:
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()





     def 							_UpperCAmelCase      (      self							)       ->							Any:
          _a    	=				torch_device

          _a    	=				DiffusionPipeline.from_pretrained('''teticio/audio-diffusion-ddim-256'''							)
          _a    	=				pipe.to(__UpperCAmelCase							)
          pipe.set_progress_bar_config(disable=__UpperCAmelCase							)

          _a    	=				torch.Generator(device=__UpperCAmelCase							).manual_seed(42							)
          _a    	=				pipe(generator=__UpperCAmelCase							)
          _a    	=				output.audios[0]
          _a    	=				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]
          _a    	=				np.frombuffer(image.tobytes()    ,				dtype='''uint8'''							)[:10]
          _a    	=				np.array([151, 167, 154, 144, 122, 134, 121, 105, 70, 26]							)

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



"""simple docstring"""






import warnings
from typing import Dict, List, Optional, Tuple

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


__snake_case		=       logging.get_logger(__name__)


class     __lowerCamelCase  (					a__				):

     '''simple docstring'''
     A_       :	Any    =		['input_ids', 'attention_mask']


     def __init__(      self    ,				__UpperCAmelCase="</s>"    ,				__UpperCAmelCase="<unk>"    ,				__UpperCAmelCase="<pad>"    ,				__UpperCAmelCase=125    ,				__UpperCAmelCase=None    ,				**__UpperCAmelCase    ,				)       ->							None:
          # Add extra_ids to the special token list
          if extra_ids > 0 and additional_special_tokens is None:
               _a    	=				[F'<extra_id_{i}>' for i in range(__UpperCAmelCase							)]
          elif extra_ids > 0 and additional_special_tokens is not None:
               # Check that we have the right number of extra_id special tokens
               _a    	=				len(set(filter(lambda __UpperCAmelCase							: bool('''extra_id''' in str(__UpperCAmelCase							)							)    ,				__UpperCAmelCase							)							)							)
               if extra_tokens != extra_ids:
                    raise ValueError(
                        F'Both extra_ids ({extra_ids}) and additional_special_tokens ({additional_special_tokens}) are'
                        ''' provided to ByT5Tokenizer. In this case the additional_special_tokens must include the'''
                        ''' extra_ids tokens'''							)

          _a    	=				AddedToken(__UpperCAmelCase    ,				lstrip=__UpperCAmelCase    ,				rstrip=__UpperCAmelCase							) if isinstance(__UpperCAmelCase    ,				__UpperCAmelCase							) else pad_token
          _a    	=				AddedToken(__UpperCAmelCase    ,				lstrip=__UpperCAmelCase    ,				rstrip=__UpperCAmelCase							) if isinstance(__UpperCAmelCase    ,				__UpperCAmelCase							) else eos_token
          _a    	=				AddedToken(__UpperCAmelCase    ,				lstrip=__UpperCAmelCase    ,				rstrip=__UpperCAmelCase							) if isinstance(__UpperCAmelCase    ,				__UpperCAmelCase							) else unk_token

          super().__init__(
              eos_token=__UpperCAmelCase    ,				unk_token=__UpperCAmelCase    ,				pad_token=__UpperCAmelCase    ,				extra_ids=__UpperCAmelCase    ,				additional_special_tokens=__UpperCAmelCase    ,				**__UpperCAmelCase    ,				)

          _a    	=				extra_ids

          _a    	=				2**8  # utf is 8 bits

          # define special tokens dict
          _a    	=				{
              self.pad_token: 0,
              self.eos_token: 1,
              self.unk_token: 2,
          }
          _a    	=				len(self.special_tokens_encoder							)
          _a    	=				len(__UpperCAmelCase							)
          for i, token in enumerate(__UpperCAmelCase							):
               _a    	=				self.vocab_size + i - n
          _a    	=				{v: k for k, v in self.special_tokens_encoder.items()}


     @property
     def 							_UpperCAmelCase      (      self							)       ->							Dict:
          return self._utf_vocab_size + self._num_special_tokens + self._extra_ids


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = False							)       ->							List[int]:
          if already_has_special_tokens:
               return super().get_special_tokens_mask(
                   token_ids_a=__UpperCAmelCase    ,				token_ids_a=__UpperCAmelCase    ,				already_has_special_tokens=__UpperCAmelCase							)

          # normal case: some special tokens
          if token_ids_a is None:
               return ([0] * len(__UpperCAmelCase							)) + [1]
          return ([0] * len(__UpperCAmelCase							)) + [1] + ([0] * len(__UpperCAmelCase							)) + [1]


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							List[int]:
          if len(__UpperCAmelCase							) > 0 and token_ids[-1] == self.eos_token_id:
               warnings.warn(
                   F'This sequence already has {self.eos_token}. In future versions this behavior may lead to duplicated'
                   ''' eos tokens being added.'''							)
               return token_ids
          else:
               return token_ids + [self.eos_token_id]


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase = None							)       ->							List[int]:
          _a    	=				[self.eos_token_id]

          if token_ids_a is None:
               return len(token_ids_a + eos							) * [0]
          return len(token_ids_a + eos + token_ids_a + eos							) * [0]


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase = None							)       ->							List[int]:
          _a    	=				self._add_eos_if_not_present(__UpperCAmelCase							)
          if token_ids_a is None:
               return token_ids_a
          else:
               _a    	=				self._add_eos_if_not_present(__UpperCAmelCase							)
               return token_ids_a + token_ids_a


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							List[str]:
          _a    	=				[chr(__UpperCAmelCase							) for i in text.encode('''utf-8'''							)]
          return tokens


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							Union[str, Any]:
          if token in self.special_tokens_encoder:
               _a    	=				self.special_tokens_encoder[token]
          elif token in self.added_tokens_encoder:
               _a    	=				self.added_tokens_encoder[token]
          elif len(__UpperCAmelCase							) != 1:
               _a    	=				self.unk_token_id
          else:
               _a    	=				ord(__UpperCAmelCase							) + self._num_special_tokens
          return token_id


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							List[Any]:
          if index in self.special_tokens_decoder:
               _a    	=				self.special_tokens_decoder[index]
          else:
               _a    	=				chr(index - self._num_special_tokens							)
          return token


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							Any:
          _a    	=				B''''''
          for token in tokens:
               if token in self.special_tokens_decoder:
                    _a    	=				self.special_tokens_decoder[token].encode('''utf-8'''							)
               elif token in self.added_tokens_decoder:
                    _a    	=				self.special_tokens_decoder[token].encode('''utf-8'''							)
               elif token in self.special_tokens_encoder:
                    _a    	=				token.encode('''utf-8'''							)
               elif token in self.added_tokens_encoder:
                    _a    	=				token.encode('''utf-8'''							)
               else:
                    _a    	=				bytes([ord(__UpperCAmelCase							)]							)
               bstring += tok_string
          _a    	=				bstring.decode('''utf-8'''    ,				errors='''ignore'''							)
          return string





     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase = None							)       ->							Tuple[str]:
          return ()