
'''simple docstring'''

import gc
import unittest

import numpy as np
import torch

from diffusers import DDPMScheduler, MidiProcessor, SpectrogramDiffusionPipeline
from diffusers.pipelines.spectrogram_diffusion import SpectrogramContEncoder, SpectrogramNotesEncoder, TaFilmDecoder
from diffusers.utils import require_torch_gpu, skip_mps, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_note_seq, require_onnxruntime

from ..pipeline_params import TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS, TOKENS_TO_AUDIO_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()


UpperCAmelCase__   : int					=		"./tests/fixtures/elise_format0.mid"






@unittest.skip('The note-seq package currently throws an error on import'    )
class    A     (	SCREAMING_SNAKE_CASE__  ,      unittest.TestCase    ):
					snake_case__      :Union[str, Any]        =     SpectrogramDiffusionPipeline
					snake_case__      :Any        =     PipelineTesterMixin.required_optional_params - {
					    'callback',
					    'latents',
					    'callback_steps',
					    'output_type',
					    'num_images_per_prompt',
					}
					snake_case__      :Optional[int]        =     False
					snake_case__      :Optional[Any]        =     False
					snake_case__      :List[Any]        =     TOKENS_TO_AUDIO_GENERATION_PARAMS
					snake_case__      :Union[str, Any]        =     TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS





					def 							__SCREAMING_SNAKE_CASE  (					self			:       Any							):


								"""simple docstring"""

								torch.manual_seed(0							)
								lowerCAmelCase__								=							SpectrogramNotesEncoder(
								    max_length=2048      ,							vocab_size=1536      ,							d_model=768      ,							dropout_rate=0.1      ,							num_layers=1      ,							num_heads=1      ,							d_kv=4      ,							d_ff=2048      ,							feed_forward_proj="gated-gelu"      ,							)

								lowerCAmelCase__								=							SpectrogramContEncoder(
								    input_dims=128      ,							targets_context_length=256      ,							d_model=768      ,							dropout_rate=0.1      ,							num_layers=1      ,							num_heads=1      ,							d_kv=4      ,							d_ff=2048      ,							feed_forward_proj="gated-gelu"      ,							)

								lowerCAmelCase__								=							TaFilmDecoder(
								    input_dims=128      ,							targets_length=256      ,							max_decoder_noise_time=20000.0      ,							d_model=768      ,							num_layers=1      ,							num_heads=1      ,							d_kv=4      ,							d_ff=2048      ,							dropout_rate=0.1      ,							)

								lowerCAmelCase__								=							DDPMScheduler()

								lowerCAmelCase__								=							{
								    "notes_encoder": notes_encoder.eval(),
								    "continuous_encoder": continuous_encoder.eval(),
								    "decoder": decoder.eval(),
								    "scheduler": scheduler,
								    "melgan": None,
								}
								return components





					def 							__SCREAMING_SNAKE_CASE  (					self			:       Dict      ,							__magic_name__			:       Any      ,							__magic_name__			:       List[str]=0							):


								"""simple docstring"""

								if str(__magic_name__							).startswith("mps"							):
											lowerCAmelCase__								=							torch.manual_seed(__magic_name__							)
								else:
											lowerCAmelCase__								=							torch.Generator(device=__magic_name__							).manual_seed(__magic_name__							)
								lowerCAmelCase__								=							{
								    "input_tokens": [
								        [1134, 90, 1135, 1133, 1080, 112, 1132, 1080, 1133, 1079, 133, 1132, 1079, 1133, 1] + [0] * 2033
								    ],
								    "generator": generator,
								    "num_inference_steps": 4,
								    "output_type": "mel",
								}
								return inputs





					def 							__SCREAMING_SNAKE_CASE  (					self			:       Optional[Any]							):


								"""simple docstring"""

								lowerCAmelCase__								=							"cpu"  # ensure determinism for the device-dependent torch.Generator
								lowerCAmelCase__								=							self.get_dummy_components()
								lowerCAmelCase__								=							SpectrogramDiffusionPipeline(**__magic_name__							)
								lowerCAmelCase__								=							pipe.to(__magic_name__							)
								pipe.set_progress_bar_config(disable=__magic_name__							)

								lowerCAmelCase__								=							self.get_dummy_inputs(__magic_name__							)
								lowerCAmelCase__								=							pipe(**__magic_name__							)
								lowerCAmelCase__								=							output.audios

								lowerCAmelCase__								=							mel[0, -3:, -3:]

								assert mel_slice.shape == (3, 3)
								lowerCAmelCase__								=							np.array(
								    [-11.51_2925, -4.78_8215, -0.4617_2905, -2.05_1715, -10.53_9147, -10.97_0963, -9.09_1634, 4.0, 4.0]							)
								assert np.abs(mel_slice.flatten() - expected_slice							).max() < 1E-2





					@skip_mps
					def 							__SCREAMING_SNAKE_CASE  (					self			:       List[str]							):


								"""simple docstring"""

								return super().test_save_load_local()





					@skip_mps
					def 							__SCREAMING_SNAKE_CASE  (					self			:       str							):


								"""simple docstring"""

								return super().test_dict_tuple_outputs_equivalent()





					@skip_mps
					def 							__SCREAMING_SNAKE_CASE  (					self			:       int							):


								"""simple docstring"""

								return super().test_save_load_optional_components()





					@skip_mps
					def 							__SCREAMING_SNAKE_CASE  (					self			:       Any							):


								"""simple docstring"""

								return super().test_attention_slicing_forward_pass()





					def 							__SCREAMING_SNAKE_CASE  (					self			:       Dict							):


								"""simple docstring"""

								pass





					def 							__SCREAMING_SNAKE_CASE  (					self			:       Optional[int]							):


								"""simple docstring"""

								pass







					@skip_mps
					def 							__SCREAMING_SNAKE_CASE  (					self			:       Any							):


								"""simple docstring"""

								return super().test_progress_bar()






