



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     :							Dict  							=				logging.get_logger(__name__)  # pylint: disable=invalid-name

__lowerCamelCase     :							Optional[int]  							=				256





class  __magic_name__   (							A__						):
		lowercase      :       Any          =['''melgan''']






		def __init__(						self      :     Dict				,		UpperCamelCase__      :     SpectrogramNotesEncoder				,		UpperCamelCase__      :     SpectrogramContEncoder				,		UpperCamelCase__      :     TaFilmDecoder				,		UpperCamelCase__      :     DDPMScheduler				,		UpperCamelCase__      :     OnnxRuntimeModel if is_onnx_available() else Any				,		)				->					None:
								'''simple docstring'''







								super().__init__()

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

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






		def SCREAMING_SNAKE_CASE_			(						self      :     Union[str, Any]				,		UpperCamelCase__      :     Any				,		UpperCamelCase__      :     List[str]=(-1.0, 1.0)				,		UpperCamelCase__      :     Optional[Any]=False      )				->					List[Any]:
								'''simple docstring'''







								UpperCAmelCase       ,		UpperCAmelCase         =     output_range
								if clip:
														UpperCAmelCase         =     torch.clip(UpperCamelCase__				,		self.min_value				,		self.max_value      )
								# Scale to [0, 1].
								UpperCAmelCase         =     (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      :     Dict				,		UpperCamelCase__      :     Optional[Any]				,		UpperCamelCase__      :     Optional[Any]=(-1.0, 1.0)				,		UpperCamelCase__      :     int=False      )				->					Tuple:
								'''simple docstring'''







								UpperCAmelCase       ,		UpperCAmelCase         =     input_range
								UpperCAmelCase         =     torch.clip(UpperCamelCase__				,		UpperCamelCase__				,		UpperCamelCase__      ) if clip else outputs
								# Scale to [0, 1].
								UpperCAmelCase         =     (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      :     Any				,		UpperCamelCase__      :     str				,		UpperCamelCase__      :     List[Any]				,		UpperCamelCase__      :     Any      )				->					List[Any]:
								'''simple docstring'''







								UpperCAmelCase         =     input_tokens > 0
								UpperCAmelCase       ,		UpperCAmelCase         =     self.notes_encoder(
								    encoder_input_tokens=UpperCamelCase__				,		encoder_inputs_mask=UpperCamelCase__      )

								UpperCAmelCase       ,		UpperCAmelCase         =     self.continuous_encoder(
								    encoder_inputs=UpperCamelCase__				,		encoder_inputs_mask=UpperCamelCase__      )

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






		def SCREAMING_SNAKE_CASE_			(						self      :     List[Any]				,		UpperCamelCase__      :     str				,		UpperCamelCase__      :     Any				,		UpperCamelCase__      :     Optional[Any]      )				->					int:
								'''simple docstring'''







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

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

								UpperCAmelCase         =     self.decoder(
								    encodings_and_masks=UpperCamelCase__				,		decoder_input_tokens=UpperCamelCase__				,		decoder_noise_time=UpperCamelCase__      )
								return logits







		@torch.no_grad()
		def __call__(						self      :     List[str]				,		UpperCamelCase__      :     List[List[int]]				,		UpperCamelCase__      :     Optional[torch.Generator] = None				,		UpperCamelCase__      :     int = 1_00				,		UpperCamelCase__      :     bool = True				,		UpperCamelCase__      :     str = "numpy"				,		UpperCamelCase__      :     Optional[Callable[[int, int, torch.FloatTensor], None]] = None				,		UpperCamelCase__      :     int = 1				,		)				->					Union[AudioPipelineOutput, Tuple]:
								'''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__      )}.'      )

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

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

														UpperCAmelCase         =     self.scale_features(
														    UpperCamelCase__				,		output_range=[-1.0, 1.0]				,		clip=UpperCamelCase__      )

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

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

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

								if not return_dict:
														return (output,)

								return AudioPipelineOutput(audios=UpperCamelCase__      )








from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
		import PIL


__lowerCamelCase     :							List[Any]  							=				logging.get_logger(__name__)





class  __magic_name__   (							A__						):
		lowercase      :       Tuple          =['''pixel_values''']






		def __init__(						self      :     Any				,		UpperCamelCase__      :     bool = True				,		UpperCamelCase__      :     Dict[str, int] = None				,		UpperCamelCase__      :     float = None				,		UpperCamelCase__      :     PILImageResampling = PILImageResampling.BILINEAR				,		UpperCamelCase__      :     bool = True				,		UpperCamelCase__      :     Union[int, float] = 1 / 2_55				,		UpperCamelCase__      :     bool = True				,		UpperCamelCase__      :     Optional[Union[float, List[float]]] = None				,		UpperCamelCase__      :     Optional[Union[float, List[float]]] = None				,		**UpperCamelCase__      :     List[str]				,		)				->					None:
								'''simple docstring'''







								super().__init__(**UpperCamelCase__      )
								UpperCAmelCase         =     size if size is not None else {"shortest_edge": 3_84}
								UpperCAmelCase         =     get_size_dict(UpperCamelCase__				,		default_to_square=UpperCamelCase__      )

								UpperCAmelCase         =     do_resize
								UpperCAmelCase         =     size
								# Default value set here for backwards compatibility where the value in config is None
								UpperCAmelCase         =     crop_pct if crop_pct is not None else 2_24 / 2_56
								UpperCAmelCase         =     resample
								UpperCAmelCase         =     do_rescale
								UpperCAmelCase         =     rescale_factor
								UpperCAmelCase         =     do_normalize
								UpperCAmelCase         =     image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
								UpperCAmelCase         =     image_std if image_std is not None else IMAGENET_STANDARD_STD






		def SCREAMING_SNAKE_CASE_			(						self      :     Any				,		UpperCamelCase__      :     np.ndarray				,		UpperCamelCase__      :     Dict[str, int]				,		UpperCamelCase__      :     float				,		UpperCamelCase__      :     PILImageResampling = PILImageResampling.BICUBIC				,		UpperCamelCase__      :     Optional[Union[str, ChannelDimension]] = None				,		**UpperCamelCase__      :     Union[str, Any]				,		)				->					np.ndarray:
								'''simple docstring'''







								UpperCAmelCase         =     get_size_dict(UpperCamelCase__				,		default_to_square=UpperCamelCase__      )
								if "shortest_edge" not in size:
														raise ValueError(F'Size dictionary must contain \'shortest_edge\' key. Got {size.keys()}'      )
								UpperCAmelCase         =     size["shortest_edge"]

								if shortest_edge < 3_84:
														# maintain same ratio, resizing shortest edge to shortest_edge/crop_pct
														UpperCAmelCase         =     int(shortest_edge / crop_pct      )
														UpperCAmelCase         =     get_resize_output_image_size(UpperCamelCase__				,		size=UpperCamelCase__				,		default_to_square=UpperCamelCase__      )
														UpperCAmelCase         =     resize(image=UpperCamelCase__				,		size=UpperCamelCase__				,		resample=UpperCamelCase__				,		data_format=UpperCamelCase__				,		**UpperCamelCase__      )
														# then crop to (shortest_edge, shortest_edge)
														return center_crop(image=UpperCamelCase__				,		size=(shortest_edge, shortest_edge)				,		data_format=UpperCamelCase__				,		**UpperCamelCase__      )
								else:
														# warping (no cropping) when evaluated at 384 or larger
														return resize(
														    UpperCamelCase__				,		size=(shortest_edge, shortest_edge)				,		resample=UpperCamelCase__				,		data_format=UpperCamelCase__				,		**UpperCamelCase__      )






		def SCREAMING_SNAKE_CASE_			(						self      :     Optional[Any]				,		UpperCamelCase__      :     np.ndarray				,		UpperCamelCase__      :     Union[int, float]				,		UpperCamelCase__      :     Optional[Union[str, ChannelDimension]] = None				,		**UpperCamelCase__      :     Optional[Any]				,		)				->					List[str]:
								'''simple docstring'''







								return rescale(UpperCamelCase__				,		scale=UpperCamelCase__				,		data_format=UpperCamelCase__				,		**UpperCamelCase__      )






		def SCREAMING_SNAKE_CASE_			(						self      :     List[str]				,		UpperCamelCase__      :     np.ndarray				,		UpperCamelCase__      :     Union[float, List[float]]				,		UpperCamelCase__      :     Union[float, List[float]]				,		UpperCamelCase__      :     Optional[Union[str, ChannelDimension]] = None				,		**UpperCamelCase__      :     Optional[Any]				,		)				->					np.ndarray:
								'''simple docstring'''







								return normalize(UpperCamelCase__				,		mean=UpperCamelCase__				,		std=UpperCamelCase__				,		data_format=UpperCamelCase__				,		**UpperCamelCase__      )







		def SCREAMING_SNAKE_CASE_			(						self      :     Tuple				,		UpperCamelCase__      :     ImageInput				,		UpperCamelCase__      :     bool = None				,		UpperCamelCase__      :     Dict[str, int] = None				,		UpperCamelCase__      :     float = None				,		UpperCamelCase__      :     PILImageResampling = None				,		UpperCamelCase__      :     bool = None				,		UpperCamelCase__      :     float = None				,		UpperCamelCase__      :     bool = None				,		UpperCamelCase__      :     Optional[Union[float, List[float]]] = None				,		UpperCamelCase__      :     Optional[Union[float, List[float]]] = None				,		UpperCamelCase__      :     Optional[Union[str, TensorType]] = None				,		UpperCamelCase__      :     ChannelDimension = ChannelDimension.FIRST				,		**UpperCamelCase__      :     Any				,		)				->					PIL.Image.Image:
								'''simple docstring'''







								UpperCAmelCase         =     do_resize if do_resize is not None else self.do_resize
								UpperCAmelCase         =     crop_pct if crop_pct is not None else self.crop_pct
								UpperCAmelCase         =     resample if resample is not None else self.resample
								UpperCAmelCase         =     do_rescale if do_rescale is not None else self.do_rescale
								UpperCAmelCase         =     rescale_factor if rescale_factor is not None else self.rescale_factor
								UpperCAmelCase         =     do_normalize if do_normalize is not None else self.do_normalize
								UpperCAmelCase         =     image_mean if image_mean is not None else self.image_mean
								UpperCAmelCase         =     image_std if image_std is not None else self.image_std

								UpperCAmelCase         =     size if size is not None else self.size
								UpperCAmelCase         =     get_size_dict(UpperCamelCase__				,		default_to_square=UpperCamelCase__      )

								UpperCAmelCase         =     make_list_of_images(UpperCamelCase__      )

								if not valid_images(UpperCamelCase__      ):
														raise ValueError(
														    "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
														    "torch.Tensor, tf.Tensor or jax.ndarray."      )

								if do_resize and size is None or resample is None:
														raise ValueError("Size and resample must be specified if do_resize is True."      )

								if do_resize and size["shortest_edge"] < 3_84 and crop_pct is None:
														raise ValueError("crop_pct must be specified if size < 384."      )

								if do_rescale and rescale_factor is None:
														raise ValueError("Rescale factor must be specified if do_rescale is True."      )

								if do_normalize and (image_mean is None or image_std is None):
														raise ValueError("Image mean and std must be specified if do_normalize is True."      )

								# All transformations expect numpy arrays.
								UpperCAmelCase         =     [to_numpy_array(UpperCamelCase__      ) for image in images]

								if do_resize:
														UpperCAmelCase         =     [self.resize(image=UpperCamelCase__				,		size=UpperCamelCase__				,		crop_pct=UpperCamelCase__				,		resample=UpperCamelCase__      ) for image in images]

								if do_rescale:
														UpperCAmelCase         =     [self.rescale(image=UpperCamelCase__				,		scale=UpperCamelCase__      ) for image in images]

								if do_normalize:
														UpperCAmelCase         =     [self.normalize(image=UpperCamelCase__				,		mean=UpperCamelCase__				,		std=UpperCamelCase__      ) for image in images]

								UpperCAmelCase         =     [to_channel_dimension_format(UpperCamelCase__				,		UpperCamelCase__      ) for image in images]

								UpperCAmelCase         =     {"pixel_values": images}
								return BatchFeature(data=UpperCamelCase__				,		tensor_type=UpperCamelCase__      )



