




import copy
import os
from typing import TYPE_CHECKING, List, Union


if TYPE_CHECKING:
     pass

from ...configuration_utils import PretrainedConfig
from ...utils import logging


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

_UpperCAmelCase		:  Tuple	=   {
    "kakaobrain/align-base": "https://huggingface.co/kakaobrain/align-base/resolve/main/config.json",
}




class       __lowerCAmelCase				(							lowerCAmelCase):
       _a           =  '''align_text_model'''






       def __init__(      self:    Any				,    _lowerCAmelCase:    Optional[int]=3_05_22				,    _lowerCAmelCase:    List[Any]=7_68				,    _lowerCAmelCase:    List[str]=12				,    _lowerCAmelCase:    Optional[Any]=12				,    _lowerCAmelCase:    Dict=30_72				,    _lowerCAmelCase:    Dict="gelu"				,    _lowerCAmelCase:    int=0.1				,    _lowerCAmelCase:    Optional[int]=0.1				,    _lowerCAmelCase:    Tuple=5_12				,    _lowerCAmelCase:    List[Any]=2				,    _lowerCAmelCase:    Union[str, Any]=0.02				,    _lowerCAmelCase:    Dict=1e-1_2				,    _lowerCAmelCase:    Any=0				,    _lowerCAmelCase:    List[str]="absolute"				,    _lowerCAmelCase:    int=True				,    **_lowerCAmelCase:    List[Any]				,    ):
         super().__init__(**_lowerCAmelCase )

         lowercase			:Optional[Any]							=  vocab_size
         lowercase			:Optional[Any]							=  hidden_size
         lowercase			:List[str]							=  num_hidden_layers
         lowercase			:int							=  num_attention_heads
         lowercase			:Optional[int]							=  hidden_act
         lowercase			:Dict							=  intermediate_size
         lowercase			:List[Any]							=  hidden_dropout_prob
         lowercase			:Optional[Any]							=  attention_probs_dropout_prob
         lowercase			:str							=  max_position_embeddings
         lowercase			:List[str]							=  type_vocab_size
         lowercase			:List[Any]							=  initializer_range
         lowercase			:int							=  layer_norm_eps
         lowercase			:Dict							=  position_embedding_type
         lowercase			:Optional[Any]							=  use_cache
         lowercase			:Union[str, Any]							=  pad_token_id






       @classmethod
       def     SCREAMING_SNAKE_CASE			(      cls:    Tuple				,    _lowerCAmelCase:    Union[str, os.PathLike]				,    **_lowerCAmelCase:    Dict ):
         cls._set_token_in_kwargs(_lowerCAmelCase )

         lowercase		,			lowercase			:Optional[int]							=  cls.get_config_dict(_lowerCAmelCase				,    **_lowerCAmelCase )

         # get the text config dict if we are loading from AlignConfig
         if config_dict.get("model_type" ) == "align":
           lowercase			:int							=  config_dict["text_config"]

         if "model_type" in config_dict and hasattr(cls				,    "model_type" ) and config_dict["model_type"] != cls.model_type:
           logger.warning(
               F"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
               F"{cls.model_type}. This is not supported for all configurations of models and can yield errors." )

         return cls.from_dict(_lowerCAmelCase				,    **_lowerCAmelCase )







