

'''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 					__a   (		_snake_case							,unittest.TestCase  ):
						__SCREAMING_SNAKE_CASE     :					Optional[int]  =      SpectrogramDiffusionPipeline
						__SCREAMING_SNAKE_CASE     :					Tuple  =      PipelineTesterMixin.required_optional_params - {
						    'callback',
						    'latents',
						    'callback_steps',
						    'output_type',
						    'num_images_per_prompt',
						}
						__SCREAMING_SNAKE_CASE     :					int  =      False
						__SCREAMING_SNAKE_CASE     :					List[str]  =      False
						__SCREAMING_SNAKE_CASE     :					str  =      TOKENS_TO_AUDIO_GENERATION_PARAMS
						__SCREAMING_SNAKE_CASE     :					List[Any]  =      TOKENS_TO_AUDIO_GENERATION_BATCH_PARAMS

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

							"""simple docstring"""



							torch.manual_seed(0)
							_lowercase							=     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"""    ,			)

							_lowercase							=     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"""    ,			)

							_lowercase							=     TaFilmDecoder(
							    input_dims=1_28    ,			targets_length=2_56    ,			max_decoder_noise_time=20000.0    ,			d_model=7_68    ,			num_layers=1    ,			num_heads=1    ,			d_kv=4    ,			d_ff=20_48    ,			dropout_rate=0.1    ,			)

							_lowercase							=     DDPMScheduler()

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

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

							"""simple docstring"""



							if str(lowercase__).startswith("""mps"""):
								_lowercase							=     torch.manual_seed(lowercase__)
							else:
								_lowercase							=     torch.Generator(device=lowercase__).manual_seed(lowercase__)
							_lowercase							=     {
							    """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 	_UpperCAmelCase							(		self	:						Tuple)						->Optional[int]:

							"""simple docstring"""



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

							_lowercase							=     self.get_dummy_inputs(lowercase__)
							_lowercase							=     pipe(**lowercase__)
							_lowercase							=     output.audios

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

							assert mel_slice.shape == (3, 3)
							_lowercase							=     np.array(
							    [-11.512925, -4.788215, -0.46172905, -2.051715, -10.539147, -10.970963, -9.091634, 4.0, 4.0])
							assert np.abs(mel_slice.flatten() - expected_slice).max() < 1e-2

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

							"""simple docstring"""



							return super().test_save_load_local()

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

							"""simple docstring"""



							return super().test_dict_tuple_outputs_equivalent()

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

							"""simple docstring"""



							return super().test_save_load_optional_components()

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

							"""simple docstring"""



							return super().test_attention_slicing_forward_pass()

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

							"""simple docstring"""



							pass

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

							"""simple docstring"""



							pass


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

							"""simple docstring"""



							return super().test_progress_bar()






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

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

							"""simple docstring"""



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

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

							"""simple docstring"""



							_lowercase							=     torch_device

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

							_lowercase							=     pipe.to(lowercase__)
							pipe.set_progress_bar_config(disable=lowercase__)

							def callback(lowercase__	:						Optional[int]    ,			lowercase__	:						Optional[Any]):
								# decode mel to audio
								_lowercase							=     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

							_lowercase							=     MidiProcessor()
							_lowercase							=     processor(lowercase__)

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

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

							"""simple docstring"""



							_lowercase							=     torch_device

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

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

							_lowercase							=     torch.manual_seed(0)
							_lowercase							=     pipe(lowercase__    ,			num_inference_steps=2    ,			generator=lowercase__)

							_lowercase							=     output.audios[0]

							assert abs(np.abs(lowercase__).sum() - 3612.841) < 1e-1


						def 	_UpperCAmelCase							(		self	:						str)						->Tuple:

							"""simple docstring"""



							_lowercase							=     torch_device

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

							_lowercase							=     MidiProcessor()

							_lowercase							=     processor(lowercase__)

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

							_lowercase							=     torch.manual_seed(0)
							_lowercase							=     pipe(lowercase__    ,			num_inference_steps=1_00    ,			generator=lowercase__)

							_lowercase							=     output.audios[0]
							assert abs(np.abs(lowercase__).sum() - 9389.1111) < 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,
							)




