



import copy
from typing import Any, Dict, List, Optional, Union

import numpy as np
import torch

from ...audio_utils import mel_filter_bank, spectrogram, window_function
from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import TensorType, logging


__UpperCAmelCase			  =			logging.get_logger(__name__)



class 							lowerCAmelCase_						(    A_				):
       UpperCAmelCase__							:   Optional[int]  	=    ["input_features", "is_longer"]
       def __init__(     self,     SCREAMING_SNAKE_CASE_=64,     SCREAMING_SNAKE_CASE_=4_8000,     SCREAMING_SNAKE_CASE_=480,     SCREAMING_SNAKE_CASE_=10,     SCREAMING_SNAKE_CASE_=1024,     SCREAMING_SNAKE_CASE_=0.0,     SCREAMING_SNAKE_CASE_=False,     SCREAMING_SNAKE_CASE_ = 0,     SCREAMING_SNAKE_CASE_ = 1_4000,     SCREAMING_SNAKE_CASE_ = None,     SCREAMING_SNAKE_CASE_ = "fusion",     SCREAMING_SNAKE_CASE_ = "repeatpad",     **SCREAMING_SNAKE_CASE_,     )	->						Dict:
          super().__init__(
              feature_size=SCREAMING_SNAKE_CASE_,     sampling_rate=SCREAMING_SNAKE_CASE_,     padding_value=SCREAMING_SNAKE_CASE_,     return_attention_mask=SCREAMING_SNAKE_CASE_,     **SCREAMING_SNAKE_CASE_,     )
          UpperCamelCase       :    List[Any]		      =      top_db
          UpperCamelCase       :    Dict		      =      truncation
          UpperCamelCase       :    str		      =      padding
          UpperCamelCase       :    Any		      =      fft_window_size
          UpperCamelCase       :    Dict		      =      (fft_window_size >> 1) + 1
          UpperCamelCase       :    Optional[Any]		      =      hop_length
          UpperCamelCase       :    Any		      =      max_length_s
          UpperCamelCase       :    List[str]		      =      max_length_s * sampling_rate
          UpperCamelCase       :    List[str]		      =      sampling_rate
          UpperCamelCase       :    List[Any]		      =      frequency_min
          UpperCamelCase       :    str		      =      frequency_max
          UpperCamelCase       :    List[Any]		      =      mel_filter_bank(
              num_frequency_bins=self.nb_frequency_bins,     num_mel_filters=SCREAMING_SNAKE_CASE_,     min_frequency=SCREAMING_SNAKE_CASE_,     max_frequency=SCREAMING_SNAKE_CASE_,     sampling_rate=SCREAMING_SNAKE_CASE_,     norm=SCREAMING_SNAKE_CASE_,     mel_scale='htk',     )
          UpperCamelCase       :    Union[str, Any]		      =      mel_filter_bank(
              num_frequency_bins=self.nb_frequency_bins,     num_mel_filters=SCREAMING_SNAKE_CASE_,     min_frequency=SCREAMING_SNAKE_CASE_,     max_frequency=SCREAMING_SNAKE_CASE_,     sampling_rate=SCREAMING_SNAKE_CASE_,     norm='slaney',     mel_scale='slaney',     )
       def    snake_case_			(     self			)	->						Dict[str, Any]:
          UpperCamelCase       :    Optional[int]		      =      copy.deepcopy(self.__dict__			)
          UpperCamelCase       :    Optional[Any]		      =      self.__class__.__name__
          if "mel_filters" in output:
             del output["mel_filters"]
          if "mel_filters_slaney" in output:
             del output["mel_filters_slaney"]
          return output
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_ = None			)	->						np.ndarray:
          UpperCamelCase       :    Optional[Any]		      =      spectrogram(
              SCREAMING_SNAKE_CASE_,     window_function(self.fft_window_size,     'hann'			),     frame_length=self.fft_window_size,     hop_length=self.hop_length,     power=2.0,     mel_filters=SCREAMING_SNAKE_CASE_,     log_mel='dB',     )
          return log_mel_spectrogram.T
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)	->						Optional[int]:
          UpperCamelCase       :    int		      =      np.array_split(list(range(0,     total_frames - chunk_frames + 1			)			),     3			)
          if len(ranges[1]			) == 0:
             # if the audio is too short, we just use the first chunk
             UpperCamelCase       :    List[Any]		      =      [0]
          if len(ranges[2]			) == 0:
             # if the audio is too short, we just use the first chunk
             UpperCamelCase       :    List[Any]		      =      [0]
          # randomly choose index for each part
          UpperCamelCase       :    List[str]		      =      np.random.choice(ranges[0]			)
          UpperCamelCase       :    Union[str, Any]		      =      np.random.choice(ranges[1]			)
          UpperCamelCase       :    Any		      =      np.random.choice(ranges[2]			)

          UpperCamelCase       :    Union[str, Any]		      =      mel[idx_front : idx_front + chunk_frames, :]
          UpperCamelCase       :    Union[str, Any]		      =      mel[idx_middle : idx_middle + chunk_frames, :]
          UpperCamelCase       :    Dict		      =      mel[idx_back : idx_back + chunk_frames, :]

          UpperCamelCase       :    Optional[int]		      =      torch.tensor(mel[None, None, :]			)
          UpperCamelCase       :    Tuple		      =      torch.nn.functional.interpolate(
              SCREAMING_SNAKE_CASE_,     size=[chunk_frames, 64],     mode='bilinear',     align_corners=SCREAMING_SNAKE_CASE_			)
          UpperCamelCase       :    Dict		      =      mel_shrink[0][0].numpy()
          UpperCamelCase       :    List[str]		      =      np.stack([mel_shrink, mel_chunk_front, mel_chunk_middle, mel_chunk_back],     axis=0			)
          return mel_fusion
       def    snake_case_			(     self,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)	->						np.array:
          if waveform.shape[0] > max_length:
             if truncation == "rand_trunc":
                UpperCamelCase       :    Dict		      =      True
                # random crop to max_length (for compatibility) -> this should be handled by self.pad
                UpperCamelCase       :    List[str]		      =      len(SCREAMING_SNAKE_CASE_			) - max_length
                UpperCamelCase       :    str		      =      np.random.randint(0,     overflow + 1			)
                UpperCamelCase       :    Any		      =      waveform[idx : idx + max_length]
                UpperCamelCase       :    Tuple		      =      self._np_extract_fbank_features(SCREAMING_SNAKE_CASE_,     self.mel_filters_slaney			)[None, :]
             elif truncation == "fusion":
                UpperCamelCase       :    str		      =      self._np_extract_fbank_features(SCREAMING_SNAKE_CASE_,     self.mel_filters			)
                UpperCamelCase       :    List[str]		      =      max_length // self.hop_length + 1  # the +1 related to how the spectrogram is computed
                UpperCamelCase       :    List[str]		      =      mel.shape[0]
                if chunk_frames == total_frames:
                   # there is a corner case where the audio length is larger than max_length but smaller than max_length+hop_length.
                   # In this case, we just use the whole audio.
                   UpperCamelCase       :    Dict		      =      np.stack([mel, mel, mel, mel],     axis=0			)
                   UpperCamelCase       :    List[Any]		      =      False
                else:
                   UpperCamelCase       :    Tuple		      =      self._random_mel_fusion(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)
                   UpperCamelCase       :    Any		      =      True
             else:
                raise NotImplementedError(F"""data_truncating {truncation} not implemented"""			)

          else:
             UpperCamelCase       :    Union[str, Any]		      =      False
             # only use repeat as a new possible value for padding. you repeat the audio before applying the usual max_length padding
             if waveform.shape[0] < max_length:
                if padding == "repeat":
                   UpperCamelCase       :    Tuple		      =      int(max_length / len(SCREAMING_SNAKE_CASE_			)			)
                   UpperCamelCase       :    Union[str, Any]		      =      np.stack(np.tile(SCREAMING_SNAKE_CASE_,     n_repeat + 1			)			)[:max_length]
                if padding == "repeatpad":
                   UpperCamelCase       :    List[Any]		      =      int(max_length / len(SCREAMING_SNAKE_CASE_			)			)
                   UpperCamelCase       :    str		      =      np.stack(np.tile(SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)			)
                UpperCamelCase       :    str		      =      np.pad(SCREAMING_SNAKE_CASE_,     (0, max_length - waveform.shape[0]),     mode='constant',     constant_values=0			)

             if truncation == "fusion":
                UpperCamelCase       :    Optional[Any]		      =      self._np_extract_fbank_features(SCREAMING_SNAKE_CASE_,     self.mel_filters			)
                UpperCamelCase       :    Optional[Any]		      =      np.stack([input_mel, input_mel, input_mel, input_mel],     axis=0			)
             else:
                UpperCamelCase       :    Tuple		      =      self._np_extract_fbank_features(SCREAMING_SNAKE_CASE_,     self.mel_filters_slaney			)[None, :]

          return input_mel, longer






       def __call__(     self,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_ = None,     SCREAMING_SNAKE_CASE_ = None,     SCREAMING_SNAKE_CASE_ = None,     SCREAMING_SNAKE_CASE_ = None,     SCREAMING_SNAKE_CASE_ = None,     **SCREAMING_SNAKE_CASE_,     )	->						BatchFeature:
          UpperCamelCase       :    Tuple		      =      truncation if truncation is not None else self.truncation
          UpperCamelCase       :    int		      =      padding if padding else self.padding

          if sampling_rate is not None:
             if sampling_rate != self.sampling_rate:
                raise ValueError(
                    F"""The model corresponding to this feature extractor: {self.__class__.__name__} was trained using a"""
                    F""" sampling rate of {self.sampling_rate}. Please make sure that the provided `raw_speech` input"""
                    F""" was sampled with {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.'			)

          UpperCamelCase       :    Dict		      =      isinstance(SCREAMING_SNAKE_CASE_,     np.ndarray			) and len(raw_speech.shape			) > 1
          if is_batched_numpy and len(raw_speech.shape			) > 2:
             raise ValueError(F"""Only mono-channel audio is supported for input to {self}"""			)
          UpperCamelCase       :    Any		      =      is_batched_numpy or (
              isinstance(SCREAMING_SNAKE_CASE_,     (list, tuple)			) and (isinstance(raw_speech[0],     (np.ndarray, tuple, list)			))
          )

          if is_batched:
             UpperCamelCase       :    Any		      =      [np.asarray(SCREAMING_SNAKE_CASE_,     dtype=np.floataa			) for speech in raw_speech]
          elif not is_batched and not isinstance(SCREAMING_SNAKE_CASE_,     np.ndarray			):
             UpperCamelCase       :    int		      =      np.asarray(SCREAMING_SNAKE_CASE_,     dtype=np.floataa			)
          elif isinstance(SCREAMING_SNAKE_CASE_,     np.ndarray			) and raw_speech.dtype is np.dtype(np.floataa			):
             UpperCamelCase       :    str		      =      raw_speech.astype(np.floataa			)

          # always return batch
          if not is_batched:
             UpperCamelCase       :    Dict		      =      [np.asarray(SCREAMING_SNAKE_CASE_			)]

          # convert to mel spectrogram, truncate and pad if needed.
          UpperCamelCase       :    int		      =      [
              self._get_input_mel(SCREAMING_SNAKE_CASE_,     max_length if max_length else self.nb_max_samples,     SCREAMING_SNAKE_CASE_,     SCREAMING_SNAKE_CASE_			)
              for waveform in raw_speech
          ]

          UpperCamelCase       :    Dict		      =      []
          UpperCamelCase       :    Optional[int]		      =      []
          for mel, longer in padded_inputs:
             input_mel.append(SCREAMING_SNAKE_CASE_			)
             is_longer.append(SCREAMING_SNAKE_CASE_			)

          if truncation == "fusion" and sum(SCREAMING_SNAKE_CASE_			) == 0:
             # if no audio is longer than 10s, then randomly select one audio to be longer
             UpperCamelCase       :    Any		      =      np.random.randint(0,     len(SCREAMING_SNAKE_CASE_			)			)
             UpperCamelCase       :    List[Any]		      =      True

          if isinstance(input_mel[0],     SCREAMING_SNAKE_CASE_			):
             UpperCamelCase       :    Any		      =      [np.asarray(SCREAMING_SNAKE_CASE_,     dtype=np.floataa			) for feature in input_mel]

          # is_longer is a list of bool
          UpperCamelCase       :    Union[str, Any]		      =      [[longer] for longer in is_longer]

          UpperCamelCase       :    Optional[int]		      =      {"""input_features""": input_mel, """is_longer""": is_longer}
          UpperCamelCase       :    Optional[int]		      =      BatchFeature(SCREAMING_SNAKE_CASE_			)

          if return_tensors is not None:
             UpperCamelCase       :    List[Any]		      =      input_features.convert_to_tensors(SCREAMING_SNAKE_CASE_			)

          return input_features







"""simple docstring"""
import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    StableDiffusionSAGPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()




class 	lowercase__				(					A_ ,A_ ,unittest.TestCase				):
   __UpperCAmelCase							       =					StableDiffusionSAGPipeline
   __UpperCAmelCase							       =					TEXT_TO_IMAGE_PARAMS
   __UpperCAmelCase							       =					TEXT_TO_IMAGE_BATCH_PARAMS
   __UpperCAmelCase							       =					TEXT_TO_IMAGE_IMAGE_PARAMS
   __UpperCAmelCase							       =					TEXT_TO_IMAGE_IMAGE_PARAMS
   __UpperCAmelCase							       =					False





   def 		UpperCamelCase_  (							self)  ->  Optional[Any]:
      torch.manual_seed(0)
      _lowerCamelCase	:					Dict   =		UNetaDConditionModel(
          block_out_channels=(32, 64)      ,   layers_per_block=2      ,   sample_size=32      ,   in_channels=4      ,   out_channels=4      ,   down_block_types=("""DownBlock2D""", """CrossAttnDownBlock2D""")      ,   up_block_types=("""CrossAttnUpBlock2D""", """UpBlock2D""")      ,   cross_attention_dim=32      ,   )
      _lowerCamelCase	:					int   =		DDIMScheduler(
          beta_start=0.0_00_85      ,   beta_end=0.0_12      ,   beta_schedule="""scaled_linear"""      ,   clip_sample=SCREAMING_SNAKE_CASE      ,   set_alpha_to_one=SCREAMING_SNAKE_CASE      ,   )
      torch.manual_seed(0)
      _lowerCamelCase	:					Tuple   =		AutoencoderKL(
          block_out_channels=[32, 64]      ,   in_channels=3      ,   out_channels=3      ,   down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]      ,   up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]      ,   latent_channels=4      ,   )
      torch.manual_seed(0)
      _lowerCamelCase	:					Tuple   =		CLIPTextConfig(
          bos_token_id=0      ,   eos_token_id=2      ,   hidden_size=32      ,   intermediate_size=37      ,   layer_norm_eps=1e-0_5      ,   num_attention_heads=4      ,   num_hidden_layers=5      ,   pad_token_id=1      ,   vocab_size=1000      ,   )
      _lowerCamelCase	:					List[Any]   =		CLIPTextModel(SCREAMING_SNAKE_CASE)
      _lowerCamelCase	:					Tuple   =		CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip""")

      _lowerCamelCase	:					List[Any]   =		{
          """unet""": unet,
          """scheduler""": scheduler,
          """vae""": vae,
          """text_encoder""": text_encoder,
          """tokenizer""": tokenizer,
          """safety_checker""": None,
          """feature_extractor""": None,
      }
      return components





   def 		UpperCamelCase_  (							self      ,   SCREAMING_SNAKE_CASE      ,   SCREAMING_SNAKE_CASE=0)  ->  List[Any]:
      if str(SCREAMING_SNAKE_CASE).startswith("""mps"""):
         _lowerCamelCase	:					List[str]   =		torch.manual_seed(SCREAMING_SNAKE_CASE)
      else:
         _lowerCamelCase	:					List[str]   =		torch.Generator(device=SCREAMING_SNAKE_CASE).manual_seed(SCREAMING_SNAKE_CASE)
      _lowerCamelCase	:					List[Any]   =		{
          """prompt""": """.""",
          """generator""": generator,
          """num_inference_steps""": 2,
          """guidance_scale""": 1.0,
          """sag_scale""": 1.0,
          """output_type""": """numpy""",
      }
      return inputs





   def 		UpperCamelCase_  (							self)  ->  Tuple:
      super().test_inference_batch_single_identical(expected_max_diff=3e-3)