@slow
@require_torch_gpu
@require_onnxruntime
@require_note_seq
class    A     (	unittest.TestCase    ):





					def 							__SCREAMING_SNAKE_CASE  (					self			:       List[str]							):


								"""simple docstring"""

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





					def 							__SCREAMING_SNAKE_CASE  (					self			:       Optional[int]							):


								"""simple docstring"""

								lowerCAmelCase__								=							torch_device

								lowerCAmelCase__								=							SpectrogramDiffusionPipeline.from_pretrained("google/music-spectrogram-diffusion"							)
								lowerCAmelCase__								=							pipe.melgan
								lowerCAmelCase__								=							None

								lowerCAmelCase__								=							pipe.to(__magic_name__							)
								pipe.set_progress_bar_config(disable=__magic_name__							)

								def callback(__magic_name__			:       Tuple      ,							__magic_name__			:       Optional[int]							):
											# decode mel to audio
											lowerCAmelCase__								=							melgan(input_features=mel_output.astype(np.floataa							)							)[0]
											assert len(audio[0]							) == 81920 * (step + 1)
											# simulate that audio is played
											return audio

								lowerCAmelCase__								=							MidiProcessor()
								lowerCAmelCase__								=							processor(__magic_name__							)

								lowerCAmelCase__								=							input_tokens[:3]
								lowerCAmelCase__								=							torch.manual_seed(0							)
								pipe(__magic_name__      ,							num_inference_steps=5      ,							generator=__magic_name__      ,							callback=__magic_name__      ,							output_type="mel"							)





					def 							__SCREAMING_SNAKE_CASE  (					self			:       int							):


								"""simple docstring"""

								lowerCAmelCase__								=							torch_device

								lowerCAmelCase__								=							SpectrogramDiffusionPipeline.from_pretrained("google/music-spectrogram-diffusion"							)
								lowerCAmelCase__								=							pipe.to(__magic_name__							)
								pipe.set_progress_bar_config(disable=__magic_name__							)
								lowerCAmelCase__								=							MidiProcessor()

								lowerCAmelCase__								=							processor(__magic_name__							)
								# just run two denoising loops
								lowerCAmelCase__								=							input_tokens[:2]

								lowerCAmelCase__								=							torch.manual_seed(0							)
								lowerCAmelCase__								=							pipe(__magic_name__      ,							num_inference_steps=2      ,							generator=__magic_name__							)

								lowerCAmelCase__								=							output.audios[0]

								assert abs(np.abs(__magic_name__							).sum() - 3612.841							) < 1E-1







					def 							__SCREAMING_SNAKE_CASE  (					self			:       Union[str, Any]							):


								"""simple docstring"""

								lowerCAmelCase__								=							torch_device

								lowerCAmelCase__								=							SpectrogramDiffusionPipeline.from_pretrained("google/music-spectrogram-diffusion"							)
								lowerCAmelCase__								=							pipe.to(__magic_name__							)
								pipe.set_progress_bar_config(disable=__magic_name__							)

								lowerCAmelCase__								=							MidiProcessor()

								lowerCAmelCase__								=							processor(__magic_name__							)

								# just run 4 denoising loops
								lowerCAmelCase__								=							input_tokens[:4]

								lowerCAmelCase__								=							torch.manual_seed(0							)
								lowerCAmelCase__								=							pipe(__magic_name__      ,							num_inference_steps=100      ,							generator=__magic_name__							)

								lowerCAmelCase__								=							output.audios[0]
								assert abs(np.abs(__magic_name__							).sum() - 9389.1111							) < 5E-2







