
"""simple docstring"""





import argparse
import json
from pathlib import Path

import torch
import torchaudio
from datasets import load_dataset
from huggingface_hub import hf_hub_download

from transformers import ASTConfig, ASTFeatureExtractor, ASTForAudioClassification
from transformers.utils import logging


logging.set_verbosity_info()
UpperCamelCase_       :    Any						      =			logging.get_logger(__name__)
def 				A_  (__a					):

					'''simple docstring'''





					A_								=    ASTConfig()

					if "10-10" in model_name:
										pass
					elif "speech-commands" in model_name:
										A_								=    128
					elif "12-12" in model_name:
										A_								=    12
										A_								=    12
					elif "14-14" in model_name:
										A_								=    14
										A_								=    14
					elif "16-16" in model_name:
										A_								=    16
										A_								=    16
					else:
										raise ValueError("Model not supported"					)

					A_								=    "huggingface/label-files"
					if "speech-commands" in model_name:
										A_								=    35
										A_								=    "speech-commands-v2-id2label.json"
					else:
										A_								=    527
										A_								=    "audioset-id2label.json"

					A_								=    json.load(open(hf_hub_download(__a					, __a					, repo_type="dataset"					)					, "r"					)					)
					A_								=    {int(__a					): v for k, v in idalabel.items()}
					A_								=    idalabel
					A_								=    {v: k for k, v in idalabel.items()}

					return config
def 				A_  (__a					):

					'''simple docstring'''





					if "module.v" in name:
										A_								=    name.replace("module.v"					, "audio_spectrogram_transformer"					)
					if "cls_token" in name:
										A_								=    name.replace("cls_token"					, "embeddings.cls_token"					)
					if "dist_token" in name:
										A_								=    name.replace("dist_token"					, "embeddings.distillation_token"					)
					if "pos_embed" in name:
										A_								=    name.replace("pos_embed"					, "embeddings.position_embeddings"					)
					if "patch_embed.proj" in name:
										A_								=    name.replace("patch_embed.proj"					, "embeddings.patch_embeddings.projection"					)
					# transformer blocks
					if "blocks" in name:
										A_								=    name.replace("blocks"					, "encoder.layer"					)
					if "attn.proj" in name:
										A_								=    name.replace("attn.proj"					, "attention.output.dense"					)
					if "attn" in name:
										A_								=    name.replace("attn"					, "attention.self"					)
					if "norm1" in name:
										A_								=    name.replace("norm1"					, "layernorm_before"					)
					if "norm2" in name:
										A_								=    name.replace("norm2"					, "layernorm_after"					)
					if "mlp.fc1" in name:
										A_								=    name.replace("mlp.fc1"					, "intermediate.dense"					)
					if "mlp.fc2" in name:
										A_								=    name.replace("mlp.fc2"					, "output.dense"					)
					# final layernorm
					if "audio_spectrogram_transformer.norm" in name:
										A_								=    name.replace("audio_spectrogram_transformer.norm"					, "audio_spectrogram_transformer.layernorm"					)
					# classifier head
					if "module.mlp_head.0" in name:
										A_								=    name.replace("module.mlp_head.0"					, "classifier.layernorm"					)
					if "module.mlp_head.1" in name:
										A_								=    name.replace("module.mlp_head.1"					, "classifier.dense"					)

					return name
def 				A_  (__a					, __a					):

					'''simple docstring'''





					for key in orig_state_dict.copy().keys():
										A_								=    orig_state_dict.pop(__a					)

										if "qkv" in key:
															A_								=    key.split("."					)
															A_								=    int(key_split[3]					)
															A_								=    config.hidden_size
															if "weight" in key:
																				A_								=    val[:dim, :]
																				A_								=    val[dim : dim * 2, :]
																				A_								=    val[-dim:, :]
															else:
																				A_								=    val[:dim]
																				A_								=    val[dim : dim * 2]
																				A_								=    val[-dim:]
										else:
															A_								=    val

					return orig_state_dict
def 				A_  (__a					):

					'''simple docstring'''





					A_								=    [
					    "module.v.head.weight",
					    "module.v.head.bias",
					    "module.v.head_dist.weight",
					    "module.v.head_dist.bias",
					]
					for k in ignore_keys:
										state_dict.pop(__a					, __a					)