@slow
@require_torch_gpu
class 	lowercase__				(					unittest.TestCase				):





   def 		UpperCamelCase_  (							self)  ->  Union[str, Any]:
      # clean up the VRAM after each test
      super().tearDown()
      gc.collect()
      torch.cuda.empty_cache()





   def 		UpperCamelCase_  (							self)  ->  Optional[Any]:
      _lowerCamelCase	:					Any   =		StableDiffusionSAGPipeline.from_pretrained("""CompVis/stable-diffusion-v1-4""")
      _lowerCamelCase	:					Union[str, Any]   =		sag_pipe.to(SCREAMING_SNAKE_CASE)
      sag_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE)

      _lowerCamelCase	:					Optional[int]   =		"""."""
      _lowerCamelCase	:					int   =		torch.manual_seed(0)
      _lowerCamelCase	:					Tuple   =		sag_pipe(
          [prompt]      ,   generator=SCREAMING_SNAKE_CASE      ,   guidance_scale=7.5      ,   sag_scale=1.0      ,   num_inference_steps=20      ,   output_type="""np""")

      _lowerCamelCase	:					Dict   =		output.images

      _lowerCamelCase	:					Tuple   =		image[0, -3:, -3:, -1]

      assert image.shape == (1, 512, 512, 3)
      _lowerCamelCase	:					Optional[Any]   =		np.array([0.15_68, 0.17_38, 0.16_95, 0.16_93, 0.15_07, 0.17_05, 0.15_47, 0.17_51, 0.19_49])

      assert np.abs(image_slice.flatten() - expected_slice).max() < 5e-2





   def 		UpperCamelCase_  (							self)  ->  List[str]:
      _lowerCamelCase	:					Optional[Any]   =		StableDiffusionSAGPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base""")
      _lowerCamelCase	:					Dict   =		sag_pipe.to(SCREAMING_SNAKE_CASE)
      sag_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE)

      _lowerCamelCase	:					Union[str, Any]   =		"""."""
      _lowerCamelCase	:					List[str]   =		torch.manual_seed(0)
      _lowerCamelCase	:					int   =		sag_pipe(
          [prompt]      ,   generator=SCREAMING_SNAKE_CASE      ,   guidance_scale=7.5      ,   sag_scale=1.0      ,   num_inference_steps=20      ,   output_type="""np""")

      _lowerCamelCase	:					Any   =		output.images

      _lowerCamelCase	:					List[str]   =		image[0, -3:, -3:, -1]

      assert image.shape == (1, 512, 512, 3)
      _lowerCamelCase	:					Any   =		np.array([0.34_59, 0.28_76, 0.25_37, 0.30_02, 0.26_71, 0.21_60, 0.30_26, 0.22_62, 0.23_71])

      assert np.abs(image_slice.flatten() - expected_slice).max() < 5e-2





   def 		UpperCamelCase_  (							self)  ->  List[str]:
      _lowerCamelCase	:					int   =		StableDiffusionSAGPipeline.from_pretrained("""stabilityai/stable-diffusion-2-1-base""")
      _lowerCamelCase	:					Optional[Any]   =		sag_pipe.to(SCREAMING_SNAKE_CASE)
      sag_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE)

      _lowerCamelCase	:					Dict   =		"""."""
      _lowerCamelCase	:					Union[str, Any]   =		torch.manual_seed(0)
      _lowerCamelCase	:					Optional[int]   =		sag_pipe(
          [prompt]      ,   width=768      ,   height=512      ,   generator=SCREAMING_SNAKE_CASE      ,   guidance_scale=7.5      ,   sag_scale=1.0      ,   num_inference_steps=20      ,   output_type="""np"""      ,   )

      _lowerCamelCase	:					Union[str, Any]   =		output.images

      assert image.shape == (1, 512, 768, 3)