'''simple docstring'''

import csv
import os

import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.csv import CsvDatasetReader, CsvDatasetWriter

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases







def   A     (      UpperCamelCase_					:     List[str]						,	UpperCamelCase_					:     Any      )					->       Optional[Any]:







			'''simple docstring'''



			assert isinstance(UpperCamelCase_						,	UpperCamelCase_      )
			assert dataset.num_rows == 4
			assert dataset.num_columns == 3
			assert dataset.column_names == ["col_1", "col_2", "col_3"]
			for feature, expected_dtype in expected_features.items():
						assert dataset.features[feature].dtype == expected_dtype







@pytest.mark.parametrize("keep_in_memory"						,	[False, True]      )
def   A     (      UpperCamelCase_					:     Optional[int]						,	UpperCamelCase_					:     Optional[Any]						,	UpperCamelCase_					:     Optional[Any]      )					->       Union[str, Any]:







			'''simple docstring'''



			lowerCAmelCase__								=							tmp_path / "cache"
			lowerCAmelCase__								=							{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
			with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
						lowerCAmelCase__								=							CsvDatasetReader(UpperCamelCase_						,	cache_dir=UpperCamelCase_						,	keep_in_memory=UpperCamelCase_      ).read()
			_check_csv_dataset(UpperCamelCase_						,	UpperCamelCase_      )







@pytest.mark.parametrize(
    "features"						,	[
        None,
        {"col_1": "string", "col_2": "int64", "col_3": "float64"},
        {"col_1": "string", "col_2": "string", "col_3": "string"},
        {"col_1": "int32", "col_2": "int32", "col_3": "int32"},
        {"col_1": "float32", "col_2": "float32", "col_3": "float32"},
    ]						,	)
def   A     (      UpperCamelCase_					:     List[Any]						,	UpperCamelCase_					:     Optional[int]						,	UpperCamelCase_					:     Tuple      )					->       Union[str, Any]:







			'''simple docstring'''



			lowerCAmelCase__								=							tmp_path / "cache"
			# CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
			lowerCAmelCase__								=							{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
			lowerCAmelCase__								=							features.copy() if features else default_expected_features
			lowerCAmelCase__								=							(
			    Features({feature: Value(UpperCamelCase_      ) for feature, dtype in features.items()}      ) if features is not None else None
			)
			lowerCAmelCase__								=							CsvDatasetReader(UpperCamelCase_						,	features=UpperCamelCase_						,	cache_dir=UpperCamelCase_      ).read()
			_check_csv_dataset(UpperCamelCase_						,	UpperCamelCase_      )







@pytest.mark.parametrize("split"						,	[None, NamedSplit("train"      ), "train", "test"]      )
def   A     (      UpperCamelCase_					:     str						,	UpperCamelCase_					:     Dict						,	UpperCamelCase_					:     str      )					->       Any:







			'''simple docstring'''



			lowerCAmelCase__								=							tmp_path / "cache"
			lowerCAmelCase__								=							{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
			lowerCAmelCase__								=							CsvDatasetReader(UpperCamelCase_						,	cache_dir=UpperCamelCase_						,	split=UpperCamelCase_      ).read()
			_check_csv_dataset(UpperCamelCase_						,	UpperCamelCase_      )
			assert dataset.split == split if split else "train"







@pytest.mark.parametrize("path_type"						,	[str, list]      )
def   A     (      UpperCamelCase_					:     Any						,	UpperCamelCase_					:     Tuple						,	UpperCamelCase_					:     List[str]      )					->       Optional[int]:







			'''simple docstring'''



			if issubclass(UpperCamelCase_						,	UpperCamelCase_      ):
						lowerCAmelCase__								=							csv_path
			elif issubclass(UpperCamelCase_						,	UpperCamelCase_      ):
						lowerCAmelCase__								=							[csv_path]
			lowerCAmelCase__								=							tmp_path / "cache"
			lowerCAmelCase__								=							{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
			lowerCAmelCase__								=							CsvDatasetReader(UpperCamelCase_						,	cache_dir=UpperCamelCase_      ).read()
			_check_csv_dataset(UpperCamelCase_						,	UpperCamelCase_      )







def   A     (      UpperCamelCase_					:     List[str]						,	UpperCamelCase_					:     Dict						,	UpperCamelCase_					:     List[str]=("train",)      )					->       Optional[int]:







			'''simple docstring'''



			assert isinstance(UpperCamelCase_						,	UpperCamelCase_      )
			for split in splits:
						lowerCAmelCase__								=							dataset_dict[split]
						assert dataset.num_rows == 4
						assert dataset.num_columns == 3
						assert dataset.column_names == ["col_1", "col_2", "col_3"]
						for feature, expected_dtype in expected_features.items():
									assert dataset.features[feature].dtype == expected_dtype







@pytest.mark.parametrize("keep_in_memory"						,	[False, True]      )
def   A     (      UpperCamelCase_					:     Optional[int]						,	UpperCamelCase_					:     List[str]						,	UpperCamelCase_					:     Optional[int]      )					->       Optional[Any]:







			'''simple docstring'''



			lowerCAmelCase__								=							tmp_path / "cache"
			lowerCAmelCase__								=							{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
			with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
						lowerCAmelCase__								=							CsvDatasetReader({"train": csv_path}						,	cache_dir=UpperCamelCase_						,	keep_in_memory=UpperCamelCase_      ).read()
			_check_csv_datasetdict(UpperCamelCase_						,	UpperCamelCase_      )







@pytest.mark.parametrize(
    "features"						,	[
        None,
        {"col_1": "string", "col_2": "int64", "col_3": "float64"},
        {"col_1": "string", "col_2": "string", "col_3": "string"},
        {"col_1": "int32", "col_2": "int32", "col_3": "int32"},
        {"col_1": "float32", "col_2": "float32", "col_3": "float32"},
    ]						,	)
def   A     (      UpperCamelCase_					:     str						,	UpperCamelCase_					:     List[Any]						,	UpperCamelCase_					:     Dict      )					->       Optional[int]:







			'''simple docstring'''



			lowerCAmelCase__								=							tmp_path / "cache"
			# CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
			lowerCAmelCase__								=							{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
			lowerCAmelCase__								=							features.copy() if features else default_expected_features
			lowerCAmelCase__								=							(
			    Features({feature: Value(UpperCamelCase_      ) for feature, dtype in features.items()}      ) if features is not None else None
			)
			lowerCAmelCase__								=							CsvDatasetReader({"train": csv_path}						,	features=UpperCamelCase_						,	cache_dir=UpperCamelCase_      ).read()
			_check_csv_datasetdict(UpperCamelCase_						,	UpperCamelCase_      )







@pytest.mark.parametrize("split"						,	[None, NamedSplit("train"      ), "train", "test"]      )
def   A     (      UpperCamelCase_					:     List[Any]						,	UpperCamelCase_					:     List[str]						,	UpperCamelCase_					:     List[Any]      )					->       Optional[int]:







			'''simple docstring'''



			if split:
						lowerCAmelCase__								=							{split: csv_path}
			else:
						lowerCAmelCase__								=							{"train": csv_path, "test": csv_path}
			lowerCAmelCase__								=							tmp_path / "cache"
			lowerCAmelCase__								=							{"col_1": "int64", "col_2": "int64", "col_3": "float64"}
			lowerCAmelCase__								=							CsvDatasetReader(UpperCamelCase_						,	cache_dir=UpperCamelCase_      ).read()
			_check_csv_datasetdict(UpperCamelCase_						,	UpperCamelCase_						,	splits=list(path.keys()      )      )
			assert all(dataset[split].split == split for split in path.keys()      )







def   A     (      UpperCamelCase_					:     str      )					->       List[str]:







			'''simple docstring'''



			with open(UpperCamelCase_						,	encoding="utf-8"      ) as csvfile:
						yield from csv.reader(UpperCamelCase_      )







def   A     (      UpperCamelCase_					:     Dict						,	UpperCamelCase_					:     int      )					->       int:







			'''simple docstring'''



			lowerCAmelCase__								=							tmp_path / "cache"
			lowerCAmelCase__								=							os.path.join(UpperCamelCase_						,	"tmp.csv"      )
			lowerCAmelCase__								=							CsvDatasetReader({"train": csv_path}						,	cache_dir=UpperCamelCase_      ).read()
			CsvDatasetWriter(dataset["train"]						,	UpperCamelCase_						,	num_proc=1      ).write()

			lowerCAmelCase__								=							iter_csv_file(UpperCamelCase_      )
			lowerCAmelCase__								=							iter_csv_file(UpperCamelCase_      )

			for rowa, rowa in zip(UpperCamelCase_						,	UpperCamelCase_      ):
						assert rowa == rowa







def   A     (      UpperCamelCase_					:     Union[str, Any]						,	UpperCamelCase_					:     Dict      )					->       List[str]:







			'''simple docstring'''



			lowerCAmelCase__								=							tmp_path / "cache"
			lowerCAmelCase__								=							os.path.join(UpperCamelCase_						,	"tmp.csv"      )
			lowerCAmelCase__								=							CsvDatasetReader({"train": csv_path}						,	cache_dir=UpperCamelCase_      ).read()
			CsvDatasetWriter(dataset["train"]						,	UpperCamelCase_						,	num_proc=2      ).write()

			lowerCAmelCase__								=							iter_csv_file(UpperCamelCase_      )
			lowerCAmelCase__								=							iter_csv_file(UpperCamelCase_      )

			for rowa, rowa in zip(UpperCamelCase_						,	UpperCamelCase_      ):
						assert rowa == rowa







def   A     (      UpperCamelCase_					:     List[str]						,	UpperCamelCase_					:     Tuple      )					->       Any:







			'''simple docstring'''



			lowerCAmelCase__								=							tmp_path / "cache"
			lowerCAmelCase__								=							os.path.join(UpperCamelCase_						,	"tmp.csv"      )
			lowerCAmelCase__								=							CsvDatasetReader({"train": csv_path}						,	cache_dir=UpperCamelCase_      ).read()
			with pytest.raises(UpperCamelCase_      ):
						CsvDatasetWriter(dataset["train"]						,	UpperCamelCase_						,	num_proc=0      )