@torch.no_grad()
def 				A_  (__a					, __a					, __a=False					):

					'''simple docstring'''





					A_								=    get_audio_spectrogram_transformer_config(__a					)

					A_								=    {
					    "ast-finetuned-audioset-10-10-0.4593": (
					        "https://www.dropbox.com/s/ca0b1v2nlxzyeb4/audioset_10_10_0.4593.pth?dl=1"
					    ),
					    "ast-finetuned-audioset-10-10-0.450": (
					        "https://www.dropbox.com/s/1tv0hovue1bxupk/audioset_10_10_0.4495.pth?dl=1"
					    ),
					    "ast-finetuned-audioset-10-10-0.448": (
					        "https://www.dropbox.com/s/6u5sikl4b9wo4u5/audioset_10_10_0.4483.pth?dl=1"
					    ),
					    "ast-finetuned-audioset-10-10-0.448-v2": (
					        "https://www.dropbox.com/s/kt6i0v9fvfm1mbq/audioset_10_10_0.4475.pth?dl=1"
					    ),
					    "ast-finetuned-audioset-12-12-0.447": (
					        "https://www.dropbox.com/s/snfhx3tizr4nuc8/audioset_12_12_0.4467.pth?dl=1"
					    ),
					    "ast-finetuned-audioset-14-14-0.443": (
					        "https://www.dropbox.com/s/z18s6pemtnxm4k7/audioset_14_14_0.4431.pth?dl=1"
					    ),
					    "ast-finetuned-audioset-16-16-0.442": (
					        "https://www.dropbox.com/s/mdsa4t1xmcimia6/audioset_16_16_0.4422.pth?dl=1"
					    ),
					    "ast-finetuned-speech-commands-v2": (
					        "https://www.dropbox.com/s/q0tbqpwv44pquwy/speechcommands_10_10_0.9812.pth?dl=1"
					    ),
					}

					# load original state_dict
					A_								=    model_name_to_url[model_name]
					A_								=    torch.hub.load_state_dict_from_url(__a					, map_location="cpu"					)
					# remove some keys
					remove_keys(__a					)
					# rename some keys
					A_								=    convert_state_dict(__a					, __a					)

					# load 🤗 model
					A_								=    ASTForAudioClassification(__a					)
					model.eval()

					model.load_state_dict(__a					)

					# verify outputs on dummy input
					# source: https://github.com/YuanGongND/ast/blob/79e873b8a54d0a3b330dd522584ff2b9926cd581/src/run.py#L62
					A_								=    -4.2677393 if "speech-commands" not in model_name else -6.845978
					A_								=    4.5689974 if "speech-commands" not in model_name else 5.5654526
					A_								=    1024 if "speech-commands" not in model_name else 128
					A_								=    ASTFeatureExtractor(mean=__a					, std=__a					, max_length=__a					)

					if "speech-commands" in model_name:
										A_								=    load_dataset("speech_commands"					, "v0.02"					, split="validation"					)
										A_								=    dataset[0]["audio"]["array"]
					else:
										A_								=    hf_hub_download(
										    repo_id="nielsr/audio-spectogram-transformer-checkpoint"					, filename="sample_audio.flac"					, repo_type="dataset"					, )

										A_			,     A_								=    torchaudio.load(__a					)
										A_								=    waveform.squeeze().numpy()

					A_								=    feature_extractor(__a					, sampling_rate=1_6000					, return_tensors="pt"					)

					# forward pass
					A_								=    model(**__a					)
					A_								=    outputs.logits

					if model_name == "ast-finetuned-audioset-10-10-0.4593":
										A_								=    torch.tensor([-0.8760, -7.0042, -8.6602]					)
					elif model_name == "ast-finetuned-audioset-10-10-0.450":
										A_								=    torch.tensor([-1.1986, -7.0903, -8.2718]					)
					elif model_name == "ast-finetuned-audioset-10-10-0.448":
										A_								=    torch.tensor([-2.6128, -8.0080, -9.4344]					)
					elif model_name == "ast-finetuned-audioset-10-10-0.448-v2":
										A_								=    torch.tensor([-1.5080, -7.4534, -8.8917]					)
					elif model_name == "ast-finetuned-audioset-12-12-0.447":
										A_								=    torch.tensor([-0.5050, -6.5833, -8.0843]					)
					elif model_name == "ast-finetuned-audioset-14-14-0.443":
										A_								=    torch.tensor([-0.3826, -7.0336, -8.2413]					)
					elif model_name == "ast-finetuned-audioset-16-16-0.442":
										A_								=    torch.tensor([-1.2113, -6.9101, -8.3470]					)
					elif model_name == "ast-finetuned-speech-commands-v2":
										A_								=    torch.tensor([6.1589, -8.0566, -8.7984]					)
					else:
										raise ValueError("Unknown model name"					)
					if not torch.allclose(logits[0, :3]					, __a					, atol=1e-4					):
										raise ValueError("Logits don't match"					)
					print("Looks ok!"					)

					if pytorch_dump_folder_path is not None:
										Path(__a					).mkdir(exist_ok=__a					)
										print(f'Saving model {model_name} to {pytorch_dump_folder_path}'					)
										model.save_pretrained(__a					)
										print(f'Saving feature extractor to {pytorch_dump_folder_path}'					)
										feature_extractor.save_pretrained(__a					)

					if push_to_hub:
										print("Pushing model and feature extractor to the hub..."					)
										model.push_to_hub(f'MIT/{model_name}'					)
										feature_extractor.push_to_hub(f'MIT/{model_name}'					)


