





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





       '''simple docstring'''







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




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

         __snake_case:						str   =						DDPMScheduler()
         __snake_case:						Optional[Any]   =						AudioDiffusionPipeline(vqvae=A							,				unet=self.dummy_unet							,				mel=A							,				scheduler=A		)
         __snake_case:						str   =						pipe.to(A		)
         pipe.set_progress_bar_config(disable=A		)

         __snake_case:						str   =						torch.Generator(device=A		).manual_seed(42		)
         __snake_case:						Dict   =						pipe(generator=A							,				steps=4		)
         __snake_case:						Optional[Any]   =						output.audios[0]
         __snake_case:						Optional[Any]   =						output.images[0]

         __snake_case:						Tuple   =						torch.Generator(device=A		).manual_seed(42		)
         __snake_case:						Dict   =						pipe(generator=A							,				steps=4							,				return_dict=A		)
         __snake_case:						Any   =						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]
         )
         __snake_case:						Any   =						np.frombuffer(image.tobytes()							,				dtype="""uint8"""		)[:10]
         __snake_case:						Optional[Any]   =						np.frombuffer(image_from_tuple.tobytes()							,				dtype="""uint8"""		)[:10]
         __snake_case:						Any   =						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

         __snake_case:						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]							,				)

         __snake_case:						Tuple   =						DDIMScheduler()
         __snake_case:						int   =						self.dummy_vqvae_and_unet
         __snake_case:						List[Any]   =						AudioDiffusionPipeline(
             vqvae=self.dummy_vqvae_and_unet[0]							,				unet=dummy_vqvae_and_unet[1]							,				mel=A							,				scheduler=A		)
         __snake_case:						Dict   =						pipe.to(A		)
         pipe.set_progress_bar_config(disable=A		)

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

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

         __snake_case:						List[Any]   =						self.dummy_unet_condition
         __snake_case:						Optional[Any]   =						AudioDiffusionPipeline(
             vqvae=self.dummy_vqvae_and_unet[0]							,				unet=A							,				mel=A							,				scheduler=A		)
         __snake_case:						Optional[int]   =						pipe.to(A		)
         pipe.set_progress_bar_config(disable=A		)

         np.random.seed(0		)
         __snake_case:						Any   =						torch.rand((1, 1, 10)		)
         __snake_case:						Optional[Any]   =						pipe(generator=A							,				encoding=A		)
         __snake_case:						Union[str, Any]   =						output.images[0]
         __snake_case:						Optional[Any]   =						np.frombuffer(image.tobytes()							,				dtype="""uint8"""		)[:10]
         __snake_case:						int   =						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     __snake_case       (							unittest.TestCase			):





       '''simple docstring'''







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






       def 							UpperCAmelCase__ (  self :    List[Any]		):
         __snake_case:						str   =						torch_device

         __snake_case:						Optional[int]   =						DiffusionPipeline.from_pretrained("""teticio/audio-diffusion-ddim-256"""		)
         __snake_case:						List[Any]   =						pipe.to(A		)
         pipe.set_progress_bar_config(disable=A		)

         __snake_case:						Dict   =						torch.Generator(device=A		).manual_seed(42		)
         __snake_case:						List[str]   =						pipe(generator=A		)
         __snake_case:						Optional[int]   =						output.audios[0]
         __snake_case:						Optional[int]   =						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]
         __snake_case:						str   =						np.frombuffer(image.tobytes()							,				dtype="""uint8"""		)[:10]
         __snake_case:						Tuple   =						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 TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_sentencepiece_available


__UpperCAmelCase       :   Union[str, Any]				       =							{}

try:
      if not is_sentencepiece_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __UpperCAmelCase       :   str				       =							["BartphoTokenizer"]

if TYPE_CHECKING:
      try:
            if not is_sentencepiece_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_bartpho import BartphoTokenizer

else:
      import sys

      __UpperCAmelCase       :   Optional[Any]				       =							_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)






