





'''simple docstring'''


import argparse
import os.path as osp
import re

import torch
from safetensors.torch import load_file, save_file


# =================#
# UNet Conversion #
# =================#

_SCREAMING_SNAKE_CASE		       =							[
    # (stable-diffusion, HF Diffusers)
    ("time_embed.0.weight", "time_embedding.linear_1.weight"),
    ("time_embed.0.bias", "time_embedding.linear_1.bias"),
    ("time_embed.2.weight", "time_embedding.linear_2.weight"),
    ("time_embed.2.bias", "time_embedding.linear_2.bias"),
    ("input_blocks.0.0.weight", "conv_in.weight"),
    ("input_blocks.0.0.bias", "conv_in.bias"),
    ("out.0.weight", "conv_norm_out.weight"),
    ("out.0.bias", "conv_norm_out.bias"),
    ("out.2.weight", "conv_out.weight"),
    ("out.2.bias", "conv_out.bias"),
]

_SCREAMING_SNAKE_CASE		       =							[
    # (stable-diffusion, HF Diffusers)
    ("in_layers.0", "norm1"),
    ("in_layers.2", "conv1"),
    ("out_layers.0", "norm2"),
    ("out_layers.3", "conv2"),
    ("emb_layers.1", "time_emb_proj"),
    ("skip_connection", "conv_shortcut"),
]

_SCREAMING_SNAKE_CASE		       =							[]
# hardcoded number of downblocks and resnets/attentions...
# would need smarter logic for other networks.
for i in range(4):
       # loop over downblocks/upblocks

       for j in range(2):
              # loop over resnets/attentions for downblocks
              _SCREAMING_SNAKE_CASE		       =							f'''down_blocks.{i}.resnets.{j}.'''
              _SCREAMING_SNAKE_CASE		       =							f'''input_blocks.{3*i + j + 1}.0.'''
              unet_conversion_map_layer.append((sd_down_res_prefix, hf_down_res_prefix))

              if i < 3:
                     # no attention layers in down_blocks.3
                     _SCREAMING_SNAKE_CASE		       =							f'''down_blocks.{i}.attentions.{j}.'''
                     _SCREAMING_SNAKE_CASE		       =							f'''input_blocks.{3*i + j + 1}.1.'''
                     unet_conversion_map_layer.append((sd_down_atn_prefix, hf_down_atn_prefix))

       for j in range(3):
              # loop over resnets/attentions for upblocks
              _SCREAMING_SNAKE_CASE		       =							f'''up_blocks.{i}.resnets.{j}.'''
              _SCREAMING_SNAKE_CASE		       =							f'''output_blocks.{3*i + j}.0.'''
              unet_conversion_map_layer.append((sd_up_res_prefix, hf_up_res_prefix))

              if i > 0:
                     # no attention layers in up_blocks.0
                     _SCREAMING_SNAKE_CASE		       =							f'''up_blocks.{i}.attentions.{j}.'''
                     _SCREAMING_SNAKE_CASE		       =							f'''output_blocks.{3*i + j}.1.'''
                     unet_conversion_map_layer.append((sd_up_atn_prefix, hf_up_atn_prefix))

       if i < 3:
              # no downsample in down_blocks.3
              _SCREAMING_SNAKE_CASE		       =							f'''down_blocks.{i}.downsamplers.0.conv.'''
              _SCREAMING_SNAKE_CASE		       =							f'''input_blocks.{3*(i+1)}.0.op.'''
              unet_conversion_map_layer.append((sd_downsample_prefix, hf_downsample_prefix))

              # no upsample in up_blocks.3
              _SCREAMING_SNAKE_CASE		       =							f'''up_blocks.{i}.upsamplers.0.'''
              _SCREAMING_SNAKE_CASE		       =							f'''output_blocks.{3*i + 2}.{1 if i == 0 else 2}.'''
              unet_conversion_map_layer.append((sd_upsample_prefix, hf_upsample_prefix))

