


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 A    (  UpperCamelCase_	):
      UpperCamelCase__				:    Tuple =['melgan']






      def __init__(   self		:	str     ,  lowercase_		:	SpectrogramNotesEncoder     ,  lowercase_		:	SpectrogramContEncoder     ,  lowercase_		:	TaFilmDecoder     ,  lowercase_		:	DDPMScheduler     ,  lowercase_		:	OnnxRuntimeModel if is_onnx_available() else Any     ,  )   ->		None:
             """simple docstring"""




             super().__init__()

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

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






      def 			lowerCamelCase			(   self		:	Optional[Any]     ,  lowercase_		:	List[str]     ,  lowercase_		:	List[Any]=(-1.0, 1.0)     ,  lowercase_		:	Any=False					)   ->		List[str]:
             """simple docstring"""




             _lowerCamelCase      ,	_lowerCamelCase			:						Tuple            =output_range
             if clip:
                    _lowerCamelCase			:						int            =torch.clip(lowercase_     ,  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 			lowerCamelCase			(   self		:	List[str]     ,  lowercase_		:	str     ,  lowercase_		:	int=(-1.0, 1.0)     ,  lowercase_		:	List[str]=False					)   ->		Union[str, Any]:
             """simple docstring"""




             _lowerCamelCase      ,	_lowerCamelCase			:						int            =input_range
             _lowerCamelCase			:						Union[str, Any]            =torch.clip(lowercase_     ,  lowercase_     ,  lowercase_					) if clip else outputs
             # Scale to [0, 1].
             _lowerCamelCase			:						Optional[int]            =(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]     ,  lowercase_		:	Optional[Any]     ,  lowercase_		:	Tuple     ,  lowercase_		:	Union[str, Any]					)   ->		List[str]:
             """simple docstring"""




             _lowerCamelCase			:						Union[str, Any]            =input_tokens > 0
             _lowerCamelCase      ,	_lowerCamelCase			:						Union[str, Any]            =self.notes_encoder(
                 encoder_input_tokens=lowercase_     ,  encoder_inputs_mask=lowercase_					)

             _lowerCamelCase      ,	_lowerCamelCase			:						str            =self.continuous_encoder(
                 encoder_inputs=lowercase_     ,  encoder_inputs_mask=lowercase_					)

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






      def 			lowerCamelCase			(   self		:	List[str]     ,  lowercase_		:	Optional[Any]     ,  lowercase_		:	str     ,  lowercase_		:	List[Any]					)   ->		Dict:
             """simple docstring"""




             _lowerCamelCase			:						str            =noise_time
             if not torch.is_tensor(lowercase_					):
                    _lowerCamelCase			:						Any            =torch.tensor([timesteps]     ,  dtype=torch.long     ,  device=input_tokens.device					)
             elif torch.is_tensor(lowercase_					) and len(timesteps.shape					) == 0:
                    _lowerCamelCase			:						Optional[int]            =timesteps[None].to(input_tokens.device					)

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

             _lowerCamelCase			:						Optional[Any]            =self.decoder(
                 encodings_and_masks=lowercase_     ,  decoder_input_tokens=lowercase_     ,  decoder_noise_time=lowercase_					)
             return logits






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




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

             _lowerCamelCase			:						List[Any]            =np.zeros([1, TARGET_FEATURE_LENGTH, self.n_dims]     ,  dtype=np.floataa					)
             _lowerCamelCase			:						Optional[Any]            =np.zeros([1, 0, self.n_dims]     ,  np.floataa					)
             _lowerCamelCase			:						Any            =torch.ones((1, TARGET_FEATURE_LENGTH)     ,  dtype=lowercase_     ,  device=self.device					)

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

                    _lowerCamelCase			:						Dict            =self.scale_features(
                        lowercase_     ,  output_range=[-1.0, 1.0]     ,  clip=lowercase_					)

                    _lowerCamelCase			:						Tuple            =self.encode(
                        input_tokens=torch.IntTensor([encoder_input_tokens]					).to(device=self.device					)     ,  continuous_inputs=lowercase_     ,  continuous_mask=lowercase_     ,  )

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

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

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

                           # Compute previous output: x_t -> x_t-1
                           _lowerCamelCase			:						Any            =self.scheduler.step(lowercase_     ,  lowercase_     ,  lowercase_     ,  generator=lowercase_					).prev_sample

                    _lowerCamelCase			:						Union[str, Any]            =self.scale_to_features(lowercase_     ,  input_range=[-1.0, 1.0]					)
                    _lowerCamelCase			:						Any            =mel[:1]
                    _lowerCamelCase			:						Tuple            =mel.cpu().float().numpy()

                    _lowerCamelCase			:						Optional[int]            =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(lowercase_     ,  lowercase_					)

                    logger.info('Generated segment'     ,  lowercase_					)

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

             if not return_dict:
                    return (output,)

             return AudioPipelineOutput(audios=lowercase_					)






