


'''simple docstring'''

import itertools
import random
import unittest

import numpy as np

from transformers import BatchFeature, SpeechTaFeatureExtractor
from transformers.testing_utils import require_torch
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_torch_available():
     import torch


a_  :				Optional[Any]						= random.Random()





def __snake_case					(      UpperCAmelCase_						:						Dict					,				UpperCAmelCase_						:						List[str]=1.0					,				UpperCAmelCase_						:						Optional[Any]=None					,				UpperCAmelCase_						:						Dict=None		):
 if rng is None:
  lowerCamelCase_      =  global_rng

 lowerCamelCase_      =  []
 for batch_idx in range(shape[0]		):
  values.append([]		)
  for _ in range(shape[1]		):
   values[-1].append(rng.random() * scale		)

 return values





@require_torch
class      snake_case    (    unittest.TestCase			):





  """simple docstring"""




  def __init__(    self			,  UpperCamelCase			,  UpperCamelCase=7			,  UpperCamelCase=400			,  UpperCamelCase=2000			,  UpperCamelCase=1			,  UpperCamelCase=0.0			,  UpperCamelCase=1_6000			,  UpperCamelCase=True			,  UpperCamelCase=80			,  UpperCamelCase=16			,  UpperCamelCase=64			,  UpperCamelCase="hann_window"			,  UpperCamelCase=80			,  UpperCamelCase=7600			,  UpperCamelCase=1e-10			,  UpperCamelCase=True			,  ):







   """simple docstring"""
   lowerCamelCase_      =  parent
   lowerCamelCase_      =  batch_size
   lowerCamelCase_      =  min_seq_length
   lowerCamelCase_      =  max_seq_length
   lowerCamelCase_      =  (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
   lowerCamelCase_      =  feature_size
   lowerCamelCase_      =  padding_value
   lowerCamelCase_      =  sampling_rate
   lowerCamelCase_      =  do_normalize
   lowerCamelCase_      =  num_mel_bins
   lowerCamelCase_      =  hop_length
   lowerCamelCase_      =  win_length
   lowerCamelCase_      =  win_function
   lowerCamelCase_      =  fmin
   lowerCamelCase_      =  fmax
   lowerCamelCase_      =  mel_floor
   lowerCamelCase_      =  return_attention_mask






  def   snake_case		(    self     ):







   """simple docstring"""
   return {
       "feature_size": self.feature_size,
       "padding_value": self.padding_value,
       "sampling_rate": self.sampling_rate,
       "do_normalize": self.do_normalize,
       "num_mel_bins": self.num_mel_bins,
       "hop_length": self.hop_length,
       "win_length": self.win_length,
       "win_function": self.win_function,
       "fmin": self.fmin,
       "fmax": self.fmax,
       "mel_floor": self.mel_floor,
       "return_attention_mask": self.return_attention_mask,
   }






  def   snake_case		(    self			,  UpperCamelCase=False			,  UpperCamelCase=False     ):







   """simple docstring"""
   def _flatten(UpperCamelCase     ):
    return list(itertools.chain(*UpperCamelCase     )     )

   if equal_length:
    lowerCamelCase_      =  floats_list((self.batch_size, self.max_seq_length)     )
   else:
    # make sure that inputs increase in size
    lowerCamelCase_      =  [
        _flatten(floats_list((x, self.feature_size)     )     )
        for x in range(self.min_seq_length			,  self.max_seq_length			,  self.seq_length_diff     )
    ]

   if numpify:
    lowerCamelCase_      =  [np.asarray(UpperCamelCase     ) for x in speech_inputs]

   return speech_inputs







  def   snake_case		(    self			,  UpperCamelCase=False			,  UpperCamelCase=False     ):







   """simple docstring"""
   if equal_length:
    lowerCamelCase_      =  [floats_list((self.max_seq_length, self.num_mel_bins)     ) for _ in range(self.batch_size     )]
   else:
    # make sure that inputs increase in size
    lowerCamelCase_      =  [
        floats_list((x, self.num_mel_bins)     )
        for x in range(self.min_seq_length			,  self.max_seq_length			,  self.seq_length_diff     )
    ]

   if numpify:
    lowerCamelCase_      =  [np.asarray(UpperCamelCase     ) for x in speech_inputs]

   return speech_inputs





@require_torch
class      snake_case    (    lowercase			,						unittest.TestCase			):





  """simple docstring"""




  _lowerCamelCase		    =  SpeechTaFeatureExtractor






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  SpeechTaFeatureExtractionTester(self     )






  def   snake_case		(    self			,  UpperCamelCase     ):







   """simple docstring"""
   self.assertTrue(np.all(np.mean(UpperCamelCase			,  axis=0     ) < 1e-3     )     )
   self.assertTrue(np.all(np.abs(np.var(UpperCamelCase			,  axis=0     ) - 1     ) < 1e-3     )     )






  def   snake_case		(    self     ):







   """simple docstring"""
   # Tests that all call wrap to encode_plus and batch_encode_plus
   lowerCamelCase_      =  self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()     )
   # create three inputs of length 800, 1000, and 1200
   lowerCamelCase_      =  [floats_list((1, x)     )[0] for x in range(800			,  1400			,  200     )]
   lowerCamelCase_      =  [np.asarray(UpperCamelCase     ) for speech_input in speech_inputs]

   # Test not batched input
   lowerCamelCase_      =  feat_extract(speech_inputs[0]			,  return_tensors="np"     ).input_values
   lowerCamelCase_      =  feat_extract(np_speech_inputs[0]			,  return_tensors="np"     ).input_values
   self.assertTrue(np.allclose(UpperCamelCase			,  UpperCamelCase			,  atol=1e-3     )     )

   # Test batched
   lowerCamelCase_      =  feat_extract(UpperCamelCase			,  return_tensors="np"     ).input_values
   lowerCamelCase_      =  feat_extract(UpperCamelCase			,  return_tensors="np"     ).input_values
   for enc_seq_a, enc_seq_a in zip(UpperCamelCase			,  UpperCamelCase     ):
    self.assertTrue(np.allclose(UpperCamelCase			,  UpperCamelCase			,  atol=1e-3     )     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()     )
   lowerCamelCase_      =  [floats_list((1, x)     )[0] for x in range(800			,  1400			,  200     )]

   lowerCamelCase_      =  ["longest", "max_length", "do_not_pad"]
   lowerCamelCase_      =  [None, 1600, None]
   for max_length, padding in zip(UpperCamelCase			,  UpperCamelCase     ):
    lowerCamelCase_      =  feat_extract(UpperCamelCase			,  padding=UpperCamelCase			,  max_length=UpperCamelCase			,  return_tensors="np"     )
    lowerCamelCase_      =  processed.input_values

    self._check_zero_mean_unit_variance(input_values[0][:800]     )
    self.assertTrue(input_values[0][800:].sum() < 1e-6     )
    self._check_zero_mean_unit_variance(input_values[1][:1000]     )
    self.assertTrue(input_values[0][1000:].sum() < 1e-6     )
    self._check_zero_mean_unit_variance(input_values[2][:1200]     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()     )
   lowerCamelCase_      =  range(800			,  1400			,  200     )
   lowerCamelCase_      =  [floats_list((1, x)     )[0] for x in lengths]

   lowerCamelCase_      =  ["longest", "max_length", "do_not_pad"]
   lowerCamelCase_      =  [None, 1600, None]

   for max_length, padding in zip(UpperCamelCase			,  UpperCamelCase     ):
    lowerCamelCase_      =  feat_extract(UpperCamelCase			,  max_length=UpperCamelCase			,  padding=UpperCamelCase     )
    lowerCamelCase_      =  processed.input_values

    self._check_zero_mean_unit_variance(input_values[0][:800]     )
    self._check_zero_mean_unit_variance(input_values[1][:1000]     )
    self._check_zero_mean_unit_variance(input_values[2][:1200]     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()     )
   lowerCamelCase_      =  [floats_list((1, x)     )[0] for x in range(800			,  1400			,  200     )]
   lowerCamelCase_      =  feat_extract(
       UpperCamelCase			,  truncation=UpperCamelCase			,  max_length=1000			,  padding="max_length"			,  return_tensors="np"     )
   lowerCamelCase_      =  processed.input_values

   self._check_zero_mean_unit_variance(input_values[0, :800]     )
   self._check_zero_mean_unit_variance(input_values[1]     )
   self._check_zero_mean_unit_variance(input_values[2]     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()     )
   lowerCamelCase_      =  [floats_list((1, x)     )[0] for x in range(800			,  1400			,  200     )]
   lowerCamelCase_      =  feat_extract(
       UpperCamelCase			,  truncation=UpperCamelCase			,  max_length=1000			,  padding="longest"			,  return_tensors="np"     )
   lowerCamelCase_      =  processed.input_values

   self._check_zero_mean_unit_variance(input_values[0, :800]     )
   self._check_zero_mean_unit_variance(input_values[1, :1000]     )
   self._check_zero_mean_unit_variance(input_values[2]     )

   # make sure that if max_length < longest -> then pad to max_length
   self.assertTrue(input_values.shape == (3, 1000)     )

   lowerCamelCase_      =  [floats_list((1, x)     )[0] for x in range(800			,  1400			,  200     )]
   lowerCamelCase_      =  feat_extract(
       UpperCamelCase			,  truncation=UpperCamelCase			,  max_length=2000			,  padding="longest"			,  return_tensors="np"     )
   lowerCamelCase_      =  processed.input_values

   self._check_zero_mean_unit_variance(input_values[0, :800]     )
   self._check_zero_mean_unit_variance(input_values[1, :1000]     )
   self._check_zero_mean_unit_variance(input_values[2]     )

   # make sure that if max_length > longest -> then pad to longest
   self.assertTrue(input_values.shape == (3, 1200)     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()     )
   lowerCamelCase_      =  np.random.rand(100     ).astype(np.floataa     )
   lowerCamelCase_      =  np_speech_inputs.tolist()

   for inputs in [py_speech_inputs, np_speech_inputs]:
    lowerCamelCase_      =  feature_extractor.pad([{"input_values": inputs}]			,  return_tensors="np"     )
    self.assertTrue(np_processed.input_values.dtype == np.floataa     )
    lowerCamelCase_      =  feature_extractor.pad([{"input_values": inputs}]			,  return_tensors="pt"     )
    self.assertTrue(pt_processed.input_values.dtype == torch.floataa     )






  def   snake_case		(    self     ):







   """simple docstring"""
   # Tests that all call wrap to encode_plus and batch_encode_plus
   lowerCamelCase_      =  self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()     )
   # create three inputs of length 800, 1000, and 1200
   lowerCamelCase_      =  [floats_list((1, x)     )[0] for x in range(800			,  1400			,  200     )]
   lowerCamelCase_      =  [np.asarray(UpperCamelCase     ) for speech_input in speech_inputs]

   # Test feature size
   lowerCamelCase_      =  feature_extractor(audio_target=UpperCamelCase			,  padding=UpperCamelCase			,  return_tensors="np"     ).input_values
   self.assertTrue(input_values.ndim == 3     )
   self.assertTrue(input_values.shape[-1] == feature_extractor.num_mel_bins     )

   # Test not batched input
   lowerCamelCase_      =  feature_extractor(speech_inputs[0]			,  return_tensors="np"     ).input_values
   lowerCamelCase_      =  feature_extractor(np_speech_inputs[0]			,  return_tensors="np"     ).input_values
   self.assertTrue(np.allclose(UpperCamelCase			,  UpperCamelCase			,  atol=1e-3     )     )

   # Test batched
   lowerCamelCase_      =  feature_extractor(UpperCamelCase			,  return_tensors="np"     ).input_values
   lowerCamelCase_      =  feature_extractor(UpperCamelCase			,  return_tensors="np"     ).input_values
   for enc_seq_a, enc_seq_a in zip(UpperCamelCase			,  UpperCamelCase     ):
    self.assertTrue(np.allclose(UpperCamelCase			,  UpperCamelCase			,  atol=1e-3     )     )

   # Test 2-D numpy arrays are batched.
   lowerCamelCase_      =  [floats_list((1, x)     )[0] for x in (800, 800, 800)]
   lowerCamelCase_      =  np.asarray(UpperCamelCase     )
   lowerCamelCase_      =  feature_extractor(UpperCamelCase			,  return_tensors="np"     ).input_values
   lowerCamelCase_      =  feature_extractor(UpperCamelCase			,  return_tensors="np"     ).input_values
   for enc_seq_a, enc_seq_a in zip(UpperCamelCase			,  UpperCamelCase     ):
    self.assertTrue(np.allclose(UpperCamelCase			,  UpperCamelCase			,  atol=1e-3     )     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.feat_extract_tester.prepare_inputs_for_target()
   lowerCamelCase_      =  self.feature_extraction_class(**self.feat_extract_dict     )
   lowerCamelCase_      =  feat_extract.model_input_names[0]

   lowerCamelCase_      =  BatchFeature({input_name: speech_inputs}     )

   self.assertTrue(all(len(UpperCamelCase     ) == len(UpperCamelCase     ) for x, y in zip(UpperCamelCase			,  processed_features[input_name]     )     )     )

   lowerCamelCase_      =  self.feat_extract_tester.prepare_inputs_for_target(equal_length=UpperCamelCase     )
   lowerCamelCase_      =  BatchFeature({input_name: speech_inputs}			,  tensor_type="np"     )

   lowerCamelCase_      =  processed_features[input_name]

   if len(batch_features_input.shape     ) < 3:
    lowerCamelCase_      =  batch_features_input[:, :, None]

   self.assertTrue(
       batch_features_input.shape
       == (self.feat_extract_tester.batch_size, len(speech_inputs[0]     ), self.feat_extract_tester.num_mel_bins)     )






  @require_torch
  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.feat_extract_tester.prepare_inputs_for_target(equal_length=UpperCamelCase     )
   lowerCamelCase_      =  self.feature_extraction_class(**self.feat_extract_dict     )
   lowerCamelCase_      =  feat_extract.model_input_names[0]

   lowerCamelCase_      =  BatchFeature({input_name: speech_inputs}			,  tensor_type="pt"     )

   lowerCamelCase_      =  processed_features[input_name]

   if len(batch_features_input.shape     ) < 3:
    lowerCamelCase_      =  batch_features_input[:, :, None]

   self.assertTrue(
       batch_features_input.shape
       == (self.feat_extract_tester.batch_size, len(speech_inputs[0]     ), self.feat_extract_tester.num_mel_bins)     )






  @require_torch
  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.feature_extraction_class(**self.feat_extract_dict     )
   lowerCamelCase_      =  self.feat_extract_tester.prepare_inputs_for_target()
   lowerCamelCase_      =  feat_extract.model_input_names[0]

   lowerCamelCase_      =  BatchFeature({input_name: speech_inputs}     )

   lowerCamelCase_      =  feat_extract.num_mel_bins  # hack!

   lowerCamelCase_      =  feat_extract.pad(UpperCamelCase			,  padding="longest"			,  return_tensors="np"     )[input_name]
   lowerCamelCase_      =  feat_extract.pad(UpperCamelCase			,  padding="longest"			,  return_tensors="pt"     )[input_name]

   self.assertTrue(abs(input_np.astype(np.floataa     ).sum() - input_pt.numpy().astype(np.floataa     ).sum()     ) < 1e-2     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.feat_extract_dict
   lowerCamelCase_      =  True
   lowerCamelCase_      =  self.feature_extraction_class(**UpperCamelCase     )
   lowerCamelCase_      =  self.feat_extract_tester.prepare_inputs_for_target()
   lowerCamelCase_      =  [len(UpperCamelCase     ) for x in speech_inputs]
   lowerCamelCase_      =  feat_extract.model_input_names[0]

   lowerCamelCase_      =  BatchFeature({input_name: speech_inputs}     )

   lowerCamelCase_      =  feat_extract.num_mel_bins  # hack!

   lowerCamelCase_      =  feat_extract.pad(UpperCamelCase			,  padding="longest"			,  return_tensors="np"     )
   self.assertIn("attention_mask"			,  UpperCamelCase     )
   self.assertListEqual(list(processed.attention_mask.shape     )			,  list(processed[input_name].shape[:2]     )     )
   self.assertListEqual(processed.attention_mask.sum(-1     ).tolist()			,  UpperCamelCase     )






  def   snake_case		(    self     ):







   """simple docstring"""
   lowerCamelCase_      =  self.feat_extract_dict
   lowerCamelCase_      =  True
   lowerCamelCase_      =  self.feature_extraction_class(**UpperCamelCase     )
   lowerCamelCase_      =  self.feat_extract_tester.prepare_inputs_for_target()
   lowerCamelCase_      =  [len(UpperCamelCase     ) for x in speech_inputs]
   lowerCamelCase_      =  feat_extract.model_input_names[0]

   lowerCamelCase_      =  BatchFeature({input_name: speech_inputs}     )
   lowerCamelCase_      =  min(UpperCamelCase     )

   lowerCamelCase_      =  feat_extract.num_mel_bins  # hack!

   lowerCamelCase_      =  feat_extract.pad(
       UpperCamelCase			,  padding="max_length"			,  max_length=UpperCamelCase			,  truncation=UpperCamelCase			,  return_tensors="np"     )
   self.assertIn("attention_mask"			,  UpperCamelCase     )
   self.assertListEqual(
       list(processed_pad.attention_mask.shape     )			,  [processed_pad[input_name].shape[0], max_length]     )
   self.assertListEqual(
       processed_pad.attention_mask[:, :max_length].sum(-1     ).tolist()			,  [max_length for x in speech_inputs]     )






  def   snake_case		(    self			,  UpperCamelCase     ):







   """simple docstring"""
   from datasets import load_dataset

   lowerCamelCase_      =  load_dataset("hf-internal-testing/librispeech_asr_dummy"			,  "clean"			,  split="validation"     )
   # automatic decoding with librispeech
   lowerCamelCase_      =  ds.sort("id"     ).select(range(UpperCamelCase     )     )[:num_samples]["audio"]

   return [x["array"] for x in speech_samples]






  def   snake_case		(    self     ):







   """simple docstring"""
   # fmt: off
   lowerCamelCase_      =  torch.tensor(
       [2.3804e-03, 2.0752e-03, 1.9836e-03, 2.1057e-03, 1.6174e-03,
        3.0518e-04, 9.1553e-05, 3.3569e-04, 9.7656e-04, 1.8311e-03,
        2.0142e-03, 2.1057e-03, 1.7395e-03, 4.5776e-04, -3.9673e-04,
        4.5776e-04, 1.0071e-03, 9.1553e-05, 4.8828e-04, 1.1597e-03,
        7.3242e-04, 9.4604e-04, 1.8005e-03, 1.8311e-03, 8.8501e-04,
        4.2725e-04, 4.8828e-04, 7.3242e-04, 1.0986e-03, 2.1057e-03]     )
   # fmt: on

   lowerCamelCase_      =  self._load_datasamples(1     )
   lowerCamelCase_      =  SpeechTaFeatureExtractor()
   lowerCamelCase_      =  feature_extractor(UpperCamelCase			,  return_tensors="pt"     ).input_values
   self.assertEquals(input_values.shape			,  (1, 9_3680)     )
   self.assertTrue(torch.allclose(input_values[0, :30]			,  UpperCamelCase			,  atol=1e-6     )     )







  def   snake_case		(    self     ):







   """simple docstring"""
   # fmt: off
   lowerCamelCase_      =  torch.tensor(
       [-2.6_870, -3.0_104, -3.1_356, -3.5_352, -3.0_044, -3.0_353, -3.4_719, -3.6_777,
        -3.1_520, -2.9_435, -2.6_553, -2.8_795, -2.9_944, -2.5_921, -3.0_279, -3.0_386,
        -3.0_864, -3.1_291, -3.2_353, -2.7_444, -2.6_831, -2.7_287, -3.1_761, -3.1_571,
        -3.2_726, -3.0_582, -3.1_007, -3.4_533, -3.4_695, -3.0_998]     )
   # fmt: on

   lowerCamelCase_      =  self._load_datasamples(1     )
   lowerCamelCase_      =  SpeechTaFeatureExtractor()
   lowerCamelCase_      =  feature_extractor(audio_target=UpperCamelCase			,  return_tensors="pt"     ).input_values
   self.assertEquals(input_values.shape			,  (1, 366, 80)     )
   self.assertTrue(torch.allclose(input_values[0, 0, :30]			,  UpperCamelCase			,  atol=1e-4     )     )







'''simple docstring'''

from typing import List, Optional, Tuple, Union

import torch

from ...models import UNetaDModel
from ...schedulers import ScoreSdeVeScheduler
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput





class      snake_case    (    lowercase			):





  """simple docstring"""




  _lowerCamelCase		    =  42
  _lowerCamelCase		    =  42






  def __init__(    self			,  UpperCamelCase			,  UpperCamelCase     ):







   """simple docstring"""
   super().__init__()
   self.register_modules(unet=UpperCamelCase			,  scheduler=UpperCamelCase     )







  @torch.no_grad()
  def __call__(    self			,  UpperCamelCase = 1			,  UpperCamelCase = 2000			,  UpperCamelCase = None			,  UpperCamelCase = "pil"			,  UpperCamelCase = True			,  **UpperCamelCase			,  ):







   """simple docstring"""
   lowerCamelCase_      =  self.unet.config.sample_size
   lowerCamelCase_      =  (batch_size, 3, img_size, img_size)

   lowerCamelCase_      =  self.unet

   lowerCamelCase_      =  randn_tensor(UpperCamelCase			,  generator=UpperCamelCase     ) * self.scheduler.init_noise_sigma
   lowerCamelCase_      =  sample.to(self.device     )

   self.scheduler.set_timesteps(UpperCamelCase     )
   self.scheduler.set_sigmas(UpperCamelCase     )

   for i, t in enumerate(self.progress_bar(self.scheduler.timesteps     )     ):
    lowerCamelCase_      =  self.scheduler.sigmas[i] * torch.ones(shape[0]			,  device=self.device     )

    # correction step
    for _ in range(self.scheduler.config.correct_steps     ):
     lowerCamelCase_      =  self.unet(UpperCamelCase			,  UpperCamelCase     ).sample
     lowerCamelCase_      =  self.scheduler.step_correct(UpperCamelCase			,  UpperCamelCase			,  generator=UpperCamelCase     ).prev_sample

    # prediction step
    lowerCamelCase_      =  model(UpperCamelCase			,  UpperCamelCase     ).sample
    lowerCamelCase_      =  self.scheduler.step_pred(UpperCamelCase			,  UpperCamelCase			,  UpperCamelCase			,  generator=UpperCamelCase     )

    lowerCamelCase_ ,lowerCamelCase_      =  output.prev_sample, output.prev_sample_mean

   lowerCamelCase_      =  sample_mean.clamp(0			,  1     )
   lowerCamelCase_      =  sample.cpu().permute(0			,  2			,  3			,  1     ).numpy()
   if output_type == "pil":
    lowerCamelCase_      =  self.numpy_to_pil(UpperCamelCase     )

   if not return_dict:
    return (sample,)

   return ImagePipelineOutput(images=UpperCamelCase     )



