




from torch import nn



def 					UpperCAmelCase__     (   lowerCamelCase      ):
  if act_fn in ["swish", "silu"]:
    return nn.SiLU()
  elif act_fn == "mish":
    return nn.Mish()
  elif act_fn == "gelu":
    return nn.GELU()
  else:
    raise ValueError(F"Unsupported activation function: {act_fn}"      )









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



