




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

UpperCamelCase_						:      Union[str, Any]       =						256

class        _a    (						__lowerCAmelCase							):
     SCREAMING_SNAKE_CASE_							:   Union[str, Any]       =						['melgan']






     def __init__(			self ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,)							->      str:
          super().__init__()

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

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






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

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

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






     def     _lowercase     (			self ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )							->      Union[str, Any]:
          _snake_case							=				noise_time
          if not torch.is_tensor(__lowerCAmelCase   ):
               _snake_case							=				torch.tensor([timesteps] ,dtype=torch.long ,device=input_tokens.device   )
          elif torch.is_tensor(__lowerCAmelCase   ) and len(timesteps.shape   ) == 0:
               _snake_case							=				timesteps[None].to(input_tokens.device   )

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

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






     @torch.no_grad()
     def __call__(			self ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE = None ,_SCREAMING_SNAKE_CASE = 100 ,_SCREAMING_SNAKE_CASE = True ,_SCREAMING_SNAKE_CASE = "numpy" ,_SCREAMING_SNAKE_CASE = None ,_SCREAMING_SNAKE_CASE = 1 ,)							->      List[Any]:



          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   )}."""   )

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

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

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

               _snake_case							=				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
               _snake_case							=				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   )   ):
                    _snake_case							=				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
                    _snake_case							=				self.scheduler.step(__lowerCAmelCase ,__lowerCAmelCase ,__lowerCAmelCase ,generator=__lowerCAmelCase   ).prev_sample

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

               _snake_case							=				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":
               _snake_case							=				self.melgan(input_features=full_pred_mel.astype(np.floataa   )   )
          else:
               _snake_case							=				full_pred_mel

          if not return_dict:
               return (output,)

          return AudioPipelineOutput(audios=__lowerCAmelCase   )








'''simple docstring'''





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCamelCase_						:      List[Any]       =						{
    '''configuration_convbert''': ['''CONVBERT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ConvBertConfig''', '''ConvBertOnnxConfig'''],
    '''tokenization_convbert''': ['''ConvBertTokenizer'''],
}

try:
   if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCamelCase_						:      Optional[Any]       =						['''ConvBertTokenizerFast''']

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCamelCase_						:      Optional[int]       =						[
       '''CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
       '''ConvBertForMaskedLM''',
       '''ConvBertForMultipleChoice''',
       '''ConvBertForQuestionAnswering''',
       '''ConvBertForSequenceClassification''',
       '''ConvBertForTokenClassification''',
       '''ConvBertLayer''',
       '''ConvBertModel''',
       '''ConvBertPreTrainedModel''',
       '''load_tf_weights_in_convbert''',
   ]


try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   UpperCamelCase_						:      Tuple       =						[
       '''TF_CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST''',
       '''TFConvBertForMaskedLM''',
       '''TFConvBertForMultipleChoice''',
       '''TFConvBertForQuestionAnswering''',
       '''TFConvBertForSequenceClassification''',
       '''TFConvBertForTokenClassification''',
       '''TFConvBertLayer''',
       '''TFConvBertModel''',
       '''TFConvBertPreTrainedModel''',
   ]


if TYPE_CHECKING:
   from .configuration_convbert import CONVBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, ConvBertConfig, ConvBertOnnxConfig
   from .tokenization_convbert import ConvBertTokenizer

   try:
      if not is_tokenizers_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .tokenization_convbert_fast import ConvBertTokenizerFast

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_convbert import (
          CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
          ConvBertForMaskedLM,
          ConvBertForMultipleChoice,
          ConvBertForQuestionAnswering,
          ConvBertForSequenceClassification,
          ConvBertForTokenClassification,
          ConvBertLayer,
          ConvBertModel,
          ConvBertPreTrainedModel,
          load_tf_weights_in_convbert,
      )

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_convbert import (
          TF_CONVBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFConvBertForMaskedLM,
          TFConvBertForMultipleChoice,
          TFConvBertForQuestionAnswering,
          TFConvBertForSequenceClassification,
          TFConvBertForTokenClassification,
          TFConvBertLayer,
          TFConvBertModel,
          TFConvBertPreTrainedModel,
      )


else:
   import sys

   UpperCamelCase_						:      Tuple       =						_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)


