


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

lowerCamelCase							:   Dict  						=	256
class   __lowerCAmelCase						(lowercase_     ):





      '''simple docstring'''






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



      def __init__(self    :    Tuple   ,						UpperCamelCase    :    SpectrogramNotesEncoder   ,						UpperCamelCase    :    SpectrogramContEncoder   ,						UpperCamelCase    :    TaFilmDecoder   ,						UpperCamelCase    :    DDPMScheduler   ,						UpperCamelCase    :    OnnxRuntimeModel if is_onnx_available() else Any   ,						):

             '''simple docstring'''







             super().__init__()

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

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



      def       UpperCamelCase__		(self    :    str   ,						UpperCamelCase    :    Dict   ,						UpperCamelCase    :    Tuple=(-1.0, 1.0)   ,						UpperCamelCase    :    int=False ):

             '''simple docstring'''







             lowercase__			,lowercase__      			= output_range
             if clip:
                    lowercase__      			= torch.clip(UpperCamelCase   ,						self.min_value   ,						self.max_value )
             # Scale to [0, 1].
             lowercase__      			= (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       UpperCamelCase__		(self    :    Union[str, Any]   ,						UpperCamelCase    :    List[str]   ,						UpperCamelCase    :    Union[str, Any]=(-1.0, 1.0)   ,						UpperCamelCase    :    Dict=False ):

             '''simple docstring'''







             lowercase__			,lowercase__      			= input_range
             lowercase__      			= torch.clip(UpperCamelCase   ,						UpperCamelCase   ,						UpperCamelCase ) if clip else outputs
             # Scale to [0, 1].
             lowercase__      			= (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       UpperCamelCase__		(self    :    Dict   ,						UpperCamelCase    :    List[Any]   ,						UpperCamelCase    :    List[str]   ,						UpperCamelCase    :    Union[str, Any] ):

             '''simple docstring'''







             lowercase__      			= input_tokens > 0
             lowercase__			,lowercase__      			= self.notes_encoder(
                 encoder_input_tokens=UpperCamelCase   ,						encoder_inputs_mask=UpperCamelCase )

             lowercase__			,lowercase__      			= self.continuous_encoder(
                 encoder_inputs=UpperCamelCase   ,						encoder_inputs_mask=UpperCamelCase )

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



      def       UpperCamelCase__		(self    :    Dict   ,						UpperCamelCase    :    Union[str, Any]   ,						UpperCamelCase    :    Optional[int]   ,						UpperCamelCase    :    Any ):

             '''simple docstring'''







             lowercase__      			= noise_time
             if not torch.is_tensor(UpperCamelCase ):
                    lowercase__      			= torch.tensor([timesteps]   ,						dtype=torch.long   ,						device=input_tokens.device )
             elif torch.is_tensor(UpperCamelCase ) and len(timesteps.shape ) == 0:
                    lowercase__      			= timesteps[None].to(input_tokens.device )

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

             lowercase__      			= 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    :    List[List[int]]   ,						UpperCamelCase    :    Optional[torch.Generator] = None   ,						UpperCamelCase    :    int = 100   ,						UpperCamelCase    :    bool = True   ,						UpperCamelCase    :    str = "numpy"   ,						UpperCamelCase    :    Optional[Callable[[int, int, torch.FloatTensor], None]] = None   ,						UpperCamelCase    :    int = 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 )}." )

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

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

                    lowercase__      			= self.scale_features(
                        UpperCamelCase   ,						output_range=[-1.0, 1.0]   ,						clip=UpperCamelCase )

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

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

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

             if not return_dict:
                    return (output,)

             return AudioPipelineOutput(audios=UpperCamelCase )










import warnings
from typing import List, Optional, Union

from ...image_utils import ImageInput
from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding, PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy
from ...utils import TensorType







