






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

  def    a					(    self						:	List[Any]			):
     """simple docstring"""

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

  @property
  def    a					(    self						:	int			):
     """simple docstring"""

     torch.manual_seed(0			)
     __UpperCamelCase	:			int					=					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    a					(    self						:	Optional[Any]			):
     """simple docstring"""

     torch.manual_seed(0			)
     __UpperCamelCase	:			List[Any]					=					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    a					(    self						:	str			):
     """simple docstring"""

     torch.manual_seed(0			)
     __UpperCamelCase	:			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""")   ,				)
     __UpperCamelCase	:			List[Any]					=					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    a					(    self						:	Any			):
     """simple docstring"""

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

     __UpperCamelCase	:			Dict					=					DDPMScheduler()
     __UpperCamelCase	:			Tuple					=					AudioDiffusionPipeline(vqvae=lowerCamelCase__   ,				unet=self.dummy_unet   ,				mel=lowerCamelCase__   ,				scheduler=lowerCamelCase__			)
     __UpperCamelCase	:			Union[str, Any]					=					pipe.to(lowerCamelCase__			)
     pipe.set_progress_bar_config(disable=lowerCamelCase__			)

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

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

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

     np.random.seed(0			)
     __UpperCamelCase	:			Any					=					np.random.uniform(-1   ,				1   ,				((dummy_vqvae_and_unet[0].config.sample_size[1] - 1) * mel.hop_length,)			)
     __UpperCamelCase	:			Optional[Any]					=					torch.Generator(device=lowerCamelCase__			).manual_seed(42			)
     __UpperCamelCase	:			List[str]					=					pipe(raw_audio=lowerCamelCase__   ,				generator=lowerCamelCase__   ,				start_step=5   ,				steps=10			)
     __UpperCamelCase	:			Optional[int]					=					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	:			Any					=					np.frombuffer(image.tobytes()   ,				dtype="""uint8"""			)[:10]
     __UpperCamelCase	:			Dict					=					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

     __UpperCamelCase	:			Union[str, Any]					=					self.dummy_unet_condition
     __UpperCamelCase	:			int					=					AudioDiffusionPipeline(
         vqvae=self.dummy_vqvae_and_unet[0]   ,				unet=lowerCamelCase__   ,				mel=lowerCamelCase__   ,				scheduler=lowerCamelCase__			)
     __UpperCamelCase	:			Union[str, Any]					=					pipe.to(lowerCamelCase__			)
     pipe.set_progress_bar_config(disable=lowerCamelCase__			)

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

  def    a					(    self						:	int			):
     """simple docstring"""

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

  def    a					(    self						:	Union[str, Any]			):
     """simple docstring"""

     __UpperCamelCase	:			str					=					torch_device

     __UpperCamelCase	:			Dict					=					DiffusionPipeline.from_pretrained("""teticio/audio-diffusion-ddim-256"""			)
     __UpperCamelCase	:			List[str]					=					pipe.to(lowerCamelCase__			)
     pipe.set_progress_bar_config(disable=lowerCamelCase__			)

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











from collections.abc import Iterable
from typing import Generic, TypeVar

UpperCamelCase    							=      TypeVar('_T')







class 		_A						( Generic[_T] ):

  def __init__(    self						:	int   ,				lowerCamelCase__						:	Iterable[_T] | None = None			):
     """simple docstring"""

     __UpperCamelCase	:			list[_T]					=					list(iterable or []			)
     __UpperCamelCase	:			list[_T]					=					[]

  def __len__(    self						:	str			):
     """simple docstring"""

     return len(self._stacka			) + len(self._stacka			)

  def __repr__(    self						:	Dict			):
     """simple docstring"""

     return f'Queue({tuple(self._stacka[::-1] + self._stacka			)})'

  def    a					(    self						:	Union[str, Any]   ,				lowerCamelCase__						:	_T			):
     """simple docstring"""

     self._stacka.append(lowerCamelCase__			)

  def    a					(    self						:	Union[str, Any]			):
     """simple docstring"""

     __UpperCamelCase	:			Any					=					self._stacka.pop
     __UpperCamelCase	:			int					=					self._stacka.append

     if not self._stacka:
        while self._stacka:
           stacka_append(stacka_pop()			)

     if not self._stacka:
        raise IndexError("""Queue is empty"""			)
     return self._stacka.pop()


if __name__ == "__main__":
       from doctest import testmod

       testmod()



