

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

_lowerCAmelCase					:			Dict  				=			256






class 	__magic_name__    (			lowerCAmelCase_      ):
			SCREAMING_SNAKE_CASE							   =						['melgan']





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




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

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

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





			def 		__magic_name__    (		self					,  __snake_case					,  __snake_case=(-1.0, 1.0)					,  __snake_case=False					)   -> Optional[Any]:




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




						'''simple docstring'''
						__a      ,    __a        =input_range
						__a        =torch.clip(__snake_case					,  __snake_case					,  __snake_case					) if clip else outputs
						# Scale to [0, 1].
						__a        =(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 		__magic_name__    (		self					,  __snake_case					,  __snake_case					,  __snake_case					)   -> Dict:




						'''simple docstring'''
						__a        =input_tokens > 0
						__a      ,    __a        =self.notes_encoder(
						    encoder_input_tokens=__snake_case					,  encoder_inputs_mask=__snake_case					)

						__a      ,    __a        =self.continuous_encoder(
						    encoder_inputs=__snake_case					,  encoder_inputs_mask=__snake_case					)

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





			def 		__magic_name__    (		self					,  __snake_case					,  __snake_case					,  __snake_case					)   -> List[Any]:




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

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

						__a        =self.decoder(
						    encodings_and_masks=__snake_case					,  decoder_input_tokens=__snake_case					,  decoder_noise_time=__snake_case					)
						return logits





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




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

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

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

									__a        =self.scale_features(
									    __snake_case					,  output_range=[-1.0, 1.0]					,  clip=__snake_case					)

									__a        =self.encode(
									    input_tokens=torch.IntTensor([encoder_input_tokens]					).to(device=self.device					)					,  continuous_inputs=__snake_case					,  continuous_mask=__snake_case					,  )

									# Sample encoder_continuous_inputs shaped gaussian noise to begin loop
									__a        =randn_tensor(
									    shape=encoder_continuous_inputs.shape					,  generator=__snake_case					,  device=self.device					,  dtype=self.decoder.dtype					,  )

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

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

												# Compute previous output: x_t -> x_t-1
												__a        =self.scheduler.step(__snake_case					,  __snake_case					,  __snake_case					,  generator=__snake_case					).prev_sample

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

									__a        =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(__snake_case					,  __snake_case					)

									logger.info('Generated segment'					,  __snake_case					)

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

						if not return_dict:
									return (output,)

						return AudioPipelineOutput(audios=__snake_case					)









import copy
import os
from typing import Union

from ...configuration_utils import PretrainedConfig
from ...models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
from ...utils import logging
from ..auto import CONFIG_MAPPING


_lowerCAmelCase					:			List[str]  				=			logging.get_logger(__name__)

_lowerCAmelCase					:			Union[str, Any]  				=			{
    "Salesforce/instruct-blip-flan-t5": "https://huggingface.co/Salesforce/instruct-blip-flan-t5/resolve/main/config.json",
}






class 	__magic_name__    (			lowerCAmelCase_      ):
			SCREAMING_SNAKE_CASE							   =						'instructblip_vision_model'





			def __init__(		self					,  __snake_case=1408					,  __snake_case=6144					,  __snake_case=39					,  __snake_case=16					,  __snake_case=224					,  __snake_case=14					,  __snake_case="gelu"					,  __snake_case=1e-6					,  __snake_case=0.0					,  __snake_case=1e-10					,  __snake_case=True					,  **__snake_case					,  )   -> str:




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

						__a        =hidden_size
						__a        =intermediate_size
						__a        =num_hidden_layers
						__a        =num_attention_heads
						__a        =patch_size
						__a        =image_size
						__a        =initializer_range
						__a        =attention_dropout
						__a        =layer_norm_eps
						__a        =hidden_act
						__a        =qkv_bias





			@classmethod
			def 		__magic_name__    (		cls					,  __snake_case					,  **__snake_case					)   -> "PretrainedConfig":




						'''simple docstring'''
						cls._set_token_in_kwargs(__snake_case					)

						__a      ,    __a        =cls.get_config_dict(__snake_case					,  **__snake_case					)

						# get the vision config dict if we are loading from InstructBlipConfig
						if config_dict.get('model_type'					) == "instructblip":
									__a        =config_dict['vision_config']

						if "model_type" in config_dict and hasattr(cls					,  'model_type'					) and config_dict["model_type"] != cls.model_type:
									logger.warning(
									    f'You are using a model of type {config_dict["model_type"]} to instantiate a model of type '
									    f'{cls.model_type}. This is not supported for all configurations of models and can yield errors.'					)

						return cls.from_dict(__snake_case					,  **__snake_case					)






class 	__magic_name__    (			lowerCAmelCase_      ):
			SCREAMING_SNAKE_CASE							   =						'instructblip_qformer'





			def __init__(		self					,  __snake_case=3_0522					,  __snake_case=768					,  __snake_case=12					,  __snake_case=12					,  __snake_case=3072					,  __snake_case="gelu"					,  __snake_case=0.1					,  __snake_case=0.1					,  __snake_case=512					,  __snake_case=0.02					,  __snake_case=1e-12					,  __snake_case=0					,  __snake_case="absolute"					,  __snake_case=2					,  __snake_case=1408					,  **__snake_case					,  )   -> List[str]:




						'''simple docstring'''
						super().__init__(pad_token_id=__snake_case					,  **__snake_case					)

						__a        =vocab_size
						__a        =hidden_size
						__a        =num_hidden_layers
						__a        =num_attention_heads
						__a        =hidden_act
						__a        =intermediate_size
						__a        =hidden_dropout_prob
						__a        =attention_probs_dropout_prob
						__a        =max_position_embeddings
						__a        =initializer_range
						__a        =layer_norm_eps
						__a        =position_embedding_type
						__a        =cross_attention_frequency
						__a        =encoder_hidden_size





			@classmethod
			def 		__magic_name__    (		cls					,  __snake_case					,  **__snake_case					)   -> "PretrainedConfig":




						'''simple docstring'''
						cls._set_token_in_kwargs(__snake_case					)

						__a      ,    __a        =cls.get_config_dict(__snake_case					,  **__snake_case					)

						# get the qformer config dict if we are loading from InstructBlipConfig
						if config_dict.get('model_type'					) == "instructblip":
									__a        =config_dict['qformer_config']

						if "model_type" in config_dict and hasattr(cls					,  'model_type'					) and config_dict["model_type"] != cls.model_type:
									logger.warning(
									    f'You are using a model of type {config_dict["model_type"]} to instantiate a model of type '
									    f'{cls.model_type}. This is not supported for all configurations of models and can yield errors.'					)

						return cls.from_dict(__snake_case					,  **__snake_case					)






class 	__magic_name__    (			lowerCAmelCase_      ):
			SCREAMING_SNAKE_CASE							   =						'instructblip'
			SCREAMING_SNAKE_CASE							   =						True





			def __init__(		self					,  __snake_case=None					,  __snake_case=None					,  __snake_case=None					,  __snake_case=32					,  **__snake_case					)   -> str:




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

						if vision_config is None:
									__a        ={}
									logger.info('vision_config is None. initializing the InstructBlipVisionConfig with default values.'					)

						if qformer_config is None:
									__a        ={}
									logger.info('qformer_config is None. Initializing the InstructBlipQFormerConfig with default values.'					)

						if text_config is None:
									__a        ={}
									logger.info('text_config is None. Initializing the text config with default values (`OPTConfig`).'					)

						__a        =InstructBlipVisionConfig(**__snake_case					)
						__a        =InstructBlipQFormerConfig(**__snake_case					)
						__a        =text_config['model_type'] if 'model_type' in text_config else 'opt'
						__a        =CONFIG_MAPPING[text_model_type](**__snake_case					)

						__a        =self.text_config.tie_word_embeddings
						__a        =self.text_config.is_encoder_decoder

						__a        =num_query_tokens
						__a        =self.vision_config.hidden_size
						__a        =self.text_config.model_type in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
						__a        =1.0
						__a        =0.02





			@classmethod
			def 		__magic_name__    (		cls					,  __snake_case					,  __snake_case					,  __snake_case					,  **__snake_case					,  )   -> Optional[Any]:




						'''simple docstring'''
						return cls(
						    vision_config=vision_config.to_dict()					,  qformer_config=qformer_config.to_dict()					,  text_config=text_config.to_dict()					,  **__snake_case					,  )





			def 		__magic_name__    (		self					)   -> List[Any]:




						'''simple docstring'''
						__a        =copy.deepcopy(self.__dict__					)
						__a        =self.vision_config.to_dict()
						__a        =self.qformer_config.to_dict()
						__a        =self.text_config.to_dict()
						__a        =self.__class__.model_type
						return output






