



"""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__              =							256






class 			__snake_case			(      _lowercase):
  snake_case__							:  Optional[Any]         =		["melgan"]




  def __init__(					self	:		List[Any]		,   __lowerCAmelCase	:		SpectrogramNotesEncoder		,   __lowerCAmelCase	:		SpectrogramContEncoder		,   __lowerCAmelCase	:		TaFilmDecoder		,   __lowerCAmelCase	:		DDPMScheduler		,   __lowerCAmelCase	:		OnnxRuntimeModel if is_onnx_available() else Any		,   ):






         """simple docstring"""






         super().__init__()

         # From MELGAN
         _lowerCamelCase		:							Tuple            =					math.log(1E-5						)  # Matches MelGAN training.
         _lowerCamelCase		:							Optional[Any]            =					4.0  # Largest value for most examples
         _lowerCamelCase		:							Tuple            =					1_2_8

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




  def 		SCREAMING_SNAKE_CASE							(					self	:		List[str]		,   __lowerCAmelCase	:		Union[str, Any]		,   __lowerCAmelCase	:		List[Any]=(-1.0, 1.0)		,   __lowerCAmelCase	:		Union[str, Any]=False						):






         """simple docstring"""






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






         """simple docstring"""






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






         """simple docstring"""






         _lowerCamelCase		:							Any            =					input_tokens > 0
         _lowerCamelCase     ,       _lowerCamelCase		:							Optional[int]            =					self.notes_encoder(
             encoder_input_tokens=__lowerCAmelCase		,   encoder_inputs_mask=__lowerCAmelCase						)

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

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




  def 		SCREAMING_SNAKE_CASE							(					self	:		Any		,   __lowerCAmelCase	:		Dict		,   __lowerCAmelCase	:		List[str]		,   __lowerCAmelCase	:		Optional[int]						):






         """simple docstring"""






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

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

         _lowerCamelCase		:							Dict            =					self.decoder(
             encodings_and_masks=__lowerCAmelCase		,   decoder_input_tokens=__lowerCAmelCase		,   decoder_noise_time=__lowerCAmelCase						)
         return logits




  @torch.no_grad()
  def __call__(					self	:		str		,   __lowerCAmelCase	:		List[List[int]]		,   __lowerCAmelCase	:		Optional[torch.Generator] = None		,   __lowerCAmelCase	:		int = 1_0_0		,   __lowerCAmelCase	:		bool = True		,   __lowerCAmelCase	:		str = "numpy"		,   __lowerCAmelCase	:		Optional[Callable[[int, int, torch.FloatTensor], None]] = None		,   __lowerCAmelCase	:		int = 1		,   ):






         """simple docstring"""






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

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

         for i, encoder_input_tokens in enumerate(__lowerCAmelCase						):
                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		:							Optional[Any]            =					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.
                       _lowerCamelCase		:							str            =					ones

                _lowerCamelCase		:							Union[str, Any]            =					self.scale_features(
                    __lowerCAmelCase		,   output_range=[-1.0, 1.0]		,   clip=__lowerCAmelCase						)

                _lowerCamelCase		:							Union[str, Any]            =					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
                _lowerCamelCase		:							Union[str, Any]            =					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						)						):
                       _lowerCamelCase		:							Tuple            =					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
                       _lowerCamelCase		:							Union[str, Any]            =					self.scheduler.step(__lowerCAmelCase		,   __lowerCAmelCase		,   __lowerCAmelCase		,   generator=__lowerCAmelCase						).prev_sample

                _lowerCamelCase		:							Optional[int]            =					self.scale_to_features(__lowerCAmelCase		,   input_range=[-1.0, 1.0]						)
                _lowerCamelCase		:							Union[str, Any]            =					mel[:1]
                _lowerCamelCase		:							int            =					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(__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":
                _lowerCamelCase		:							Union[str, Any]            =					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=__lowerCAmelCase						)





"""simple docstring"""




import unittest

import numpy as np

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
						import torch

if is_vision_available():
						from PIL import Image

						from transformers import MobileViTImageProcessor





class  lowerCamelCase    (unittest.TestCase      ):

			def __init__(       self					:    Tuple     ,					__UpperCAmelCase					:    str     ,					__UpperCAmelCase					:    Any=7     ,					__UpperCAmelCase					:    Union[str, Any]=3     ,					__UpperCAmelCase					:    Any=1_8     ,					__UpperCAmelCase					:    Tuple=3_0     ,					__UpperCAmelCase					:    List[Any]=4_0_0     ,					__UpperCAmelCase					:    Dict=True     ,					__UpperCAmelCase					:    int=None     ,					__UpperCAmelCase					:    List[Any]=True     ,					__UpperCAmelCase					:    Any=None     ,					__UpperCAmelCase					:    Any=True     ,					) ->     Tuple:
						SCREAMING_SNAKE_CASE__										=   size if size is not None else {"""shortest_edge""": 2_0}
						SCREAMING_SNAKE_CASE__										=   crop_size if crop_size is not None else {"""height""": 1_8, """width""": 1_8}
						SCREAMING_SNAKE_CASE__										=   parent
						SCREAMING_SNAKE_CASE__										=   batch_size
						SCREAMING_SNAKE_CASE__										=   num_channels
						SCREAMING_SNAKE_CASE__										=   image_size
						SCREAMING_SNAKE_CASE__										=   min_resolution
						SCREAMING_SNAKE_CASE__										=   max_resolution
						SCREAMING_SNAKE_CASE__										=   do_resize
						SCREAMING_SNAKE_CASE__										=   size
						SCREAMING_SNAKE_CASE__										=   do_center_crop
						SCREAMING_SNAKE_CASE__										=   crop_size
						SCREAMING_SNAKE_CASE__										=   do_flip_channel_order





			def        SCREAMING_SNAKE_CASE					(       self					:    Tuple							) ->     Any:
						return {
						    "do_resize": self.do_resize,
						    "size": self.size,
						    "do_center_crop": self.do_center_crop,
						    "crop_size": self.crop_size,
						    "do_flip_channel_order": self.do_flip_channel_order,
						}





@require_torch
@require_vision
class  lowerCamelCase    (A__      ,unittest.TestCase      ):
			lowerCamelCase__ :		Tuple       		=				MobileViTImageProcessor if is_vision_available() else None

			def        SCREAMING_SNAKE_CASE					(       self					:    Dict							) ->     Any:
						SCREAMING_SNAKE_CASE__										=   MobileViTImageProcessingTester(self							)

			@property
			def        SCREAMING_SNAKE_CASE					(       self					:    List[str]							) ->     int:
						return self.image_processor_tester.prepare_image_processor_dict()

			def        SCREAMING_SNAKE_CASE					(       self					:    Tuple							) ->     Optional[Any]:
						SCREAMING_SNAKE_CASE__										=   self.image_processing_class(**self.image_processor_dict							)
						self.assertTrue(hasattr(__UpperCAmelCase     ,					"""do_resize"""							)							)
						self.assertTrue(hasattr(__UpperCAmelCase     ,					"""size"""							)							)
						self.assertTrue(hasattr(__UpperCAmelCase     ,					"""do_center_crop"""							)							)
						self.assertTrue(hasattr(__UpperCAmelCase     ,					"""center_crop"""							)							)
						self.assertTrue(hasattr(__UpperCAmelCase     ,					"""do_flip_channel_order"""							)							)

			def        SCREAMING_SNAKE_CASE					(       self					:    Tuple							) ->     List[str]:
						SCREAMING_SNAKE_CASE__										=   self.image_processing_class.from_dict(self.image_processor_dict							)
						self.assertEqual(image_processor.size     ,					{"""shortest_edge""": 2_0}							)
						self.assertEqual(image_processor.crop_size     ,					{"""height""": 1_8, """width""": 1_8}							)

						SCREAMING_SNAKE_CASE__										=   self.image_processing_class.from_dict(self.image_processor_dict     ,					size=4_2     ,					crop_size=8_4							)
						self.assertEqual(image_processor.size     ,					{"""shortest_edge""": 4_2}							)
						self.assertEqual(image_processor.crop_size     ,					{"""height""": 8_4, """width""": 8_4}							)

			def        SCREAMING_SNAKE_CASE					(       self					:    Dict							) ->     Dict:
						pass

			def        SCREAMING_SNAKE_CASE					(       self					:    Any							) ->     List[str]:
						# Initialize image_processing
						SCREAMING_SNAKE_CASE__										=   self.image_processing_class(**self.image_processor_dict							)
						# create random PIL images
						SCREAMING_SNAKE_CASE__										=   prepare_image_inputs(self.image_processor_tester     ,					equal_resolution=__UpperCAmelCase							)
						for image in image_inputs:
									self.assertIsInstance(__UpperCAmelCase     ,					Image.Image							)

						# Test not batched input
						SCREAMING_SNAKE_CASE__										=   image_processing(image_inputs[0]     ,					return_tensors="""pt"""							).pixel_values
						self.assertEqual(
						    encoded_images.shape     ,					(
						        1,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size["""height"""],
						        self.image_processor_tester.crop_size["""width"""],
						    )     ,					)

						# Test batched
						SCREAMING_SNAKE_CASE__										=   image_processing(__UpperCAmelCase     ,					return_tensors="""pt"""							).pixel_values
						self.assertEqual(
						    encoded_images.shape     ,					(
						        self.image_processor_tester.batch_size,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size["""height"""],
						        self.image_processor_tester.crop_size["""width"""],
						    )     ,					)

			def        SCREAMING_SNAKE_CASE					(       self					:    List[Any]							) ->     Optional[int]:
						# Initialize image_processing
						SCREAMING_SNAKE_CASE__										=   self.image_processing_class(**self.image_processor_dict							)
						# create random numpy tensors
						SCREAMING_SNAKE_CASE__										=   prepare_image_inputs(self.image_processor_tester     ,					equal_resolution=__UpperCAmelCase     ,					numpify=__UpperCAmelCase							)
						for image in image_inputs:
									self.assertIsInstance(__UpperCAmelCase     ,					np.ndarray							)

						# Test not batched input
						SCREAMING_SNAKE_CASE__										=   image_processing(image_inputs[0]     ,					return_tensors="""pt"""							).pixel_values
						self.assertEqual(
						    encoded_images.shape     ,					(
						        1,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size["""height"""],
						        self.image_processor_tester.crop_size["""width"""],
						    )     ,					)

						# Test batched
						SCREAMING_SNAKE_CASE__										=   image_processing(__UpperCAmelCase     ,					return_tensors="""pt"""							).pixel_values
						self.assertEqual(
						    encoded_images.shape     ,					(
						        self.image_processor_tester.batch_size,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size["""height"""],
						        self.image_processor_tester.crop_size["""width"""],
						    )     ,					)





			def        SCREAMING_SNAKE_CASE					(       self					:    List[str]							) ->     str:
						# Initialize image_processing
						SCREAMING_SNAKE_CASE__										=   self.image_processing_class(**self.image_processor_dict							)
						# create random PyTorch tensors
						SCREAMING_SNAKE_CASE__										=   prepare_image_inputs(self.image_processor_tester     ,					equal_resolution=__UpperCAmelCase     ,					torchify=__UpperCAmelCase							)
						for image in image_inputs:
									self.assertIsInstance(__UpperCAmelCase     ,					torch.Tensor							)

						# Test not batched input
						SCREAMING_SNAKE_CASE__										=   image_processing(image_inputs[0]     ,					return_tensors="""pt"""							).pixel_values
						self.assertEqual(
						    encoded_images.shape     ,					(
						        1,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size["""height"""],
						        self.image_processor_tester.crop_size["""width"""],
						    )     ,					)

						# Test batched
						SCREAMING_SNAKE_CASE__										=   image_processing(__UpperCAmelCase     ,					return_tensors="""pt"""							).pixel_values
						self.assertEqual(
						    encoded_images.shape     ,					(
						        self.image_processor_tester.batch_size,
						        self.image_processor_tester.num_channels,
						        self.image_processor_tester.crop_size["""height"""],
						        self.image_processor_tester.crop_size["""width"""],
						    )     ,					)


