





import copy
import os
from typing import TYPE_CHECKING, List, Union


if TYPE_CHECKING:
       pass

from ...configuration_utils import PretrainedConfig
from ...utils import logging


UpperCamelCase__      =logging.get_logger(__name__)

UpperCamelCase__      ={
    'kakaobrain/align-base': 'https://huggingface.co/kakaobrain/align-base/resolve/main/config.json',
}





class 			lowerCAmelCase__( __lowercase    ):






       '''simple docstring'''


       __snake_case 			=    'align_text_model'




       def __init__(      self    ,						__lowerCamelCase=3_0_5_2_2    ,						__lowerCamelCase=7_6_8    ,						__lowerCamelCase=1_2    ,						__lowerCamelCase=1_2    ,						__lowerCamelCase=3_0_7_2    ,						__lowerCamelCase="gelu"    ,						__lowerCamelCase=0.1    ,						__lowerCamelCase=0.1    ,						__lowerCamelCase=5_1_2    ,						__lowerCamelCase=2    ,						__lowerCamelCase=0.02    ,						__lowerCamelCase=1E-12    ,						__lowerCamelCase=0    ,						__lowerCamelCase="absolute"    ,						__lowerCamelCase=True    ,						**__lowerCamelCase    ,						)			->     List[Any]:
            super().__init__(**__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	Tuple						    =							vocab_size
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							hidden_size
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							num_hidden_layers
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							num_attention_heads
            _SCREAMING_SNAKE_CASE       :	int						    =							hidden_act
            _SCREAMING_SNAKE_CASE       :	Any						    =							intermediate_size
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							hidden_dropout_prob
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							attention_probs_dropout_prob
            _SCREAMING_SNAKE_CASE       :	Dict						    =							max_position_embeddings
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							type_vocab_size
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							initializer_range
            _SCREAMING_SNAKE_CASE       :	Dict						    =							layer_norm_eps
            _SCREAMING_SNAKE_CASE       :	str						    =							position_embedding_type
            _SCREAMING_SNAKE_CASE       :	Dict						    =							use_cache
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							pad_token_id







       @classmethod
       def     UpperCamelCase_			(      cls    ,						__lowerCamelCase    ,						**__lowerCamelCase							)			->     "PretrainedConfig":
            cls._set_token_in_kwargs(__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE       :	List[Any]						    =							cls.get_config_dict(__lowerCamelCase    ,						**__lowerCamelCase							)

            # get the text config dict if we are loading from AlignConfig
            if config_dict.get("model_type"							) == "align":
                 _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							config_dict["text_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(__lowerCamelCase    ,						**__lowerCamelCase							)





class 			lowerCAmelCase__( __lowercase    ):






       '''simple docstring'''


       __snake_case 			=    'align_vision_model'




       def __init__(      self    ,						__lowerCamelCase = 3    ,						__lowerCamelCase = 6_0_0    ,						__lowerCamelCase = 2.0    ,						__lowerCamelCase = 3.1    ,						__lowerCamelCase = 8    ,						__lowerCamelCase = [3, 3, 5, 3, 5, 5, 3]    ,						__lowerCamelCase = [3_2, 1_6, 2_4, 4_0, 8_0, 1_1_2, 1_9_2]    ,						__lowerCamelCase = [1_6, 2_4, 4_0, 8_0, 1_1_2, 1_9_2, 3_2_0]    ,						__lowerCamelCase = []    ,						__lowerCamelCase = [1, 2, 2, 2, 1, 2, 1]    ,						__lowerCamelCase = [1, 2, 2, 3, 3, 4, 1]    ,						__lowerCamelCase = [1, 6, 6, 6, 6, 6, 6]    ,						__lowerCamelCase = 0.25    ,						__lowerCamelCase = "swish"    ,						__lowerCamelCase = 2_5_6_0    ,						__lowerCamelCase = "mean"    ,						__lowerCamelCase = 0.02    ,						__lowerCamelCase = 0.001    ,						__lowerCamelCase = 0.99    ,						__lowerCamelCase = 0.2    ,						**__lowerCamelCase    ,						)			->     Dict:
            super().__init__(**__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	int						    =							num_channels
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							image_size
            _SCREAMING_SNAKE_CASE       :	Dict						    =							width_coefficient
            _SCREAMING_SNAKE_CASE       :	str						    =							depth_coefficient
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							depth_divisor
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							kernel_sizes
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							in_channels
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							out_channels
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							depthwise_padding
            _SCREAMING_SNAKE_CASE       :	str						    =							strides
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							num_block_repeats
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							expand_ratios
            _SCREAMING_SNAKE_CASE       :	int						    =							squeeze_expansion_ratio
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							hidden_act
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							hidden_dim
            _SCREAMING_SNAKE_CASE       :	Dict						    =							pooling_type
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							initializer_range
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							batch_norm_eps
            _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							batch_norm_momentum
            _SCREAMING_SNAKE_CASE       :	int						    =							drop_connect_rate
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							sum(__lowerCamelCase							) * 4







       @classmethod
       def     UpperCamelCase_			(      cls    ,						__lowerCamelCase    ,						**__lowerCamelCase							)			->     "PretrainedConfig":
            cls._set_token_in_kwargs(__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE    ,			_SCREAMING_SNAKE_CASE       :	int						    =							cls.get_config_dict(__lowerCamelCase    ,						**__lowerCamelCase							)

            # get the vision config dict if we are loading from AlignConfig
            if config_dict.get("model_type"							) == "align":
                 _SCREAMING_SNAKE_CASE       :	int						    =							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(__lowerCamelCase    ,						**__lowerCamelCase							)





class 			lowerCAmelCase__( __lowercase    ):






       '''simple docstring'''


       __snake_case 			=    'align'
       __snake_case 			=    True




       def __init__(      self    ,						__lowerCamelCase=None    ,						__lowerCamelCase=None    ,						__lowerCamelCase=6_4_0    ,						__lowerCamelCase=1.0    ,						__lowerCamelCase=0.02    ,						**__lowerCamelCase    ,						)			->     List[Any]:
            super().__init__(**__lowerCamelCase							)

            if text_config is None:
                 _SCREAMING_SNAKE_CASE       :	List[Any]						    =							{}
                 logger.info("text_config is None. Initializing the AlignTextConfig with default values."							)

            if vision_config is None:
                 _SCREAMING_SNAKE_CASE       :	List[str]						    =							{}
                 logger.info("vision_config is None. Initializing the AlignVisionConfig with default values."							)

            _SCREAMING_SNAKE_CASE       :	Dict						    =							AlignTextConfig(**__lowerCamelCase							)
            _SCREAMING_SNAKE_CASE       :	Any						    =							AlignVisionConfig(**__lowerCamelCase							)

            _SCREAMING_SNAKE_CASE       :	int						    =							projection_dim
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							temperature_init_value
            _SCREAMING_SNAKE_CASE       :	Any						    =							initializer_range




       @classmethod
       def     UpperCamelCase_			(      cls    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						**__lowerCamelCase							)			->     List[str]:
            return cls(text_config=text_config.to_dict()    ,						vision_config=vision_config.to_dict()    ,						**__lowerCamelCase							)







       def     UpperCamelCase_			(      self							)			->     Union[str, Any]:
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							copy.deepcopy(self.__dict__							)
            _SCREAMING_SNAKE_CASE       :	Any						    =							self.text_config.to_dict()
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							self.vision_config.to_dict()
            _SCREAMING_SNAKE_CASE       :	Dict						    =							self.__class__.model_type
            return output






from math import acos, sin
from typing import List, Tuple, Union

import numpy as np
import torch
from PIL import Image

from ...models import AutoencoderKL, UNetaDConditionModel
from ...schedulers import DDIMScheduler, DDPMScheduler
from ...utils import randn_tensor
from ..pipeline_utils import AudioPipelineOutput, BaseOutput, DiffusionPipeline, ImagePipelineOutput
from .mel import Mel





class 			lowerCAmelCase__( __lowercase    ):






       '''simple docstring'''


       __snake_case 			=    ['vqvae']




       def __init__(      self    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase    ,						)			->     Optional[Any]:
            super().__init__()
            self.register_modules(unet=__lowerCamelCase    ,						scheduler=__lowerCamelCase    ,						mel=__lowerCamelCase    ,						vqvae=__lowerCamelCase							)




       def     UpperCamelCase_			(      self							)			->     int:
            return 5_0 if isinstance(self.scheduler    ,						__lowerCamelCase							) else 1_0_0_0




       @torch.no_grad()
       def __call__(      self    ,						__lowerCamelCase = 1    ,						__lowerCamelCase = None    ,						__lowerCamelCase = None    ,						__lowerCamelCase = 0    ,						__lowerCamelCase = 0    ,						__lowerCamelCase = None    ,						__lowerCamelCase = None    ,						__lowerCamelCase = 0    ,						__lowerCamelCase = 0    ,						__lowerCamelCase = None    ,						__lowerCamelCase = 0    ,						__lowerCamelCase = None    ,						__lowerCamelCase = None    ,						__lowerCamelCase=True    ,						)			->     Union[
           Union[AudioPipelineOutput, ImagePipelineOutput],
           Tuple[List[Image.Image], Tuple[int, List[np.ndarray]]],
       ]:
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							steps or self.get_default_steps()
            self.scheduler.set_timesteps(__lowerCamelCase							)
            _SCREAMING_SNAKE_CASE       :	List[Any]						    =							step_generator or generator
            # For backwards compatibility
            if type(self.unet.config.sample_size							) == int:
                 _SCREAMING_SNAKE_CASE       :	Dict						    =							(self.unet.config.sample_size, self.unet.config.sample_size)
            if noise is None:
                 _SCREAMING_SNAKE_CASE       :	Tuple						    =							randn_tensor(
                     (
                         batch_size,
                         self.unet.config.in_channels,
                         self.unet.config.sample_size[0],
                         self.unet.config.sample_size[1],
                     )    ,						generator=__lowerCamelCase    ,						device=self.device    ,						)
            _SCREAMING_SNAKE_CASE       :	Any						    =							noise
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							None

            if audio_file is not None or raw_audio is not None:
                 self.mel.load_audio(__lowerCamelCase    ,						__lowerCamelCase							)
                 _SCREAMING_SNAKE_CASE       :	Tuple						    =							self.mel.audio_slice_to_image(__lowerCamelCase							)
                 _SCREAMING_SNAKE_CASE       :	str						    =							np.frombuffer(input_image.tobytes()    ,						dtype="uint8"							).reshape(
                     (input_image.height, input_image.width)							)
                 _SCREAMING_SNAKE_CASE       :	List[Any]						    =							(input_image / 2_5_5) * 2 - 1
                 _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							torch.tensor(input_image[np.newaxis, :, :]    ,						dtype=torch.float							).to(self.device							)

                 if self.vqvae is not None:
                      _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							self.vqvae.encode(torch.unsqueeze(__lowerCamelCase    ,						0							)							).latent_dist.sample(
                          generator=__lowerCamelCase							)[0]
                      _SCREAMING_SNAKE_CASE       :	Dict						    =							self.vqvae.config.scaling_factor * input_images

                 if start_step > 0:
                      _SCREAMING_SNAKE_CASE       :	List[str]						    =							self.scheduler.add_noise(__lowerCamelCase    ,						__lowerCamelCase    ,						self.scheduler.timesteps[start_step - 1]							)

                 _SCREAMING_SNAKE_CASE       :	Dict						    =							(
                     self.unet.config.sample_size[1] * self.mel.get_sample_rate() / self.mel.x_res / self.mel.hop_length
                 )
                 _SCREAMING_SNAKE_CASE       :	Dict						    =							int(mask_start_secs * pixels_per_second							)
                 _SCREAMING_SNAKE_CASE       :	str						    =							int(mask_end_secs * pixels_per_second							)
                 _SCREAMING_SNAKE_CASE       :	List[str]						    =							self.scheduler.add_noise(__lowerCamelCase    ,						__lowerCamelCase    ,						torch.tensor(self.scheduler.timesteps[start_step:]							)							)

            for step, t in enumerate(self.progress_bar(self.scheduler.timesteps[start_step:]							)							):
                 if isinstance(self.unet    ,						__lowerCamelCase							):
                      _SCREAMING_SNAKE_CASE       :	str						    =							self.unet(__lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase							)["sample"]
                 else:
                      _SCREAMING_SNAKE_CASE       :	List[str]						    =							self.unet(__lowerCamelCase    ,						__lowerCamelCase							)["sample"]

                 if isinstance(self.scheduler    ,						__lowerCamelCase							):
                      _SCREAMING_SNAKE_CASE       :	Tuple						    =							self.scheduler.step(
                          model_output=__lowerCamelCase    ,						timestep=__lowerCamelCase    ,						sample=__lowerCamelCase    ,						eta=__lowerCamelCase    ,						generator=__lowerCamelCase    ,						)["prev_sample"]
                 else:
                      _SCREAMING_SNAKE_CASE       :	Tuple						    =							self.scheduler.step(
                          model_output=__lowerCamelCase    ,						timestep=__lowerCamelCase    ,						sample=__lowerCamelCase    ,						generator=__lowerCamelCase    ,						)["prev_sample"]

                 if mask is not None:
                      if mask_start > 0:
                           _SCREAMING_SNAKE_CASE       :	Tuple						    =							mask[:, step, :, :mask_start]
                      if mask_end > 0:
                           _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							mask[:, step, :, -mask_end:]

            if self.vqvae is not None:
                 # 0.18215 was scaling factor used in training to ensure unit variance
                 _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							1 / self.vqvae.config.scaling_factor * images
                 _SCREAMING_SNAKE_CASE       :	Tuple						    =							self.vqvae.decode(__lowerCamelCase							)["sample"]

            _SCREAMING_SNAKE_CASE       :	Tuple						    =							(images / 2 + 0.5).clamp(0    ,						1							)
            _SCREAMING_SNAKE_CASE       :	List[str]						    =							images.cpu().permute(0    ,						2    ,						3    ,						1							).numpy()
            _SCREAMING_SNAKE_CASE       :	int						    =							(images * 2_5_5).round().astype("uint8"							)
            _SCREAMING_SNAKE_CASE       :	Dict						    =							list(
                (Image.fromarray(_[:, :, 0]							) for _ in images)
                if images.shape[3] == 1
                else (Image.fromarray(__lowerCamelCase    ,						mode="RGB"							).convert("L"							) for _ in images)							)

            _SCREAMING_SNAKE_CASE       :	Dict						    =							[self.mel.image_to_audio(__lowerCamelCase							) for _ in images]
            if not return_dict:
                 return images, (self.mel.get_sample_rate(), audios)

            return BaseOutput(**AudioPipelineOutput(np.array(__lowerCamelCase							)[:, np.newaxis, :]							)    ,						**ImagePipelineOutput(__lowerCamelCase							)							)




       @torch.no_grad()
       def     UpperCamelCase_			(      self    ,						__lowerCamelCase    ,						__lowerCamelCase = 5_0							)			->     np.ndarray:
            assert isinstance(self.scheduler    ,						__lowerCamelCase							)
            self.scheduler.set_timesteps(__lowerCamelCase							)
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							np.array(
                [np.frombuffer(image.tobytes()    ,						dtype="uint8"							).reshape((1, image.height, image.width)							) for image in images]							)
            _SCREAMING_SNAKE_CASE       :	Optional[Any]						    =							(sample / 2_5_5) * 2 - 1
            _SCREAMING_SNAKE_CASE       :	Optional[int]						    =							torch.Tensor(__lowerCamelCase							).to(self.device							)

            for t in self.progress_bar(torch.flip(self.scheduler.timesteps    ,						(0,)							)							):
                 _SCREAMING_SNAKE_CASE       :	Any						    =							t - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
                 _SCREAMING_SNAKE_CASE       :	int						    =							self.scheduler.alphas_cumprod[t]
                 _SCREAMING_SNAKE_CASE       :	List[str]						    =							(
                     self.scheduler.alphas_cumprod[prev_timestep]
                     if prev_timestep >= 0
                     else self.scheduler.final_alpha_cumprod
                 )
                 _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							1 - alpha_prod_t
                 _SCREAMING_SNAKE_CASE       :	List[Any]						    =							self.unet(__lowerCamelCase    ,						__lowerCamelCase							)["sample"]
                 _SCREAMING_SNAKE_CASE       :	Dict						    =							(1 - alpha_prod_t_prev) ** 0.5 * model_output
                 _SCREAMING_SNAKE_CASE       :	Any						    =							(sample - pred_sample_direction) * alpha_prod_t_prev ** (-0.5)
                 _SCREAMING_SNAKE_CASE       :	Union[str, Any]						    =							sample * alpha_prod_t ** 0.5 + beta_prod_t ** 0.5 * model_output

            return sample







       @staticmethod
       def     UpperCamelCase_			(      __lowerCamelCase    ,						__lowerCamelCase    ,						__lowerCamelCase							)			->     torch.Tensor:
            _SCREAMING_SNAKE_CASE       :	Tuple						    =							acos(torch.dot(torch.flatten(__lowerCamelCase							)    ,						torch.flatten(__lowerCamelCase							)							) / torch.norm(__lowerCamelCase							) / torch.norm(__lowerCamelCase							)							)
            return sin((1 - alpha) * theta							) * xa / sin(__lowerCamelCase							) + sin(alpha * theta							) * xa / sin(__lowerCamelCase							)