




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


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

__A								=       2_56







class        __lowerCAmelCase					(			__magic_name__					):

			"""simple docstring"""







			snake_case_  						=			['''melgan''']


			def __init__( self     ,		lowerCamelCase__     ,		lowerCamelCase__     ,		lowerCamelCase__     ,		lowerCamelCase__     ,		lowerCamelCase__     ,		)       -> None:



							'''simple docstring'''
							super().__init__()

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

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


			def  lowercase_       ( self     ,		lowerCamelCase__     ,		lowerCamelCase__=(-1.0, 1.0)     ,		lowerCamelCase__=False				)       -> Optional[Any]:



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



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



							'''simple docstring'''
							__lowerCamelCase									=       input_tokens > 0
							__lowerCamelCase			,		__lowerCamelCase									=       self.notes_encoder(
							    encoder_input_tokens=lowerCamelCase__     ,		encoder_inputs_mask=lowerCamelCase__				)

							__lowerCamelCase			,		__lowerCamelCase									=       self.continuous_encoder(
							    encoder_inputs=lowerCamelCase__     ,		encoder_inputs_mask=lowerCamelCase__				)

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


			def  lowercase_       ( self     ,		lowerCamelCase__     ,		lowerCamelCase__     ,		lowerCamelCase__				)       -> int:



							'''simple docstring'''
							__lowerCamelCase									=       noise_time
							if not torch.is_tensor(lowerCamelCase__				):
											__lowerCamelCase									=       torch.tensor([timesteps]     ,		dtype=torch.long     ,		device=input_tokens.device				)
							elif torch.is_tensor(lowerCamelCase__				) and len(timesteps.shape				) == 0:
											__lowerCamelCase									=       timesteps[None].to(input_tokens.device				)

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

							__lowerCamelCase									=       self.decoder(
							    encodings_and_masks=lowerCamelCase__     ,		decoder_input_tokens=lowerCamelCase__     ,		decoder_noise_time=lowerCamelCase__				)
							return logits






			@torch.no_grad()
			def __call__( self     ,		lowerCamelCase__     ,		lowerCamelCase__ = None     ,		lowerCamelCase__ = 100     ,		lowerCamelCase__ = True     ,		lowerCamelCase__ = "numpy"     ,		lowerCamelCase__ = None     ,		lowerCamelCase__ = 1     ,		)       -> Union[AudioPipelineOutput, Tuple]:



							'''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									=       np.zeros([1, TARGET_FEATURE_LENGTH, self.n_dims]     ,		dtype=np.floataa				)
							__lowerCamelCase									=       np.zeros([1, 0, self.n_dims]     ,		np.floataa				)
							__lowerCamelCase									=       torch.ones((1, TARGET_FEATURE_LENGTH)     ,		dtype=lowerCamelCase__     ,		device=self.device				)

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

											__lowerCamelCase									=       self.scale_features(
											    lowerCamelCase__     ,		output_range=[-1.0, 1.0]     ,		clip=lowerCamelCase__				)

											__lowerCamelCase									=       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									=       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									=       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									=       self.scheduler.step(lowerCamelCase__     ,		lowerCamelCase__     ,		lowerCamelCase__     ,		generator=lowerCamelCase__				).prev_sample

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

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

							if not return_dict:
											return (output,)

							return AudioPipelineOutput(audios=lowerCamelCase__				)








import warnings
from typing import List

import numpy as np

from ...processing_utils import ProcessorMixin
from ...tokenization_utils_base import BatchEncoding
from ...utils import is_flax_available, is_tf_available, is_torch_available







class        __lowerCAmelCase					(			__magic_name__					):

			"""simple docstring"""







			snake_case_  						=			['''image_processor''', '''tokenizer''']
			snake_case_  						=			'''OwlViTImageProcessor'''
			snake_case_  						=			('''CLIPTokenizer''', '''CLIPTokenizerFast''')


			def __init__( self     ,		lowerCamelCase__=None     ,		lowerCamelCase__=None     ,		**lowerCamelCase__				)       -> Tuple:



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

							__lowerCamelCase									=       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__(lowerCamelCase__     ,		lowerCamelCase__				)


			def __call__( self     ,		lowerCamelCase__=None     ,		lowerCamelCase__=None     ,		lowerCamelCase__=None     ,		lowerCamelCase__="max_length"     ,		lowerCamelCase__="np"     ,		**lowerCamelCase__				)       -> Optional[int]:



							'''simple docstring'''

							if text is None and query_images is None and images is None:
											raise ValueError(
											    'You have to specify at least one text or query image or image. All three cannot be none.'				)

							if text is not None:
											if isinstance(lowerCamelCase__     ,		lowerCamelCase__				) or (isinstance(lowerCamelCase__     ,		lowerCamelCase__				) and not isinstance(text[0]     ,		lowerCamelCase__				)):
															__lowerCamelCase									=       [self.tokenizer(lowerCamelCase__     ,		padding=lowerCamelCase__     ,		return_tensors=lowerCamelCase__     ,		**lowerCamelCase__				)]

											elif isinstance(lowerCamelCase__     ,		lowerCamelCase__				) and isinstance(text[0]     ,		lowerCamelCase__				):
															__lowerCamelCase									=       []

															# Maximum number of queries across batch
															__lowerCamelCase									=       max([len(lowerCamelCase__				) for t in text]				)

															# Pad all batch samples to max number of text queries
															for t in text:
																			if len(lowerCamelCase__				) != max_num_queries:
																							__lowerCamelCase									=       t + [' '] * (max_num_queries - len(lowerCamelCase__				))

																			__lowerCamelCase									=       self.tokenizer(lowerCamelCase__     ,		padding=lowerCamelCase__     ,		return_tensors=lowerCamelCase__     ,		**lowerCamelCase__				)
																			encodings.append(lowerCamelCase__				)
											else:
															raise TypeError('Input text should be a string, a list of strings or a nested list of strings'				)

											if return_tensors == "np":
															__lowerCamelCase									=       np.concatenate([encoding['input_ids'] for encoding in encodings]     ,		axis=0				)
															__lowerCamelCase									=       np.concatenate([encoding['attention_mask'] for encoding in encodings]     ,		axis=0				)

											elif return_tensors == "jax" and is_flax_available():
															import jax.numpy as jnp

															__lowerCamelCase									=       jnp.concatenate([encoding['input_ids'] for encoding in encodings]     ,		axis=0				)
															__lowerCamelCase									=       jnp.concatenate([encoding['attention_mask'] for encoding in encodings]     ,		axis=0				)

											elif return_tensors == "pt" and is_torch_available():
															import torch

															__lowerCamelCase									=       torch.cat([encoding['input_ids'] for encoding in encodings]     ,		dim=0				)
															__lowerCamelCase									=       torch.cat([encoding['attention_mask'] for encoding in encodings]     ,		dim=0				)

											elif return_tensors == "tf" and is_tf_available():
															import tensorflow as tf

															__lowerCamelCase									=       tf.stack([encoding['input_ids'] for encoding in encodings]     ,		axis=0				)
															__lowerCamelCase									=       tf.stack([encoding['attention_mask'] for encoding in encodings]     ,		axis=0				)

											else:
															raise ValueError('Target return tensor type could not be returned'				)

											__lowerCamelCase									=       BatchEncoding()
											__lowerCamelCase									=       input_ids
											__lowerCamelCase									=       attention_mask

							if query_images is not None:
											__lowerCamelCase									=       BatchEncoding()
											__lowerCamelCase									=       self.image_processor(
											    lowerCamelCase__     ,		return_tensors=lowerCamelCase__     ,		**lowerCamelCase__				).pixel_values
											__lowerCamelCase									=       query_pixel_values

							if images is not None:
											__lowerCamelCase									=       self.image_processor(lowerCamelCase__     ,		return_tensors=lowerCamelCase__     ,		**lowerCamelCase__				)

							if text is not None and images is not None:
											__lowerCamelCase									=       image_features.pixel_values
											return encoding
							elif query_images is not None and images is not None:
											__lowerCamelCase									=       image_features.pixel_values
											return encoding
							elif text is not None or query_images is not None:
											return encoding
							else:
											return BatchEncoding(data=dict(**lowerCamelCase__				)     ,		tensor_type=lowerCamelCase__				)


			def  lowercase_       ( self     ,		*lowerCamelCase__     ,		**lowerCamelCase__				)       -> Tuple:



							'''simple docstring'''
							return self.image_processor.post_process(*lowerCamelCase__     ,		**lowerCamelCase__				)


			def  lowercase_       ( self     ,		*lowerCamelCase__     ,		**lowerCamelCase__				)       -> Optional[Any]:



							'''simple docstring'''
							return self.image_processor.post_process_object_detection(*lowerCamelCase__     ,		**lowerCamelCase__				)


			def  lowercase_       ( self     ,		*lowerCamelCase__     ,		**lowerCamelCase__				)       -> Optional[int]:



							'''simple docstring'''
							return self.image_processor.post_process_image_guided_detection(*lowerCamelCase__     ,		**lowerCamelCase__				)


			def  lowercase_       ( self     ,		*lowerCamelCase__     ,		**lowerCamelCase__				)       -> Optional[int]:



							'''simple docstring'''
							return self.tokenizer.batch_decode(*lowerCamelCase__     ,		**lowerCamelCase__				)


			def  lowercase_       ( self     ,		*lowerCamelCase__     ,		**lowerCamelCase__				)       -> Optional[int]:



							'''simple docstring'''
							return self.tokenizer.decode(*lowerCamelCase__     ,		**lowerCamelCase__				)


			@property
			def  lowercase_       ( self				)       -> Dict:



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






			@property
			def  lowercase_       ( self				)       -> Union[str, Any]:



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


