



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 		_lowerCamelCase      (		self							:					List[str]):


           '''simple docstring'''

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



       @property
       def 		_lowerCamelCase      (		self							:					int):


           '''simple docstring'''

           torch.manual_seed(0)
           __a							=    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							:					Tuple):


           '''simple docstring'''

           torch.manual_seed(0)
           __a							=    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							:					Union[str, Any]):


           '''simple docstring'''

           torch.manual_seed(0)
           __a							=    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''')					,		)
           __a							=    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							:					List[str]):


           '''simple docstring'''

           __a							=    '''cpu'''  # ensure determinism for the device-dependent torch.Generator
           __a							=    Mel(
               x_res=self.dummy_unet.config.sample_size[1]					,		y_res=self.dummy_unet.config.sample_size[0]					,		)

           __a							=    DDPMScheduler()
           __a							=    AudioDiffusionPipeline(vqvae=__SCREAMING_SNAKE_CASE					,		unet=self.dummy_unet					,		mel=__SCREAMING_SNAKE_CASE					,		scheduler=__SCREAMING_SNAKE_CASE)
           __a							=    pipe.to(__SCREAMING_SNAKE_CASE)
           pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE)

           __a							=    torch.Generator(device=__SCREAMING_SNAKE_CASE).manual_seed(42)
           __a							=    pipe(generator=__SCREAMING_SNAKE_CASE					,		steps=4)
           __a							=    output.audios[0]
           __a							=    output.images[0]

           __a							=    torch.Generator(device=__SCREAMING_SNAKE_CASE).manual_seed(42)
           __a							=    pipe(generator=__SCREAMING_SNAKE_CASE					,		steps=4					,		return_dict=__SCREAMING_SNAKE_CASE)
           __a							=    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]
           )
           __a							=    np.frombuffer(image.tobytes()					,		dtype='''uint8''')[:10]
           __a							=    np.frombuffer(image_from_tuple.tobytes()					,		dtype='''uint8''')[:10]
           __a							=    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

           __a							=    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]					,		)

           __a							=    DDIMScheduler()
           __a							=    self.dummy_vqvae_and_unet
           __a							=    AudioDiffusionPipeline(
               vqvae=self.dummy_vqvae_and_unet[0]					,		unet=dummy_vqvae_and_unet[1]					,		mel=__SCREAMING_SNAKE_CASE					,		scheduler=__SCREAMING_SNAKE_CASE)
           __a							=    pipe.to(__SCREAMING_SNAKE_CASE)
           pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE)

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

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

           __a							=    self.dummy_unet_condition
           __a							=    AudioDiffusionPipeline(
               vqvae=self.dummy_vqvae_and_unet[0]					,		unet=__SCREAMING_SNAKE_CASE					,		mel=__SCREAMING_SNAKE_CASE					,		scheduler=__SCREAMING_SNAKE_CASE)
           __a							=    pipe.to(__SCREAMING_SNAKE_CASE)
           pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE)

           np.random.seed(0)
           __a							=    torch.rand((1, 1, 10))
           __a							=    pipe(generator=__SCREAMING_SNAKE_CASE					,		encoding=__SCREAMING_SNAKE_CASE)
           __a							=    output.images[0]
           __a							=    np.frombuffer(image.tobytes()					,		dtype='''uint8''')[:10]
           __a							=    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    _A	(     unittest.TestCase   ):



       def 		_lowerCamelCase      (		self							:					str):


           '''simple docstring'''

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





       def 		_lowerCamelCase      (		self							:					str):


           '''simple docstring'''

           __a							=    torch_device

           __a							=    DiffusionPipeline.from_pretrained('''teticio/audio-diffusion-ddim-256''')
           __a							=    pipe.to(__SCREAMING_SNAKE_CASE)
           pipe.set_progress_bar_config(disable=__SCREAMING_SNAKE_CASE)

           __a							=    torch.Generator(device=__SCREAMING_SNAKE_CASE).manual_seed(42)
           __a							=    pipe(generator=__SCREAMING_SNAKE_CASE)
           __a							=    output.audios[0]
           __a							=    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]
           __a							=    np.frombuffer(image.tobytes()					,		dtype='''uint8''')[:10]
           __a							=    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_tokenizers_available, is_torch_available


__snake_case						:Tuple						=					{
    '''configuration_bloom''': ['''BLOOM_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''BloomConfig''', '''BloomOnnxConfig'''],
}
try:
 if not is_tokenizers_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __snake_case						:List[Any]						=					['''BloomTokenizerFast''']

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __snake_case						:Any						=					[
     '''BLOOM_PRETRAINED_MODEL_ARCHIVE_LIST''',
     '''BloomForCausalLM''',
     '''BloomModel''',
     '''BloomPreTrainedModel''',
     '''BloomForSequenceClassification''',
     '''BloomForTokenClassification''',
     '''BloomForQuestionAnswering''',
 ]

if TYPE_CHECKING:
 from .configuration_bloom import BLOOM_PRETRAINED_CONFIG_ARCHIVE_MAP, BloomConfig, BloomOnnxConfig

 try:
  if not is_tokenizers_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tokenization_bloom_fast import BloomTokenizerFast

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_bloom import (
      BLOOM_PRETRAINED_MODEL_ARCHIVE_LIST,
      BloomForCausalLM,
      BloomForQuestionAnswering,
      BloomForSequenceClassification,
      BloomForTokenClassification,
      BloomModel,
      BloomPreTrainedModel,
  )

else:
 import sys

 __snake_case						:Any						=					_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)





