mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-19 07:54:56 +00:00
Brandon Thomas
133 lines
3.7 KiB
Python
133 lines
3.7 KiB
Python
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from typing import Optional
|
|
|
|
import torch
|
|
from packaging import version
|
|
|
|
from nemo.utils import logging
|
|
|
|
# Library version globals
|
|
TORCH_VERSION = None
|
|
TORCH_VERSION_MIN = version.Version('1.7')
|
|
|
|
|
|
def stft_patch(
|
|
input: torch.Tensor,
|
|
n_fft: int,
|
|
hop_length: Optional[int] = None,
|
|
win_length: Optional[int] = None,
|
|
window: Optional[torch.Tensor] = None,
|
|
center: bool = True,
|
|
pad_mode: str = 'reflect',
|
|
normalized: bool = False,
|
|
onesided: Optional[bool] = None,
|
|
return_complex: Optional[bool] = None,
|
|
):
|
|
"""
|
|
Patch over torch.stft for PyTorch <= 1.6.
|
|
Arguments are same as torch.stft().
|
|
|
|
# TODO: Remove once PyTorch 1.7+ is a requirement.
|
|
"""
|
|
global TORCH_VERSION
|
|
if TORCH_VERSION is None:
|
|
TORCH_VERSION = version.parse(torch.__version__)
|
|
|
|
logging.warning(
|
|
"torch.stft() signature has been updated for PyTorch 1.7+\n"
|
|
"Please update PyTorch to remain compatible with later versions of NeMo."
|
|
)
|
|
|
|
if TORCH_VERSION < TORCH_VERSION_MIN:
|
|
return torch.stft(
|
|
input=input,
|
|
n_fft=n_fft,
|
|
hop_length=hop_length,
|
|
win_length=win_length,
|
|
window=window,
|
|
center=center,
|
|
pad_mode=pad_mode,
|
|
normalized=normalized,
|
|
onesided=True,
|
|
)
|
|
else:
|
|
return torch.stft(
|
|
input=input,
|
|
n_fft=n_fft,
|
|
hop_length=hop_length,
|
|
win_length=win_length,
|
|
window=window,
|
|
center=center,
|
|
pad_mode=pad_mode,
|
|
normalized=normalized,
|
|
onesided=onesided,
|
|
return_complex=return_complex,
|
|
)
|
|
|
|
|
|
def istft_patch(
|
|
input: torch.Tensor,
|
|
n_fft: int,
|
|
hop_length: Optional[int] = None,
|
|
win_length: Optional[int] = None,
|
|
window: Optional[torch.Tensor] = None,
|
|
center: bool = True,
|
|
normalized: bool = False,
|
|
onesided: Optional[bool] = None,
|
|
length: int = None,
|
|
return_complex: Optional[bool] = False,
|
|
):
|
|
"""
|
|
Patch over torch.stft for PyTorch <= 1.6.
|
|
Arguments are same as torch.stft().
|
|
|
|
# TODO: Remove once PyTorch 1.7+ is a requirement.
|
|
"""
|
|
global TORCH_VERSION
|
|
if TORCH_VERSION is None:
|
|
TORCH_VERSION = version.parse(torch.__version__)
|
|
|
|
logging.warning(
|
|
"torch.stft() signature has been updated for PyTorch 1.7+\n"
|
|
"Please update PyTorch to remain compatible with later versions of NeMo."
|
|
)
|
|
|
|
if TORCH_VERSION < TORCH_VERSION_MIN:
|
|
return torch.istft(
|
|
input=input,
|
|
n_fft=n_fft,
|
|
hop_length=hop_length,
|
|
win_length=win_length,
|
|
window=window,
|
|
center=center,
|
|
length=length,
|
|
normalized=normalized,
|
|
onesided=True,
|
|
)
|
|
else:
|
|
return torch.istft(
|
|
input=input,
|
|
n_fft=n_fft,
|
|
hop_length=hop_length,
|
|
win_length=win_length,
|
|
window=window,
|
|
center=center,
|
|
length=length,
|
|
normalized=normalized,
|
|
onesided=onesided,
|
|
return_complex=return_complex,
|
|
)
|