




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 			_UpperCamelCase      (				_UpperCAmelCase		):
      UpperCAmelCase_        =		['vqvae']
      def __init__(			self   :str ,      lowerCamelCase   :AutoencoderKL ,      lowerCamelCase   :UNetaDConditionModel ,      lowerCamelCase   :Mel ,      lowerCamelCase   :Union[DDIMScheduler, DDPMScheduler] ,      )     ->			int:
       super().__init__()
       self.register_modules(unet=__UpperCamelCase ,      scheduler=__UpperCamelCase ,      mel=__UpperCamelCase ,      vqvae=__UpperCamelCase      )
      def 	UpperCAmelCase_						(			self   :Optional[Any]      )     ->			int:
       return 50 if isinstance(self.scheduler ,      __UpperCamelCase      ) else 1000
      @torch.no_grad()
      def __call__(			self   :Optional[Any] ,      lowerCamelCase   :int = 1 ,      lowerCamelCase   :str = None ,      lowerCamelCase   :np.ndarray = None ,      lowerCamelCase   :int = 0 ,      lowerCamelCase   :int = 0 ,      lowerCamelCase   :int = None ,      lowerCamelCase   :torch.Generator = None ,      lowerCamelCase   :float = 0 ,      lowerCamelCase   :float = 0 ,      lowerCamelCase   :torch.Generator = None ,      lowerCamelCase   :float = 0 ,      lowerCamelCase   :torch.Tensor = None ,      lowerCamelCase   :torch.Tensor = None ,      lowerCamelCase   :Any=True ,      )     ->			Union[
          Union[AudioPipelineOutput, ImagePipelineOutput],
          Tuple[List[Image.Image], Tuple[int, List[np.ndarray]]],
      ]:
       UpperCAmelCase__							 =						steps or self.get_default_steps()
       self.scheduler.set_timesteps(__UpperCamelCase      )
       UpperCAmelCase__							 =						step_generator or generator
       # For backwards compatibility
       if type(self.unet.config.sample_size      ) == int:
        UpperCAmelCase__							 =						(self.unet.config.sample_size, self.unet.config.sample_size)
       if noise is None:
        UpperCAmelCase__							 =						randn_tensor(
            (
                batch_size,
                self.unet.config.in_channels,
                self.unet.config.sample_size[0],
                self.unet.config.sample_size[1],
            ) ,      generator=__UpperCamelCase ,      device=self.device ,      )
       UpperCAmelCase__							 =						noise
       UpperCAmelCase__							 =						None

       if audio_file is not None or raw_audio is not None:
        self.mel.load_audio(__UpperCamelCase ,      __UpperCamelCase      )
        UpperCAmelCase__							 =						self.mel.audio_slice_to_image(__UpperCamelCase      )
        UpperCAmelCase__							 =						np.frombuffer(input_image.tobytes() ,      dtype="uint8"      ).reshape(
            (input_image.height, input_image.width)      )
        UpperCAmelCase__							 =						(input_image / 255) * 2 - 1
        UpperCAmelCase__							 =						torch.tensor(input_image[np.newaxis, :, :] ,      dtype=torch.float      ).to(self.device      )

        if self.vqvae is not None:
         UpperCAmelCase__							 =						self.vqvae.encode(torch.unsqueeze(__UpperCamelCase ,      0      )      ).latent_dist.sample(
             generator=__UpperCamelCase      )[0]
         UpperCAmelCase__							 =						self.vqvae.config.scaling_factor * input_images

        if start_step > 0:
         UpperCAmelCase__							 =						self.scheduler.add_noise(__UpperCamelCase ,      __UpperCamelCase ,      self.scheduler.timesteps[start_step - 1]      )

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

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

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

        if mask is not None:
         if mask_start > 0:
          UpperCAmelCase__							 =						mask[:, step, :, :mask_start]
         if mask_end > 0:
          UpperCAmelCase__							 =						mask[:, step, :, -mask_end:]

       if self.vqvae is not None:
        # 0.18215 was scaling factor used in training to ensure unit variance
        UpperCAmelCase__							 =						1 / self.vqvae.config.scaling_factor * images
        UpperCAmelCase__							 =						self.vqvae.decode(__UpperCamelCase      )["sample"]

       UpperCAmelCase__							 =						(images / 2 + 0.5).clamp(0 ,      1      )
       UpperCAmelCase__							 =						images.cpu().permute(0 ,      2 ,      3 ,      1      ).numpy()
       UpperCAmelCase__							 =						(images * 255).round().astype("uint8"      )
       UpperCAmelCase__							 =						list(
           (Image.fromarray(_[:, :, 0]      ) for _ in images)
           if images.shape[3] == 1
           else (Image.fromarray(__UpperCamelCase ,      mode="RGB"      ).convert("L"      ) for _ in images)      )

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

       return BaseOutput(**AudioPipelineOutput(np.array(__UpperCamelCase      )[:, np.newaxis, :]      ) ,      **ImagePipelineOutput(__UpperCamelCase      )      )
      @torch.no_grad()
      def 	UpperCAmelCase_						(			self   :Optional[Any] ,      lowerCamelCase   :List[Image.Image] ,      lowerCamelCase   :int = 50      )     ->			np.ndarray:
       assert isinstance(self.scheduler ,      __UpperCamelCase      )
       self.scheduler.set_timesteps(__UpperCamelCase      )
       UpperCAmelCase__							 =						np.array(
           [np.frombuffer(image.tobytes() ,      dtype="uint8"      ).reshape((1, image.height, image.width)      ) for image in images]      )
       UpperCAmelCase__							 =						(sample / 255) * 2 - 1
       UpperCAmelCase__							 =						torch.Tensor(__UpperCamelCase      ).to(self.device      )

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

       return sample







      @staticmethod
      def 	UpperCAmelCase_						(			lowerCamelCase   :torch.Tensor ,      lowerCamelCase   :torch.Tensor ,      lowerCamelCase   :float      )     ->			torch.Tensor:
       UpperCAmelCase__							 =						acos(torch.dot(torch.flatten(__UpperCamelCase      ) ,      torch.flatten(__UpperCamelCase      )      ) / torch.norm(__UpperCamelCase      ) / torch.norm(__UpperCamelCase      )      )
       return sin((1 - alpha) * theta      ) * xa / sin(__UpperCamelCase      ) + sin(alpha * theta      ) * xa / sin(__UpperCamelCase      )











