






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






  """simple docstring"""





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




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

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

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

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

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

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

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

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

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

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

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

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

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








from typing import Dict, List, Optional

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


_lowerCamelCase      :					Any					    =      logging.get_logger(__name__)


_lowerCamelCase      :					Tuple					    =      {
    '''nielsr/canine-s''': 2048,
}

# Unicode defines 1,114,112 total “codepoints”
_lowerCamelCase      :					List[str]					    =      111_4112

# Below: Constants defining canonical codepoints for special, pseudo-characters.
# Copied from https://github.com/google-research/language/blob/master/language/canine/special_codepoints.py
_lowerCamelCase      :					List[str]					    =      0

_lowerCamelCase      :					Any					    =      0xe_0_0_0
_lowerCamelCase      :					Union[str, Any]					    =      0xe_0_0_1
_lowerCamelCase      :					Any					    =      0xe_0_0_2
_lowerCamelCase      :					List[str]					    =      0xe_0_0_3
_lowerCamelCase      :					Any					    =      0xe_0_0_4

# Maps special codepoints to human-readable names.
_lowerCamelCase      :					Dict[int, str]					    =      {
    # Special symbols are represented using codepoints values that are valid,
    # but designated as "Private Use", meaning that they will never be assigned
    # characters by the Unicode Consortium, and are thus safe for use here.
    #
    # NOTE: Do *NOT* add any sort of [UNK_CHAR] here. They are explicitly
    # excluded and should fail with a hard error.
    CLS: "[CLS]",
    SEP: "[SEP]",
    BOS: "[BOS]",
    MASK: "[MASK]",
    PAD: "[PAD]",
    RESERVED: "[RESERVED]",
}

# Maps special codepoint human-readable names to their codepoint values.
_lowerCamelCase      :					Dict[str, int]					    =      {name: codepoint for codepoint, name in SPECIAL_CODEPOINTS.items()}





class        lowerCAmelCase__  (					__magic_name__			):



       '''simple docstring'''







       lowercase_						=    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES

       def __init__(							self ,   lowercase__=chr(lowercase__     ) ,   lowercase__=chr(lowercase__     ) ,   lowercase__=chr(lowercase__     ) ,   lowercase__=chr(lowercase__     ) ,   lowercase__=chr(lowercase__     ) ,   lowercase__=chr(lowercase__     ) ,   lowercase__=False ,   lowercase__=2_0_4_8 ,   **lowercase__ ,   ):







             '''simple docstring'''



             __A    =AddedToken(lowercase__ ,   lstrip=lowercase__ ,   rstrip=lowercase__     ) if isinstance(lowercase__ ,   lowercase__     ) else bos_token
             __A    =AddedToken(lowercase__ ,   lstrip=lowercase__ ,   rstrip=lowercase__     ) if isinstance(lowercase__ ,   lowercase__     ) else eos_token
             __A    =AddedToken(lowercase__ ,   lstrip=lowercase__ ,   rstrip=lowercase__     ) if isinstance(lowercase__ ,   lowercase__     ) else sep_token
             __A    =AddedToken(lowercase__ ,   lstrip=lowercase__ ,   rstrip=lowercase__     ) if isinstance(lowercase__ ,   lowercase__     ) else cls_token
             __A    =AddedToken(lowercase__ ,   lstrip=lowercase__ ,   rstrip=lowercase__     ) if isinstance(lowercase__ ,   lowercase__     ) else pad_token

             # Mask token behave like a normal word, i.e. include the space before it
             __A    =AddedToken(lowercase__ ,   lstrip=lowercase__ ,   rstrip=lowercase__     ) if isinstance(lowercase__ ,   lowercase__     ) else mask_token

             super().__init__(
                 bos_token=lowercase__ ,   eos_token=lowercase__ ,   sep_token=lowercase__ ,   cls_token=lowercase__ ,   pad_token=lowercase__ ,   mask_token=lowercase__ ,   add_prefix_space=lowercase__ ,   model_max_length=lowercase__ ,   **lowercase__ ,   )

             # Creates a mapping for looking up the IDs of special symbols.
             __A    ={}
             for codepoint, name in SPECIAL_CODEPOINTS.items():
                   __A    =codepoint

             # Creates a mapping for looking up the string forms of special symbol IDs.
             __A    ={
                 codepoint: name for name, codepoint in self._special_codepoints.items()
             }

             __A    =UNICODE_VOCAB_SIZE
             __A    =len(self._special_codepoints     )

       @property
       def     __UpperCamelCase     (							self     ):







             '''simple docstring'''



             return self._unicode_vocab_size

       def     __UpperCamelCase     (							self ,   lowercase__     ):







             '''simple docstring'''



             return list(lowercase__     )

       def     __UpperCamelCase     (							self ,   lowercase__     ):







             '''simple docstring'''



             try:
                   return ord(lowercase__     )
             except TypeError:
                   raise ValueError(f'''invalid token: \'{token}\''''     )

       def     __UpperCamelCase     (							self ,   lowercase__     ):







             '''simple docstring'''



             try:
                   if index in SPECIAL_CODEPOINTS:
                         return SPECIAL_CODEPOINTS[index]
                   return chr(lowercase__     )
             except TypeError:
                   raise ValueError(f'''invalid id: {index}'''     )

       def     __UpperCamelCase     (							self ,   lowercase__     ):







             '''simple docstring'''



             return "".join(lowercase__     )

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ = None     ):







             '''simple docstring'''



             __A    =[self.sep_token_id]
             __A    =[self.cls_token_id]

             __A    =cls + token_ids_a + sep
             if token_ids_a is not None:
                   result += token_ids_a + sep
             return result

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ = None ,   lowercase__ = False     ):







             '''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__     )

             __A    =[1] + ([0] * len(lowercase__     )) + [1]
             if token_ids_a is not None:
                   result += ([0] * len(lowercase__     )) + [1]
             return result

       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ = None     ):







             '''simple docstring'''



             __A    =[self.sep_token_id]
             __A    =[self.cls_token_id]

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


       def     __UpperCamelCase     (							self ,   lowercase__ ,   lowercase__ = None     ):







             '''simple docstring'''



             return ()


