

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


lowerCAmelCase							:  Optional[int]            =						logging.get_logger(__name__)  # pylint: disable=invalid-name

lowerCAmelCase							:  Tuple            =						256



class   __lowercase    (			UpperCAmelCase_     ):




			"""simple docstring"""







			_UpperCAmelCase       :							List[str]      							=					['''melgan''']



			def __init__(			self  :					int       ,						lowerCAmelCase__  :					SpectrogramNotesEncoder       ,						lowerCAmelCase__  :					SpectrogramContEncoder       ,						lowerCAmelCase__  :					TaFilmDecoder       ,						lowerCAmelCase__  :					DDPMScheduler       ,						lowerCAmelCase__  :					OnnxRuntimeModel if is_onnx_available() else Any       ,						):
									super().__init__()

									# From MELGAN
									SCREAMING_SNAKE_CASE_:			Dict  		=	math.log(1E-5)  # Matches MelGAN training.
									SCREAMING_SNAKE_CASE_:			str  		=	4.0  # Largest value for most examples
									SCREAMING_SNAKE_CASE_:			Any  		=	128

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



			def 	_SCREAMING_SNAKE_CASE      (			self  :					int       ,						lowerCAmelCase__  :					Dict       ,						lowerCAmelCase__  :					str=(-1.0, 1.0)       ,						lowerCAmelCase__  :					Tuple=False):
									SCREAMING_SNAKE_CASE_		,   SCREAMING_SNAKE_CASE_:			List[Any]  		=	output_range
									if clip:
															SCREAMING_SNAKE_CASE_:			Any  		=	torch.clip(lowerCAmelCase__       ,						self.min_value       ,						self.max_value)
									# Scale to [0, 1].
									SCREAMING_SNAKE_CASE_:			List[Any]  		=	(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 	_SCREAMING_SNAKE_CASE      (			self  :					Optional[int]       ,						lowerCAmelCase__  :					Optional[int]       ,						lowerCAmelCase__  :					List[Any]=(-1.0, 1.0)       ,						lowerCAmelCase__  :					Tuple=False):
									SCREAMING_SNAKE_CASE_		,   SCREAMING_SNAKE_CASE_:			Tuple  		=	input_range
									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	torch.clip(lowerCAmelCase__       ,						lowerCAmelCase__       ,						lowerCAmelCase__) if clip else outputs
									# Scale to [0, 1].
									SCREAMING_SNAKE_CASE_:			Tuple  		=	(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 	_SCREAMING_SNAKE_CASE      (			self  :					Optional[Any]       ,						lowerCAmelCase__  :					Any       ,						lowerCAmelCase__  :					Dict       ,						lowerCAmelCase__  :					List[Any]):
									SCREAMING_SNAKE_CASE_:			int  		=	input_tokens > 0
									SCREAMING_SNAKE_CASE_		,   SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	self.notes_encoder(
									    encoder_input_tokens=lowerCAmelCase__       ,						encoder_inputs_mask=lowerCAmelCase__)

									SCREAMING_SNAKE_CASE_		,   SCREAMING_SNAKE_CASE_:			Any  		=	self.continuous_encoder(
									    encoder_inputs=lowerCAmelCase__       ,						encoder_inputs_mask=lowerCAmelCase__)

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



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Any       ,						lowerCAmelCase__  :					str       ,						lowerCAmelCase__  :					Tuple       ,						lowerCAmelCase__  :					str):
									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	noise_time
									if not torch.is_tensor(lowerCAmelCase__):
															SCREAMING_SNAKE_CASE_:			Any  		=	torch.tensor([timesteps]       ,						dtype=torch.long       ,						device=input_tokens.device)
									elif torch.is_tensor(lowerCAmelCase__) and len(timesteps.shape) == 0:
															SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	timesteps[None].to(input_tokens.device)

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

									SCREAMING_SNAKE_CASE_:			str  		=	self.decoder(
									    encodings_and_masks=lowerCAmelCase__       ,						decoder_input_tokens=lowerCAmelCase__       ,						decoder_noise_time=lowerCAmelCase__)
									return logits




			@torch.no_grad()
			def __call__(			self  :					Dict       ,						lowerCAmelCase__  :					List[List[int]]       ,						lowerCAmelCase__  :					Optional[torch.Generator] = None       ,						lowerCAmelCase__  :					int = 100       ,						lowerCAmelCase__  :					bool = True       ,						lowerCAmelCase__  :					str = "numpy"       ,						lowerCAmelCase__  :					Optional[Callable[[int, int, torch.FloatTensor], None]] = None       ,						lowerCAmelCase__  :					int = 1       ,						):
									if (callback_steps is None) or (
									    callback_steps is not None and (not isinstance(lowerCAmelCase__       ,						lowerCAmelCase__) or callback_steps <= 0)
									):
															raise ValueError(
															    F"`callback_steps` has to be a positive integer but is {callback_steps} of type"
															    F" {type(lowerCAmelCase__)}.")

									SCREAMING_SNAKE_CASE_:			List[Any]  		=	np.zeros([1, TARGET_FEATURE_LENGTH, self.n_dims]       ,						dtype=np.floataa)
									SCREAMING_SNAKE_CASE_:			Tuple  		=	np.zeros([1, 0, self.n_dims]       ,						np.floataa)
									SCREAMING_SNAKE_CASE_:			List[str]  		=	torch.ones((1, TARGET_FEATURE_LENGTH)       ,						dtype=lowerCAmelCase__       ,						device=self.device)

									for i, encoder_input_tokens in enumerate(lowerCAmelCase__):
															if i == 0:
																					SCREAMING_SNAKE_CASE_:			Any  		=	torch.from_numpy(pred_mel[:1].copy()).to(
																					    device=self.device       ,						dtype=self.decoder.dtype)
																					# The first chunk has no previous context.
																					SCREAMING_SNAKE_CASE_:			int  		=	torch.zeros((1, TARGET_FEATURE_LENGTH)       ,						dtype=lowerCAmelCase__       ,						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.
																					SCREAMING_SNAKE_CASE_:			Optional[int]  		=	ones

															SCREAMING_SNAKE_CASE_:			str  		=	self.scale_features(
															    lowerCAmelCase__       ,						output_range=[-1.0, 1.0]       ,						clip=lowerCAmelCase__)

															SCREAMING_SNAKE_CASE_:			List[Any]  		=	self.encode(
															    input_tokens=torch.IntTensor([encoder_input_tokens]).to(device=self.device)       ,						continuous_inputs=lowerCAmelCase__       ,						continuous_mask=lowerCAmelCase__       ,						)

															# Sample encoder_continuous_inputs shaped gaussian noise to begin loop
															SCREAMING_SNAKE_CASE_:			Any  		=	randn_tensor(
															    shape=encoder_continuous_inputs.shape       ,						generator=lowerCAmelCase__       ,						device=self.device       ,						dtype=self.decoder.dtype       ,						)

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

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

																					# Compute previous output: x_t -> x_t-1
																					SCREAMING_SNAKE_CASE_:			int  		=	self.scheduler.step(lowerCAmelCase__       ,						lowerCAmelCase__       ,						lowerCAmelCase__       ,						generator=lowerCAmelCase__).prev_sample

															SCREAMING_SNAKE_CASE_:			str  		=	self.scale_to_features(lowerCAmelCase__       ,						input_range=[-1.0, 1.0])
															SCREAMING_SNAKE_CASE_:			Tuple  		=	mel[:1]
															SCREAMING_SNAKE_CASE_:			str  		=	mel.cpu().float().numpy()

															SCREAMING_SNAKE_CASE_:			List[str]  		=	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(lowerCAmelCase__       ,						lowerCAmelCase__)

															logger.info("Generated segment"       ,						lowerCAmelCase__)

									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":
															SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	self.melgan(input_features=full_pred_mel.astype(np.floataa))
									else:
															SCREAMING_SNAKE_CASE_:			int  		=	full_pred_mel

									if not return_dict:
															return (output,)

									return AudioPipelineOutput(audios=lowerCAmelCase__)









import importlib
import json
import os
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
import transformers.models.auto
from transformers.models.auto.configuration_auto import CONFIG_MAPPING, AutoConfig
from transformers.models.bert.configuration_bert import BertConfig
from transformers.models.roberta.configuration_roberta import RobertaConfig
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER, get_tests_dir


sys.path.append(str(Path(__file__).parent.parent.parent.parent / """utils"""))

from test_module.custom_configuration import CustomConfig  # noqa E402


lowerCAmelCase							:  Union[str, Any]            =						get_tests_dir("""fixtures/dummy-config.json""")



class   __lowercase    (			unittest.TestCase     ):




			"""simple docstring"""







			def 	_SCREAMING_SNAKE_CASE      (			self  :					Tuple):
									SCREAMING_SNAKE_CASE_:			Dict  		=	0



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Any):
									self.assertIsNotNone(transformers.models.auto.__spec__)
									self.assertIsNotNone(importlib.util.find_spec("transformers.models.auto"))



			def 	_SCREAMING_SNAKE_CASE      (			self  :					List[Any]):
									SCREAMING_SNAKE_CASE_:			List[Any]  		=	AutoConfig.from_pretrained("bert-base-uncased")
									self.assertIsInstance(lowerCAmelCase__       ,						lowerCAmelCase__)



			def 	_SCREAMING_SNAKE_CASE      (			self  :					str):
									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	AutoConfig.from_pretrained(lowerCAmelCase__)
									self.assertIsInstance(lowerCAmelCase__       ,						lowerCAmelCase__)



			def 	_SCREAMING_SNAKE_CASE      (			self  :					int):
									SCREAMING_SNAKE_CASE_:			Tuple  		=	AutoConfig.from_pretrained(lowerCAmelCase__)
									self.assertIsInstance(lowerCAmelCase__       ,						lowerCAmelCase__)



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Union[str, Any]):
									SCREAMING_SNAKE_CASE_:			List[Any]  		=	AutoConfig.for_model("roberta")
									self.assertIsInstance(lowerCAmelCase__       ,						lowerCAmelCase__)



			def 	_SCREAMING_SNAKE_CASE      (			self  :					List[Any]):
									with tempfile.TemporaryDirectory() as tmp_dir:
															# This model name contains bert and roberta, but roberta ends up being picked.
															SCREAMING_SNAKE_CASE_:			int  		=	os.path.join(lowerCAmelCase__       ,						"fake-roberta")
															os.makedirs(lowerCAmelCase__       ,						exist_ok=lowerCAmelCase__)
															with open(os.path.join(lowerCAmelCase__       ,						"config.json")       ,						"w") as f:
																					f.write(json.dumps({}))
															SCREAMING_SNAKE_CASE_:			Any  		=	AutoConfig.from_pretrained(lowerCAmelCase__)
															self.assertEqual(type(lowerCAmelCase__)       ,						lowerCAmelCase__)



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Optional[int]):
									try:
															AutoConfig.register("custom"       ,						lowerCAmelCase__)
															# Wrong model type will raise an error
															with self.assertRaises(lowerCAmelCase__):
																					AutoConfig.register("model"       ,						lowerCAmelCase__)
															# Trying to register something existing in the Transformers library will raise an error
															with self.assertRaises(lowerCAmelCase__):
																					AutoConfig.register("bert"       ,						lowerCAmelCase__)

															# Now that the config is registered, it can be used as any other config with the auto-API
															SCREAMING_SNAKE_CASE_:			List[Any]  		=	CustomConfig()
															with tempfile.TemporaryDirectory() as tmp_dir:
																					config.save_pretrained(lowerCAmelCase__)
																					SCREAMING_SNAKE_CASE_:			Tuple  		=	AutoConfig.from_pretrained(lowerCAmelCase__)
																					self.assertIsInstance(lowerCAmelCase__       ,						lowerCAmelCase__)

									finally:
															if "custom" in CONFIG_MAPPING._extra_content:
																					del CONFIG_MAPPING._extra_content["custom"]



			def 	_SCREAMING_SNAKE_CASE      (			self  :					List[str]):
									with self.assertRaisesRegex(
									    lowerCAmelCase__       ,						"bert-base is not a local folder and is not a valid model identifier"):
															SCREAMING_SNAKE_CASE_:			List[str]  		=	AutoConfig.from_pretrained("bert-base")



			def 	_SCREAMING_SNAKE_CASE      (			self  :					List[Any]):
									with self.assertRaisesRegex(
									    lowerCAmelCase__       ,						R"aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)"):
															SCREAMING_SNAKE_CASE_:			str  		=	AutoConfig.from_pretrained(lowerCAmelCase__       ,						revision="aaaaaa")



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Optional[int]):
									with self.assertRaisesRegex(
									    lowerCAmelCase__       ,						"hf-internal-testing/no-config-test-repo does not appear to have a file named config.json."       ,						):
															SCREAMING_SNAKE_CASE_:			int  		=	AutoConfig.from_pretrained("hf-internal-testing/no-config-test-repo")



			def 	_SCREAMING_SNAKE_CASE      (			self  :					List[str]):
									# If remote code is not set, we will time out when asking whether to load the model.
									with self.assertRaises(lowerCAmelCase__):
															SCREAMING_SNAKE_CASE_:			str  		=	AutoConfig.from_pretrained("hf-internal-testing/test_dynamic_model")
									# If remote code is disabled, we can't load this config.
									with self.assertRaises(lowerCAmelCase__):
															SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	AutoConfig.from_pretrained("hf-internal-testing/test_dynamic_model"       ,						trust_remote_code=lowerCAmelCase__)

									SCREAMING_SNAKE_CASE_:			str  		=	AutoConfig.from_pretrained("hf-internal-testing/test_dynamic_model"       ,						trust_remote_code=lowerCAmelCase__)
									self.assertEqual(config.__class__.__name__       ,						"NewModelConfig")

									# Test config can be reloaded.
									with tempfile.TemporaryDirectory() as tmp_dir:
															config.save_pretrained(lowerCAmelCase__)
															SCREAMING_SNAKE_CASE_:			Any  		=	AutoConfig.from_pretrained(lowerCAmelCase__       ,						trust_remote_code=lowerCAmelCase__)
									self.assertEqual(reloaded_config.__class__.__name__       ,						"NewModelConfig")




			def 	_SCREAMING_SNAKE_CASE      (			self  :					List[Any]):



									class   __lowercase    (			UpperCAmelCase_     ):




												"""simple docstring"""







												_UpperCAmelCase       :							int      							=					'''new-model'''

									try:
															AutoConfig.register("new-model"       ,						lowerCAmelCase__)
															# If remote code is not set, the default is to use local
															SCREAMING_SNAKE_CASE_:			int  		=	AutoConfig.from_pretrained("hf-internal-testing/test_dynamic_model")
															self.assertEqual(config.__class__.__name__       ,						"NewModelConfigLocal")

															# If remote code is disabled, we load the local one.
															SCREAMING_SNAKE_CASE_:			Optional[int]  		=	AutoConfig.from_pretrained("hf-internal-testing/test_dynamic_model"       ,						trust_remote_code=lowerCAmelCase__)
															self.assertEqual(config.__class__.__name__       ,						"NewModelConfigLocal")

															# If remote is enabled, we load from the Hub
															SCREAMING_SNAKE_CASE_:			int  		=	AutoConfig.from_pretrained("hf-internal-testing/test_dynamic_model"       ,						trust_remote_code=lowerCAmelCase__)
															self.assertEqual(config.__class__.__name__       ,						"NewModelConfig")

									finally:
															if "new-model" in CONFIG_MAPPING._extra_content:
																					del CONFIG_MAPPING._extra_content["new-model"]






