# Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. ############################################################################### """ _torch_cvt.py: Data processing graph converted from PyTorch """ import io import onnx import torch import numpy as np from onnx import numpy_helper from ._ortapi2 import make_onnx_model from ._cuops import SingleOpGraph from ._hf_cvt import HFTokenizerConverter from .util import remove_unused_initializers, mel_filterbank class _WhisperHParams: SAMPLE_RATE = 16000 N_FFT = 400 N_MELS = 80 HOP_LENGTH = 160 CHUNK_LENGTH = 30 N_SAMPLES = CHUNK_LENGTH * SAMPLE_RATE # 480000 samples in a 30-second chunk N_FRAMES = N_SAMPLES // HOP_LENGTH class CustomOpStftNorm(torch.autograd.Function): @staticmethod def symbolic(g, self, n_fft, hop_length, window): t_n_fft = g.op('Constant', value_t=torch.tensor( n_fft, dtype=torch.int64)) t_hop_length = g.op('Constant', value_t=torch.tensor( hop_length, dtype=torch.int64)) t_frame_size = g.op( 'Constant', value_t=torch.tensor(n_fft, dtype=torch.int64)) return g.op("ai.onnx.contrib::StftNorm", self, t_n_fft, t_hop_length, window, t_frame_size) @staticmethod def forward(ctx, audio, n_fft, hop_length, window): win_length = window.shape[0] stft = torch.stft(audio, n_fft, hop_length, win_length, window, center=True, pad_mode="reflect", normalized=False, onesided=True, return_complex=True) return stft.abs() ** 2 class WhisperPrePipeline(torch.nn.Module): def __init__(self, sr=_WhisperHParams.SAMPLE_RATE, n_fft=_WhisperHParams.N_FFT, hop_length=_WhisperHParams.HOP_LENGTH, n_mels=_WhisperHParams.N_MELS, n_samples=_WhisperHParams.N_SAMPLES): super().__init__() self.n_samples = n_samples self.hop_length = hop_length self.n_fft = n_fft self.window = torch.hann_window(n_fft) self.mel_filters = torch.from_numpy( mel_filterbank(sr=sr, n_fft=n_fft, n_mels=n_mels)) def forward(self, audio_pcm: torch.Tensor): stft_norm = CustomOpStftNorm.apply(audio_pcm, self.n_fft, self.hop_length, self.window) magnitudes = stft_norm[:, :, :-1] mel_spec = self.mel_filters @ magnitudes log_spec = torch.clamp(mel_spec, min=1e-10).log10() spec_min = log_spec.max() - 8.0 log_spec = torch.maximum(log_spec, spec_min) spec_shape = log_spec.shape padding_spec = torch.ones(spec_shape[0], spec_shape[1], self.n_samples // self.hop_length - spec_shape[2], dtype=torch.float) padding_spec *= spec_min log_spec = torch.cat((log_spec, padding_spec), dim=2) log_spec = (log_spec + 4.0) / 4.0 return log_spec def _to_onnx_stft(onnx_model, n_fft): """Convert custom-op STFT-Norm to ONNX STFT""" node_idx = 0 new_stft_nodes = [] stft_norm_node = None for node in onnx_model.graph.node: if node.op_type == "StftNorm": stft_norm_node = node break node_idx += 1 if stft_norm_node is None: raise RuntimeError("Cannot find STFTNorm node in the graph") make_node = onnx.helper.make_node replaced_nodes = [ make_node('Constant', inputs=[], outputs=['const_minus_1_output_0'], name='const_minus_1', value=numpy_helper.from_array(np.array([-1], dtype='int64'))), make_node('Constant', inputs=[], outputs=['const_14_output_0'], name='const_14', value=numpy_helper.from_array(np.array([0, n_fft // 2, 0, n_fft // 2], dtype='int64'), name='const_14')), make_node('Pad', inputs=[stft_norm_node.input[0], 'const_14_output_0'], outputs=['pad_1_output_0'], mode='reflect'), make_node('Unsqueeze', inputs=['pad_1_output_0', 'const_minus_1_output_0'], outputs=['unsqueeze_1_output_0'], name='unsqueeze_1'), make_node('STFT', inputs=['unsqueeze_1_output_0', stft_norm_node.input[2], stft_norm_node.input[3], stft_norm_node.input[4]], outputs=['stft_output_0'], name='stft', onesided=1), make_node('Transpose', inputs=['stft_output_0'], outputs=['transpose_1_output_0'], name='transpose_1', perm=[0, 2, 1, 3]), make_node('Constant', inputs=[], outputs=['const_17_output_0'], name='const_17', value=numpy_helper.from_array(np.array([2], dtype='int64'), name='')), make_node('Constant', inputs=[], outputs=['const_18_output_0'], name='const_18', value=numpy_helper.from_array(np.array([0], dtype='int64'), name='')), make_node('Constant', inputs=[], outputs=['const_20_output_0'], name='const_20', value=numpy_helper.from_array(np.array([1], dtype='int64'), name='')), make_node('Slice', inputs=['transpose_1_output_0', 'const_18_output_0', 'const_minus_1_output_0', 'const_17_output_0', 'const_20_output_0'], outputs=['slice_1_output_0'], name='slice_1'), make_node('Constant', inputs=[], outputs=[ 'const0_output_0'], name='const0', value_int=0), make_node('Constant', inputs=[], outputs=[ 'const1_output_0'], name='const1', value_int=1), make_node('Gather', inputs=['slice_1_output_0', 'const0_output_0'], outputs=['gather_4_output_0'], name='gather_4', axis=3), make_node('Gather', inputs=['slice_1_output_0', 'const1_output_0'], outputs=['gather_5_output_0'], name='gather_5', axis=3), make_node('Mul', inputs=['gather_4_output_0', 'gather_4_output_0'], outputs=[ 'mul_output_0'], name='mul0'), make_node('Mul', inputs=['gather_5_output_0', 'gather_5_output_0'], outputs=[ 'mul_1_output_0'], name='mul1'), make_node('Add', inputs=['mul_output_0', 'mul_1_output_0'], outputs=[ stft_norm_node.output[0]], name='add0'), ] new_stft_nodes.extend(onnx_model.graph.node[:node_idx]) new_stft_nodes.extend(replaced_nodes) new_stft_nodes.extend(onnx_model.graph.node[node_idx + 1:]) del onnx_model.graph.node[:] onnx_model.graph.node.extend(new_stft_nodes) onnx.checker.check_model(onnx_model) return onnx_model def _torch_export(*arg, **kwargs): with io.BytesIO() as f: torch.onnx.export(*arg, f, **kwargs) return onnx.load_from_string(f.getvalue()) class WhisperDataProcGraph: def __init__(self, processor, **kwargs): self.hf_processor = processor _opset = kwargs.pop('opset', 17) self.opset_version = _opset if _opset else 17 def pre_processing(self, **kwargs): use_audio_decoder = kwargs.pop('USE_AUDIO_DECODER', True) use_onnx_stft = kwargs.pop('USE_ONNX_STFT', True) feature_extractor = self.hf_processor.feature_extractor whisper_processing = WhisperPrePipeline( feature_extractor.sampling_rate, feature_extractor.n_fft, feature_extractor.hop_length, feature_extractor.feature_size, feature_extractor.n_samples) audio_pcm = torch.rand((1, 32000), dtype=torch.float32) model_args = (audio_pcm,) pre_model = _torch_export( whisper_processing, model_args, input_names=["audio_pcm"], output_names=["log_mel"], do_constant_folding=True, export_params=True, opset_version=self.opset_version, dynamic_axes={ "audio_pcm": {1: "sample_len"}, } ) if use_onnx_stft: pre_model = _to_onnx_stft(pre_model, feature_extractor.n_fft) remove_unused_initializers(pre_model.graph) pre_full = pre_model if use_audio_decoder: audecoder_g = SingleOpGraph.build_graph( "AudioDecoder", downsampling_rate=feature_extractor.sampling_rate, stereo_to_mono=1) audecoder_m = make_onnx_model(audecoder_g) pre_full = onnx.compose.merge_models( audecoder_m, pre_model, io_map=[("floatPCM", "audio_pcm")]) return pre_full def post_processing(self, **kwargs): skip_special_tokens = kwargs.get('skip_special_tokens', True) g = SingleOpGraph.build_graph( "BpeDecoder", cvt=HFTokenizerConverter(self.hf_processor.tokenizer).bpe_decoder, skip_special_tokens=skip_special_tokens) bpenode = g.node[0] bpenode.input[0] = "generated_ids" nodes = [onnx.helper.make_node('Cast', ['sequences'], ["generated_ids"], to=onnx.TensorProto.INT64), bpenode] del g.node[:] g.node.extend(nodes) inputs = [onnx.helper.make_tensor_value_info( "sequences", onnx.TensorProto.INT32, ['N', 'seq_len', 'ids'])] del g.input[:] g.input.extend(inputs) g.output[0].type.CopyFrom(onnx.helper.make_tensor_type_proto( onnx.TensorProto.STRING, ['N', 'text'])) return make_onnx_model(g, opset_version=self.opset_version)