



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        _snake_case    (self				):
						# clean up the VRAM after each test
						super().tearDown()
						gc.collect()
						torch.cuda.empty_cache()





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

						lowerCamelCase__ :      str								=		DDPMScheduler()
						lowerCamelCase__ :      Any								=		AudioDiffusionPipeline(vqvae=__magic_name__							,    unet=self.dummy_unet							,    mel=__magic_name__							,    scheduler=__magic_name__				)
						lowerCamelCase__ :      List[Any]								=		pipe.to(__magic_name__				)
						pipe.set_progress_bar_config(disable=__magic_name__				)

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

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

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

						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__ :      str								=		torch.Generator(device=__magic_name__				).manual_seed(42				)
						lowerCamelCase__ :      str								=		pipe(raw_audio=__magic_name__							,    generator=__magic_name__							,    start_step=5							,    steps=10				)
						lowerCamelCase__ :      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__ :      str								=		np.frombuffer(image.tobytes()							,    dtype="""uint8"""				)[:10]
						lowerCamelCase__ :      List[str]								=		np.array([120, 117, 110, 109, 138, 167, 138, 148, 132, 121]				)

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

						lowerCamelCase__ :      Tuple								=		self.dummy_unet_condition
						lowerCamelCase__ :      Tuple								=		AudioDiffusionPipeline(
						    vqvae=self.dummy_vqvae_and_unet[0]							,    unet=__magic_name__							,    mel=__magic_name__							,    scheduler=__magic_name__				)
						lowerCamelCase__ :      List[Any]								=		pipe.to(__magic_name__				)
						pipe.set_progress_bar_config(disable=__magic_name__				)

						np.random.seed(0				)
						lowerCamelCase__ :      Tuple								=		torch.rand((1, 1, 10)				)
						lowerCamelCase__ :      Any								=		pipe(generator=__magic_name__							,    encoding=__magic_name__				)
						lowerCamelCase__ :      int								=		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 				__A					( unittest.TestCase      ):





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





			def        _snake_case    (self				):
						lowerCamelCase__ :      Union[str, Any]								=		torch_device

						lowerCamelCase__ :      int								=		DiffusionPipeline.from_pretrained("""teticio/audio-diffusion-ddim-256"""				)
						lowerCamelCase__ :      Any								=		pipe.to(__magic_name__				)
						pipe.set_progress_bar_config(disable=__magic_name__				)

						lowerCamelCase__ :      Union[str, Any]								=		torch.Generator(device=__magic_name__				).manual_seed(42				)
						lowerCamelCase__ :      Tuple								=		pipe(generator=__magic_name__				)
						lowerCamelCase__ :      Any								=		output.audios[0]
						lowerCamelCase__ :      Dict								=		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__ :      Union[str, Any]								=		np.frombuffer(image.tobytes()							,    dtype="""uint8"""				)[:10]
						lowerCamelCase__ :      str								=		np.array([151, 167, 154, 144, 122, 134, 121, 105, 70, 26]				)

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








import copy
from dataclasses import dataclass, field
from typing import ClassVar, Dict

from ..features import Audio, Features, Value
from .base import TaskTemplate
@dataclass(frozen=A_      )
class 				__A					( A_      ):
			UpperCamelCase					:str     =			field(default='''automatic-speech-recognition'''					, metadata={'''include_in_asdict_even_if_is_default''': True}      )
			UpperCamelCase					:ClassVar[Features]     =			Features({'''audio''': Audio()}      )
			UpperCamelCase					:ClassVar[Features]     =			Features({'''transcription''': Value('''string'''      )}      )
			UpperCamelCase					:str     =			"audio"
			UpperCamelCase					:str     =			"transcription"





			def        _snake_case    (self							,    __magic_name__				):
						if self.audio_column not in features:
									raise ValueError(f"Column {self.audio_column} is not present in features."				)
						if not isinstance(features[self.audio_column]							,    __magic_name__				):
									raise ValueError(f"Column {self.audio_column} is not an Audio type."				)
						lowerCamelCase__ :      int								=		copy.deepcopy(self				)
						lowerCamelCase__ :      Tuple								=		self.input_schema.copy()
						lowerCamelCase__ :      Tuple								=		features[self.audio_column]
						lowerCamelCase__ :      List[str]								=		input_schema
						return task_template





			@property
			def        _snake_case    (self				):
						return {self.audio_column: "audio", self.transcription_column: "transcription"}



