


'''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


__UpperCAmelCase       							=				random.Random()


def    _snake_case  (	A					,       A=1.0					,       A=None					,       A=None			)		->       Optional[Any]:
							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 						a__							(				unittest.TestCase    ):





							'''simple docstring'''






							def __init__(						self   ,       lowerCamelCase_   ,       lowerCamelCase_=7   ,       lowerCamelCase_=4_00   ,       lowerCamelCase_=20_00   ,       lowerCamelCase_=1   ,       lowerCamelCase_=0.0   ,       lowerCamelCase_=1_60_00   ,       lowerCamelCase_=True   ,       lowerCamelCase_=80   ,       lowerCamelCase_=16   ,       lowerCamelCase_=64   ,       lowerCamelCase_="hann_window"   ,       lowerCamelCase_=80   ,       lowerCamelCase_=76_00   ,       lowerCamelCase_=1e-10   ,       lowerCamelCase_=True   ,       )			->				List[Any]:
														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     __SCREAMING_SNAKE_CASE				(						self  )			->				Optional[Any]:
														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     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_=False   ,       lowerCamelCase_=False  )			->				str:
														def _flatten(lowerCamelCase_  ):
																					return list(itertools.chain(*lowerCamelCase_  )  )

														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(lowerCamelCase_  ) for x in speech_inputs]

														return speech_inputs







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_=False   ,       lowerCamelCase_=False  )			->				Union[str, Any]:
														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(lowerCamelCase_  ) for x in speech_inputs]

														return speech_inputs




@require_torch
class 						a__							(				a__							,							unittest.TestCase    ):





							'''simple docstring'''






							lowercase__     :  str      		= SpeechTaFeatureExtractor







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				List[str]:
														lowerCAmelCase__			  =      SpeechTaFeatureExtractionTester(self  )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_  )			->				Tuple:
														self.assertTrue(np.all(np.mean(lowerCamelCase_   ,       axis=0  ) < 1e-3  )  )
														self.assertTrue(np.all(np.abs(np.var(lowerCamelCase_   ,       axis=0  ) - 1  ) < 1e-3  )  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Union[str, Any]:
														# 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(8_00   ,       14_00   ,       2_00  )]
														lowerCAmelCase__			  =      [np.asarray(lowerCamelCase_  ) 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(lowerCamelCase_   ,       lowerCamelCase_   ,       atol=1e-3  )  )

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







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				int:
														lowerCAmelCase__			  =      self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()  )
														lowerCAmelCase__			  =      [floats_list((1, x)  )[0] for x in range(8_00   ,       14_00   ,       2_00  )]

														lowerCAmelCase__			  =      ['''longest''', '''max_length''', '''do_not_pad''']
														lowerCAmelCase__			  =      [None, 16_00, None]
														for max_length, padding in zip(lowerCamelCase_   ,       lowerCamelCase_  ):
																					lowerCAmelCase__			  =      feat_extract(lowerCamelCase_   ,       padding=lowerCamelCase_   ,       max_length=lowerCamelCase_   ,       return_tensors='''np'''  )
																					lowerCAmelCase__			  =      processed.input_values

																					self._check_zero_mean_unit_variance(input_values[0][:8_00]  )
																					self.assertTrue(input_values[0][8_00:].sum() < 1e-6  )
																					self._check_zero_mean_unit_variance(input_values[1][:10_00]  )
																					self.assertTrue(input_values[0][10_00:].sum() < 1e-6  )
																					self._check_zero_mean_unit_variance(input_values[2][:12_00]  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				int:
														lowerCAmelCase__			  =      self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()  )
														lowerCAmelCase__			  =      range(8_00   ,       14_00   ,       2_00  )
														lowerCAmelCase__			  =      [floats_list((1, x)  )[0] for x in lengths]

														lowerCAmelCase__			  =      ['''longest''', '''max_length''', '''do_not_pad''']
														lowerCAmelCase__			  =      [None, 16_00, None]

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

																					self._check_zero_mean_unit_variance(input_values[0][:8_00]  )
																					self._check_zero_mean_unit_variance(input_values[1][:10_00]  )
																					self._check_zero_mean_unit_variance(input_values[2][:12_00]  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Tuple:
														lowerCAmelCase__			  =      self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()  )
														lowerCAmelCase__			  =      [floats_list((1, x)  )[0] for x in range(8_00   ,       14_00   ,       2_00  )]
														lowerCAmelCase__			  =      feat_extract(
														    lowerCamelCase_   ,       truncation=lowerCamelCase_   ,       max_length=10_00   ,       padding='''max_length'''   ,       return_tensors='''np'''  )
														lowerCAmelCase__			  =      processed.input_values

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







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				str:
														lowerCAmelCase__			  =      self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()  )
														lowerCAmelCase__			  =      [floats_list((1, x)  )[0] for x in range(8_00   ,       14_00   ,       2_00  )]
														lowerCAmelCase__			  =      feat_extract(
														    lowerCamelCase_   ,       truncation=lowerCamelCase_   ,       max_length=10_00   ,       padding='''longest'''   ,       return_tensors='''np'''  )
														lowerCAmelCase__			  =      processed.input_values

														self._check_zero_mean_unit_variance(input_values[0, :8_00]  )
														self._check_zero_mean_unit_variance(input_values[1, :10_00]  )
														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, 10_00)  )

														lowerCAmelCase__			  =      [floats_list((1, x)  )[0] for x in range(8_00   ,       14_00   ,       2_00  )]
														lowerCAmelCase__			  =      feat_extract(
														    lowerCamelCase_   ,       truncation=lowerCamelCase_   ,       max_length=20_00   ,       padding='''longest'''   ,       return_tensors='''np'''  )
														lowerCAmelCase__			  =      processed.input_values

														self._check_zero_mean_unit_variance(input_values[0, :8_00]  )
														self._check_zero_mean_unit_variance(input_values[1, :10_00]  )
														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, 12_00)  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				int:
														lowerCAmelCase__			  =      self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()  )
														lowerCAmelCase__			  =      np.random.rand(1_00  ).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     __SCREAMING_SNAKE_CASE				(						self  )			->				int:
														# 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(8_00   ,       14_00   ,       2_00  )]
														lowerCAmelCase__			  =      [np.asarray(lowerCamelCase_  ) for speech_input in speech_inputs]

														# Test feature size
														lowerCAmelCase__			  =      feature_extractor(audio_target=lowerCamelCase_   ,       padding=lowerCamelCase_   ,       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(lowerCamelCase_   ,       lowerCamelCase_   ,       atol=1e-3  )  )

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

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







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Dict:
														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(lowerCamelCase_  ) == len(lowerCamelCase_  ) for x, y in zip(lowerCamelCase_   ,       processed_features[input_name]  )  )  )

														lowerCAmelCase__			  =      self.feat_extract_tester.prepare_inputs_for_target(equal_length=lowerCamelCase_  )
														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     __SCREAMING_SNAKE_CASE				(						self  )			->				List[Any]:
														lowerCAmelCase__			  =      self.feat_extract_tester.prepare_inputs_for_target(equal_length=lowerCamelCase_  )
														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     __SCREAMING_SNAKE_CASE				(						self  )			->				Dict:
														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(lowerCamelCase_   ,       padding='''longest'''   ,       return_tensors='''np'''  )[input_name]
														lowerCAmelCase__			  =      feat_extract.pad(lowerCamelCase_   ,       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     __SCREAMING_SNAKE_CASE				(						self  )			->				List[Any]:
														lowerCAmelCase__			  =      self.feat_extract_dict
														lowerCAmelCase__			  =      True
														lowerCAmelCase__			  =      self.feature_extraction_class(**lowerCamelCase_  )
														lowerCAmelCase__			  =      self.feat_extract_tester.prepare_inputs_for_target()
														lowerCAmelCase__			  =      [len(lowerCamelCase_  ) 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(lowerCamelCase_   ,       padding='''longest'''   ,       return_tensors='''np'''  )
														self.assertIn('''attention_mask'''   ,       lowerCamelCase_  )
														self.assertListEqual(list(processed.attention_mask.shape  )   ,       list(processed[input_name].shape[:2]  )  )
														self.assertListEqual(processed.attention_mask.sum(-1  ).tolist()   ,       lowerCamelCase_  )







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				str:
														lowerCAmelCase__			  =      self.feat_extract_dict
														lowerCAmelCase__			  =      True
														lowerCAmelCase__			  =      self.feature_extraction_class(**lowerCamelCase_  )
														lowerCAmelCase__			  =      self.feat_extract_tester.prepare_inputs_for_target()
														lowerCAmelCase__			  =      [len(lowerCamelCase_  ) for x in speech_inputs]
														lowerCAmelCase__			  =      feat_extract.model_input_names[0]

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

														lowerCAmelCase__			  =      feat_extract.num_mel_bins  # hack!

														lowerCAmelCase__			  =      feat_extract.pad(
														    lowerCamelCase_   ,       padding='''max_length'''   ,       max_length=lowerCamelCase_   ,       truncation=lowerCamelCase_   ,       return_tensors='''np'''  )
														self.assertIn('''attention_mask'''   ,       lowerCamelCase_  )
														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     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_  )			->				Dict:
														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(lowerCamelCase_  )  )[:num_samples]['''audio''']

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







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Tuple:
														# fmt: off
														lowerCAmelCase__			  =      torch.tensor(
														    [2.38_04e-03, 2.07_52e-03, 1.98_36e-03, 2.10_57e-03, 1.61_74e-03,
														     3.05_18e-04, 9.15_53e-05, 3.35_69e-04, 9.76_56e-04, 1.83_11e-03,
														     2.01_42e-03, 2.10_57e-03, 1.73_95e-03, 4.57_76e-04, -3.96_73e-04,
														     4.57_76e-04, 1.00_71e-03, 9.15_53e-05, 4.88_28e-04, 1.15_97e-03,
														     7.32_42e-04, 9.46_04e-04, 1.80_05e-03, 1.83_11e-03, 8.85_01e-04,
														     4.27_25e-04, 4.88_28e-04, 7.32_42e-04, 1.09_86e-03, 2.10_57e-03]  )
														# fmt: on

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







							def     __SCREAMING_SNAKE_CASE				(						self  )			->				Dict:
														# 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=lowerCamelCase_   ,       return_tensors='''pt'''  ).input_values
														self.assertEquals(input_values.shape   ,       (1, 3_66, 80)  )
														self.assertTrue(torch.allclose(input_values[0, 0, :30]   ,       lowerCamelCase_   ,       atol=1e-4  )  )



'''simple docstring'''






from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput, apply_forward_hook
from .modeling_utils import ModelMixin
from .vae import Decoder, DecoderOutput, Encoder, VectorQuantizer




@dataclass
class 						a__							(				a__    ):





							'''simple docstring'''






							lowercase__     :  torch.FloatTensor




class 						a__							(				a__							,							a__    ):





							'''simple docstring'''






							@register_to_config
							def __init__(						self   ,       lowerCamelCase_ = 3   ,       lowerCamelCase_ = 3   ,       lowerCamelCase_ = ("DownEncoderBlock2D",)   ,       lowerCamelCase_ = ("UpDecoderBlock2D",)   ,       lowerCamelCase_ = (64,)   ,       lowerCamelCase_ = 1   ,       lowerCamelCase_ = "silu"   ,       lowerCamelCase_ = 3   ,       lowerCamelCase_ = 32   ,       lowerCamelCase_ = 2_56   ,       lowerCamelCase_ = 32   ,       lowerCamelCase_ = None   ,       lowerCamelCase_ = 0.18_215   ,       lowerCamelCase_ = "group"   ,       )			->				Union[str, Any]:
														super().__init__()

														# pass init params to Encoder
														lowerCAmelCase__			  =      Encoder(
														    in_channels=lowerCamelCase_   ,       out_channels=lowerCamelCase_   ,       down_block_types=lowerCamelCase_   ,       block_out_channels=lowerCamelCase_   ,       layers_per_block=lowerCamelCase_   ,       act_fn=lowerCamelCase_   ,       norm_num_groups=lowerCamelCase_   ,       double_z=lowerCamelCase_   ,       )

														lowerCAmelCase__			  =      vq_embed_dim if vq_embed_dim is not None else latent_channels

														lowerCAmelCase__			  =      nn.Convad(lowerCamelCase_   ,       lowerCamelCase_   ,       1  )
														lowerCAmelCase__			  =      VectorQuantizer(lowerCamelCase_   ,       lowerCamelCase_   ,       beta=0.25   ,       remap=lowerCamelCase_   ,       sane_index_shape=lowerCamelCase_  )
														lowerCAmelCase__			  =      nn.Convad(lowerCamelCase_   ,       lowerCamelCase_   ,       1  )

														# pass init params to Decoder
														lowerCAmelCase__			  =      Decoder(
														    in_channels=lowerCamelCase_   ,       out_channels=lowerCamelCase_   ,       up_block_types=lowerCamelCase_   ,       block_out_channels=lowerCamelCase_   ,       layers_per_block=lowerCamelCase_   ,       act_fn=lowerCamelCase_   ,       norm_num_groups=lowerCamelCase_   ,       norm_type=lowerCamelCase_   ,       )







							@apply_forward_hook
							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_ = True  )			->				VQEncoderOutput:
														lowerCAmelCase__			  =      self.encoder(lowerCamelCase_  )
														lowerCAmelCase__			  =      self.quant_conv(lowerCamelCase_  )

														if not return_dict:
																					return (h,)

														return VQEncoderOutput(latents=lowerCamelCase_  )







							@apply_forward_hook
							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_ = False   ,       lowerCamelCase_ = True  )			->				Union[DecoderOutput, torch.FloatTensor]:
														# also go through quantization layer
														if not force_not_quantize:
																					lowerCAmelCase__  ,   lowerCAmelCase__  ,   lowerCAmelCase__			  =      self.quantize(lowerCamelCase_  )
														else:
																					lowerCAmelCase__			  =      h
														lowerCAmelCase__			  =      self.post_quant_conv(lowerCamelCase_  )
														lowerCAmelCase__			  =      self.decoder(lowerCamelCase_   ,       quant if self.config.norm_type == '''spatial''' else None  )

														if not return_dict:
																					return (dec,)

														return DecoderOutput(sample=lowerCamelCase_  )







							def     __SCREAMING_SNAKE_CASE				(						self   ,       lowerCamelCase_   ,       lowerCamelCase_ = True  )			->				Union[DecoderOutput, torch.FloatTensor]:
														lowerCAmelCase__			  =      sample
														lowerCAmelCase__			  =      self.encode(lowerCamelCase_  ).latents
														lowerCAmelCase__			  =      self.decode(lowerCamelCase_  ).sample

														if not return_dict:
																					return (dec,)

														return DecoderOutput(sample=lowerCamelCase_  )