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 	_SCREAMING_SNAKE_CASE    (					self     :				List[Any])				->List[Any]:


          """simple docstring"""


          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()



      @property
      def 	_SCREAMING_SNAKE_CASE    (					self     :				Optional[Any])				->str:


          """simple docstring"""


          torch.manual_seed(0)
          _lowerCamelCase     :		List[str] 						=						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 	_SCREAMING_SNAKE_CASE    (					self     :				Optional[int])				->Union[str, Any]:


          """simple docstring"""


          torch.manual_seed(0)
          _lowerCamelCase     :		Tuple 						=						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 	_SCREAMING_SNAKE_CASE    (					self     :				Optional[Any])				->Any:


          """simple docstring"""


          torch.manual_seed(0)
          _lowerCamelCase     :		Optional[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""")    ,  )
          _lowerCamelCase     :		Tuple 						=						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 	_SCREAMING_SNAKE_CASE    (					self     :				Optional[Any])				->List[Any]:


          """simple docstring"""


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

          _lowerCamelCase     :		int 						=						DDPMScheduler()
          _lowerCamelCase     :		str 						=						AudioDiffusionPipeline(vqvae=_UpperCamelCase    ,  unet=self.dummy_unet    ,  mel=_UpperCamelCase    ,  scheduler=_UpperCamelCase)
          _lowerCamelCase     :		Optional[int] 						=						pipe.to(_UpperCamelCase)
          pipe.set_progress_bar_config(disable=_UpperCamelCase)

          _lowerCamelCase     :		Any 						=						torch.Generator(device=_UpperCamelCase).manual_seed(42)
          _lowerCamelCase     :		Union[str, Any] 						=						pipe(generator=_UpperCamelCase    ,  steps=4)
          _lowerCamelCase     :		List[str] 						=						output.audios[0]
          _lowerCamelCase     :		Tuple 						=						output.images[0]

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

          _lowerCamelCase     :		List[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]    ,  )

          _lowerCamelCase     :		str 						=						DDIMScheduler()
          _lowerCamelCase     :		Any 						=						self.dummy_vqvae_and_unet
          _lowerCamelCase     :		Optional[Any] 						=						AudioDiffusionPipeline(
              vqvae=self.dummy_vqvae_and_unet[0]    ,  unet=dummy_vqvae_and_unet[1]    ,  mel=_UpperCamelCase    ,  scheduler=_UpperCamelCase)
          _lowerCamelCase     :		Optional[Any] 						=						pipe.to(_UpperCamelCase)
          pipe.set_progress_bar_config(disable=_UpperCamelCase)

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

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

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

          np.random.seed(0)
          _lowerCamelCase     :		Dict 						=						torch.rand((1, 1, 10))
          _lowerCamelCase     :		List[Any] 						=						pipe(generator=_UpperCamelCase    ,  encoding=_UpperCamelCase)
          _lowerCamelCase     :		Tuple 						=						output.images[0]
          _lowerCamelCase     :		int 						=						np.frombuffer(image.tobytes()    ,  dtype="""uint8""")[:10]
          _lowerCamelCase     :		List[str] 						=						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 	_SCREAMING_SNAKE_CASE    (					self     :				str)				->Dict:


          """simple docstring"""


          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()




      def 	_SCREAMING_SNAKE_CASE    (					self     :				Union[str, Any])				->List[Any]:


          """simple docstring"""


          _lowerCamelCase     :		Dict 						=						torch_device

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

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

from .. import Features, NamedSplit
from ..packaged_modules.text.text import Text
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader




class 	__snake_case	( __lowerCAmelCase		):





      '''simple docstring'''






      def __init__(					self     :				Dict    ,  _UpperCamelCase     :				NestedDataStructureLike[PathLike]    ,  _UpperCamelCase     :				Optional[NamedSplit] = None    ,  _UpperCamelCase     :				Optional[Features] = None    ,  _UpperCamelCase     :				str = None    ,  _UpperCamelCase     :				bool = False    ,  _UpperCamelCase     :				bool = False    ,  _UpperCamelCase     :				Optional[int] = None    ,  **_UpperCamelCase     :				Tuple    ,  )				->Union[str, Any]:


          """simple docstring"""


          super().__init__(
              _UpperCamelCase    ,  split=_UpperCamelCase    ,  features=_UpperCamelCase    ,  cache_dir=_UpperCamelCase    ,  keep_in_memory=_UpperCamelCase    ,  streaming=_UpperCamelCase    ,  num_proc=_UpperCamelCase    ,  **_UpperCamelCase    ,  )
          _lowerCamelCase     :		List[Any] 						=						path_or_paths if isinstance(_UpperCamelCase    ,  _UpperCamelCase) else {self.split: path_or_paths}
          _lowerCamelCase     :		Any 						=						Text(
              cache_dir=_UpperCamelCase    ,  data_files=_UpperCamelCase    ,  features=_UpperCamelCase    ,  **_UpperCamelCase    ,  )




      def 	_SCREAMING_SNAKE_CASE    (					self     :				Optional[Any])				->Optional[Any]:


          """simple docstring"""


          if self.streaming:
              _lowerCamelCase     :		Tuple 						=						self.builder.as_streaming_dataset(split=self.split)
          # Build regular (map-style) dataset
          else:
              _lowerCamelCase     :		List[Any] 						=						None
              _lowerCamelCase     :		Any 						=						None
              _lowerCamelCase     :		List[str] 						=						None
              _lowerCamelCase     :		Dict 						=						None

              self.builder.download_and_prepare(
                  download_config=_UpperCamelCase    ,  download_mode=_UpperCamelCase    ,  verification_mode=_UpperCamelCase    ,  base_path=_UpperCamelCase    ,  num_proc=self.num_proc    ,  )
              _lowerCamelCase     :		Optional[int] 						=						self.builder.as_dataset(
                  split=self.split    ,  verification_mode=_UpperCamelCase    ,  in_memory=self.keep_in_memory)
          return dataset




