


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			:  Union[str, Any]     =logging.get_logger(__name__)  # pylint: disable=invalid-name

_UpperCAmelCase			:  Optional[Any]     =256







class 				snake_case__(							UpperCAmelCase__							):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__      :  List[Any]					=							["""melgan"""]




						def __init__(				self   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						__lowercase   ,						)					->    None:
										super().__init__()

										# From MELGAN
										lowerCAmelCase_ :      Optional[Any]    				=				math.log(1e-5							)  # Matches MelGAN training.
										lowerCAmelCase_ :      List[str]    				=				4.0  # Largest value for most examples
										lowerCAmelCase_ :      Union[str, Any]    				=				1_2_8

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




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

										lowerCAmelCase_				,  lowerCAmelCase_ :      Union[str, Any]    				=				self.continuous_encoder(
										    encoder_inputs=_SCREAMING_SNAKE_CASE   ,						encoder_inputs_mask=_SCREAMING_SNAKE_CASE							)

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




						def 						lowercase_ (				self   ,						__lowercase   ,						__lowercase   ,						__lowercase							)					->    str:
										lowerCAmelCase_ :      Optional[int]    				=				noise_time
										if not torch.is_tensor(_SCREAMING_SNAKE_CASE							):
														lowerCAmelCase_ :      Dict    				=				torch.tensor([timesteps]   ,						dtype=torch.long   ,						device=input_tokens.device							)
										elif torch.is_tensor(_SCREAMING_SNAKE_CASE							) and len(timesteps.shape							) == 0:
														lowerCAmelCase_ :      str    				=				timesteps[None].to(input_tokens.device							)

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

										lowerCAmelCase_ :      Dict    				=				self.decoder(
										    encodings_and_masks=_SCREAMING_SNAKE_CASE   ,						decoder_input_tokens=_SCREAMING_SNAKE_CASE   ,						decoder_noise_time=_SCREAMING_SNAKE_CASE							)
										return logits




						@torch.no_grad()
						def __call__(				self   ,						__lowercase   ,						__lowercase = None   ,						__lowercase = 1_0_0   ,						__lowercase = True   ,						__lowercase = "numpy"   ,						__lowercase = None   ,						__lowercase = 1   ,						)					->    Union[AudioPipelineOutput, Tuple]:






										if (callback_steps is None) or (
										    callback_steps is not None and (not isinstance(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_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(_SCREAMING_SNAKE_CASE							)}."""							)

										lowerCAmelCase_ :      Any    				=				np.zeros([1, TARGET_FEATURE_LENGTH, self.n_dims]   ,						dtype=np.floataa							)
										lowerCAmelCase_ :      Any    				=				np.zeros([1, 0, self.n_dims]   ,						np.floataa							)
										lowerCAmelCase_ :      int    				=				torch.ones((1, TARGET_FEATURE_LENGTH)   ,						dtype=_SCREAMING_SNAKE_CASE   ,						device=self.device							)

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

														lowerCAmelCase_ :      Any    				=				self.scale_features(
														    _SCREAMING_SNAKE_CASE   ,						output_range=[-1.0, 1.0]   ,						clip=_SCREAMING_SNAKE_CASE							)

														lowerCAmelCase_ :      Optional[Any]    				=				self.encode(
														    input_tokens=torch.IntTensor([encoder_input_tokens]							).to(device=self.device							)   ,						continuous_inputs=_SCREAMING_SNAKE_CASE   ,						continuous_mask=_SCREAMING_SNAKE_CASE   ,						)

														# Sample encoder_continuous_inputs shaped gaussian noise to begin loop
														lowerCAmelCase_ :      Dict    				=				randn_tensor(
														    shape=encoder_continuous_inputs.shape   ,						generator=_SCREAMING_SNAKE_CASE   ,						device=self.device   ,						dtype=self.decoder.dtype   ,						)

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

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

																		# Compute previous output: x_t -> x_t-1
																		lowerCAmelCase_ :      List[str]    				=				self.scheduler.step(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						generator=_SCREAMING_SNAKE_CASE							).prev_sample

														lowerCAmelCase_ :      Optional[int]    				=				self.scale_to_features(_SCREAMING_SNAKE_CASE   ,						input_range=[-1.0, 1.0]							)
														lowerCAmelCase_ :      Any    				=				mel[:1]
														lowerCAmelCase_ :      int    				=				mel.cpu().float().numpy()

														lowerCAmelCase_ :      Any    				=				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(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE							)

														logger.info('''Generated segment'''   ,						_SCREAMING_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":
														lowerCAmelCase_ :      Tuple    				=				self.melgan(input_features=full_pred_mel.astype(np.floataa							)							)
										else:
														lowerCAmelCase_ :      str    				=				full_pred_mel

										if not return_dict:
														return (output,)

										return AudioPipelineOutput(audios=_SCREAMING_SNAKE_CASE							)



import math
from enum import Enum
from typing import Optional, Union

from torch.optim import Optimizer
from torch.optim.lr_scheduler import LambdaLR

from .utils import logging


_UpperCAmelCase			:  Tuple     =logging.get_logger(__name__)







class 				snake_case__(							UpperCAmelCase__							):







						'''simple docstring'''




						SCREAMING_SNAKE_CASE__      :  List[str]					=							"""linear"""
						SCREAMING_SNAKE_CASE__      :  Union[str, Any]					=							"""cosine"""
						SCREAMING_SNAKE_CASE__      :  Dict					=							"""cosine_with_restarts"""
						SCREAMING_SNAKE_CASE__      :  List[str]					=							"""polynomial"""
						SCREAMING_SNAKE_CASE__      :  Dict					=							"""constant"""
						SCREAMING_SNAKE_CASE__      :  List[str]					=							"""constant_with_warmup"""
						SCREAMING_SNAKE_CASE__      :  str					=							"""piecewise_constant"""






def 			lowerCAmelCase (		lowerCAmelCase_ ,   lowerCAmelCase_ = -1 )-> Tuple:
				return LambdaLR(lowerCAmelCase_ ,   lambda lowerCAmelCase_ : 1 ,   last_epoch=lowerCAmelCase_ )



def 			lowerCAmelCase (		lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ = -1 )-> List[Any]:

				def lr_lambda(lowerCAmelCase_ ):
								if current_step < num_warmup_steps:
												return float(lowerCAmelCase_ ) / float(max(1.0 ,   lowerCAmelCase_ ) )
								return 1.0

				return LambdaLR(lowerCAmelCase_ ,   lowerCAmelCase_ ,   last_epoch=lowerCAmelCase_ )



def 			lowerCAmelCase (		lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ = -1 )-> int:
				lowerCAmelCase_ :      Optional[int]    				=				{}
				lowerCAmelCase_ :      Union[str, Any]    				=				step_rules.split(''',''' )
				for rule_str in rule_list[:-1]:
								lowerCAmelCase_				,  lowerCAmelCase_ :      str    				=				rule_str.split(''':''' )
								lowerCAmelCase_ :      int    				=				int(lowerCAmelCase_ )
								lowerCAmelCase_ :      str    				=				float(lowerCAmelCase_ )
								lowerCAmelCase_ :      List[Any]    				=				value
				lowerCAmelCase_ :      int    				=				float(rule_list[-1] )

				def create_rules_function(lowerCAmelCase_ ,   lowerCAmelCase_ ):
								def rule_func(lowerCAmelCase_ ) -> float:
												lowerCAmelCase_ :      Tuple    				=				sorted(rules_dict.keys() )
												for i, sorted_step in enumerate(lowerCAmelCase_ ):
																if steps < sorted_step:
																				return rules_dict[sorted_steps[i]]
												return last_lr_multiple

								return rule_func

				lowerCAmelCase_ :      Tuple    				=				create_rules_function(lowerCAmelCase_ ,   lowerCAmelCase_ )

				return LambdaLR(lowerCAmelCase_ ,   lowerCAmelCase_ ,   last_epoch=lowerCAmelCase_ )



def 			lowerCAmelCase (		lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_=-1 )-> Optional[int]:

				def lr_lambda(lowerCAmelCase_ ):
								if current_step < num_warmup_steps:
												return float(lowerCAmelCase_ ) / float(max(1 ,   lowerCAmelCase_ ) )
								return max(
								    0.0 ,   float(num_training_steps - current_step ) / float(max(1 ,   num_training_steps - num_warmup_steps ) ) )

				return LambdaLR(lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ )



def 			lowerCAmelCase (		lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ = 0.5 ,   lowerCAmelCase_ = -1 )-> List[Any]:

				def lr_lambda(lowerCAmelCase_ ):
								if current_step < num_warmup_steps:
												return float(lowerCAmelCase_ ) / float(max(1 ,   lowerCAmelCase_ ) )
								lowerCAmelCase_ :      Tuple    				=				float(current_step - num_warmup_steps ) / float(max(1 ,   num_training_steps - num_warmup_steps ) )
								return max(0.0 ,   0.5 * (1.0 + math.cos(math.pi * float(lowerCAmelCase_ ) * 2.0 * progress )) )

				return LambdaLR(lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ )



def 			lowerCAmelCase (		lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ = 1 ,   lowerCAmelCase_ = -1 )-> Dict:

				def lr_lambda(lowerCAmelCase_ ):
								if current_step < num_warmup_steps:
												return float(lowerCAmelCase_ ) / float(max(1 ,   lowerCAmelCase_ ) )
								lowerCAmelCase_ :      List[Any]    				=				float(current_step - num_warmup_steps ) / float(max(1 ,   num_training_steps - num_warmup_steps ) )
								if progress >= 1.0:
												return 0.0
								return max(0.0 ,   0.5 * (1.0 + math.cos(math.pi * ((float(lowerCAmelCase_ ) * progress) % 1.0) )) )

				return LambdaLR(lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ )



def 			lowerCAmelCase (		lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_=1e-7 ,   lowerCAmelCase_=1.0 ,   lowerCAmelCase_=-1 )-> Any:
				lowerCAmelCase_ :      Dict    				=				optimizer.defaults['''lr''']
				if not (lr_init > lr_end):
								raise ValueError(f"""lr_end ({lr_end}) must be be smaller than initial lr ({lr_init})""" )

				def lr_lambda(lowerCAmelCase_ ):
								if current_step < num_warmup_steps:
												return float(lowerCAmelCase_ ) / float(max(1 ,   lowerCAmelCase_ ) )
								elif current_step > num_training_steps:
												return lr_end / lr_init  # as LambdaLR multiplies by lr_init
								else:
												lowerCAmelCase_ :      List[Any]    				=				lr_init - lr_end
												lowerCAmelCase_ :      Optional[Any]    				=				num_training_steps - num_warmup_steps
												lowerCAmelCase_ :      Any    				=				1 - (current_step - num_warmup_steps) / decay_steps
												lowerCAmelCase_ :      List[Any]    				=				lr_range * pct_remaining**power + lr_end
												return decay / lr_init  # as LambdaLR multiplies by lr_init

				return LambdaLR(lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ )


_UpperCAmelCase			:  Union[str, Any]     ={
    SchedulerType.LINEAR: get_linear_schedule_with_warmup,
    SchedulerType.COSINE: get_cosine_schedule_with_warmup,
    SchedulerType.COSINE_WITH_RESTARTS: get_cosine_with_hard_restarts_schedule_with_warmup,
    SchedulerType.POLYNOMIAL: get_polynomial_decay_schedule_with_warmup,
    SchedulerType.CONSTANT: get_constant_schedule,
    SchedulerType.CONSTANT_WITH_WARMUP: get_constant_schedule_with_warmup,
    SchedulerType.PIECEWISE_CONSTANT: get_piecewise_constant_schedule,
}




def 			lowerCAmelCase (		lowerCAmelCase_ ,   lowerCAmelCase_ ,   lowerCAmelCase_ = None ,   lowerCAmelCase_ = None ,   lowerCAmelCase_ = None ,   lowerCAmelCase_ = 1 ,   lowerCAmelCase_ = 1.0 ,   lowerCAmelCase_ = -1 ,   )-> Optional[int]:
				lowerCAmelCase_ :      Union[str, Any]    				=				SchedulerType(lowerCAmelCase_ )
				lowerCAmelCase_ :      Optional[int]    				=				TYPE_TO_SCHEDULER_FUNCTION[name]
				if name == SchedulerType.CONSTANT:
								return schedule_func(lowerCAmelCase_ ,   last_epoch=lowerCAmelCase_ )

				if name == SchedulerType.PIECEWISE_CONSTANT:
								return schedule_func(lowerCAmelCase_ ,   step_rules=lowerCAmelCase_ ,   last_epoch=lowerCAmelCase_ )

				# All other schedulers require `num_warmup_steps`
				if num_warmup_steps is None:
								raise ValueError(f"""{name} requires `num_warmup_steps`, please provide that argument.""" )

				if name == SchedulerType.CONSTANT_WITH_WARMUP:
								return schedule_func(lowerCAmelCase_ ,   num_warmup_steps=lowerCAmelCase_ ,   last_epoch=lowerCAmelCase_ )

				# All other schedulers require `num_training_steps`
				if num_training_steps is None:
								raise ValueError(f"""{name} requires `num_training_steps`, please provide that argument.""" )

				if name == SchedulerType.COSINE_WITH_RESTARTS:
								return schedule_func(
								    lowerCAmelCase_ ,   num_warmup_steps=lowerCAmelCase_ ,   num_training_steps=lowerCAmelCase_ ,   num_cycles=lowerCAmelCase_ ,   last_epoch=lowerCAmelCase_ ,   )

				if name == SchedulerType.POLYNOMIAL:
								return schedule_func(
								    lowerCAmelCase_ ,   num_warmup_steps=lowerCAmelCase_ ,   num_training_steps=lowerCAmelCase_ ,   power=lowerCAmelCase_ ,   last_epoch=lowerCAmelCase_ ,   )

				return schedule_func(
				    lowerCAmelCase_ ,   num_warmup_steps=lowerCAmelCase_ ,   num_training_steps=lowerCAmelCase_ ,   last_epoch=lowerCAmelCase_ )