import tempfile

import torch

from diffusers import PNDMScheduler

from .test_schedulers import SchedulerCommonTest




class A    (  UpperCamelCase_	):
      UpperCamelCase__				:    List[str] =(PNDMScheduler,)
      UpperCamelCase__				:    Dict =(('num_inference_steps', 50),)






      def 			lowerCamelCase			(   self		:	Dict     ,  **lowercase_		:	Dict					)   ->		str:
             """simple docstring"""




             _lowerCamelCase			:						List[Any]            ={
                 'num_train_timesteps': 1000,
                 'beta_start': 0.0001,
                 'beta_end': 0.02,
                 'beta_schedule': 'linear',
             }

             config.update(**lowercase_					)
             return config






      def 			lowerCamelCase			(   self		:	Any     ,  lowercase_		:	str=0     ,  **lowercase_		:	str					)   ->		Optional[int]:
             """simple docstring"""




             _lowerCamelCase			:						Optional[Any]            =dict(self.forward_default_kwargs					)
             _lowerCamelCase			:						Optional[Any]            =kwargs.pop('num_inference_steps'     ,  lowercase_					)
             _lowerCamelCase			:						Tuple            =self.dummy_sample
             _lowerCamelCase			:						int            =0.1 * sample
             _lowerCamelCase			:						int            =[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

             for scheduler_class in self.scheduler_classes:
                    _lowerCamelCase			:						str            =self.get_scheduler_config(**lowercase_					)
                    _lowerCamelCase			:						str            =scheduler_class(**lowercase_					)
                    scheduler.set_timesteps(lowercase_					)
                    # copy over dummy past residuals
                    _lowerCamelCase			:						Any            =dummy_past_residuals[:]

                    with tempfile.TemporaryDirectory() as tmpdirname:
                           scheduler.save_config(lowercase_					)
                           _lowerCamelCase			:						Optional[Any]            =scheduler_class.from_pretrained(lowercase_					)
                           new_scheduler.set_timesteps(lowercase_					)
                           # copy over dummy past residuals
                           _lowerCamelCase			:						Any            =dummy_past_residuals[:]

                    _lowerCamelCase			:						Union[str, Any]            =scheduler.step_prk(lowercase_     ,  lowercase_     ,  lowercase_     ,  **lowercase_					).prev_sample
                    _lowerCamelCase			:						List[str]            =new_scheduler.step_prk(lowercase_     ,  lowercase_     ,  lowercase_     ,  **lowercase_					).prev_sample

                    assert torch.sum(torch.abs(output - new_output					)					) < 1E-5, "Scheduler outputs are not identical"

                    _lowerCamelCase			:						Optional[int]            =scheduler.step_plms(lowercase_     ,  lowercase_     ,  lowercase_     ,  **lowercase_					).prev_sample
                    _lowerCamelCase			:						Optional[Any]            =new_scheduler.step_plms(lowercase_     ,  lowercase_     ,  lowercase_     ,  **lowercase_					).prev_sample

                    assert torch.sum(torch.abs(output - new_output					)					) < 1E-5, "Scheduler outputs are not identical"






      def 			lowerCamelCase			(   self		:	Any					)   ->		Any:
             """simple docstring"""




             pass






      def 			lowerCamelCase			(   self		:	Union[str, Any]     ,  lowercase_		:	List[str]=0     ,  **lowercase_		:	int					)   ->		int:
             """simple docstring"""




             _lowerCamelCase			:						Any            =dict(self.forward_default_kwargs					)
             _lowerCamelCase			:						Dict            =kwargs.pop('num_inference_steps'     ,  lowercase_					)
             _lowerCamelCase			:						Union[str, Any]            =self.dummy_sample
             _lowerCamelCase			:						Optional[int]            =0.1 * sample
             _lowerCamelCase			:						List[str]            =[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]

             for scheduler_class in self.scheduler_classes:
                    _lowerCamelCase			:						Any            =self.get_scheduler_config()
                    _lowerCamelCase			:						int            =scheduler_class(**lowercase_					)
                    scheduler.set_timesteps(lowercase_					)

                    # copy over dummy past residuals (must be after setting timesteps)
                    _lowerCamelCase			:						Any            =dummy_past_residuals[:]

                    with tempfile.TemporaryDirectory() as tmpdirname:
                           scheduler.save_config(lowercase_					)
                           _lowerCamelCase			:						Optional[Any]            =scheduler_class.from_pretrained(lowercase_					)
                           # copy over dummy past residuals
                           new_scheduler.set_timesteps(lowercase_					)

                           # copy over dummy past residual (must be after setting timesteps)
                           _lowerCamelCase			:						Optional[int]            =dummy_past_residuals[:]

                    _lowerCamelCase			:						List[str]            =scheduler.step_prk(lowercase_     ,  lowercase_     ,  lowercase_     ,  **lowercase_					).prev_sample
                    _lowerCamelCase			:						str            =new_scheduler.step_prk(lowercase_     ,  lowercase_     ,  lowercase_     ,  **lowercase_					).prev_sample

                    assert torch.sum(torch.abs(output - new_output					)					) < 1E-5, "Scheduler outputs are not identical"

                    _lowerCamelCase			:						Optional[Any]            =scheduler.step_plms(lowercase_     ,  lowercase_     ,  lowercase_     ,  **lowercase_					).prev_sample
                    _lowerCamelCase			:						List[str]            =new_scheduler.step_plms(lowercase_     ,  lowercase_     ,  lowercase_     ,  **lowercase_					).prev_sample

                    assert torch.sum(torch.abs(output - new_output					)					) < 1E-5, "Scheduler outputs are not identical"






      def 			lowerCamelCase			(   self		:	Optional[Any]     ,  **lowercase_		:	Any					)   ->		Tuple:
             """simple docstring"""




             _lowerCamelCase			:						Union[str, Any]            =self.scheduler_classes[0]
             _lowerCamelCase			:						Optional[Any]            =self.get_scheduler_config(**lowercase_					)
             _lowerCamelCase			:						Any            =scheduler_class(**lowercase_					)

             _lowerCamelCase			:						Union[str, Any]            =10
             _lowerCamelCase			:						str            =self.dummy_model()
             _lowerCamelCase			:						Union[str, Any]            =self.dummy_sample_deter
             scheduler.set_timesteps(lowercase_					)

             for i, t in enumerate(scheduler.prk_timesteps					):
                    _lowerCamelCase			:						Union[str, Any]            =model(lowercase_     ,  lowercase_					)
                    _lowerCamelCase			:						Any            =scheduler.step_prk(lowercase_     ,  lowercase_     ,  lowercase_					).prev_sample

             for i, t in enumerate(scheduler.plms_timesteps					):
                    _lowerCamelCase			:						List[Any]            =model(lowercase_     ,  lowercase_					)
                    _lowerCamelCase			:						int            =scheduler.step_plms(lowercase_     ,  lowercase_     ,  lowercase_					).prev_sample

             return sample






      def 			lowerCamelCase			(   self		:	Dict					)   ->		List[Any]:
             """simple docstring"""




             _lowerCamelCase			:						Any            =dict(self.forward_default_kwargs					)

             _lowerCamelCase			:						Optional[int]            =kwargs.pop('num_inference_steps'     ,  lowercase_					)

             for scheduler_class in self.scheduler_classes:
                    _lowerCamelCase			:						List[str]            =self.get_scheduler_config()
                    _lowerCamelCase			:						List[Any]            =scheduler_class(**lowercase_					)

                    _lowerCamelCase			:						Union[str, Any]            =self.dummy_sample
                    _lowerCamelCase			:						Optional[int]            =0.1 * sample

                    if num_inference_steps is not None and hasattr(lowercase_     ,  'set_timesteps'					):
                           scheduler.set_timesteps(lowercase_					)
                    elif num_inference_steps is not None and not hasattr(lowercase_     ,  'set_timesteps'					):
                           _lowerCamelCase			:						Tuple            =num_inference_steps

                    # copy over dummy past residuals (must be done after set_timesteps)
                    _lowerCamelCase			:						List[Any]            =[residual + 0.2, residual + 0.15, residual + 0.1, residual + 0.05]
                    _lowerCamelCase			:						Union[str, Any]            =dummy_past_residuals[:]

                    _lowerCamelCase			:						List[Any]            =scheduler.step_prk(lowercase_     ,  0     ,  lowercase_     ,  **lowercase_					).prev_sample
                    _lowerCamelCase			:						str            =scheduler.step_prk(lowercase_     ,  1     ,  lowercase_     ,  **lowercase_					).prev_sample

                    self.assertEqual(output_a.shape     ,  sample.shape					)
                    self.assertEqual(output_a.shape     ,  output_a.shape					)

                    _lowerCamelCase			:						Tuple            =scheduler.step_plms(lowercase_     ,  0     ,  lowercase_     ,  **lowercase_					).prev_sample
                    _lowerCamelCase			:						Optional[Any]            =scheduler.step_plms(lowercase_     ,  1     ,  lowercase_     ,  **lowercase_					).prev_sample

                    self.assertEqual(output_a.shape     ,  sample.shape					)
                    self.assertEqual(output_a.shape     ,  output_a.shape					)






      def 			lowerCamelCase			(   self		:	Dict					)   ->		Any:
             """simple docstring"""




             for timesteps in [100, 1000]:
                    self.check_over_configs(num_train_timesteps=lowercase_					)






      def 			lowerCamelCase			(   self		:	Optional[int]					)   ->		Optional[Any]:
             """simple docstring"""




             for steps_offset in [0, 1]:
                    self.check_over_configs(steps_offset=lowercase_					)

             _lowerCamelCase			:						Optional[int]            =self.scheduler_classes[0]
             _lowerCamelCase			:						Dict            =self.get_scheduler_config(steps_offset=1					)
             _lowerCamelCase			:						List[Any]            =scheduler_class(**lowercase_					)
             scheduler.set_timesteps(10					)
             assert torch.equal(
                 scheduler.timesteps     ,  torch.LongTensor(
                     [901, 851, 851, 801, 801, 751, 751, 701, 701, 651, 651, 601, 601, 501, 401, 301, 201, 101, 1]					)     ,  )






      def 			lowerCamelCase			(   self		:	List[str]					)   ->		str:
             """simple docstring"""




             for beta_start, beta_end in zip([0.0001, 0.001]     ,  [0.002, 0.02]					):
                    self.check_over_configs(beta_start=lowercase_     ,  beta_end=lowercase_					)






      def 			lowerCamelCase			(   self		:	Tuple					)   ->		Union[str, Any]:
             """simple docstring"""




             for schedule in ["linear", "squaredcos_cap_v2"]:
                    self.check_over_configs(beta_schedule=lowercase_					)






      def 			lowerCamelCase			(   self		:	Tuple					)   ->		Any:
             """simple docstring"""




             for prediction_type in ["epsilon", "v_prediction"]:
                    self.check_over_configs(prediction_type=lowercase_					)






      def 			lowerCamelCase			(   self		:	Optional[Any]					)   ->		Union[str, Any]:
             """simple docstring"""




             for t in [1, 5, 10]:
                    self.check_over_forward(time_step=lowercase_					)






      def 			lowerCamelCase			(   self		:	List[str]					)   ->		List[Any]:
             """simple docstring"""




             for t, num_inference_steps in zip([1, 5, 10]     ,  [10, 50, 100]					):
                    self.check_over_forward(num_inference_steps=lowercase_					)






      def 			lowerCamelCase			(   self		:	int					)   ->		Optional[Any]:
             """simple docstring"""




             _lowerCamelCase			:						Optional[int]            =27

             for scheduler_class in self.scheduler_classes:
                    _lowerCamelCase			:						Dict            =self.dummy_sample
                    _lowerCamelCase			:						List[Any]            =0.1 * sample

                    _lowerCamelCase			:						List[Any]            =self.get_scheduler_config()
                    _lowerCamelCase			:						List[str]            =scheduler_class(**lowercase_					)

                    scheduler.set_timesteps(lowercase_					)

                    # before power of 3 fix, would error on first step, so we only need to do two
                    for i, t in enumerate(scheduler.prk_timesteps[:2]					):
                           _lowerCamelCase			:						Tuple            =scheduler.step_prk(lowercase_     ,  lowercase_     ,  lowercase_					).prev_sample






      def 			lowerCamelCase			(   self		:	List[str]					)   ->		Optional[int]:
             """simple docstring"""




             with self.assertRaises(lowercase_					):
                    _lowerCamelCase			:						Dict            =self.scheduler_classes[0]
                    _lowerCamelCase			:						Optional[Any]            =self.get_scheduler_config()
                    _lowerCamelCase			:						List[Any]            =scheduler_class(**lowercase_					)

                    scheduler.step_plms(self.dummy_sample     ,  1     ,  self.dummy_sample					).prev_sample






      def 			lowerCamelCase			(   self		:	List[str]					)   ->		List[str]:
             """simple docstring"""




             _lowerCamelCase			:						Tuple            =self.full_loop()
             _lowerCamelCase			:						Optional[Any]            =torch.sum(torch.abs(lowercase_					)					)
             _lowerCamelCase			:						Union[str, Any]            =torch.mean(torch.abs(lowercase_					)					)

             assert abs(result_sum.item() - 198.1318					) < 1E-2
             assert abs(result_mean.item() - 0.2580					) < 1E-3






      def 			lowerCamelCase			(   self		:	Any					)   ->		Optional[int]:
             """simple docstring"""




             _lowerCamelCase			:						int            =self.full_loop(prediction_type='v_prediction'					)
             _lowerCamelCase			:						Dict            =torch.sum(torch.abs(lowercase_					)					)
             _lowerCamelCase			:						str            =torch.mean(torch.abs(lowercase_					)					)

             assert abs(result_sum.item() - 67.3986					) < 1E-2
             assert abs(result_mean.item() - 0.0878					) < 1E-3






      def 			lowerCamelCase			(   self		:	Tuple					)   ->		Optional[Any]:
             """simple docstring"""




             _lowerCamelCase			:						Tuple            =self.full_loop(set_alpha_to_one=lowercase_     ,  beta_start=0.01					)
             _lowerCamelCase			:						Optional[Any]            =torch.sum(torch.abs(lowercase_					)					)
             _lowerCamelCase			:						Union[str, Any]            =torch.mean(torch.abs(lowercase_					)					)

             assert abs(result_sum.item() - 230.0399					) < 1E-2
             assert abs(result_mean.item() - 0.2995					) < 1E-3






      def 			lowerCamelCase			(   self		:	Optional[Any]					)   ->		Any:
             """simple docstring"""




             _lowerCamelCase			:						str            =self.full_loop(set_alpha_to_one=lowercase_     ,  beta_start=0.01					)
             _lowerCamelCase			:						Union[str, Any]            =torch.sum(torch.abs(lowercase_					)					)
             _lowerCamelCase			:						str            =torch.mean(torch.abs(lowercase_					)					)

             assert abs(result_sum.item() - 186.9482					) < 1E-2
             assert abs(result_mean.item() - 0.2434					) < 1E-3


