






'''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


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

lowerCamelCase_				   =   2_56

class   _UpperCAmelCase     (	snake_case__		):




   """simple docstring"""

   snake_case       							=		['''melgan''']

   def __init__(						self					:      List[Any] , __UpperCAmelCase					:      List[Any] , __UpperCAmelCase					:      Union[str, Any] , __UpperCAmelCase					:      Dict , __UpperCAmelCase					:      Optional[Any] , __UpperCAmelCase					:      List[Any] , ):





         '''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=_UpperCamelCase , continuous_encoder=_UpperCamelCase , decoder=_UpperCamelCase , scheduler=_UpperCamelCase , melgan=_UpperCamelCase , )

   def 		lowerCAmelCase					(						self					:      Dict , __UpperCAmelCase					:      Optional[int] , __UpperCAmelCase					:      Optional[Any]=(-1.0, 1.0) , __UpperCAmelCase					:      List[str]=False						):





         '''simple docstring'''





         _A	,      _A     	=      output_range
         if clip:
               _A     	=      torch.clip(_UpperCamelCase , 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					:      List[str] , __UpperCAmelCase					:      Optional[int] , __UpperCAmelCase					:      Tuple=(-1.0, 1.0) , __UpperCAmelCase					:      int=False						):





         '''simple docstring'''





         _A	,      _A     	=      input_range
         _A     	=      torch.clip(_UpperCamelCase , _UpperCamelCase , _UpperCamelCase						) 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] , __UpperCAmelCase					:      Optional[int] , __UpperCAmelCase					:      List[Any] , __UpperCAmelCase					:      Dict						):





         '''simple docstring'''





         _A     	=      input_tokens > 0
         _A	,      _A     	=      self.notes_encoder(
             encoder_input_tokens=_UpperCamelCase , encoder_inputs_mask=_UpperCamelCase						)

         _A	,      _A     	=      self.continuous_encoder(
             encoder_inputs=_UpperCamelCase , encoder_inputs_mask=_UpperCamelCase						)

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

   def 		lowerCAmelCase					(						self					:      int , __UpperCAmelCase					:      Dict , __UpperCAmelCase					:      Tuple , __UpperCAmelCase					:      Optional[Any]						):





         '''simple docstring'''





         _A     	=      noise_time
         if not torch.is_tensor(_UpperCamelCase						):
               _A     	=      torch.tensor([timesteps] , dtype=torch.long , device=input_tokens.device						)
         elif torch.is_tensor(_UpperCamelCase						) 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=_UpperCamelCase , decoder_input_tokens=_UpperCamelCase , decoder_noise_time=_UpperCamelCase						)
         return logits

   @torch.no_grad()
   def __call__(						self					:      List[Any] , __UpperCAmelCase					:      str , __UpperCAmelCase					:      int = None , __UpperCAmelCase					:      Union[str, Any] = 100 , __UpperCAmelCase					:      Optional[Any] = True , __UpperCAmelCase					:      Dict = "numpy" , __UpperCAmelCase					:      Optional[int] = None , __UpperCAmelCase					:      List[str] = 1 , ):





         '''simple docstring'''





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

         _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=_UpperCamelCase , device=self.device						)

         for i, encoder_input_tokens in enumerate(_UpperCamelCase						):
               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=_UpperCamelCase , 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(
                   _UpperCamelCase , output_range=[-1.0, 1.0] , clip=_UpperCamelCase						)

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

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

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

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

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

               _A     	=      self.scale_to_features(_UpperCamelCase , 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(_UpperCamelCase , _UpperCamelCase						)

               logger.info("Generated segment" , _UpperCamelCase						)

         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=_UpperCamelCase						)








# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import copy
import importlib.metadata
import json
import os
from dataclasses import dataclass
from typing import Any, Dict, Union

from packaging import version

from ..utils import is_torch_available, logging


if is_torch_available():
    import torch


UpperCamelCase__		 =     logging.get_logger(__name__)