class       __lowerCAmelCase				(							lowerCAmelCase):
       _a           =  '''align_vision_model'''






       def __init__(      self:    List[str]				,    _lowerCAmelCase:    int = 3				,    _lowerCAmelCase:    int = 6_00				,    _lowerCAmelCase:    float = 2.0				,    _lowerCAmelCase:    float = 3.1				,    _lowerCAmelCase:    int = 8				,    _lowerCAmelCase:    List[int] = [3, 3, 5, 3, 5, 5, 3]				,    _lowerCAmelCase:    List[int] = [32, 16, 24, 40, 80, 1_12, 1_92]				,    _lowerCAmelCase:    List[int] = [16, 24, 40, 80, 1_12, 1_92, 3_20]				,    _lowerCAmelCase:    List[int] = []				,    _lowerCAmelCase:    List[int] = [1, 2, 2, 2, 1, 2, 1]				,    _lowerCAmelCase:    List[int] = [1, 2, 2, 3, 3, 4, 1]				,    _lowerCAmelCase:    List[int] = [1, 6, 6, 6, 6, 6, 6]				,    _lowerCAmelCase:    float = 0.25				,    _lowerCAmelCase:    str = "swish"				,    _lowerCAmelCase:    int = 25_60				,    _lowerCAmelCase:    str = "mean"				,    _lowerCAmelCase:    float = 0.02				,    _lowerCAmelCase:    float = 0.0_01				,    _lowerCAmelCase:    float = 0.99				,    _lowerCAmelCase:    float = 0.2				,    **_lowerCAmelCase:    List[Any]				,    ):
         super().__init__(**_lowerCAmelCase )

         lowercase			:Optional[Any]							=  num_channels
         lowercase			:Optional[int]							=  image_size
         lowercase			:Optional[int]							=  width_coefficient
         lowercase			:Dict							=  depth_coefficient
         lowercase			:List[str]							=  depth_divisor
         lowercase			:List[Any]							=  kernel_sizes
         lowercase			:Any							=  in_channels
         lowercase			:int							=  out_channels
         lowercase			:Optional[int]							=  depthwise_padding
         lowercase			:str							=  strides
         lowercase			:Optional[int]							=  num_block_repeats
         lowercase			:Optional[Any]							=  expand_ratios
         lowercase			:Tuple							=  squeeze_expansion_ratio
         lowercase			:Union[str, Any]							=  hidden_act
         lowercase			:Union[str, Any]							=  hidden_dim
         lowercase			:Optional[int]							=  pooling_type
         lowercase			:List[str]							=  initializer_range
         lowercase			:int							=  batch_norm_eps
         lowercase			:Any							=  batch_norm_momentum
         lowercase			:Optional[Any]							=  drop_connect_rate
         lowercase			:Optional[Any]							=  sum(_lowerCAmelCase ) * 4






       @classmethod
       def     SCREAMING_SNAKE_CASE			(      cls:    Optional[Any]				,    _lowerCAmelCase:    Union[str, os.PathLike]				,    **_lowerCAmelCase:    str ):
         cls._set_token_in_kwargs(_lowerCAmelCase )

         lowercase		,			lowercase			:List[str]							=  cls.get_config_dict(_lowerCAmelCase				,    **_lowerCAmelCase )

         # get the vision config dict if we are loading from AlignConfig
         if config_dict.get("model_type" ) == "align":
           lowercase			:Tuple							=  config_dict["vision_config"]

         if "model_type" in config_dict and hasattr(cls				,    "model_type" ) and config_dict["model_type"] != cls.model_type:
           logger.warning(
               F"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
               F"{cls.model_type}. This is not supported for all configurations of models and can yield errors." )

         return cls.from_dict(_lowerCAmelCase				,    **_lowerCAmelCase )