_SCREAMING_SNAKE_CASE		       =							"mid_block.attentions.0."
_SCREAMING_SNAKE_CASE		       =							"middle_block.1."
unet_conversion_map_layer.append((sd_mid_atn_prefix, hf_mid_atn_prefix))

for j in range(2):
       _SCREAMING_SNAKE_CASE		       =							f'''mid_block.resnets.{j}.'''
       _SCREAMING_SNAKE_CASE		       =							f'''middle_block.{2*j}.'''
       unet_conversion_map_layer.append((sd_mid_res_prefix, hf_mid_res_prefix))


def 		__a(SCREAMING_SNAKE_CASE_     :    Union[str, Any]			):



    '''simple docstring'''





    _lowerCAmelCase			=					{k: k for k in unet_state_dict.keys()}
    for sd_name, hf_name in unet_conversion_map:
        _lowerCAmelCase			=					sd_name
    for k, v in mapping.items():
        if "resnets" in k:
            for sd_part, hf_part in unet_conversion_map_resnet:
                _lowerCAmelCase			=					v.replace(__snake_case						,				__snake_case			)
            _lowerCAmelCase			=					v
    for k, v in mapping.items():
        for sd_part, hf_part in unet_conversion_map_layer:
            _lowerCAmelCase			=					v.replace(__snake_case						,				__snake_case			)
        _lowerCAmelCase			=					v
    _lowerCAmelCase			=					{v: unet_state_dict[k] for k, v in mapping.items()}
    return new_state_dict


# ================#
# VAE Conversion #
# ================#

_SCREAMING_SNAKE_CASE		       =							[
    # (stable-diffusion, HF Diffusers)
    ("nin_shortcut", "conv_shortcut"),
    ("norm_out", "conv_norm_out"),
    ("mid.attn_1.", "mid_block.attentions.0."),
]

for i in range(4):
       # down_blocks have two resnets
       for j in range(2):
              _SCREAMING_SNAKE_CASE		       =							f'''encoder.down_blocks.{i}.resnets.{j}.'''
              _SCREAMING_SNAKE_CASE		       =							f'''encoder.down.{i}.block.{j}.'''
              vae_conversion_map.append((sd_down_prefix, hf_down_prefix))

       if i < 3:
              _SCREAMING_SNAKE_CASE		       =							f'''down_blocks.{i}.downsamplers.0.'''
              _SCREAMING_SNAKE_CASE		       =							f'''down.{i}.downsample.'''
              vae_conversion_map.append((sd_downsample_prefix, hf_downsample_prefix))

              _SCREAMING_SNAKE_CASE		       =							f'''up_blocks.{i}.upsamplers.0.'''
              _SCREAMING_SNAKE_CASE		       =							f'''up.{3-i}.upsample.'''
              vae_conversion_map.append((sd_upsample_prefix, hf_upsample_prefix))

       # up_blocks have three resnets
       # also, up blocks in hf are numbered in reverse from sd
       for j in range(3):
              _SCREAMING_SNAKE_CASE		       =							f'''decoder.up_blocks.{i}.resnets.{j}.'''
              _SCREAMING_SNAKE_CASE		       =							f'''decoder.up.{3-i}.block.{j}.'''
              vae_conversion_map.append((sd_up_prefix, hf_up_prefix))

# this part accounts for mid blocks in both the encoder and the decoder
for i in range(2):
       _SCREAMING_SNAKE_CASE		       =							f'''mid_block.resnets.{i}.'''
       _SCREAMING_SNAKE_CASE		       =							f'''mid.block_{i+1}.'''
       vae_conversion_map.append((sd_mid_res_prefix, hf_mid_res_prefix))


_SCREAMING_SNAKE_CASE		       =							[
    # (stable-diffusion, HF Diffusers)
    ("norm.", "group_norm."),
    ("q.", "query."),
    ("k.", "key."),
    ("v.", "value."),
    ("proj_out.", "proj_attn."),
]


def 		__a(SCREAMING_SNAKE_CASE_     :    Optional[int]			):



    '''simple docstring'''





    return w.reshape(*w.shape						,				1						,				1			)