@dataclass
class     __snake_case    :
       """simple docstring"""







       def __init__(		self      , _UpperCamelCase=False      , _UpperCamelCase=False      , _UpperCamelCase=6.0      , _UpperCamelCase=None      , _UpperCamelCase=False      , _UpperCamelCase=False      , _UpperCamelCase=None      , _UpperCamelCase="fp4"      , _UpperCamelCase=False      , **_UpperCamelCase      , )      ->							int:






              """simple docstring"""


              __snake_case       	=							load_in_abit
              __snake_case       	=							load_in_abit
              __snake_case       	=							llm_inta_threshold
              __snake_case       	=							llm_inta_skip_modules
              __snake_case       	=							llm_inta_enable_fpaa_cpu_offload
              __snake_case       	=							llm_inta_has_fpaa_weight
              __snake_case       	=							bnb_abit_quant_type
              __snake_case       	=							bnb_abit_use_double_quant

              if bnb_abit_compute_dtype is None:
                     __snake_case       	=							torch.floataa
              elif isinstance(_UpperCamelCase      , _UpperCamelCase   ):
                     __snake_case       	=							getattr(_UpperCamelCase      , _UpperCamelCase   )
              elif isinstance(_UpperCamelCase      , torch.dtype   ):
                     __snake_case       	=							bnb_abit_compute_dtype
              else:
                     raise ValueError("""bnb_4bit_compute_dtype must be a string or a torch.dtype"""   )

              self.post_init()


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






              """simple docstring"""


              if not isinstance(self.llm_inta_threshold      , _UpperCamelCase   ):
                     raise ValueError("""llm_int8_threshold must be a float"""   )

              if self.llm_inta_skip_modules is not None and not isinstance(self.llm_inta_skip_modules      , _UpperCamelCase   ):
                     raise ValueError("""llm_int8_skip_modules must be a list of strings"""   )
              if not isinstance(self.llm_inta_enable_fpaa_cpu_offload      , _UpperCamelCase   ):
                     raise ValueError("""llm_int8_enable_fp32_cpu_offload must be a boolean"""   )

              if not isinstance(self.llm_inta_has_fpaa_weight      , _UpperCamelCase   ):
                     raise ValueError("""llm_int8_has_fp16_weight must be a boolean"""   )

              if self.bnb_abit_compute_dtype is not None and not isinstance(self.bnb_abit_compute_dtype      , torch.dtype   ):
                     raise ValueError("""bnb_4bit_compute_dtype must be torch.dtype"""   )

              if not isinstance(self.bnb_abit_quant_type      , _UpperCamelCase   ):
                     raise ValueError("""bnb_4bit_quant_type must be a string"""   )

              if not isinstance(self.bnb_abit_use_double_quant      , _UpperCamelCase   ):
                     raise ValueError("""bnb_4bit_use_double_quant must be a boolean"""   )

              if self.load_in_abit and not version.parse(importlib.metadata.version("""bitsandbytes"""   )   ) >= version.parse(
                  """0.39.0"""   ):
                     raise ValueError(
                         """4 bit quantization requires bitsandbytes>=0.39.0 - please upgrade your bitsandbytes version"""   )


       def 			a		(		self   )      ->							int:






              """simple docstring"""


              return self.load_in_abit or self.load_in_abit


       def 			a		(		self   )      ->							int:






              """simple docstring"""


              if self.load_in_abit:
                     return "llm_int8"
              elif self.load_in_abit and self.bnb_abit_quant_type == "fp4":
                     return "fp4"
              elif self.load_in_abit and self.bnb_abit_quant_type == "nf4":
                     return "nf4"
              else:
                     return None


       @classmethod
       def 			a		(		cls      , _UpperCamelCase      , _UpperCamelCase      , **_UpperCamelCase   )      ->							Union[str, Any]:






              """simple docstring"""


              __snake_case       	=							cls(**_UpperCamelCase   )

              __snake_case       	=							[]
              for key, value in kwargs.items():
                     if hasattr(_UpperCamelCase      , _UpperCamelCase   ):
                            setattr(_UpperCamelCase      , _UpperCamelCase      , _UpperCamelCase   )
                            to_remove.append(_UpperCamelCase   )
              for key in to_remove:
                     kwargs.pop(_UpperCamelCase      , _UpperCamelCase   )

              if return_unused_kwargs:
                     return config, kwargs
              else:
                     return config


       def 			a		(		self      , _UpperCamelCase   )      ->							int:






              """simple docstring"""


              with open(_UpperCamelCase      , """w"""      , encoding="""utf-8"""   ) as writer:
                     __snake_case       	=							self.to_dict()
                     __snake_case       	=							json.dumps(_UpperCamelCase      , indent=2      , sort_keys=_UpperCamelCase   ) + """\n"""

                     writer.write(_UpperCamelCase   )


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






              """simple docstring"""


              __snake_case       	=							copy.deepcopy(self.__dict__   )
              __snake_case       	=							str(output["""bnb_4bit_compute_dtype"""]   ).split("""."""   )[1]

              return output


       def __repr__(		self   )      ->							Dict:






              """simple docstring"""


              return F'{self.__class__.__name__} {self.to_json_string()}'


       def 			a		(		self      , _UpperCamelCase = True   )      ->							str:






              """simple docstring"""


              if use_diff is True:
                     __snake_case       	=							self.to_diff_dict()
              else:
                     __snake_case       	=							self.to_dict()
              return json.dumps(_UpperCamelCase      , indent=2      , sort_keys=_UpperCamelCase   ) + "\n"






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






              """simple docstring"""


              __snake_case       	=							self.to_dict()

              # get the default config dict
              __snake_case       	=							BitsAndBytesConfig().to_dict()

              __snake_case       	=							{}

              # only serialize values that differ from the default config
              for key, value in config_dict.items():
                     if value != default_config_dict[key]:
                            __snake_case       	=							value

              return serializable_config_dict