from typing import List, Optional, Union

import numpy as np

from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import PaddingStrategy, TensorType, logging


_lowerCAmelCase						:   Any		=   logging.get_logger(__name__)
class 			_UpperCamelCase      (				lowerCAmelCase		):
      UpperCAmelCase_        =		["""input_values""", """padding_mask"""]
      def __init__(			self   :Dict ,      lowerCamelCase   :int = 1 ,      lowerCamelCase   :int = 2_4000 ,      lowerCamelCase   :float = 0.0 ,      lowerCamelCase   :float = None ,      lowerCamelCase   :float = None ,      **lowerCamelCase   :Optional[Any] ,      )     ->			str:
       super().__init__(feature_size=lowerCamelCase ,      sampling_rate=lowerCamelCase ,      padding_value=lowerCamelCase ,      **lowerCamelCase      )
       UpperCAmelCase__							 =						chunk_length_s
       UpperCAmelCase__							 =						overlap
      @property
      def 	UpperCAmelCase_						(			self   :Optional[Any]      )     ->			Optional[int]:
       if self.chunk_length_s is None:
        return None
       else:
        return int(self.chunk_length_s * self.sampling_rate      )
      @property
      def 	UpperCAmelCase_						(			self   :str      )     ->			Optional[int]:
       if self.chunk_length_s is None or self.overlap is None:
        return None
       else:
        return max(1 ,      int((1.0 - self.overlap) * self.chunk_length      )      )







      def __call__(			self   :List[Any] ,      lowerCamelCase   :Union[np.ndarray, List[float], List[np.ndarray], List[List[float]]] ,      lowerCamelCase   :Optional[Union[bool, str, PaddingStrategy]] = None ,      lowerCamelCase   :Optional[bool] = False ,      lowerCamelCase   :Optional[int] = None ,      lowerCamelCase   :Optional[Union[str, TensorType]] = None ,      lowerCamelCase   :Optional[int] = None ,      )     ->			BatchFeature:
       if sampling_rate is not None:
        if sampling_rate != self.sampling_rate:
         raise ValueError(
             f'''The model corresponding to this feature extractor: {self} was trained using a sampling rate of'''
             f''' {self.sampling_rate}. Please make sure that the provided audio input was sampled with'''
             f''' {self.sampling_rate} and not {sampling_rate}.'''      )
       else:
        logger.warning(
            "It is strongly recommended to pass the `sampling_rate` argument to this function. "
            "Failing to do so can result in silent errors that might be hard to debug."      )

       if padding and truncation:
        raise ValueError("Both padding and truncation were set. Make sure you only set one."      )
       elif padding is None:
        # by default let's pad the inputs
        UpperCAmelCase__							 =						True

       UpperCAmelCase__							 =						bool(
           isinstance(lowerCamelCase ,      (list, tuple)      ) and (isinstance(raw_audio[0] ,      (np.ndarray, tuple, list)      ))      )

       if is_batched:
        UpperCAmelCase__							 =						[np.asarray(lowerCamelCase ,      dtype=np.floataa      ).T for audio in raw_audio]
       elif not is_batched and not isinstance(lowerCamelCase ,      np.ndarray      ):
        UpperCAmelCase__							 =						np.asarray(lowerCamelCase ,      dtype=np.floataa      )
       elif isinstance(lowerCamelCase ,      np.ndarray      ) and raw_audio.dtype is np.dtype(np.floataa      ):
        UpperCAmelCase__							 =						raw_audio.astype(np.floataa      )

       # always return batch
       if not is_batched:
        UpperCAmelCase__							 =						[np.asarray(lowerCamelCase      ).T]

       # verify inputs are valid
       for idx, example in enumerate(lowerCamelCase      ):
        if example.ndim > 2:
         raise ValueError(f'''Expected input shape (channels, length) but got shape {example.shape}'''      )
        if self.feature_size == 1 and example.ndim != 1:
         raise ValueError(f'''Expected mono audio but example has {example.shape[-1]} channels'''      )
        if self.feature_size == 2 and example.shape[-1] != 2:
         raise ValueError(f'''Expected stereo audio but example has {example.shape[-1]} channels'''      )

       UpperCAmelCase__							 =						None
       UpperCAmelCase__							 =						BatchFeature({"input_values": raw_audio}      )
       if self.chunk_stride is not None and self.chunk_length is not None and max_length is None:
        if truncation:
         UpperCAmelCase__							 =						min(array.shape[0] for array in raw_audio      )
         UpperCAmelCase__							 =						int(np.floor(max_length / self.chunk_stride      )      )
         UpperCAmelCase__							 =						(nb_step - 1) * self.chunk_stride + self.chunk_length
        elif padding:
         UpperCAmelCase__							 =						max(array.shape[0] for array in raw_audio      )
         UpperCAmelCase__							 =						int(np.ceil(max_length / self.chunk_stride      )      )
         UpperCAmelCase__							 =						(nb_step - 1) * self.chunk_stride + self.chunk_length
         UpperCAmelCase__							 =						"max_length"
        else:
         UpperCAmelCase__							 =						input_values

        # normal padding on batch
       if padded_inputs is None:
        UpperCAmelCase__							 =						self.pad(
            lowerCamelCase ,      max_length=lowerCamelCase ,      truncation=lowerCamelCase ,      padding=lowerCamelCase ,      return_attention_mask=lowerCamelCase ,      )
        if padding:
         UpperCAmelCase__							 =						padded_inputs.pop("attention_mask"      )

       UpperCAmelCase__							 =						[]
       for example in padded_inputs.pop("input_values"      ):
        if self.feature_size == 1:
         UpperCAmelCase__							 =						example[..., None]
        input_values.append(example.T      )

       UpperCAmelCase__							 =						input_values
       if return_tensors is not None:
        UpperCAmelCase__							 =						padded_inputs.convert_to_tensors(lowerCamelCase      )

       return padded_inputs