if __name__ == "__main__":
		UpperCamelCase_       :    Dict						      =			argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '''--model_name''',
		    default='''ast-finetuned-audioset-10-10-0.4593''',
		    type=str,
		    help='''Name of the Audio Spectrogram Transformer model you\'d like to convert.''',
		)
		parser.add_argument(
		    '''--pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model directory.'''
		)
		parser.add_argument(
		    '''--push_to_hub''', action='''store_true''', help='''Whether or not to push the converted model to the 🤗 hub.'''
		)

		UpperCamelCase_       :    List[str]						      =			parser.parse_args()
		convert_audio_spectrogram_transformer_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)






"""simple docstring"""





import math
from typing import Any, Callable, List, Optional, Tuple, Union

import numpy as np
import torch

from ...models import TaFilmDecoder
from ...schedulers import DDPMScheduler
from ...utils import is_onnx_available, logging, randn_tensor


if is_onnx_available():
		from ..onnx_utils import OnnxRuntimeModel

from ..pipeline_utils import AudioPipelineOutput, DiffusionPipeline
from .continous_encoder import SpectrogramContEncoder
from .notes_encoder import SpectrogramNotesEncoder


UpperCamelCase_       :    Any						      =			logging.get_logger(__name__)  # pylint: disable=invalid-name

UpperCamelCase_       :    List[Any]						      =			256



