'''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  lowerCAmelCase__ (      self			:       List[Any]  )				->Tuple:
        # clean up the VRAM after each test
        super().tearDown()
        gc.collect()
        torch.cuda.empty_cache()






  @property
  def  lowerCAmelCase__ (      self			:       Any  )				->str:
        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  lowerCAmelCase__ (      self			:       str  )				->Union[str, Any]:
        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  lowerCAmelCase__ (      self			:       List[Any]  )				->Any:
        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  lowerCAmelCase__ (      self			:       Optional[Any]  )				->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 lowerCamelCase							(	unittest.TestCase					):




  '''simple docstring'''



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






  def  lowerCAmelCase__ (      self			:       Dict  )				->Union[str, Any]:
        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






'''simple docstring'''
from typing import List, Optional, Union

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


lowercase__  :      int        =					logging.get_logger(__name__)

lowercase__  :      List[str]        =					{
    "huggingface/time-series-transformer-tourism-monthly": (
        "https://huggingface.co/huggingface/time-series-transformer-tourism-monthly/resolve/main/config.json"
    ),
    # See all TimeSeriesTransformer models at https://huggingface.co/models?filter=time_series_transformer
}







class lowerCamelCase							(	lowerCamelCase					):




  '''simple docstring'''



  lowerCAmelCase__										=	'''time_series_transformer'''
  lowerCAmelCase__										=	{
      '''hidden_size''': '''d_model''',
      '''num_attention_heads''': '''encoder_attention_heads''',
      '''num_hidden_layers''': '''encoder_layers''',
  }






  def __init__(      self			:       Optional[int]  ,							UpperCAmelCase__			:       Optional[int] = None  ,							UpperCAmelCase__			:       Optional[int] = None  ,							UpperCAmelCase__			:       str = "student_t"  ,							UpperCAmelCase__			:       str = "nll"  ,							UpperCAmelCase__			:       int = 1  ,							UpperCAmelCase__			:       List[int] = [1, 2, 3, 4, 5, 6, 7]  ,							UpperCAmelCase__			:       Optional[Union[str, bool]] = "mean"  ,							UpperCAmelCase__			:       int = 0  ,							UpperCAmelCase__			:       int = 0  ,							UpperCAmelCase__			:       int = 0  ,							UpperCAmelCase__			:       int = 0  ,							UpperCAmelCase__			:       Optional[List[int]] = None  ,							UpperCAmelCase__			:       Optional[List[int]] = None  ,							UpperCAmelCase__			:       int = 32  ,							UpperCAmelCase__			:       int = 32  ,							UpperCAmelCase__			:       int = 2  ,							UpperCAmelCase__			:       int = 2  ,							UpperCAmelCase__			:       int = 2  ,							UpperCAmelCase__			:       int = 2  ,							UpperCAmelCase__			:       bool = True  ,							UpperCAmelCase__			:       str = "gelu"  ,							UpperCAmelCase__			:       int = 64  ,							UpperCAmelCase__			:       float = 0.1  ,							UpperCAmelCase__			:       float = 0.1  ,							UpperCAmelCase__			:       float = 0.1  ,							UpperCAmelCase__			:       float = 0.1  ,							UpperCAmelCase__			:       float = 0.1  ,							UpperCAmelCase__			:       int = 100  ,							UpperCAmelCase__			:       float = 0.02  ,							UpperCAmelCase__			:       Union[str, Any]=True  ,							**UpperCAmelCase__			:       Tuple  ,							)				->Optional[int]:
        # time series specific configuration
        UpperCAmelCase_            =  prediction_length
        UpperCAmelCase_            =  context_length or prediction_length
        UpperCAmelCase_            =  distribution_output
        UpperCAmelCase_            =  loss
        UpperCAmelCase_            =  input_size
        UpperCAmelCase_            =  num_time_features
        UpperCAmelCase_            =  lags_sequence
        UpperCAmelCase_            =  scaling
        UpperCAmelCase_            =  num_dynamic_real_features
        UpperCAmelCase_            =  num_static_real_features
        UpperCAmelCase_            =  num_static_categorical_features
        if cardinality and num_static_categorical_features > 0:
              if len(UpperCAmelCase__  ) != num_static_categorical_features:
                    raise ValueError(
                        '''The cardinality should be a list of the same length as `num_static_categorical_features`'''  )
              UpperCAmelCase_            =  cardinality
        else:
              UpperCAmelCase_            =  [0]
        if embedding_dimension and num_static_categorical_features > 0:
              if len(UpperCAmelCase__  ) != num_static_categorical_features:
                    raise ValueError(
                        '''The embedding dimension should be a list of the same length as `num_static_categorical_features`'''  )
              UpperCAmelCase_            =  embedding_dimension
        else:
              UpperCAmelCase_            =  [min(50  ,							(cat + 1) // 2  ) for cat in self.cardinality]
        UpperCAmelCase_            =  num_parallel_samples

        # Transformer architecture configuration
        UpperCAmelCase_            =  input_size * len(UpperCAmelCase__  ) + self._number_of_features
        UpperCAmelCase_            =  d_model
        UpperCAmelCase_            =  encoder_attention_heads
        UpperCAmelCase_            =  decoder_attention_heads
        UpperCAmelCase_            =  encoder_ffn_dim
        UpperCAmelCase_            =  decoder_ffn_dim
        UpperCAmelCase_            =  encoder_layers
        UpperCAmelCase_            =  decoder_layers

        UpperCAmelCase_            =  dropout
        UpperCAmelCase_            =  attention_dropout
        UpperCAmelCase_            =  activation_dropout
        UpperCAmelCase_            =  encoder_layerdrop
        UpperCAmelCase_            =  decoder_layerdrop

        UpperCAmelCase_            =  activation_function
        UpperCAmelCase_            =  init_std

        UpperCAmelCase_            =  use_cache

        super().__init__(is_encoder_decoder=UpperCAmelCase__  ,							**UpperCAmelCase__  )






  @property
  def  lowerCAmelCase__ (      self			:       List[str]  )				->int:
        return (
            sum(self.embedding_dimension  )
            + self.num_dynamic_real_features
            + self.num_time_features
            + self.num_static_real_features
            + self.input_size * 2  # the log1p(abs(loc)) and log(scale) features
        )