def 		__a(SCREAMING_SNAKE_CASE_     :    int			):



    '''simple docstring'''





    _lowerCAmelCase			=					{k: k for k in vae_state_dict.keys()}
    for k, v in mapping.items():
        for sd_part, hf_part in vae_conversion_map:
            _lowerCAmelCase			=					v.replace(__snake_case						,				__snake_case			)
        _lowerCAmelCase			=					v
    for k, v in mapping.items():
        if "attentions" in k:
            for sd_part, hf_part in vae_conversion_map_attn:
                _lowerCAmelCase			=					v.replace(__snake_case						,				__snake_case			)
            _lowerCAmelCase			=					v
    _lowerCAmelCase			=					{v: vae_state_dict[k] for k, v in mapping.items()}
    _lowerCAmelCase			=					["q", "k", "v", "proj_out"]
    for k, v in new_state_dict.items():
        for weight_name in weights_to_convert:
            if F'''mid.attn_1.{weight_name}.weight''' in k:
                print(F'''Reshaping {k} for SD format'''			)
                _lowerCAmelCase			=					reshape_weight_for_sd(__snake_case			)
    return new_state_dict


# =========================#
# Text Encoder Conversion #
# =========================#


_SCREAMING_SNAKE_CASE		       =							[
    # (stable-diffusion, HF Diffusers)
    ("resblocks.", "text_model.encoder.layers."),
    ("ln_1", "layer_norm1"),
    ("ln_2", "layer_norm2"),
    (".c_fc.", ".fc1."),
    (".c_proj.", ".fc2."),
    (".attn", ".self_attn"),
    ("ln_final.", "transformer.text_model.final_layer_norm."),
    ("token_embedding.weight", "transformer.text_model.embeddings.token_embedding.weight"),
    ("positional_embedding", "transformer.text_model.embeddings.position_embedding.weight"),
]
_SCREAMING_SNAKE_CASE		       =							{re.escape(x[1]): x[0] for x in textenc_conversion_lst}
_SCREAMING_SNAKE_CASE		       =							re.compile("|".join(protected.keys()))

# Ordering is from https://github.com/pytorch/pytorch/blob/master/test/cpp/api/modules.cpp
_SCREAMING_SNAKE_CASE		       =							{"q": 0, "k": 1, "v": 2}


def 		__a(SCREAMING_SNAKE_CASE_     :    List[Any]			):



    '''simple docstring'''





    _lowerCAmelCase			=					{}
    _lowerCAmelCase			=					{}
    _lowerCAmelCase			=					{}
    for k, v in text_enc_dict.items():
        if (
            k.endswith(".self_attn.q_proj.weight"			)
            or k.endswith(".self_attn.k_proj.weight"			)
            or k.endswith(".self_attn.v_proj.weight"			)
        ):
            _lowerCAmelCase			=					k[: -len(".q_proj.weight"			)]
            _lowerCAmelCase			=					k[-len("q_proj.weight"			)]
            if k_pre not in capture_qkv_weight:
                _lowerCAmelCase			=					[None, None, None]
            _lowerCAmelCase			=					v
            continue

        if (
            k.endswith(".self_attn.q_proj.bias"			)
            or k.endswith(".self_attn.k_proj.bias"			)
            or k.endswith(".self_attn.v_proj.bias"			)
        ):
            _lowerCAmelCase			=					k[: -len(".q_proj.bias"			)]
            _lowerCAmelCase			=					k[-len("q_proj.bias"			)]
            if k_pre not in capture_qkv_bias:
                _lowerCAmelCase			=					[None, None, None]
            _lowerCAmelCase			=					v
            continue

        _lowerCAmelCase			=					textenc_pattern.sub(lambda SCREAMING_SNAKE_CASE_			: protected[re.escape(m.group(0			)			)]						,				__snake_case			)
        _lowerCAmelCase			=					v

    for k_pre, tensors in capture_qkv_weight.items():
        if None in tensors:
            raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing"			)
        _lowerCAmelCase			=					textenc_pattern.sub(lambda SCREAMING_SNAKE_CASE_			: protected[re.escape(m.group(0			)			)]						,				__snake_case			)
        _lowerCAmelCase			=					torch.cat(__snake_case			)

    for k_pre, tensors in capture_qkv_bias.items():
        if None in tensors:
            raise Exception("CORRUPTED MODEL: one of the q-k-v values for the text encoder was missing"			)
        _lowerCAmelCase			=					textenc_pattern.sub(lambda SCREAMING_SNAKE_CASE_			: protected[re.escape(m.group(0			)			)]						,				__snake_case			)
        _lowerCAmelCase			=					torch.cat(__snake_case			)

    return new_state_dict