class __lowerCAmelCase					(       _lowercase   ):




			"""simple docstring"""







			snake_case						   =					["melgan"]




			def __init__(					self    :       Dict							, _snake_case    :       SpectrogramNotesEncoder							, _snake_case    :       SpectrogramContEncoder							, _snake_case    :       TaFilmDecoder							, _snake_case    :       DDPMScheduler							, _snake_case    :       OnnxRuntimeModel if is_onnx_available() else Any							, )   ->			None:







								"""simple docstring"""
								super().__init__()

								# From MELGAN
								A_								=    math.log(1e-5						)  # Matches MelGAN training.
								A_								=    4.0  # Largest value for most examples
								A_								=    128

								self.register_modules(
								    notes_encoder=_snake_case							, continuous_encoder=_snake_case							, decoder=_snake_case							, scheduler=_snake_case							, melgan=_snake_case							, )




			def   lowerCamelCase__  (					self    :       Union[str, Any]							, _snake_case    :       Tuple							, _snake_case    :       str=(-1.0, 1.0)							, _snake_case    :       int=False						)   ->			str:







								"""simple docstring"""
								A_			,     A_								=    output_range
								if clip:
													A_								=    torch.clip(_snake_case							, self.min_value							, self.max_value						)
								# Scale to [0, 1].
								A_								=    (features - self.min_value) / (self.max_value - self.min_value)
								# Scale to [min_out, max_out].
								return zero_one * (max_out - min_out) + min_out




			def   lowerCamelCase__  (					self    :       Dict							, _snake_case    :       Tuple							, _snake_case    :       Optional[Any]=(-1.0, 1.0)							, _snake_case    :       List[str]=False						)   ->			List[str]:







								"""simple docstring"""
								A_			,     A_								=    input_range
								A_								=    torch.clip(_snake_case							, _snake_case							, _snake_case						) if clip else outputs
								# Scale to [0, 1].
								A_								=    (outputs - min_out) / (max_out - min_out)
								# Scale to [self.min_value, self.max_value].
								return zero_one * (self.max_value - self.min_value) + self.min_value




			def   lowerCamelCase__  (					self    :       Union[str, Any]							, _snake_case    :       str							, _snake_case    :       List[Any]							, _snake_case    :       Union[str, Any]						)   ->			List[str]:







								"""simple docstring"""
								A_								=    input_tokens > 0
								A_			,     A_								=    self.notes_encoder(
								    encoder_input_tokens=_snake_case							, encoder_inputs_mask=_snake_case						)

								A_			,     A_								=    self.continuous_encoder(
								    encoder_inputs=_snake_case							, encoder_inputs_mask=_snake_case						)

								return [(tokens_encoded, tokens_mask), (continuous_encoded, continuous_mask)]




			def   lowerCamelCase__  (					self    :       List[Any]							, _snake_case    :       List[Any]							, _snake_case    :       int							, _snake_case    :       Tuple						)   ->			Optional[int]:







								"""simple docstring"""
								A_								=    noise_time
								if not torch.is_tensor(_snake_case						):
													A_								=    torch.tensor([timesteps]							, dtype=torch.long							, device=input_tokens.device						)
								elif torch.is_tensor(_snake_case						) and len(timesteps.shape						) == 0:
													A_								=    timesteps[None].to(input_tokens.device						)

								# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
								A_								=    timesteps * torch.ones(input_tokens.shape[0]							, dtype=timesteps.dtype							, device=timesteps.device						)

								A_								=    self.decoder(
								    encodings_and_masks=_snake_case							, decoder_input_tokens=_snake_case							, decoder_noise_time=_snake_case						)
								return logits






			@torch.no_grad()
			def __call__(					self    :       List[Any]							, _snake_case    :       List[List[int]]							, _snake_case    :       Optional[torch.Generator] = None							, _snake_case    :       int = 100							, _snake_case    :       bool = True							, _snake_case    :       str = "numpy"							, _snake_case    :       Optional[Callable[[int, int, torch.FloatTensor], None]] = None							, _snake_case    :       int = 1							, )   ->			Union[AudioPipelineOutput, Tuple]:







								"""simple docstring"""
								if (callback_steps is None) or (
								    callback_steps is not None and (not isinstance(_snake_case							, _snake_case						) or callback_steps <= 0)
								):
													raise ValueError(
													    F'`callback_steps` has to be a positive integer but is {callback_steps} of type'
													    F' {type(_snake_case						)}.'						)

								A_								=    np.zeros([1, TARGET_FEATURE_LENGTH, self.n_dims]							, dtype=np.floataa						)
								A_								=    np.zeros([1, 0, self.n_dims]							, np.floataa						)
								A_								=    torch.ones((1, TARGET_FEATURE_LENGTH)							, dtype=_snake_case							, device=self.device						)

								for i, encoder_input_tokens in enumerate(_snake_case						):
													if i == 0:
																		A_								=    torch.from_numpy(pred_mel[:1].copy()						).to(
																		    device=self.device							, dtype=self.decoder.dtype						)
																		# The first chunk has no previous context.
																		A_								=    torch.zeros((1, TARGET_FEATURE_LENGTH)							, dtype=_snake_case							, device=self.device						)
													else:
																		# The full song pipeline does not feed in a context feature, so the mask
																		# will be all 0s after the feature converter. Because we know we're
																		# feeding in a full context chunk from the previous prediction, set it
																		# to all 1s.
																		A_								=    ones

													A_								=    self.scale_features(
													    _snake_case							, output_range=[-1.0, 1.0]							, clip=_snake_case						)

													A_								=    self.encode(
													    input_tokens=torch.IntTensor([encoder_input_tokens]						).to(device=self.device						)							, continuous_inputs=_snake_case							, continuous_mask=_snake_case							, )

													# Sample encoder_continuous_inputs shaped gaussian noise to begin loop
													A_								=    randn_tensor(
													    shape=encoder_continuous_inputs.shape							, generator=_snake_case							, device=self.device							, dtype=self.decoder.dtype							, )

													# set step values
													self.scheduler.set_timesteps(_snake_case						)

													# Denoising diffusion loop
													for j, t in enumerate(self.progress_bar(self.scheduler.timesteps						)						):
																		A_								=    self.decode(
																		    encodings_and_masks=_snake_case							, input_tokens=_snake_case							, noise_time=t / self.scheduler.config.num_train_timesteps							, )

																		# Compute previous output: x_t -> x_t-1
																		A_								=    self.scheduler.step(_snake_case							, _snake_case							, _snake_case							, generator=_snake_case						).prev_sample

													A_								=    self.scale_to_features(_snake_case							, input_range=[-1.0, 1.0]						)
													A_								=    mel[:1]
													A_								=    mel.cpu().float().numpy()

													A_								=    np.concatenate([full_pred_mel, pred_mel[:1]]							, axis=1						)

													# call the callback, if provided
													if callback is not None and i % callback_steps == 0:
																		callback(_snake_case							, _snake_case						)

													logger.info("Generated segment"							, _snake_case						)

								if output_type == "numpy" and not is_onnx_available():
													raise ValueError(
													    "Cannot return output in 'np' format if ONNX is not available. Make sure to have ONNX installed or set 'output_type' to 'mel'."						)
								elif output_type == "numpy" and self.melgan is None:
													raise ValueError(
													    "Cannot return output in 'np' format if melgan component is not defined. Make sure to define `self.melgan` or set 'output_type' to 'mel'."						)

								if output_type == "numpy":
													A_								=    self.melgan(input_features=full_pred_mel.astype(np.floataa						)						)
								else:
													A_								=    full_pred_mel

								if not return_dict:
													return (output,)

								return AudioPipelineOutput(audios=_snake_case						)




