


'''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()


lowerCamelCase__     						=  './tests/fixtures/elise_format0.mid'






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




			def 				lowerCAmelCase__     (		self	:  Dict			) ->Union[str, Any]:



					'''simple docstring'''







					torch.manual_seed(0			)
					_UpperCAmelCase		:							Optional[int]             =   SpectrogramNotesEncoder(
					    max_length=20_48						,	vocab_size=15_36						,	d_model=7_68						,	dropout_rate=0.1						,	num_layers=1						,	num_heads=1						,	d_kv=4						,	d_ff=20_48						,	feed_forward_proj="gated-gelu"						,	)

					_UpperCAmelCase		:							Optional[Any]             =   SpectrogramContEncoder(
					    input_dims=1_28						,	targets_context_length=2_56						,	d_model=7_68						,	dropout_rate=0.1						,	num_layers=1						,	num_heads=1						,	d_kv=4						,	d_ff=20_48						,	feed_forward_proj="gated-gelu"						,	)

					_UpperCAmelCase		:							Union[str, Any]             =   TaFilmDecoder(
					    input_dims=1_28						,	targets_length=2_56						,	max_decoder_noise_time=2_0_0_0_0.0						,	d_model=7_68						,	num_layers=1						,	num_heads=1						,	d_kv=4						,	d_ff=20_48						,	dropout_rate=0.1						,	)

					_UpperCAmelCase		:							Optional[Any]             =   DDPMScheduler()

					_UpperCAmelCase		:							Optional[Any]             =   {
					    "notes_encoder": notes_encoder.eval(),
					    "continuous_encoder": continuous_encoder.eval(),
					    "decoder": decoder.eval(),
					    "scheduler": scheduler,
					    "melgan": None,
					}
					return components




			def 				lowerCAmelCase__     (		self	:  int						,	lowerCamelCase__	:  Dict						,	lowerCamelCase__	:  List[str]=0			) ->Dict:



					'''simple docstring'''







					if str(lowercase__			).startswith("mps"			):
							_UpperCAmelCase		:							Any             =   torch.manual_seed(lowercase__			)
					else:
							_UpperCAmelCase		:							List[str]             =   torch.Generator(device=lowercase__			).manual_seed(lowercase__			)
					_UpperCAmelCase		:							int             =   {
					    "input_tokens": [
					        [11_34, 90, 11_35, 11_33, 10_80, 1_12, 11_32, 10_80, 11_33, 10_79, 1_33, 11_32, 10_79, 11_33, 1] + [0] * 20_33
					    ],
					    "generator": generator,
					    "num_inference_steps": 4,
					    "output_type": "mel",
					}
					return inputs




			def 				lowerCAmelCase__     (		self	:  Tuple			) ->Optional[int]:



					'''simple docstring'''







					_UpperCAmelCase		:							Optional[Any]             =   "cpu"  # ensure determinism for the device-dependent torch.Generator
					_UpperCAmelCase		:							int             =   self.get_dummy_components()
					_UpperCAmelCase		:							List[str]             =   SpectrogramDiffusionPipeline(**lowercase__			)
					_UpperCAmelCase		:							List[Any]             =   pipe.to(lowercase__			)
					pipe.set_progress_bar_config(disable=lowercase__			)

					_UpperCAmelCase		:							Union[str, Any]             =   self.get_dummy_inputs(lowercase__			)
					_UpperCAmelCase		:							Union[str, Any]             =   pipe(**lowercase__			)
					_UpperCAmelCase		:							Tuple             =   output.audios

					_UpperCAmelCase		:							Tuple             =   mel[0, -3:, -3:]

					assert mel_slice.shape == (3, 3)
					_UpperCAmelCase		:							Optional[Any]             =   np.array(
					    [-1_1.5_1_2_9_2_5, -4.7_8_8_2_1_5, -0.4_6_1_7_2_9_0_5, -2.0_5_1_7_1_5, -1_0.5_3_9_1_4_7, -1_0.9_7_0_9_6_3, -9.0_9_1_6_3_4, 4.0, 4.0]			)
					assert np.abs(mel_slice.flatten() - expected_slice			).max() < 1E-2




			@skip_mps
			def 				lowerCAmelCase__     (		self	:  Tuple			) ->Tuple:



					'''simple docstring'''







					return super().test_save_load_local()




			@skip_mps
			def 				lowerCAmelCase__     (		self	:  Any			) ->Dict:



					'''simple docstring'''







					return super().test_dict_tuple_outputs_equivalent()




			@skip_mps
			def 				lowerCAmelCase__     (		self	:  List[str]			) ->Union[str, Any]:



					'''simple docstring'''







					return super().test_save_load_optional_components()




			@skip_mps
			def 				lowerCAmelCase__     (		self	:  Optional[Any]			) ->Optional[int]:



					'''simple docstring'''







					return super().test_attention_slicing_forward_pass()




			def 				lowerCAmelCase__     (		self	:  int			) ->Optional[Any]:



					'''simple docstring'''







					pass




			def 				lowerCAmelCase__     (		self	:  Tuple			) ->Optional[Any]:



					'''simple docstring'''







					pass






			@skip_mps
			def 				lowerCAmelCase__     (		self	:  Optional[int]			) ->Any:



					'''simple docstring'''







					return super().test_progress_bar()






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




			def 				lowerCAmelCase__     (		self	:  Optional[Any]			) ->Optional[int]:



					'''simple docstring'''







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




			def 				lowerCAmelCase__     (		self	:  str			) ->List[str]:



					'''simple docstring'''







					_UpperCAmelCase		:							str             =   torch_device

					_UpperCAmelCase		:							Optional[int]             =   SpectrogramDiffusionPipeline.from_pretrained("google/music-spectrogram-diffusion"			)
					_UpperCAmelCase		:							List[str]             =   pipe.melgan
					_UpperCAmelCase		:							int             =   None

					_UpperCAmelCase		:							Any             =   pipe.to(lowercase__			)
					pipe.set_progress_bar_config(disable=lowercase__			)

					def callback(lowerCamelCase__	:  Optional[int]						,	lowerCamelCase__	:  Optional[Any]			):
							# decode mel to audio
							_UpperCAmelCase		:							Optional[int]             =   melgan(input_features=mel_output.astype(np.floataa			)			)[0]
							assert len(audio[0]			) == 8_19_20 * (step + 1)
							# simulate that audio is played
							return audio

					_UpperCAmelCase		:							Tuple             =   MidiProcessor()
					_UpperCAmelCase		:							Union[str, Any]             =   processor(lowercase__			)

					_UpperCAmelCase		:							List[str]             =   input_tokens[:3]
					_UpperCAmelCase		:							Any             =   torch.manual_seed(0			)
					pipe(lowercase__						,	num_inference_steps=5						,	generator=lowercase__						,	callback=lowercase__						,	output_type="mel"			)




			def 				lowerCAmelCase__     (		self	:  List[Any]			) ->Tuple:



					'''simple docstring'''







					_UpperCAmelCase		:							Optional[int]             =   torch_device

					_UpperCAmelCase		:							List[str]             =   SpectrogramDiffusionPipeline.from_pretrained("google/music-spectrogram-diffusion"			)
					_UpperCAmelCase		:							Dict             =   pipe.to(lowercase__			)
					pipe.set_progress_bar_config(disable=lowercase__			)
					_UpperCAmelCase		:							Any             =   MidiProcessor()

					_UpperCAmelCase		:							str             =   processor(lowercase__			)
					# just run two denoising loops
					_UpperCAmelCase		:							Dict             =   input_tokens[:2]

					_UpperCAmelCase		:							Tuple             =   torch.manual_seed(0			)
					_UpperCAmelCase		:							Union[str, Any]             =   pipe(lowercase__						,	num_inference_steps=2						,	generator=lowercase__			)

					_UpperCAmelCase		:							str             =   output.audios[0]

					assert abs(np.abs(lowercase__			).sum() - 3_6_1_2.8_4_1			) < 1E-1






			def 				lowerCAmelCase__     (		self	:  str			) ->Tuple:



					'''simple docstring'''







					_UpperCAmelCase		:							Dict             =   torch_device

					_UpperCAmelCase		:							Tuple             =   SpectrogramDiffusionPipeline.from_pretrained("google/music-spectrogram-diffusion"			)
					_UpperCAmelCase		:							Optional[int]             =   pipe.to(lowercase__			)
					pipe.set_progress_bar_config(disable=lowercase__			)

					_UpperCAmelCase		:							str             =   MidiProcessor()

					_UpperCAmelCase		:							List[str]             =   processor(lowercase__			)

					# just run 4 denoising loops
					_UpperCAmelCase		:							List[Any]             =   input_tokens[:4]

					_UpperCAmelCase		:							Optional[int]             =   torch.manual_seed(0			)
					_UpperCAmelCase		:							int             =   pipe(lowercase__						,	num_inference_steps=1_00						,	generator=lowercase__			)

					_UpperCAmelCase		:							int             =   output.audios[0]
					assert abs(np.abs(lowercase__			).sum() - 9_3_8_9.1_1_1_1			) < 5E-2






'''simple docstring'''







import argparse
import copy
import json
import multiprocessing
import os
import time

from create_dummy_models import COMPOSITE_MODELS, create_tiny_models
from huggingface_hub import ModelFilter, hf_api

import transformers
from transformers import AutoFeatureExtractor, AutoImageProcessor, AutoTokenizer
from transformers.image_processing_utils import BaseImageProcessor






def 					_SCREAMING_SNAKE_CASE		(						):
	_lowercase							=     set()
	# Each auto modeling files contains multiple mappings. Let's get them in a dynamic way.
	for module_name in ["modeling_auto", "modeling_tf_auto", "modeling_flax_auto"]:
		_lowercase							=     getattr(transformers.models.auto							,       snake_case_							,       snake_case_	)
		if module is None:
			continue
		# all mappings in a single auto modeling file
		_lowercase							=     [
		    x
		    for x in dir(snake_case_	)
		    if x.endswith("""_MAPPING_NAMES"""	)
		    and (x.startswith("""MODEL_"""	) or x.startswith("""TF_MODEL_"""	) or x.startswith("""FLAX_MODEL_"""	))
		]
		for name in mapping_names:
			_lowercase							=     getattr(snake_case_							,       snake_case_	)
			if mapping is not None:
				for v in mapping.values():
					if isinstance(snake_case_							,       (list, tuple)	):
						model_names.update(snake_case_	)
					elif isinstance(snake_case_							,       snake_case_	):
						model_names.add(snake_case_	)

	return sorted(snake_case_	)






def 					_SCREAMING_SNAKE_CASE		(						):
	# All model names defined in auto mappings
	_lowercase							=     set(get_all_model_names()	)

	with open("""tests/utils/tiny_model_summary.json"""	) as fp:
		_lowercase							=     json.load(snake_case_	)
	_lowercase							=     set()
	for model_base_name in tiny_model_info:
		tiny_models_names.update(tiny_model_info[model_base_name]["""model_classes"""]	)

	# Remove a tiny model name if one of its framework implementation hasn't yet a tiny version on the Hub.
	_lowercase							=     model_names.difference(snake_case_	)
	for model_name in copy.copy(snake_case_	):
		if not model_name.startswith("""TF"""	) and F"""TF{model_name}""" in not_on_hub:
			tiny_models_names.remove(snake_case_	)
		elif model_name.startswith("""TF"""	) and model_name[2:] in not_on_hub:
			tiny_models_names.remove(snake_case_	)

	return sorted(snake_case_	)






def 					_SCREAMING_SNAKE_CASE		(						snake_case_	):
	_lowercase							=     COMPOSITE_MODELS.values()

	# All tiny model base names on Hub
	_lowercase							=     get_all_model_names()
	_lowercase							=     hf_api.list_models(
	    filter=ModelFilter(
	        author="""hf-internal-testing"""							,       )	)
	_lowercase							=     set()
	for x in models:
		_lowercase							=     x.modelId
		_lowercase		,		_lowercase							=     model.split("""/"""	)
		if not model.startswith("""tiny-random-"""	):
			continue
		_lowercase							=     model.replace("""tiny-random-"""							,       """"""	)
		if not model[0].isupper():
			continue
		if model not in model_names and model not in special_models:
			continue
		_models.add(snake_case_	)

	_lowercase							=     sorted(_models	)
	# All tiny model names on Hub
	_lowercase							=     {}
	for model in models:
		_lowercase							=     F"""hf-internal-testing/tiny-random-{model}"""
		_lowercase							=     model.split("""-"""	)[0]
		try:
			_lowercase							=     hf_api.repo_info(snake_case_	)
			_lowercase							=     {
			    """tokenizer_classes""": set(),
			    """processor_classes""": set(),
			    """model_classes""": set(),
			    """sha""": repo_info.sha,
			}
		except Exception:
			continue
		try:
			time.sleep(1	)
			_lowercase							=     AutoTokenizer.from_pretrained(snake_case_	)
			content["tokenizer_classes"].add(tokenizer_fast.__class__.__name__	)
		except Exception:
			pass
		try:
			time.sleep(1	)
			_lowercase							=     AutoTokenizer.from_pretrained(snake_case_							,       use_fast=snake_case_	)
			content["tokenizer_classes"].add(tokenizer_slow.__class__.__name__	)
		except Exception:
			pass
		try:
			time.sleep(1	)
			_lowercase							=     AutoImageProcessor.from_pretrained(snake_case_	)
			content["processor_classes"].add(img_p.__class__.__name__	)
		except Exception:
			pass
		try:
			time.sleep(1	)
			_lowercase							=     AutoFeatureExtractor.from_pretrained(snake_case_	)
			if not isinstance(snake_case_							,       snake_case_	):
				content["processor_classes"].add(feat_p.__class__.__name__	)
		except Exception:
			pass
		try:
			time.sleep(1	)
			_lowercase							=     getattr(snake_case_							,       snake_case_	)
			_lowercase							=     model_class.from_pretrained(snake_case_	)
			content["model_classes"].add(m.__class__.__name__	)
		except Exception:
			pass
		try:
			time.sleep(1	)
			_lowercase							=     getattr(snake_case_							,       F"""TF{model}"""	)
			_lowercase							=     model_class.from_pretrained(snake_case_	)
			content["model_classes"].add(m.__class__.__name__	)
		except Exception:
			pass

		_lowercase							=     sorted(content["""tokenizer_classes"""]	)
		_lowercase							=     sorted(content["""processor_classes"""]	)
		_lowercase							=     sorted(content["""model_classes"""]	)

		_lowercase							=     content
		with open(os.path.join(snake_case_							,       """hub_tiny_model_summary.json"""	)							,       """w"""	) as fp:
			json.dump(snake_case_							,       snake_case_							,       ensure_ascii=snake_case_							,       indent=4	)


if __name__ == "__main__":
							_lowerCamelCase    =  argparse.ArgumentParser()
							parser.add_argument('--num_workers', default=1, type=int, help='The number of workers to run.')
							_lowerCamelCase    =  parser.parse_args()

							# This has to be `spawn` to avoid hanging forever!
							multiprocessing.set_start_method('spawn')

							_lowerCamelCase    =  'tiny_models'
							_lowerCamelCase    =  True
							_lowerCamelCase    =  None
							_lowerCamelCase    =  get_tiny_model_names_from_repo()
							_lowerCamelCase    =  True
							_lowerCamelCase    =  True
							_lowerCamelCase    =  'hf-internal-testing'

							create_tiny_models(
							    output_path,
							    all,
							    model_types,
							    models_to_skip,
							    no_check,
							    upload,
							    organization,
							    token=os.environ.get('TOKEN', None),
							    num_workers=args.num_workers,
							)