class __snake_case			(			a		):
   UpperCAmelCase__  :    Dict          =  ['''image_processor''', '''tokenizer''']
   UpperCAmelCase__  :    Dict          =  '''FlavaImageProcessor'''
   UpperCAmelCase__  :    Dict          =  ('''BertTokenizer''', '''BertTokenizerFast''')







   def __init__(      self		:						Union[str, Any]					,						_snake_case		:						List[str]=None					,						_snake_case		:						str=None					,						**_snake_case		:						int):

        """simple docstring"""


        UpperCAmelCase_          =							None
        if "feature_extractor" in kwargs:
             warnings.warn(
                 '''The `feature_extractor` argument is deprecated and will be removed in v5, use `image_processor`'''
                 ''' instead.'''					,						_snake_case					,						)
             UpperCAmelCase_          =							kwargs.pop('''feature_extractor''')

        UpperCAmelCase_          =							image_processor if image_processor is not None else feature_extractor
        if image_processor is None:
             raise ValueError('''You need to specify an `image_processor`.''')
        if tokenizer is None:
             raise ValueError('''You need to specify a `tokenizer`.''')

        super().__init__(_snake_case					,						_snake_case)
        UpperCAmelCase_          =							self.image_processor







   def __call__(      self		:						List[Any]					,						_snake_case		:						Optional[ImageInput] = None					,						_snake_case		:						Optional[Union[TextInput, PreTokenizedInput, List[TextInput], List[PreTokenizedInput]]] = None					,						_snake_case		:						bool = True					,						_snake_case		:						Union[bool, str, PaddingStrategy] = False					,						_snake_case		:						Union[bool, str, TruncationStrategy] = False					,						_snake_case		:						Optional[int] = None					,						_snake_case		:						int = 0					,						_snake_case		:						Optional[int] = None					,						_snake_case		:						Optional[bool] = None					,						_snake_case		:						Optional[bool] = None					,						_snake_case		:						Optional[bool] = None					,						_snake_case		:						Optional[bool] = None					,						_snake_case		:						bool = False					,						_snake_case		:						bool = False					,						_snake_case		:						bool = False					,						_snake_case		:						bool = False					,						_snake_case		:						bool = True					,						_snake_case		:						Optional[Union[str, TensorType]] = None					,						**_snake_case		:						Any					,						):

        """simple docstring"""


        if text is None and images is None:
             raise ValueError('''You have to specify either text or images. Both cannot be none.''')

        if text is not None:
             UpperCAmelCase_          =							self.tokenizer(
                 text=_snake_case					,						add_special_tokens=_snake_case					,						padding=_snake_case					,						truncation=_snake_case					,						max_length=_snake_case					,						stride=_snake_case					,						pad_to_multiple_of=_snake_case					,						return_token_type_ids=_snake_case					,						return_attention_mask=_snake_case					,						return_overflowing_tokens=_snake_case					,						return_special_tokens_mask=_snake_case					,						return_offsets_mapping=_snake_case					,						return_length=_snake_case					,						verbose=_snake_case					,						return_tensors=_snake_case					,						**_snake_case					,						)
        if images is not None:
             UpperCAmelCase_          =							self.image_processor(
                 _snake_case					,						return_image_mask=_snake_case					,						return_codebook_pixels=_snake_case					,						return_tensors=_snake_case					,						**_snake_case					,						)

        if text is not None and images is not None:
             encoding.update(_snake_case)
             return encoding
        elif text is not None:
             return encoding
        else:
             return BatchEncoding(data=dict(**_snake_case)					,						tensor_type=_snake_case)







   def lowerCamelCase						(      self		:						Any					,						*_snake_case		:						Optional[Any]					,						**_snake_case		:						int):

        """simple docstring"""


        return self.tokenizer.batch_decode(*_snake_case					,						**_snake_case)







   def lowerCamelCase						(      self		:						Optional[int]					,						*_snake_case		:						int					,						**_snake_case		:						Dict):

        """simple docstring"""


        return self.tokenizer.decode(*_snake_case					,						**_snake_case)







   @property
   def lowerCamelCase						(      self		:						str):

        """simple docstring"""


        UpperCAmelCase_          =							self.tokenizer.model_input_names
        UpperCAmelCase_          =							self.image_processor.model_input_names
        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))







   @property
   def lowerCamelCase						(      self		:						str):

        """simple docstring"""


        warnings.warn(
            '''`feature_extractor_class` is deprecated and will be removed in v5. Use `image_processor_class` instead.'''					,						_snake_case					,						)
        return self.image_processor_class







   @property
   def lowerCamelCase						(      self		:						Any):

        """simple docstring"""


        warnings.warn(
            '''`feature_extractor` is deprecated and will be removed in v5. Use `image_processor` instead.'''					,						_snake_case					,						)
        return self.image_processor