def 		__a(SCREAMING_SNAKE_CASE_     :    List[str]			):



    '''simple docstring'''





    return text_enc_dict


if __name__ == "__main__":
       _SCREAMING_SNAKE_CASE		       =							argparse.ArgumentParser()

       parser.add_argument("--model_path", default=None, type=str, required=True, help="Path to the model to convert.")
       parser.add_argument("--checkpoint_path", default=None, type=str, required=True, help="Path to the output model.")
       parser.add_argument("--half", action="store_true", help="Save weights in half precision.")
       parser.add_argument(
           "--use_safetensors", action="store_true", help="Save weights use safetensors, default is ckpt."
       )

       _SCREAMING_SNAKE_CASE		       =							parser.parse_args()

       assert args.model_path is not None, "Must provide a model path!"

       assert args.checkpoint_path is not None, "Must provide a checkpoint path!"

       # Path for safetensors
       _SCREAMING_SNAKE_CASE		       =							osp.join(args.model_path, "unet", "diffusion_pytorch_model.safetensors")
       _SCREAMING_SNAKE_CASE		       =							osp.join(args.model_path, "vae", "diffusion_pytorch_model.safetensors")
       _SCREAMING_SNAKE_CASE		       =							osp.join(args.model_path, "text_encoder", "model.safetensors")

       # Load models from safetensors if it exists, if it doesn't pytorch
       if osp.exists(unet_path):
              _SCREAMING_SNAKE_CASE		       =							load_file(unet_path, device="cpu")
       else:
              _SCREAMING_SNAKE_CASE		       =							osp.join(args.model_path, "unet", "diffusion_pytorch_model.bin")
              _SCREAMING_SNAKE_CASE		       =							torch.load(unet_path, map_location="cpu")

       if osp.exists(vae_path):
              _SCREAMING_SNAKE_CASE		       =							load_file(vae_path, device="cpu")
       else:
              _SCREAMING_SNAKE_CASE		       =							osp.join(args.model_path, "vae", "diffusion_pytorch_model.bin")
              _SCREAMING_SNAKE_CASE		       =							torch.load(vae_path, map_location="cpu")

       if osp.exists(text_enc_path):
              _SCREAMING_SNAKE_CASE		       =							load_file(text_enc_path, device="cpu")
       else:
              _SCREAMING_SNAKE_CASE		       =							osp.join(args.model_path, "text_encoder", "pytorch_model.bin")
              _SCREAMING_SNAKE_CASE		       =							torch.load(text_enc_path, map_location="cpu")

       # Convert the UNet model
       _SCREAMING_SNAKE_CASE		       =							convert_unet_state_dict(unet_state_dict)
       _SCREAMING_SNAKE_CASE		       =							{"model.diffusion_model." + k: v for k, v in unet_state_dict.items()}

       # Convert the VAE model
       _SCREAMING_SNAKE_CASE		       =							convert_vae_state_dict(vae_state_dict)
       _SCREAMING_SNAKE_CASE		       =							{"first_stage_model." + k: v for k, v in vae_state_dict.items()}

       # Easiest way to identify v2.0 model seems to be that the text encoder (OpenCLIP) is deeper
       _SCREAMING_SNAKE_CASE		       =							"text_model.encoder.layers.22.layer_norm2.bias" in text_enc_dict

       if is_vaa_model:
              # Need to add the tag 'transformer' in advance so we can knock it out from the final layer-norm
              _SCREAMING_SNAKE_CASE		       =							{"transformer." + k: v for k, v in text_enc_dict.items()}
              _SCREAMING_SNAKE_CASE		       =							convert_text_enc_state_dict_vaa(text_enc_dict)
              _SCREAMING_SNAKE_CASE		       =							{"cond_stage_model.model." + k: v for k, v in text_enc_dict.items()}
       else:
              _SCREAMING_SNAKE_CASE		       =							convert_text_enc_state_dict(text_enc_dict)
              _SCREAMING_SNAKE_CASE		       =							{"cond_stage_model.transformer." + k: v for k, v in text_enc_dict.items()}

       # Put together new checkpoint
       _SCREAMING_SNAKE_CASE		       =							{**unet_state_dict, **vae_state_dict, **text_enc_dict}
       if args.half:
              _SCREAMING_SNAKE_CASE		       =							{k: v.half() for k, v in state_dict.items()}

       if args.use_safetensors:
              save_file(state_dict, args.checkpoint_path)
       else:
              _SCREAMING_SNAKE_CASE		       =							{"state_dict": state_dict}
              torch.save(state_dict, args.checkpoint_path)