class       __lowerCAmelCase				(							lowerCAmelCase):
       _a           =  '''align'''
       _a           =  True






       def __init__(      self:    List[Any]				,    _lowerCAmelCase:    int=None				,    _lowerCAmelCase:    str=None				,    _lowerCAmelCase:    Optional[int]=6_40				,    _lowerCAmelCase:    Optional[int]=1.0				,    _lowerCAmelCase:    str=0.02				,    **_lowerCAmelCase:    Any				,    ):
         super().__init__(**_lowerCAmelCase )

         if text_config is None:
           lowercase			:Optional[Any]							=  {}
           logger.info("text_config is None. Initializing the AlignTextConfig with default values." )

         if vision_config is None:
           lowercase			:List[Any]							=  {}
           logger.info("vision_config is None. Initializing the AlignVisionConfig with default values." )

         lowercase			:List[Any]							=  AlignTextConfig(**_lowerCAmelCase )
         lowercase			:str							=  AlignVisionConfig(**_lowerCAmelCase )

         lowercase			:Optional[Any]							=  projection_dim
         lowercase			:Optional[int]							=  temperature_init_value
         lowercase			:List[str]							=  initializer_range






       @classmethod
       def     SCREAMING_SNAKE_CASE			(      cls:    int				,    _lowerCAmelCase:    AlignTextConfig				,    _lowerCAmelCase:    AlignVisionConfig				,    **_lowerCAmelCase:    Any ):
         return cls(text_config=text_config.to_dict()				,    vision_config=vision_config.to_dict()				,    **_lowerCAmelCase )






       def     SCREAMING_SNAKE_CASE			(      self:    Any ):
         lowercase			:Union[str, Any]							=  copy.deepcopy(self.__dict__ )
         lowercase			:List[Any]							=  self.text_config.to_dict()
         lowercase			:str							=  self.vision_config.to_dict()
         lowercase			:List[Any]							=  self.__class__.model_type
         return output









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):






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






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






       @property
       def     SCREAMING_SNAKE_CASE			(      self:    Any ):
         torch.manual_seed(0 )
         lowercase			:Dict							=  UNetaDConditionModel(
             sample_size=(64, 32)				,    in_channels=1				,    out_channels=1				,    layers_per_block=2				,    block_out_channels=(1_28, 1_28)				,    down_block_types=("CrossAttnDownBlock2D", "DownBlock2D")				,    up_block_types=("UpBlock2D", "CrossAttnUpBlock2D")				,    cross_attention_dim=10				,    )
         return model






       @property
       def     SCREAMING_SNAKE_CASE			(      self:    Union[str, Any] ):
         torch.manual_seed(0 )
         lowercase			:List[str]							=  AutoencoderKL(
             sample_size=(1_28, 64)				,    in_channels=1				,    out_channels=1				,    latent_channels=1				,    layers_per_block=2				,    block_out_channels=(1_28, 1_28)				,    down_block_types=("DownEncoderBlock2D", "DownEncoderBlock2D")				,    up_block_types=("UpDecoderBlock2D", "UpDecoderBlock2D")				,    )
         lowercase			:List[str]							=  UNetaDModel(
             sample_size=(64, 32)				,    in_channels=1				,    out_channels=1				,    layers_per_block=2				,    block_out_channels=(1_28, 1_28)				,    down_block_types=("AttnDownBlock2D", "DownBlock2D")				,    up_block_types=("UpBlock2D", "AttnUpBlock2D")				,    )
         return vqvae, unet






       @slow
       def     SCREAMING_SNAKE_CASE			(      self:    str ):
         lowercase			:Optional[Any]							=  "cpu"  # ensure determinism for the device-dependent torch.Generator
         lowercase			:Optional[int]							=  Mel(
             x_res=self.dummy_unet.config.sample_size[1]				,    y_res=self.dummy_unet.config.sample_size[0]				,    )

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

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

         lowercase			:List[str]							=  torch.Generator(device=_lowerCAmelCase ).manual_seed(42 )
         lowercase			:Optional[int]							=  pipe(generator=_lowerCAmelCase				,    steps=4				,    return_dict=_lowerCAmelCase )
         lowercase			:int							=  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			:Any							=  np.frombuffer(image.tobytes()				,    dtype="uint8" )[:10]
         lowercase			:Dict							=  np.frombuffer(image_from_tuple.tobytes()				,    dtype="uint8" )[:10]
         lowercase			:List[Any]							=  np.array([69, 2_55, 2_55, 2_55, 0, 0, 77, 1_81, 12, 1_27] )

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

         lowercase			:Optional[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			:List[str]							=  DDIMScheduler()
         lowercase			:Tuple							=  self.dummy_vqvae_and_unet
         lowercase			:str							=  AudioDiffusionPipeline(
             vqvae=self.dummy_vqvae_and_unet[0]				,    unet=dummy_vqvae_and_unet[1]				,    mel=_lowerCAmelCase				,    scheduler=_lowerCAmelCase )
         lowercase			:Union[str, Any]							=  pipe.to(_lowerCAmelCase )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase )

         np.random.seed(0 )
         lowercase			:Dict							=  np.random.uniform(-1				,    1				,    ((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,) )
         lowercase			:List[Any]							=  torch.Generator(device=_lowerCAmelCase ).manual_seed(42 )
         lowercase			:Dict							=  pipe(raw_audio=_lowerCAmelCase				,    generator=_lowerCAmelCase				,    start_step=5				,    steps=10 )
         lowercase			:str							=  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			:Optional[Any]							=  np.frombuffer(image.tobytes()				,    dtype="uint8" )[:10]
         lowercase			:List[str]							=  np.array([1_20, 1_17, 1_10, 1_09, 1_38, 1_67, 1_38, 1_48, 1_32, 1_21] )

         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=_lowerCAmelCase				,    mel=_lowerCAmelCase				,    scheduler=_lowerCAmelCase )
         lowercase			:int							=  pipe.to(_lowerCAmelCase )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase )

         np.random.seed(0 )
         lowercase			:List[str]							=  torch.rand((1, 1, 10) )
         lowercase			:Union[str, Any]							=  pipe(generator=_lowerCAmelCase				,    encoding=_lowerCAmelCase )
         lowercase			:Tuple							=  output.images[0]
         lowercase			:Optional[Any]							=  np.frombuffer(image.tobytes()				,    dtype="uint8" )[:10]
         lowercase			:Any							=  np.array([1_07, 1_03, 1_20, 1_27, 1_42, 1_22, 1_13, 1_22, 97, 1_11] )

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







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






       def     SCREAMING_SNAKE_CASE			(      self:    Any ):
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()
         torch.cuda.empty_cache()






       def     SCREAMING_SNAKE_CASE			(      self:    str ):
         lowercase			:Tuple							=  torch_device

         lowercase			:List[Any]							=  DiffusionPipeline.from_pretrained("teticio/audio-diffusion-ddim-256" )
         lowercase			:Dict							=  pipe.to(_lowerCAmelCase )
         pipe.set_progress_bar_config(disable=_lowerCAmelCase )

         lowercase			:Any							=  torch.Generator(device=_lowerCAmelCase ).manual_seed(42 )
         lowercase			:Optional[int]							=  pipe(generator=_lowerCAmelCase )
         lowercase			:List[str]							=  output.audios[0]
         lowercase			: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			:Dict							=  np.frombuffer(image.tobytes()				,    dtype="uint8" )[:10]
         lowercase			:List[str]							=  np.array([1_51, 1_67, 1_54, 1_44, 1_22, 1_34, 1_21, 1_05, 70, 26] )

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