import itertools
import random
import unittest

import numpy as np

from transformers import is_speech_available
from transformers.testing_utils import require_torch, require_torchaudio

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_speech_available():
			from transformers import SpeechaTextFeatureExtractor

_a 						=      random.Random()





def lowerCAmelCase__(__snake_case       ,__snake_case=1.0       ,__snake_case=None       ,__snake_case=None							)				->      List[Any]:







	'''simple docstring'''







	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
@require_torchaudio
class 					__A    (    unittest.TestCase ):

			'''simple docstring'''







			def __init__(    self       ,  __lowerCAmelCase       ,  __lowerCAmelCase=7       ,  __lowerCAmelCase=4_0_0       ,  __lowerCAmelCase=2_0_0_0       ,  __lowerCAmelCase=2_4       ,  __lowerCAmelCase=2_4       ,  __lowerCAmelCase=0.0       ,  __lowerCAmelCase=1_6_0_0_0       ,  __lowerCAmelCase=True       ,  __lowerCAmelCase=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__								=		num_mel_bins
				lowerCamelCase__								=		padding_value
				lowerCamelCase__								=		sampling_rate
				lowerCamelCase__								=		return_attention_mask
				lowerCamelCase__								=		do_normalize





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				return {
				    "feature_size": self.feature_size,
				    "num_mel_bins": self.num_mel_bins,
				    "padding_value": self.padding_value,
				    "sampling_rate": self.sampling_rate,
				    "return_attention_mask": self.return_attention_mask,
				    "do_normalize": self.do_normalize,
				}





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase=False       ,  __lowerCAmelCase=False							):



				'''simple docstring'''







				def _flatten(__lowerCAmelCase							):
					return list(itertools.chain(*__lowerCAmelCase							)							)

				if equal_length:
					lowerCamelCase__								=		[floats_list((self.max_seq_length, self.feature_size)							) for _ in range(self.batch_size							)]
				else:
					# make sure that inputs increase in size
					lowerCamelCase__								=		[
					    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






@require_torch
@require_torchaudio
class 					__A    (    lowerCAmelCase			,				unittest.TestCase ):

			'''simple docstring'''







			lowerCAmelCase_							     =						SpeechaTextFeatureExtractor if is_speech_available() else None





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		SpeechaTextFeatureExtractionTester(self							)





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase							):



				'''simple docstring'''







				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 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				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_0_0       ,  1_4_0_0       ,  2_0_0							)]
				lowerCamelCase__								=		[np.asarray(__lowerCAmelCase							) for speech_input in speech_inputs]

				# Test feature size
				lowerCamelCase__								=		feature_extractor(__lowerCAmelCase       ,  padding=__lowerCAmelCase       ,  return_tensors='''np'''							).input_features
				self.assertTrue(input_features.ndim == 3							)
				self.assertTrue(input_features.shape[-1] == feature_extractor.feature_size							)

				# Test not batched input
				lowerCamelCase__								=		feature_extractor(speech_inputs[0]       ,  return_tensors='''np'''							).input_features
				lowerCamelCase__								=		feature_extractor(np_speech_inputs[0]       ,  return_tensors='''np'''							).input_features
				self.assertTrue(np.allclose(__lowerCAmelCase       ,  __lowerCAmelCase       ,  atol=1E-3							)							)

				# Test batched
				lowerCamelCase__								=		feature_extractor(__lowerCAmelCase       ,  return_tensors='''np'''							).input_features
				lowerCamelCase__								=		feature_extractor(__lowerCAmelCase       ,  return_tensors='''np'''							).input_features
				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_0_0, 8_0_0, 8_0_0)]
				lowerCamelCase__								=		np.asarray(__lowerCAmelCase							)
				lowerCamelCase__								=		feature_extractor(__lowerCAmelCase       ,  return_tensors='''np'''							).input_features
				lowerCamelCase__								=		feature_extractor(__lowerCAmelCase       ,  return_tensors='''np'''							).input_features
				for enc_seq_a, enc_seq_a in zip(__lowerCAmelCase       ,  __lowerCAmelCase							):
					self.assertTrue(np.allclose(__lowerCAmelCase       ,  __lowerCAmelCase       ,  atol=1E-3							)							)





			def 		__lowerCamelCase				(    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(8_0_0       ,  1_4_0_0       ,  2_0_0							)]

				lowerCamelCase__								=		['''longest''', '''max_length''', '''do_not_pad''']
				lowerCamelCase__								=		[None, 1_6, None]
				for max_length, padding in zip(__lowerCAmelCase       ,  __lowerCAmelCase							):
					lowerCamelCase__								=		feature_extractor(
					    __lowerCAmelCase       ,  padding=__lowerCAmelCase       ,  max_length=__lowerCAmelCase       ,  return_attention_mask=__lowerCAmelCase							)
					lowerCamelCase__								=		inputs.input_features
					lowerCamelCase__								=		inputs.attention_mask
					lowerCamelCase__								=		[np.sum(__lowerCAmelCase							) for x in attention_mask]

					self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]							)
					self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]							)
					self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]							)





			def 		__lowerCamelCase				(    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(8_0_0       ,  1_4_0_0       ,  2_0_0							)]

				lowerCamelCase__								=		['''longest''', '''max_length''', '''do_not_pad''']
				lowerCamelCase__								=		[None, 1_6, None]
				for max_length, padding in zip(__lowerCAmelCase       ,  __lowerCAmelCase							):
					lowerCamelCase__								=		feature_extractor(
					    __lowerCAmelCase       ,  max_length=__lowerCAmelCase       ,  padding=__lowerCAmelCase       ,  return_tensors='''np'''       ,  return_attention_mask=__lowerCAmelCase							)
					lowerCamelCase__								=		inputs.input_features
					lowerCamelCase__								=		inputs.attention_mask
					lowerCamelCase__								=		[np.sum(__lowerCAmelCase							) for x in attention_mask]

					self._check_zero_mean_unit_variance(input_features[0][: fbank_feat_lengths[0]]							)
					self.assertTrue(input_features[0][fbank_feat_lengths[0] :].sum() < 1E-6							)
					self._check_zero_mean_unit_variance(input_features[1][: fbank_feat_lengths[1]]							)
					self.assertTrue(input_features[0][fbank_feat_lengths[1] :].sum() < 1E-6							)
					self._check_zero_mean_unit_variance(input_features[2][: fbank_feat_lengths[2]]							)





			def 		__lowerCamelCase				(    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(8_0_0       ,  1_4_0_0       ,  2_0_0							)]
				lowerCamelCase__								=		feature_extractor(
				    __lowerCAmelCase       ,  padding='''max_length'''       ,  max_length=4       ,  truncation=__lowerCAmelCase       ,  return_tensors='''np'''       ,  return_attention_mask=__lowerCAmelCase       ,  )
				lowerCamelCase__								=		inputs.input_features
				lowerCamelCase__								=		inputs.attention_mask
				lowerCamelCase__								=		np.sum(attention_mask == 1       ,  axis=1							)

				self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]							)
				self._check_zero_mean_unit_variance(input_features[1]							)
				self._check_zero_mean_unit_variance(input_features[2]							)





			def 		__lowerCamelCase				(    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(8_0_0       ,  1_4_0_0       ,  2_0_0							)]
				lowerCamelCase__								=		feature_extractor(
				    __lowerCAmelCase       ,  padding='''longest'''       ,  max_length=4       ,  truncation=__lowerCAmelCase       ,  return_tensors='''np'''       ,  return_attention_mask=__lowerCAmelCase       ,  )
				lowerCamelCase__								=		inputs.input_features
				lowerCamelCase__								=		inputs.attention_mask
				lowerCamelCase__								=		np.sum(attention_mask == 1       ,  axis=1							)

				self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]							)
				self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]							)
				self._check_zero_mean_unit_variance(input_features[2]							)

				# make sure that if max_length < longest -> then pad to max_length
				self.assertEqual(input_features.shape       ,  (3, 4, 2_4)							)

				lowerCamelCase__								=		[floats_list((1, x)							)[0] for x in range(8_0_0       ,  1_4_0_0       ,  2_0_0							)]
				lowerCamelCase__								=		feature_extractor(
				    __lowerCAmelCase       ,  padding='''longest'''       ,  max_length=1_6       ,  truncation=__lowerCAmelCase       ,  return_tensors='''np'''       ,  return_attention_mask=__lowerCAmelCase       ,  )
				lowerCamelCase__								=		inputs.input_features
				lowerCamelCase__								=		inputs.attention_mask
				lowerCamelCase__								=		np.sum(attention_mask == 1       ,  axis=1							)

				self._check_zero_mean_unit_variance(input_features[0, : fbank_feat_lengths[0]]							)
				self._check_zero_mean_unit_variance(input_features[1, : fbank_feat_lengths[1]]							)
				self._check_zero_mean_unit_variance(input_features[2]							)

				# make sure that if max_length < longest -> then pad to max_length
				self.assertEqual(input_features.shape       ,  (3, 6, 2_4)							)





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				import torch

				lowerCamelCase__								=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
				lowerCamelCase__								=		np.random.rand(1_0_0       ,  3_2							).astype(np.floataa							)
				lowerCamelCase__								=		np_speech_inputs.tolist()

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





			def 		__lowerCamelCase				(    self       ,  __lowerCAmelCase							):



				'''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(__lowerCAmelCase							)							)[:num_samples]['''audio''']

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





			def 		__lowerCamelCase				(    self							):



				'''simple docstring'''







				lowerCamelCase__								=		np.array([
				    -1.5745, -1.7713, -1.7020, -1.6069, -1.2250, -1.1105, -0.9072, -0.8241,
				    -1.2310, -0.8098, -0.3320, -0.4101, -0.7985, -0.4996, -0.8213, -0.9128,
				    -1.0420, -1.1286, -1.0440, -0.7999, -0.8405, -1.2275, -1.5443, -1.4625,
				]							)
				# fmt: on

				lowerCamelCase__								=		self._load_datasamples(1							)
				lowerCamelCase__								=		self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()							)
				lowerCamelCase__								=		feature_extractor(__lowerCAmelCase       ,  return_tensors='''pt'''							).input_features
				self.assertEquals(input_features.shape       ,  (1, 5_8_4, 2_4)							)
				self.assertTrue(np.allclose(input_features[0, 0, :3_0]       ,  __lowerCAmelCase       ,  atol=1E-4							)							)



