mirror of
https://github.com/RVC-Boss/GPT-SoVITS.git
synced 2026-09-04 15:00:00 +08:00
Compare commits
5 Commits
bcf10cbb8a
...
abac19ee6a
| Author | SHA1 | Date | |
|---|---|---|---|
|
abac19ee6a | ||
|
48b1a0169a | ||
|
d523079fc0 | ||
|
a47b87bb7b | ||
|
|
735b2e3554 |
@ -1,69 +0,0 @@
|
||||
# Copyright (c) 2024 NVIDIA CORPORATION.
|
||||
# Licensed under the MIT license.
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from alias_free_activation.torch.resample import UpSample1d, DownSample1d
|
||||
|
||||
# load fused CUDA kernel: this enables importing anti_alias_activation_cuda
|
||||
from alias_free_activation.cuda import load
|
||||
|
||||
anti_alias_activation_cuda = load.load()
|
||||
|
||||
|
||||
class FusedAntiAliasActivation(torch.autograd.Function):
|
||||
"""
|
||||
Assumes filter size 12, replication padding on upsampling/downsampling, and logscale alpha/beta parameters as inputs.
|
||||
The hyperparameters are hard-coded in the kernel to maximize speed.
|
||||
NOTE: The fused kenrel is incorrect for Activation1d with different hyperparameters.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx, inputs, up_ftr, down_ftr, alpha, beta):
|
||||
activation_results = anti_alias_activation_cuda.forward(inputs, up_ftr, down_ftr, alpha, beta)
|
||||
|
||||
return activation_results
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, output_grads):
|
||||
raise NotImplementedError
|
||||
return output_grads, None, None
|
||||
|
||||
|
||||
class Activation1d(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
activation,
|
||||
up_ratio: int = 2,
|
||||
down_ratio: int = 2,
|
||||
up_kernel_size: int = 12,
|
||||
down_kernel_size: int = 12,
|
||||
fused: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
self.up_ratio = up_ratio
|
||||
self.down_ratio = down_ratio
|
||||
self.act = activation
|
||||
self.upsample = UpSample1d(up_ratio, up_kernel_size)
|
||||
self.downsample = DownSample1d(down_ratio, down_kernel_size)
|
||||
|
||||
self.fused = fused # Whether to use fused CUDA kernel or not
|
||||
|
||||
def forward(self, x):
|
||||
if not self.fused:
|
||||
x = self.upsample(x)
|
||||
x = self.act(x)
|
||||
x = self.downsample(x)
|
||||
return x
|
||||
else:
|
||||
if self.act.__class__.__name__ == "Snake":
|
||||
beta = self.act.alpha.data # Snake uses same params for alpha and beta
|
||||
else:
|
||||
beta = self.act.beta.data # Snakebeta uses different params for alpha and beta
|
||||
alpha = self.act.alpha.data
|
||||
if not self.act.alpha_logscale: # Exp baked into cuda kernel, cancel it out with a log
|
||||
alpha = torch.log(alpha)
|
||||
beta = torch.log(beta)
|
||||
|
||||
x = FusedAntiAliasActivation.apply(x, self.upsample.filter, self.downsample.lowpass.filter, alpha, beta)
|
||||
return x
|
||||
@ -1,23 +0,0 @@
|
||||
/* coding=utf-8
|
||||
* Copyright (c) 2024, 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.
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
|
||||
extern "C" torch::Tensor fwd_cuda(torch::Tensor const &input, torch::Tensor const &up_filter, torch::Tensor const &down_filter, torch::Tensor const &alpha, torch::Tensor const &beta);
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("forward", &fwd_cuda, "Anti-Alias Activation forward (CUDA)");
|
||||
}
|
||||
@ -1,246 +0,0 @@
|
||||
/* coding=utf-8
|
||||
* Copyright (c) 2024, 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.
|
||||
*/
|
||||
|
||||
#include <ATen/ATen.h>
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_profiler_api.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <torch/extension.h>
|
||||
#include "type_shim.h"
|
||||
#include <assert.h>
|
||||
#include <cfloat>
|
||||
#include <limits>
|
||||
#include <stdint.h>
|
||||
#include <c10/macros/Macros.h>
|
||||
|
||||
namespace
|
||||
{
|
||||
// Hard-coded hyperparameters
|
||||
// WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and
|
||||
constexpr int ELEMENTS_PER_LDG_STG = 1; //(WARP_ITERATIONS < 4) ? 1 : 4;
|
||||
constexpr int BUFFER_SIZE = 32;
|
||||
constexpr int FILTER_SIZE = 12;
|
||||
constexpr int HALF_FILTER_SIZE = 6;
|
||||
constexpr int UPSAMPLE_REPLICATION_PAD = 5; // 5 on each side, matching torch impl
|
||||
constexpr int DOWNSAMPLE_REPLICATION_PAD_LEFT = 5; // matching torch impl
|
||||
constexpr int DOWNSAMPLE_REPLICATION_PAD_RIGHT = 6; // matching torch impl
|
||||
|
||||
template <typename input_t, typename output_t, typename acc_t>
|
||||
__global__ void anti_alias_activation_forward(
|
||||
output_t *dst,
|
||||
const input_t *src,
|
||||
const input_t *up_ftr,
|
||||
const input_t *down_ftr,
|
||||
const input_t *alpha,
|
||||
const input_t *beta,
|
||||
int batch_size,
|
||||
int channels,
|
||||
int seq_len)
|
||||
{
|
||||
// Up and downsample filters
|
||||
input_t up_filter[FILTER_SIZE];
|
||||
input_t down_filter[FILTER_SIZE];
|
||||
|
||||
// Load data from global memory including extra indices reserved for replication paddings
|
||||
input_t elements[2 * FILTER_SIZE + 2 * BUFFER_SIZE + 2 * UPSAMPLE_REPLICATION_PAD] = {0};
|
||||
input_t intermediates[2 * FILTER_SIZE + 2 * BUFFER_SIZE + DOWNSAMPLE_REPLICATION_PAD_LEFT + DOWNSAMPLE_REPLICATION_PAD_RIGHT] = {0};
|
||||
|
||||
// Output stores downsampled output before writing to dst
|
||||
output_t output[BUFFER_SIZE];
|
||||
|
||||
// blockDim/threadIdx = (128, 1, 1)
|
||||
// gridDim/blockIdx = (seq_blocks, channels, batches)
|
||||
int block_offset = (blockIdx.x * 128 * BUFFER_SIZE + seq_len * (blockIdx.y + gridDim.y * blockIdx.z));
|
||||
int local_offset = threadIdx.x * BUFFER_SIZE;
|
||||
int seq_offset = blockIdx.x * 128 * BUFFER_SIZE + local_offset;
|
||||
|
||||
// intermediate have double the seq_len
|
||||
int intermediate_local_offset = threadIdx.x * BUFFER_SIZE * 2;
|
||||
int intermediate_seq_offset = blockIdx.x * 128 * BUFFER_SIZE * 2 + intermediate_local_offset;
|
||||
|
||||
// Get values needed for replication padding before moving pointer
|
||||
const input_t *right_most_pntr = src + (seq_len * (blockIdx.y + gridDim.y * blockIdx.z));
|
||||
input_t seq_left_most_value = right_most_pntr[0];
|
||||
input_t seq_right_most_value = right_most_pntr[seq_len - 1];
|
||||
|
||||
// Move src and dst pointers
|
||||
src += block_offset + local_offset;
|
||||
dst += block_offset + local_offset;
|
||||
|
||||
// Alpha and beta values for snake activatons. Applies exp by default
|
||||
alpha = alpha + blockIdx.y;
|
||||
input_t alpha_val = expf(alpha[0]);
|
||||
beta = beta + blockIdx.y;
|
||||
input_t beta_val = expf(beta[0]);
|
||||
|
||||
#pragma unroll
|
||||
for (int it = 0; it < FILTER_SIZE; it += 1)
|
||||
{
|
||||
up_filter[it] = up_ftr[it];
|
||||
down_filter[it] = down_ftr[it];
|
||||
}
|
||||
|
||||
// Apply replication padding for upsampling, matching torch impl
|
||||
#pragma unroll
|
||||
for (int it = -HALF_FILTER_SIZE; it < BUFFER_SIZE + HALF_FILTER_SIZE; it += 1)
|
||||
{
|
||||
int element_index = seq_offset + it; // index for element
|
||||
if ((element_index < 0) && (element_index >= -UPSAMPLE_REPLICATION_PAD))
|
||||
{
|
||||
elements[2 * (HALF_FILTER_SIZE + it)] = 2 * seq_left_most_value;
|
||||
}
|
||||
if ((element_index >= seq_len) && (element_index < seq_len + UPSAMPLE_REPLICATION_PAD))
|
||||
{
|
||||
elements[2 * (HALF_FILTER_SIZE + it)] = 2 * seq_right_most_value;
|
||||
}
|
||||
if ((element_index >= 0) && (element_index < seq_len))
|
||||
{
|
||||
elements[2 * (HALF_FILTER_SIZE + it)] = 2 * src[it];
|
||||
}
|
||||
}
|
||||
|
||||
// Apply upsampling strided convolution and write to intermediates. It reserves DOWNSAMPLE_REPLICATION_PAD_LEFT for replication padding of the downsampilng conv later
|
||||
#pragma unroll
|
||||
for (int it = 0; it < (2 * BUFFER_SIZE + 2 * FILTER_SIZE); it += 1)
|
||||
{
|
||||
input_t acc = 0.0;
|
||||
int element_index = intermediate_seq_offset + it; // index for intermediate
|
||||
#pragma unroll
|
||||
for (int f_idx = 0; f_idx < FILTER_SIZE; f_idx += 1)
|
||||
{
|
||||
if ((element_index + f_idx) >= 0)
|
||||
{
|
||||
acc += up_filter[f_idx] * elements[it + f_idx];
|
||||
}
|
||||
}
|
||||
intermediates[it + DOWNSAMPLE_REPLICATION_PAD_LEFT] = acc;
|
||||
}
|
||||
|
||||
// Apply activation function. It reserves DOWNSAMPLE_REPLICATION_PAD_LEFT and DOWNSAMPLE_REPLICATION_PAD_RIGHT for replication padding of the downsampilng conv later
|
||||
double no_div_by_zero = 0.000000001;
|
||||
#pragma unroll
|
||||
for (int it = 0; it < 2 * BUFFER_SIZE + 2 * FILTER_SIZE; it += 1)
|
||||
{
|
||||
intermediates[it + DOWNSAMPLE_REPLICATION_PAD_LEFT] += (1.0 / (beta_val + no_div_by_zero)) * sinf(intermediates[it + DOWNSAMPLE_REPLICATION_PAD_LEFT] * alpha_val) * sinf(intermediates[it + DOWNSAMPLE_REPLICATION_PAD_LEFT] * alpha_val);
|
||||
}
|
||||
|
||||
// Apply replication padding before downsampling conv from intermediates
|
||||
#pragma unroll
|
||||
for (int it = 0; it < DOWNSAMPLE_REPLICATION_PAD_LEFT; it += 1)
|
||||
{
|
||||
intermediates[it] = intermediates[DOWNSAMPLE_REPLICATION_PAD_LEFT];
|
||||
}
|
||||
#pragma unroll
|
||||
for (int it = DOWNSAMPLE_REPLICATION_PAD_LEFT + 2 * BUFFER_SIZE + 2 * FILTER_SIZE; it < DOWNSAMPLE_REPLICATION_PAD_LEFT + 2 * BUFFER_SIZE + 2 * FILTER_SIZE + DOWNSAMPLE_REPLICATION_PAD_RIGHT; it += 1)
|
||||
{
|
||||
intermediates[it] = intermediates[DOWNSAMPLE_REPLICATION_PAD_LEFT + 2 * BUFFER_SIZE + 2 * FILTER_SIZE - 1];
|
||||
}
|
||||
|
||||
// Apply downsample strided convolution (assuming stride=2) from intermediates
|
||||
#pragma unroll
|
||||
for (int it = 0; it < BUFFER_SIZE; it += 1)
|
||||
{
|
||||
input_t acc = 0.0;
|
||||
#pragma unroll
|
||||
for (int f_idx = 0; f_idx < FILTER_SIZE; f_idx += 1)
|
||||
{
|
||||
// Add constant DOWNSAMPLE_REPLICATION_PAD_RIGHT to match torch implementation
|
||||
acc += down_filter[f_idx] * intermediates[it * 2 + f_idx + DOWNSAMPLE_REPLICATION_PAD_RIGHT];
|
||||
}
|
||||
output[it] = acc;
|
||||
}
|
||||
|
||||
// Write output to dst
|
||||
#pragma unroll
|
||||
for (int it = 0; it < BUFFER_SIZE; it += ELEMENTS_PER_LDG_STG)
|
||||
{
|
||||
int element_index = seq_offset + it;
|
||||
if (element_index < seq_len)
|
||||
{
|
||||
dst[it] = output[it];
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
template <typename input_t, typename output_t, typename acc_t>
|
||||
void dispatch_anti_alias_activation_forward(
|
||||
output_t *dst,
|
||||
const input_t *src,
|
||||
const input_t *up_ftr,
|
||||
const input_t *down_ftr,
|
||||
const input_t *alpha,
|
||||
const input_t *beta,
|
||||
int batch_size,
|
||||
int channels,
|
||||
int seq_len)
|
||||
{
|
||||
if (seq_len == 0)
|
||||
{
|
||||
return;
|
||||
}
|
||||
else
|
||||
{
|
||||
// Use 128 threads per block to maximimize gpu utilization
|
||||
constexpr int threads_per_block = 128;
|
||||
constexpr int seq_len_per_block = 4096;
|
||||
int blocks_per_seq_len = (seq_len + seq_len_per_block - 1) / seq_len_per_block;
|
||||
dim3 blocks(blocks_per_seq_len, channels, batch_size);
|
||||
dim3 threads(threads_per_block, 1, 1);
|
||||
|
||||
anti_alias_activation_forward<input_t, output_t, acc_t>
|
||||
<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, up_ftr, down_ftr, alpha, beta, batch_size, channels, seq_len);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
extern "C" torch::Tensor fwd_cuda(torch::Tensor const &input, torch::Tensor const &up_filter, torch::Tensor const &down_filter, torch::Tensor const &alpha, torch::Tensor const &beta)
|
||||
{
|
||||
// Input is a 3d tensor with dimensions [batches, channels, seq_len]
|
||||
const int batches = input.size(0);
|
||||
const int channels = input.size(1);
|
||||
const int seq_len = input.size(2);
|
||||
|
||||
// Output
|
||||
auto act_options = input.options().requires_grad(false);
|
||||
|
||||
torch::Tensor anti_alias_activation_results =
|
||||
torch::empty({batches, channels, seq_len}, act_options);
|
||||
|
||||
void *input_ptr = static_cast<void *>(input.data_ptr());
|
||||
void *up_filter_ptr = static_cast<void *>(up_filter.data_ptr());
|
||||
void *down_filter_ptr = static_cast<void *>(down_filter.data_ptr());
|
||||
void *alpha_ptr = static_cast<void *>(alpha.data_ptr());
|
||||
void *beta_ptr = static_cast<void *>(beta.data_ptr());
|
||||
void *anti_alias_activation_results_ptr = static_cast<void *>(anti_alias_activation_results.data_ptr());
|
||||
|
||||
DISPATCH_FLOAT_HALF_AND_BFLOAT(
|
||||
input.scalar_type(),
|
||||
"dispatch anti alias activation_forward",
|
||||
dispatch_anti_alias_activation_forward<scalar_t, scalar_t, float>(
|
||||
reinterpret_cast<scalar_t *>(anti_alias_activation_results_ptr),
|
||||
reinterpret_cast<const scalar_t *>(input_ptr),
|
||||
reinterpret_cast<const scalar_t *>(up_filter_ptr),
|
||||
reinterpret_cast<const scalar_t *>(down_filter_ptr),
|
||||
reinterpret_cast<const scalar_t *>(alpha_ptr),
|
||||
reinterpret_cast<const scalar_t *>(beta_ptr),
|
||||
batches,
|
||||
channels,
|
||||
seq_len););
|
||||
return anti_alias_activation_results;
|
||||
}
|
||||
@ -1 +0,0 @@
|
||||
|
||||
@ -1,29 +0,0 @@
|
||||
/* coding=utf-8
|
||||
* 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.
|
||||
*/
|
||||
|
||||
/*This code is copied fron NVIDIA apex:
|
||||
* https://github.com/NVIDIA/apex
|
||||
* with minor changes. */
|
||||
|
||||
#ifndef TORCH_CHECK
|
||||
#define TORCH_CHECK AT_CHECK
|
||||
#endif
|
||||
|
||||
#ifdef VERSION_GE_1_3
|
||||
#define DATA_PTR data_ptr
|
||||
#else
|
||||
#define DATA_PTR data
|
||||
#endif
|
||||
@ -1,82 +0,0 @@
|
||||
# Copyright (c) 2024 NVIDIA CORPORATION.
|
||||
# Licensed under the MIT license.
|
||||
|
||||
import os
|
||||
import pathlib
|
||||
import subprocess
|
||||
|
||||
from torch.utils import cpp_extension
|
||||
|
||||
"""
|
||||
Setting this param to a list has a problem of generating different compilation commands (with diferent order of architectures) and leading to recompilation of fused kernels.
|
||||
Set it to empty stringo avoid recompilation and assign arch flags explicity in extra_cuda_cflags below
|
||||
"""
|
||||
os.environ["TORCH_CUDA_ARCH_LIST"] = ""
|
||||
|
||||
|
||||
def load():
|
||||
# Check if cuda 11 is installed for compute capability 8.0
|
||||
cc_flag = []
|
||||
_, bare_metal_major, _ = _get_cuda_bare_metal_version(cpp_extension.CUDA_HOME)
|
||||
if int(bare_metal_major) >= 11:
|
||||
cc_flag.append("-gencode")
|
||||
cc_flag.append("arch=compute_80,code=sm_80")
|
||||
|
||||
# Build path
|
||||
srcpath = pathlib.Path(__file__).parent.absolute()
|
||||
buildpath = srcpath / "build"
|
||||
_create_build_dir(buildpath)
|
||||
|
||||
# Helper function to build the kernels.
|
||||
def _cpp_extention_load_helper(name, sources, extra_cuda_flags):
|
||||
return cpp_extension.load(
|
||||
name=name,
|
||||
sources=sources,
|
||||
build_directory=buildpath,
|
||||
extra_cflags=[
|
||||
"-O3",
|
||||
],
|
||||
extra_cuda_cflags=[
|
||||
"-O3",
|
||||
"-gencode",
|
||||
"arch=compute_70,code=sm_70",
|
||||
"--use_fast_math",
|
||||
]
|
||||
+ extra_cuda_flags
|
||||
+ cc_flag,
|
||||
verbose=True,
|
||||
)
|
||||
|
||||
extra_cuda_flags = [
|
||||
"-U__CUDA_NO_HALF_OPERATORS__",
|
||||
"-U__CUDA_NO_HALF_CONVERSIONS__",
|
||||
"--expt-relaxed-constexpr",
|
||||
"--expt-extended-lambda",
|
||||
]
|
||||
|
||||
sources = [
|
||||
srcpath / "anti_alias_activation.cpp",
|
||||
srcpath / "anti_alias_activation_cuda.cu",
|
||||
]
|
||||
anti_alias_activation_cuda = _cpp_extention_load_helper("anti_alias_activation_cuda", sources, extra_cuda_flags)
|
||||
|
||||
return anti_alias_activation_cuda
|
||||
|
||||
|
||||
def _get_cuda_bare_metal_version(cuda_dir):
|
||||
raw_output = subprocess.check_output([cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True)
|
||||
output = raw_output.split()
|
||||
release_idx = output.index("release") + 1
|
||||
release = output[release_idx].split(".")
|
||||
bare_metal_major = release[0]
|
||||
bare_metal_minor = release[1][0]
|
||||
|
||||
return raw_output, bare_metal_major, bare_metal_minor
|
||||
|
||||
|
||||
def _create_build_dir(buildpath):
|
||||
try:
|
||||
os.mkdir(buildpath)
|
||||
except OSError:
|
||||
if not os.path.isdir(buildpath):
|
||||
print(f"Creation of the build directory {buildpath} failed")
|
||||
@ -1,92 +0,0 @@
|
||||
/* coding=utf-8
|
||||
* 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.
|
||||
*/
|
||||
|
||||
#include <ATen/ATen.h>
|
||||
#include "compat.h"
|
||||
|
||||
#define DISPATCH_FLOAT_HALF_AND_BFLOAT(TYPE, NAME, ...) \
|
||||
switch (TYPE) \
|
||||
{ \
|
||||
case at::ScalarType::Float: \
|
||||
{ \
|
||||
using scalar_t = float; \
|
||||
__VA_ARGS__; \
|
||||
break; \
|
||||
} \
|
||||
case at::ScalarType::Half: \
|
||||
{ \
|
||||
using scalar_t = at::Half; \
|
||||
__VA_ARGS__; \
|
||||
break; \
|
||||
} \
|
||||
case at::ScalarType::BFloat16: \
|
||||
{ \
|
||||
using scalar_t = at::BFloat16; \
|
||||
__VA_ARGS__; \
|
||||
break; \
|
||||
} \
|
||||
default: \
|
||||
AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \
|
||||
}
|
||||
|
||||
#define DISPATCH_FLOAT_HALF_AND_BFLOAT_INOUT_TYPES(TYPEIN, TYPEOUT, NAME, ...) \
|
||||
switch (TYPEIN) \
|
||||
{ \
|
||||
case at::ScalarType::Float: \
|
||||
{ \
|
||||
using scalar_t_in = float; \
|
||||
switch (TYPEOUT) \
|
||||
{ \
|
||||
case at::ScalarType::Float: \
|
||||
{ \
|
||||
using scalar_t_out = float; \
|
||||
__VA_ARGS__; \
|
||||
break; \
|
||||
} \
|
||||
case at::ScalarType::Half: \
|
||||
{ \
|
||||
using scalar_t_out = at::Half; \
|
||||
__VA_ARGS__; \
|
||||
break; \
|
||||
} \
|
||||
case at::ScalarType::BFloat16: \
|
||||
{ \
|
||||
using scalar_t_out = at::BFloat16; \
|
||||
__VA_ARGS__; \
|
||||
break; \
|
||||
} \
|
||||
default: \
|
||||
AT_ERROR(#NAME, " not implemented for '", toString(TYPEOUT), "'"); \
|
||||
} \
|
||||
break; \
|
||||
} \
|
||||
case at::ScalarType::Half: \
|
||||
{ \
|
||||
using scalar_t_in = at::Half; \
|
||||
using scalar_t_out = at::Half; \
|
||||
__VA_ARGS__; \
|
||||
break; \
|
||||
} \
|
||||
case at::ScalarType::BFloat16: \
|
||||
{ \
|
||||
using scalar_t_in = at::BFloat16; \
|
||||
using scalar_t_out = at::BFloat16; \
|
||||
__VA_ARGS__; \
|
||||
break; \
|
||||
} \
|
||||
default: \
|
||||
AT_ERROR(#NAME, " not implemented for '", toString(TYPEIN), "'"); \
|
||||
}
|
||||
@ -89,28 +89,18 @@ class AMPBlock1(torch.nn.Module):
|
||||
|
||||
self.num_layers = len(self.convs1) + len(self.convs2) # Total number of conv layers
|
||||
|
||||
# Select which Activation1d, lazy-load cuda version to ensure backward compatibility
|
||||
if self.h.get("use_cuda_kernel", False):
|
||||
from .alias_free_activation.cuda.activation1d import (
|
||||
Activation1d as CudaActivation1d,
|
||||
)
|
||||
|
||||
Activation1d = CudaActivation1d
|
||||
else:
|
||||
Activation1d = TorchActivation1d
|
||||
|
||||
# Activation functions
|
||||
if activation == "snake":
|
||||
self.activations = nn.ModuleList(
|
||||
[
|
||||
Activation1d(activation=activations.Snake(channels, alpha_logscale=h.snake_logscale))
|
||||
TorchActivation1d(activation=activations.Snake(channels, alpha_logscale=h.snake_logscale))
|
||||
for _ in range(self.num_layers)
|
||||
]
|
||||
)
|
||||
elif activation == "snakebeta":
|
||||
self.activations = nn.ModuleList(
|
||||
[
|
||||
Activation1d(activation=activations.SnakeBeta(channels, alpha_logscale=h.snake_logscale))
|
||||
TorchActivation1d(activation=activations.SnakeBeta(channels, alpha_logscale=h.snake_logscale))
|
||||
for _ in range(self.num_layers)
|
||||
]
|
||||
)
|
||||
@ -181,28 +171,18 @@ class AMPBlock2(torch.nn.Module):
|
||||
|
||||
self.num_layers = len(self.convs) # Total number of conv layers
|
||||
|
||||
# Select which Activation1d, lazy-load cuda version to ensure backward compatibility
|
||||
if self.h.get("use_cuda_kernel", False):
|
||||
from .alias_free_activation.cuda.activation1d import (
|
||||
Activation1d as CudaActivation1d,
|
||||
)
|
||||
|
||||
Activation1d = CudaActivation1d
|
||||
else:
|
||||
Activation1d = TorchActivation1d
|
||||
|
||||
# Activation functions
|
||||
if activation == "snake":
|
||||
self.activations = nn.ModuleList(
|
||||
[
|
||||
Activation1d(activation=activations.Snake(channels, alpha_logscale=h.snake_logscale))
|
||||
TorchActivation1d(activation=activations.Snake(channels, alpha_logscale=h.snake_logscale))
|
||||
for _ in range(self.num_layers)
|
||||
]
|
||||
)
|
||||
elif activation == "snakebeta":
|
||||
self.activations = nn.ModuleList(
|
||||
[
|
||||
Activation1d(activation=activations.SnakeBeta(channels, alpha_logscale=h.snake_logscale))
|
||||
TorchActivation1d(activation=activations.SnakeBeta(channels, alpha_logscale=h.snake_logscale))
|
||||
for _ in range(self.num_layers)
|
||||
]
|
||||
)
|
||||
@ -235,31 +215,17 @@ class BigVGAN(
|
||||
):
|
||||
"""
|
||||
BigVGAN is a neural vocoder model that applies anti-aliased periodic activation for residual blocks (resblocks).
|
||||
New in BigVGAN-v2: it can optionally use optimized CUDA kernels for AMP (anti-aliased multi-periodicity) blocks.
|
||||
|
||||
Args:
|
||||
h (AttrDict): Hyperparameters.
|
||||
use_cuda_kernel (bool): If set to True, loads optimized CUDA kernels for AMP. This should be used for inference only, as training is not supported with CUDA kernels.
|
||||
|
||||
Note:
|
||||
- The `use_cuda_kernel` parameter should be used for inference only, as training with CUDA kernels is not supported.
|
||||
- Ensure that the activation function is correctly specified in the hyperparameters (h.activation).
|
||||
"""
|
||||
|
||||
def __init__(self, h: AttrDict, use_cuda_kernel: bool = False):
|
||||
def __init__(self, h: AttrDict):
|
||||
super().__init__()
|
||||
self.h = h
|
||||
self.h["use_cuda_kernel"] = use_cuda_kernel
|
||||
|
||||
# Select which Activation1d, lazy-load cuda version to ensure backward compatibility
|
||||
if self.h.get("use_cuda_kernel", False):
|
||||
from .alias_free_activation.cuda.activation1d import (
|
||||
Activation1d as CudaActivation1d,
|
||||
)
|
||||
|
||||
Activation1d = CudaActivation1d
|
||||
else:
|
||||
Activation1d = TorchActivation1d
|
||||
|
||||
self.num_kernels = len(h.resblock_kernel_sizes)
|
||||
self.num_upsamples = len(h.upsample_rates)
|
||||
@ -312,7 +278,7 @@ class BigVGAN(
|
||||
"activation incorrectly specified. check the config file and look for 'activation'."
|
||||
)
|
||||
|
||||
self.activation_post = Activation1d(activation=activation_post)
|
||||
self.activation_post = TorchActivation1d(activation=activation_post)
|
||||
|
||||
# Whether to use bias for the final conv_post. Default to True for backward compatibility
|
||||
self.use_bias_at_final = h.get("use_bias_at_final", True)
|
||||
@ -393,7 +359,6 @@ class BigVGAN(
|
||||
token: Union[str, bool, None],
|
||||
map_location: str = "cpu", # Additional argument
|
||||
strict: bool = False, # Additional argument
|
||||
use_cuda_kernel: bool = False,
|
||||
**model_kwargs,
|
||||
):
|
||||
"""Load Pytorch pretrained weights and return the loaded model."""
|
||||
@ -417,17 +382,7 @@ class BigVGAN(
|
||||
h = load_hparams_from_json(config_file)
|
||||
|
||||
# instantiate BigVGAN using h
|
||||
if use_cuda_kernel:
|
||||
print(
|
||||
"[WARNING] You have specified use_cuda_kernel=True during BigVGAN.from_pretrained(). Only inference is supported (training is not implemented)!"
|
||||
)
|
||||
print(
|
||||
"[WARNING] You need nvcc and ninja installed in your system that matches your PyTorch build is using to build the kernel. If not, the model will fail to initialize or generate incorrect waveform!"
|
||||
)
|
||||
print(
|
||||
"[WARNING] For detail, see the official GitHub repository: https://github.com/NVIDIA/BigVGAN?tab=readme-ov-file#using-custom-cuda-kernel-for-synthesis"
|
||||
)
|
||||
model = cls(h, use_cuda_kernel=use_cuda_kernel)
|
||||
model = cls(h)
|
||||
|
||||
# Download and load pretrained generator weight
|
||||
if os.path.isdir(model_id):
|
||||
|
||||
@ -1,45 +0,0 @@
|
||||
{
|
||||
"resblock": "1",
|
||||
"num_gpus": 0,
|
||||
"batch_size": 32,
|
||||
"learning_rate": 0.0001,
|
||||
"adam_b1": 0.8,
|
||||
"adam_b2": 0.99,
|
||||
"lr_decay": 0.9999996,
|
||||
"seed": 1234,
|
||||
|
||||
"upsample_rates": [4,4,2,2,2,2],
|
||||
"upsample_kernel_sizes": [8,8,4,4,4,4],
|
||||
"upsample_initial_channel": 1536,
|
||||
"resblock_kernel_sizes": [3,7,11],
|
||||
"resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
|
||||
|
||||
"activation": "snakebeta",
|
||||
"snake_logscale": true,
|
||||
|
||||
"resolutions": [[1024, 120, 600], [2048, 240, 1200], [512, 50, 240]],
|
||||
"mpd_reshapes": [2, 3, 5, 7, 11],
|
||||
"use_spectral_norm": false,
|
||||
"discriminator_channel_mult": 1,
|
||||
|
||||
"segment_size": 8192,
|
||||
"num_mels": 80,
|
||||
"num_freq": 1025,
|
||||
"n_fft": 1024,
|
||||
"hop_size": 256,
|
||||
"win_size": 1024,
|
||||
|
||||
"sampling_rate": 22050,
|
||||
|
||||
"fmin": 0,
|
||||
"fmax": 8000,
|
||||
"fmax_for_loss": null,
|
||||
|
||||
"num_workers": 4,
|
||||
|
||||
"dist_config": {
|
||||
"dist_backend": "nccl",
|
||||
"dist_url": "tcp://localhost:54321",
|
||||
"world_size": 1
|
||||
}
|
||||
}
|
||||
@ -1,45 +0,0 @@
|
||||
{
|
||||
"resblock": "1",
|
||||
"num_gpus": 0,
|
||||
"batch_size": 32,
|
||||
"learning_rate": 0.0001,
|
||||
"adam_b1": 0.8,
|
||||
"adam_b2": 0.99,
|
||||
"lr_decay": 0.9999996,
|
||||
"seed": 1234,
|
||||
|
||||
"upsample_rates": [4,4,2,2,2,2],
|
||||
"upsample_kernel_sizes": [8,8,4,4,4,4],
|
||||
"upsample_initial_channel": 1536,
|
||||
"resblock_kernel_sizes": [3,7,11],
|
||||
"resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
|
||||
|
||||
"activation": "snakebeta",
|
||||
"snake_logscale": true,
|
||||
|
||||
"resolutions": [[1024, 120, 600], [2048, 240, 1200], [512, 50, 240]],
|
||||
"mpd_reshapes": [2, 3, 5, 7, 11],
|
||||
"use_spectral_norm": false,
|
||||
"discriminator_channel_mult": 1,
|
||||
|
||||
"segment_size": 8192,
|
||||
"num_mels": 100,
|
||||
"num_freq": 1025,
|
||||
"n_fft": 1024,
|
||||
"hop_size": 256,
|
||||
"win_size": 1024,
|
||||
|
||||
"sampling_rate": 24000,
|
||||
|
||||
"fmin": 0,
|
||||
"fmax": 12000,
|
||||
"fmax_for_loss": null,
|
||||
|
||||
"num_workers": 4,
|
||||
|
||||
"dist_config": {
|
||||
"dist_backend": "nccl",
|
||||
"dist_url": "tcp://localhost:54321",
|
||||
"world_size": 1
|
||||
}
|
||||
}
|
||||
@ -1,45 +0,0 @@
|
||||
{
|
||||
"resblock": "1",
|
||||
"num_gpus": 0,
|
||||
"batch_size": 32,
|
||||
"learning_rate": 0.0001,
|
||||
"adam_b1": 0.8,
|
||||
"adam_b2": 0.99,
|
||||
"lr_decay": 0.9999996,
|
||||
"seed": 1234,
|
||||
|
||||
"upsample_rates": [8,8,2,2],
|
||||
"upsample_kernel_sizes": [16,16,4,4],
|
||||
"upsample_initial_channel": 512,
|
||||
"resblock_kernel_sizes": [3,7,11],
|
||||
"resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
|
||||
|
||||
"activation": "snakebeta",
|
||||
"snake_logscale": true,
|
||||
|
||||
"resolutions": [[1024, 120, 600], [2048, 240, 1200], [512, 50, 240]],
|
||||
"mpd_reshapes": [2, 3, 5, 7, 11],
|
||||
"use_spectral_norm": false,
|
||||
"discriminator_channel_mult": 1,
|
||||
|
||||
"segment_size": 8192,
|
||||
"num_mels": 80,
|
||||
"num_freq": 1025,
|
||||
"n_fft": 1024,
|
||||
"hop_size": 256,
|
||||
"win_size": 1024,
|
||||
|
||||
"sampling_rate": 22050,
|
||||
|
||||
"fmin": 0,
|
||||
"fmax": 8000,
|
||||
"fmax_for_loss": null,
|
||||
|
||||
"num_workers": 4,
|
||||
|
||||
"dist_config": {
|
||||
"dist_backend": "nccl",
|
||||
"dist_url": "tcp://localhost:54321",
|
||||
"world_size": 1
|
||||
}
|
||||
}
|
||||
@ -1,45 +0,0 @@
|
||||
{
|
||||
"resblock": "1",
|
||||
"num_gpus": 0,
|
||||
"batch_size": 32,
|
||||
"learning_rate": 0.0001,
|
||||
"adam_b1": 0.8,
|
||||
"adam_b2": 0.99,
|
||||
"lr_decay": 0.9999996,
|
||||
"seed": 1234,
|
||||
|
||||
"upsample_rates": [8,8,2,2],
|
||||
"upsample_kernel_sizes": [16,16,4,4],
|
||||
"upsample_initial_channel": 512,
|
||||
"resblock_kernel_sizes": [3,7,11],
|
||||
"resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
|
||||
|
||||
"activation": "snakebeta",
|
||||
"snake_logscale": true,
|
||||
|
||||
"resolutions": [[1024, 120, 600], [2048, 240, 1200], [512, 50, 240]],
|
||||
"mpd_reshapes": [2, 3, 5, 7, 11],
|
||||
"use_spectral_norm": false,
|
||||
"discriminator_channel_mult": 1,
|
||||
|
||||
"segment_size": 8192,
|
||||
"num_mels": 100,
|
||||
"num_freq": 1025,
|
||||
"n_fft": 1024,
|
||||
"hop_size": 256,
|
||||
"win_size": 1024,
|
||||
|
||||
"sampling_rate": 24000,
|
||||
|
||||
"fmin": 0,
|
||||
"fmax": 12000,
|
||||
"fmax_for_loss": null,
|
||||
|
||||
"num_workers": 4,
|
||||
|
||||
"dist_config": {
|
||||
"dist_backend": "nccl",
|
||||
"dist_url": "tcp://localhost:54321",
|
||||
"world_size": 1
|
||||
}
|
||||
}
|
||||
@ -1,61 +0,0 @@
|
||||
{
|
||||
"resblock": "1",
|
||||
"num_gpus": 0,
|
||||
"batch_size": 4,
|
||||
"learning_rate": 0.0001,
|
||||
"adam_b1": 0.8,
|
||||
"adam_b2": 0.99,
|
||||
"lr_decay": 0.9999996,
|
||||
"seed": 1234,
|
||||
|
||||
"upsample_rates": [4,4,2,2,2,2],
|
||||
"upsample_kernel_sizes": [8,8,4,4,4,4],
|
||||
"upsample_initial_channel": 1536,
|
||||
"resblock_kernel_sizes": [3,7,11],
|
||||
"resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
|
||||
|
||||
"use_tanh_at_final": false,
|
||||
"use_bias_at_final": false,
|
||||
|
||||
"activation": "snakebeta",
|
||||
"snake_logscale": true,
|
||||
|
||||
"use_cqtd_instead_of_mrd": true,
|
||||
"cqtd_filters": 128,
|
||||
"cqtd_max_filters": 1024,
|
||||
"cqtd_filters_scale": 1,
|
||||
"cqtd_dilations": [1, 2, 4],
|
||||
"cqtd_hop_lengths": [512, 256, 256],
|
||||
"cqtd_n_octaves": [9, 9, 9],
|
||||
"cqtd_bins_per_octaves": [24, 36, 48],
|
||||
|
||||
"mpd_reshapes": [2, 3, 5, 7, 11],
|
||||
"use_spectral_norm": false,
|
||||
"discriminator_channel_mult": 1,
|
||||
|
||||
"use_multiscale_melloss": true,
|
||||
"lambda_melloss": 15,
|
||||
|
||||
"clip_grad_norm": 500,
|
||||
|
||||
"segment_size": 65536,
|
||||
"num_mels": 80,
|
||||
"num_freq": 1025,
|
||||
"n_fft": 1024,
|
||||
"hop_size": 256,
|
||||
"win_size": 1024,
|
||||
|
||||
"sampling_rate": 22050,
|
||||
|
||||
"fmin": 0,
|
||||
"fmax": null,
|
||||
"fmax_for_loss": null,
|
||||
|
||||
"num_workers": 4,
|
||||
|
||||
"dist_config": {
|
||||
"dist_backend": "nccl",
|
||||
"dist_url": "tcp://localhost:54321",
|
||||
"world_size": 1
|
||||
}
|
||||
}
|
||||
@ -1,61 +0,0 @@
|
||||
{
|
||||
"resblock": "1",
|
||||
"num_gpus": 0,
|
||||
"batch_size": 4,
|
||||
"learning_rate": 0.0001,
|
||||
"adam_b1": 0.8,
|
||||
"adam_b2": 0.99,
|
||||
"lr_decay": 0.9999996,
|
||||
"seed": 1234,
|
||||
|
||||
"upsample_rates": [4,4,2,2,2,2],
|
||||
"upsample_kernel_sizes": [8,8,4,4,4,4],
|
||||
"upsample_initial_channel": 1536,
|
||||
"resblock_kernel_sizes": [3,7,11],
|
||||
"resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
|
||||
|
||||
"use_tanh_at_final": false,
|
||||
"use_bias_at_final": false,
|
||||
|
||||
"activation": "snakebeta",
|
||||
"snake_logscale": true,
|
||||
|
||||
"use_cqtd_instead_of_mrd": true,
|
||||
"cqtd_filters": 128,
|
||||
"cqtd_max_filters": 1024,
|
||||
"cqtd_filters_scale": 1,
|
||||
"cqtd_dilations": [1, 2, 4],
|
||||
"cqtd_hop_lengths": [512, 256, 256],
|
||||
"cqtd_n_octaves": [9, 9, 9],
|
||||
"cqtd_bins_per_octaves": [24, 36, 48],
|
||||
|
||||
"mpd_reshapes": [2, 3, 5, 7, 11],
|
||||
"use_spectral_norm": false,
|
||||
"discriminator_channel_mult": 1,
|
||||
|
||||
"use_multiscale_melloss": true,
|
||||
"lambda_melloss": 15,
|
||||
|
||||
"clip_grad_norm": 500,
|
||||
|
||||
"segment_size": 65536,
|
||||
"num_mels": 80,
|
||||
"num_freq": 1025,
|
||||
"n_fft": 1024,
|
||||
"hop_size": 256,
|
||||
"win_size": 1024,
|
||||
|
||||
"sampling_rate": 22050,
|
||||
|
||||
"fmin": 0,
|
||||
"fmax": 8000,
|
||||
"fmax_for_loss": null,
|
||||
|
||||
"num_workers": 4,
|
||||
|
||||
"dist_config": {
|
||||
"dist_backend": "nccl",
|
||||
"dist_url": "tcp://localhost:54321",
|
||||
"world_size": 1
|
||||
}
|
||||
}
|
||||
@ -1,61 +0,0 @@
|
||||
{
|
||||
"resblock": "1",
|
||||
"num_gpus": 0,
|
||||
"batch_size": 4,
|
||||
"learning_rate": 0.0001,
|
||||
"adam_b1": 0.8,
|
||||
"adam_b2": 0.99,
|
||||
"lr_decay": 0.9999996,
|
||||
"seed": 1234,
|
||||
|
||||
"upsample_rates": [4,4,2,2,2,2],
|
||||
"upsample_kernel_sizes": [8,8,4,4,4,4],
|
||||
"upsample_initial_channel": 1536,
|
||||
"resblock_kernel_sizes": [3,7,11],
|
||||
"resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
|
||||
|
||||
"use_tanh_at_final": false,
|
||||
"use_bias_at_final": false,
|
||||
|
||||
"activation": "snakebeta",
|
||||
"snake_logscale": true,
|
||||
|
||||
"use_cqtd_instead_of_mrd": true,
|
||||
"cqtd_filters": 128,
|
||||
"cqtd_max_filters": 1024,
|
||||
"cqtd_filters_scale": 1,
|
||||
"cqtd_dilations": [1, 2, 4],
|
||||
"cqtd_hop_lengths": [512, 256, 256],
|
||||
"cqtd_n_octaves": [9, 9, 9],
|
||||
"cqtd_bins_per_octaves": [24, 36, 48],
|
||||
|
||||
"mpd_reshapes": [2, 3, 5, 7, 11],
|
||||
"use_spectral_norm": false,
|
||||
"discriminator_channel_mult": 1,
|
||||
|
||||
"use_multiscale_melloss": true,
|
||||
"lambda_melloss": 15,
|
||||
|
||||
"clip_grad_norm": 500,
|
||||
|
||||
"segment_size": 65536,
|
||||
"num_mels": 100,
|
||||
"num_freq": 1025,
|
||||
"n_fft": 1024,
|
||||
"hop_size": 256,
|
||||
"win_size": 1024,
|
||||
|
||||
"sampling_rate": 24000,
|
||||
|
||||
"fmin": 0,
|
||||
"fmax": null,
|
||||
"fmax_for_loss": null,
|
||||
|
||||
"num_workers": 4,
|
||||
|
||||
"dist_config": {
|
||||
"dist_backend": "nccl",
|
||||
"dist_url": "tcp://localhost:54321",
|
||||
"world_size": 1
|
||||
}
|
||||
}
|
||||
@ -1,61 +0,0 @@
|
||||
{
|
||||
"resblock": "1",
|
||||
"num_gpus": 0,
|
||||
"batch_size": 4,
|
||||
"learning_rate": 0.0001,
|
||||
"adam_b1": 0.8,
|
||||
"adam_b2": 0.99,
|
||||
"lr_decay": 0.9999996,
|
||||
"seed": 1234,
|
||||
|
||||
"upsample_rates": [4,4,2,2,2,2],
|
||||
"upsample_kernel_sizes": [8,8,4,4,4,4],
|
||||
"upsample_initial_channel": 1536,
|
||||
"resblock_kernel_sizes": [3,7,11],
|
||||
"resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
|
||||
|
||||
"use_tanh_at_final": false,
|
||||
"use_bias_at_final": false,
|
||||
|
||||
"activation": "snakebeta",
|
||||
"snake_logscale": true,
|
||||
|
||||
"use_cqtd_instead_of_mrd": true,
|
||||
"cqtd_filters": 128,
|
||||
"cqtd_max_filters": 1024,
|
||||
"cqtd_filters_scale": 1,
|
||||
"cqtd_dilations": [1, 2, 4],
|
||||
"cqtd_hop_lengths": [512, 256, 256],
|
||||
"cqtd_n_octaves": [9, 9, 9],
|
||||
"cqtd_bins_per_octaves": [24, 36, 48],
|
||||
|
||||
"mpd_reshapes": [2, 3, 5, 7, 11],
|
||||
"use_spectral_norm": false,
|
||||
"discriminator_channel_mult": 1,
|
||||
|
||||
"use_multiscale_melloss": true,
|
||||
"lambda_melloss": 15,
|
||||
|
||||
"clip_grad_norm": 500,
|
||||
|
||||
"segment_size": 65536,
|
||||
"num_mels": 128,
|
||||
"num_freq": 1025,
|
||||
"n_fft": 1024,
|
||||
"hop_size": 256,
|
||||
"win_size": 1024,
|
||||
|
||||
"sampling_rate": 44100,
|
||||
|
||||
"fmin": 0,
|
||||
"fmax": null,
|
||||
"fmax_for_loss": null,
|
||||
|
||||
"num_workers": 4,
|
||||
|
||||
"dist_config": {
|
||||
"dist_backend": "nccl",
|
||||
"dist_url": "tcp://localhost:54321",
|
||||
"world_size": 1
|
||||
}
|
||||
}
|
||||
@ -1,61 +0,0 @@
|
||||
{
|
||||
"resblock": "1",
|
||||
"num_gpus": 0,
|
||||
"batch_size": 4,
|
||||
"learning_rate": 0.0001,
|
||||
"adam_b1": 0.8,
|
||||
"adam_b2": 0.99,
|
||||
"lr_decay": 0.9999996,
|
||||
"seed": 1234,
|
||||
|
||||
"upsample_rates": [8,4,2,2,2,2],
|
||||
"upsample_kernel_sizes": [16,8,4,4,4,4],
|
||||
"upsample_initial_channel": 1536,
|
||||
"resblock_kernel_sizes": [3,7,11],
|
||||
"resblock_dilation_sizes": [[1,3,5], [1,3,5], [1,3,5]],
|
||||
|
||||
"use_tanh_at_final": false,
|
||||
"use_bias_at_final": false,
|
||||
|
||||
"activation": "snakebeta",
|
||||
"snake_logscale": true,
|
||||
|
||||
"use_cqtd_instead_of_mrd": true,
|
||||
"cqtd_filters": 128,
|
||||
"cqtd_max_filters": 1024,
|
||||
"cqtd_filters_scale": 1,
|
||||
"cqtd_dilations": [1, 2, 4],
|
||||
"cqtd_hop_lengths": [512, 256, 256],
|
||||
"cqtd_n_octaves": [9, 9, 9],
|
||||
"cqtd_bins_per_octaves": [24, 36, 48],
|
||||
|
||||
"mpd_reshapes": [2, 3, 5, 7, 11],
|
||||
"use_spectral_norm": false,
|
||||
"discriminator_channel_mult": 1,
|
||||
|
||||
"use_multiscale_melloss": true,
|
||||
"lambda_melloss": 15,
|
||||
|
||||
"clip_grad_norm": 500,
|
||||
|
||||
"segment_size": 65536,
|
||||
"num_mels": 128,
|
||||
"num_freq": 2049,
|
||||
"n_fft": 2048,
|
||||
"hop_size": 512,
|
||||
"win_size": 2048,
|
||||
|
||||
"sampling_rate": 44100,
|
||||
|
||||
"fmin": 0,
|
||||
"fmax": null,
|
||||
"fmax_for_loss": null,
|
||||
|
||||
"num_workers": 4,
|
||||
|
||||
"dist_config": {
|
||||
"dist_backend": "nccl",
|
||||
"dist_url": "tcp://localhost:54321",
|
||||
"world_size": 1
|
||||
}
|
||||
}
|
||||
@ -1,625 +0,0 @@
|
||||
# Copyright (c) 2024 NVIDIA CORPORATION.
|
||||
# Licensed under the MIT license.
|
||||
|
||||
# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.nn as nn
|
||||
from torch.nn import Conv2d
|
||||
from torch.nn.utils import weight_norm, spectral_norm
|
||||
from torchaudio.transforms import Spectrogram, Resample
|
||||
|
||||
from env import AttrDict
|
||||
from utils import get_padding
|
||||
import typing
|
||||
from typing import List, Tuple
|
||||
|
||||
|
||||
class DiscriminatorP(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
h: AttrDict,
|
||||
period: List[int],
|
||||
kernel_size: int = 5,
|
||||
stride: int = 3,
|
||||
use_spectral_norm: bool = False,
|
||||
):
|
||||
super().__init__()
|
||||
self.period = period
|
||||
self.d_mult = h.discriminator_channel_mult
|
||||
norm_f = weight_norm if not use_spectral_norm else spectral_norm
|
||||
|
||||
self.convs = nn.ModuleList(
|
||||
[
|
||||
norm_f(
|
||||
Conv2d(
|
||||
1,
|
||||
int(32 * self.d_mult),
|
||||
(kernel_size, 1),
|
||||
(stride, 1),
|
||||
padding=(get_padding(5, 1), 0),
|
||||
)
|
||||
),
|
||||
norm_f(
|
||||
Conv2d(
|
||||
int(32 * self.d_mult),
|
||||
int(128 * self.d_mult),
|
||||
(kernel_size, 1),
|
||||
(stride, 1),
|
||||
padding=(get_padding(5, 1), 0),
|
||||
)
|
||||
),
|
||||
norm_f(
|
||||
Conv2d(
|
||||
int(128 * self.d_mult),
|
||||
int(512 * self.d_mult),
|
||||
(kernel_size, 1),
|
||||
(stride, 1),
|
||||
padding=(get_padding(5, 1), 0),
|
||||
)
|
||||
),
|
||||
norm_f(
|
||||
Conv2d(
|
||||
int(512 * self.d_mult),
|
||||
int(1024 * self.d_mult),
|
||||
(kernel_size, 1),
|
||||
(stride, 1),
|
||||
padding=(get_padding(5, 1), 0),
|
||||
)
|
||||
),
|
||||
norm_f(
|
||||
Conv2d(
|
||||
int(1024 * self.d_mult),
|
||||
int(1024 * self.d_mult),
|
||||
(kernel_size, 1),
|
||||
1,
|
||||
padding=(2, 0),
|
||||
)
|
||||
),
|
||||
]
|
||||
)
|
||||
self.conv_post = norm_f(Conv2d(int(1024 * self.d_mult), 1, (3, 1), 1, padding=(1, 0)))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
|
||||
fmap = []
|
||||
|
||||
# 1d to 2d
|
||||
b, c, t = x.shape
|
||||
if t % self.period != 0: # pad first
|
||||
n_pad = self.period - (t % self.period)
|
||||
x = F.pad(x, (0, n_pad), "reflect")
|
||||
t = t + n_pad
|
||||
x = x.view(b, c, t // self.period, self.period)
|
||||
|
||||
for l in self.convs:
|
||||
x = l(x)
|
||||
x = F.leaky_relu(x, 0.1)
|
||||
fmap.append(x)
|
||||
x = self.conv_post(x)
|
||||
fmap.append(x)
|
||||
x = torch.flatten(x, 1, -1)
|
||||
|
||||
return x, fmap
|
||||
|
||||
|
||||
class MultiPeriodDiscriminator(torch.nn.Module):
|
||||
def __init__(self, h: AttrDict):
|
||||
super().__init__()
|
||||
self.mpd_reshapes = h.mpd_reshapes
|
||||
print(f"mpd_reshapes: {self.mpd_reshapes}")
|
||||
self.discriminators = nn.ModuleList(
|
||||
[DiscriminatorP(h, rs, use_spectral_norm=h.use_spectral_norm) for rs in self.mpd_reshapes]
|
||||
)
|
||||
|
||||
def forward(
|
||||
self, y: torch.Tensor, y_hat: torch.Tensor
|
||||
) -> Tuple[
|
||||
List[torch.Tensor],
|
||||
List[torch.Tensor],
|
||||
List[List[torch.Tensor]],
|
||||
List[List[torch.Tensor]],
|
||||
]:
|
||||
y_d_rs = []
|
||||
y_d_gs = []
|
||||
fmap_rs = []
|
||||
fmap_gs = []
|
||||
for i, d in enumerate(self.discriminators):
|
||||
y_d_r, fmap_r = d(y)
|
||||
y_d_g, fmap_g = d(y_hat)
|
||||
y_d_rs.append(y_d_r)
|
||||
fmap_rs.append(fmap_r)
|
||||
y_d_gs.append(y_d_g)
|
||||
fmap_gs.append(fmap_g)
|
||||
|
||||
return y_d_rs, y_d_gs, fmap_rs, fmap_gs
|
||||
|
||||
|
||||
class DiscriminatorR(nn.Module):
|
||||
def __init__(self, cfg: AttrDict, resolution: List[List[int]]):
|
||||
super().__init__()
|
||||
|
||||
self.resolution = resolution
|
||||
assert len(self.resolution) == 3, f"MRD layer requires list with len=3, got {self.resolution}"
|
||||
self.lrelu_slope = 0.1
|
||||
|
||||
norm_f = weight_norm if cfg.use_spectral_norm == False else spectral_norm
|
||||
if hasattr(cfg, "mrd_use_spectral_norm"):
|
||||
print(f"[INFO] overriding MRD use_spectral_norm as {cfg.mrd_use_spectral_norm}")
|
||||
norm_f = weight_norm if cfg.mrd_use_spectral_norm == False else spectral_norm
|
||||
self.d_mult = cfg.discriminator_channel_mult
|
||||
if hasattr(cfg, "mrd_channel_mult"):
|
||||
print(f"[INFO] overriding mrd channel multiplier as {cfg.mrd_channel_mult}")
|
||||
self.d_mult = cfg.mrd_channel_mult
|
||||
|
||||
self.convs = nn.ModuleList(
|
||||
[
|
||||
norm_f(nn.Conv2d(1, int(32 * self.d_mult), (3, 9), padding=(1, 4))),
|
||||
norm_f(
|
||||
nn.Conv2d(
|
||||
int(32 * self.d_mult),
|
||||
int(32 * self.d_mult),
|
||||
(3, 9),
|
||||
stride=(1, 2),
|
||||
padding=(1, 4),
|
||||
)
|
||||
),
|
||||
norm_f(
|
||||
nn.Conv2d(
|
||||
int(32 * self.d_mult),
|
||||
int(32 * self.d_mult),
|
||||
(3, 9),
|
||||
stride=(1, 2),
|
||||
padding=(1, 4),
|
||||
)
|
||||
),
|
||||
norm_f(
|
||||
nn.Conv2d(
|
||||
int(32 * self.d_mult),
|
||||
int(32 * self.d_mult),
|
||||
(3, 9),
|
||||
stride=(1, 2),
|
||||
padding=(1, 4),
|
||||
)
|
||||
),
|
||||
norm_f(
|
||||
nn.Conv2d(
|
||||
int(32 * self.d_mult),
|
||||
int(32 * self.d_mult),
|
||||
(3, 3),
|
||||
padding=(1, 1),
|
||||
)
|
||||
),
|
||||
]
|
||||
)
|
||||
self.conv_post = norm_f(nn.Conv2d(int(32 * self.d_mult), 1, (3, 3), padding=(1, 1)))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
|
||||
fmap = []
|
||||
|
||||
x = self.spectrogram(x)
|
||||
x = x.unsqueeze(1)
|
||||
for l in self.convs:
|
||||
x = l(x)
|
||||
x = F.leaky_relu(x, self.lrelu_slope)
|
||||
fmap.append(x)
|
||||
x = self.conv_post(x)
|
||||
fmap.append(x)
|
||||
x = torch.flatten(x, 1, -1)
|
||||
|
||||
return x, fmap
|
||||
|
||||
def spectrogram(self, x: torch.Tensor) -> torch.Tensor:
|
||||
n_fft, hop_length, win_length = self.resolution
|
||||
x = F.pad(
|
||||
x,
|
||||
(int((n_fft - hop_length) / 2), int((n_fft - hop_length) / 2)),
|
||||
mode="reflect",
|
||||
)
|
||||
x = x.squeeze(1)
|
||||
x = torch.stft(
|
||||
x,
|
||||
n_fft=n_fft,
|
||||
hop_length=hop_length,
|
||||
win_length=win_length,
|
||||
center=False,
|
||||
return_complex=True,
|
||||
)
|
||||
x = torch.view_as_real(x) # [B, F, TT, 2]
|
||||
mag = torch.norm(x, p=2, dim=-1) # [B, F, TT]
|
||||
|
||||
return mag
|
||||
|
||||
|
||||
class MultiResolutionDiscriminator(nn.Module):
|
||||
def __init__(self, cfg, debug=False):
|
||||
super().__init__()
|
||||
self.resolutions = cfg.resolutions
|
||||
assert len(self.resolutions) == 3, (
|
||||
f"MRD requires list of list with len=3, each element having a list with len=3. Got {self.resolutions}"
|
||||
)
|
||||
self.discriminators = nn.ModuleList([DiscriminatorR(cfg, resolution) for resolution in self.resolutions])
|
||||
|
||||
def forward(
|
||||
self, y: torch.Tensor, y_hat: torch.Tensor
|
||||
) -> Tuple[
|
||||
List[torch.Tensor],
|
||||
List[torch.Tensor],
|
||||
List[List[torch.Tensor]],
|
||||
List[List[torch.Tensor]],
|
||||
]:
|
||||
y_d_rs = []
|
||||
y_d_gs = []
|
||||
fmap_rs = []
|
||||
fmap_gs = []
|
||||
|
||||
for i, d in enumerate(self.discriminators):
|
||||
y_d_r, fmap_r = d(x=y)
|
||||
y_d_g, fmap_g = d(x=y_hat)
|
||||
y_d_rs.append(y_d_r)
|
||||
fmap_rs.append(fmap_r)
|
||||
y_d_gs.append(y_d_g)
|
||||
fmap_gs.append(fmap_g)
|
||||
|
||||
return y_d_rs, y_d_gs, fmap_rs, fmap_gs
|
||||
|
||||
|
||||
# Method based on descript-audio-codec: https://github.com/descriptinc/descript-audio-codec
|
||||
# Modified code adapted from https://github.com/gemelo-ai/vocos under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
class DiscriminatorB(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
window_length: int,
|
||||
channels: int = 32,
|
||||
hop_factor: float = 0.25,
|
||||
bands: Tuple[Tuple[float, float], ...] = (
|
||||
(0.0, 0.1),
|
||||
(0.1, 0.25),
|
||||
(0.25, 0.5),
|
||||
(0.5, 0.75),
|
||||
(0.75, 1.0),
|
||||
),
|
||||
):
|
||||
super().__init__()
|
||||
self.window_length = window_length
|
||||
self.hop_factor = hop_factor
|
||||
self.spec_fn = Spectrogram(
|
||||
n_fft=window_length,
|
||||
hop_length=int(window_length * hop_factor),
|
||||
win_length=window_length,
|
||||
power=None,
|
||||
)
|
||||
n_fft = window_length // 2 + 1
|
||||
bands = [(int(b[0] * n_fft), int(b[1] * n_fft)) for b in bands]
|
||||
self.bands = bands
|
||||
convs = lambda: nn.ModuleList(
|
||||
[
|
||||
weight_norm(nn.Conv2d(2, channels, (3, 9), (1, 1), padding=(1, 4))),
|
||||
weight_norm(nn.Conv2d(channels, channels, (3, 9), (1, 2), padding=(1, 4))),
|
||||
weight_norm(nn.Conv2d(channels, channels, (3, 9), (1, 2), padding=(1, 4))),
|
||||
weight_norm(nn.Conv2d(channels, channels, (3, 9), (1, 2), padding=(1, 4))),
|
||||
weight_norm(nn.Conv2d(channels, channels, (3, 3), (1, 1), padding=(1, 1))),
|
||||
]
|
||||
)
|
||||
self.band_convs = nn.ModuleList([convs() for _ in range(len(self.bands))])
|
||||
|
||||
self.conv_post = weight_norm(nn.Conv2d(channels, 1, (3, 3), (1, 1), padding=(1, 1)))
|
||||
|
||||
def spectrogram(self, x: torch.Tensor) -> List[torch.Tensor]:
|
||||
# Remove DC offset
|
||||
x = x - x.mean(dim=-1, keepdims=True)
|
||||
# Peak normalize the volume of input audio
|
||||
x = 0.8 * x / (x.abs().max(dim=-1, keepdim=True)[0] + 1e-9)
|
||||
x = self.spec_fn(x)
|
||||
x = torch.view_as_real(x)
|
||||
x = x.permute(0, 3, 2, 1) # [B, F, T, C] -> [B, C, T, F]
|
||||
# Split into bands
|
||||
x_bands = [x[..., b[0] : b[1]] for b in self.bands]
|
||||
return x_bands
|
||||
|
||||
def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
|
||||
x_bands = self.spectrogram(x.squeeze(1))
|
||||
fmap = []
|
||||
x = []
|
||||
|
||||
for band, stack in zip(x_bands, self.band_convs):
|
||||
for i, layer in enumerate(stack):
|
||||
band = layer(band)
|
||||
band = torch.nn.functional.leaky_relu(band, 0.1)
|
||||
if i > 0:
|
||||
fmap.append(band)
|
||||
x.append(band)
|
||||
|
||||
x = torch.cat(x, dim=-1)
|
||||
x = self.conv_post(x)
|
||||
fmap.append(x)
|
||||
|
||||
return x, fmap
|
||||
|
||||
|
||||
# Method based on descript-audio-codec: https://github.com/descriptinc/descript-audio-codec
|
||||
# Modified code adapted from https://github.com/gemelo-ai/vocos under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
class MultiBandDiscriminator(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
h,
|
||||
):
|
||||
"""
|
||||
Multi-band multi-scale STFT discriminator, with the architecture based on https://github.com/descriptinc/descript-audio-codec.
|
||||
and the modified code adapted from https://github.com/gemelo-ai/vocos.
|
||||
"""
|
||||
super().__init__()
|
||||
# fft_sizes (list[int]): Tuple of window lengths for FFT. Defaults to [2048, 1024, 512] if not set in h.
|
||||
self.fft_sizes = h.get("mbd_fft_sizes", [2048, 1024, 512])
|
||||
self.discriminators = nn.ModuleList([DiscriminatorB(window_length=w) for w in self.fft_sizes])
|
||||
|
||||
def forward(
|
||||
self, y: torch.Tensor, y_hat: torch.Tensor
|
||||
) -> Tuple[
|
||||
List[torch.Tensor],
|
||||
List[torch.Tensor],
|
||||
List[List[torch.Tensor]],
|
||||
List[List[torch.Tensor]],
|
||||
]:
|
||||
y_d_rs = []
|
||||
y_d_gs = []
|
||||
fmap_rs = []
|
||||
fmap_gs = []
|
||||
|
||||
for d in self.discriminators:
|
||||
y_d_r, fmap_r = d(x=y)
|
||||
y_d_g, fmap_g = d(x=y_hat)
|
||||
y_d_rs.append(y_d_r)
|
||||
fmap_rs.append(fmap_r)
|
||||
y_d_gs.append(y_d_g)
|
||||
fmap_gs.append(fmap_g)
|
||||
|
||||
return y_d_rs, y_d_gs, fmap_rs, fmap_gs
|
||||
|
||||
|
||||
# Adapted from https://github.com/open-mmlab/Amphion/blob/main/models/vocoders/gan/discriminator/mssbcqtd.py under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
class DiscriminatorCQT(nn.Module):
|
||||
def __init__(self, cfg: AttrDict, hop_length: int, n_octaves: int, bins_per_octave: int):
|
||||
super().__init__()
|
||||
self.cfg = cfg
|
||||
|
||||
self.filters = cfg["cqtd_filters"]
|
||||
self.max_filters = cfg["cqtd_max_filters"]
|
||||
self.filters_scale = cfg["cqtd_filters_scale"]
|
||||
self.kernel_size = (3, 9)
|
||||
self.dilations = cfg["cqtd_dilations"]
|
||||
self.stride = (1, 2)
|
||||
|
||||
self.in_channels = cfg["cqtd_in_channels"]
|
||||
self.out_channels = cfg["cqtd_out_channels"]
|
||||
self.fs = cfg["sampling_rate"]
|
||||
self.hop_length = hop_length
|
||||
self.n_octaves = n_octaves
|
||||
self.bins_per_octave = bins_per_octave
|
||||
|
||||
# Lazy-load
|
||||
from nnAudio import features
|
||||
|
||||
self.cqt_transform = features.cqt.CQT2010v2(
|
||||
sr=self.fs * 2,
|
||||
hop_length=self.hop_length,
|
||||
n_bins=self.bins_per_octave * self.n_octaves,
|
||||
bins_per_octave=self.bins_per_octave,
|
||||
output_format="Complex",
|
||||
pad_mode="constant",
|
||||
)
|
||||
|
||||
self.conv_pres = nn.ModuleList()
|
||||
for _ in range(self.n_octaves):
|
||||
self.conv_pres.append(
|
||||
nn.Conv2d(
|
||||
self.in_channels * 2,
|
||||
self.in_channels * 2,
|
||||
kernel_size=self.kernel_size,
|
||||
padding=self.get_2d_padding(self.kernel_size),
|
||||
)
|
||||
)
|
||||
|
||||
self.convs = nn.ModuleList()
|
||||
|
||||
self.convs.append(
|
||||
nn.Conv2d(
|
||||
self.in_channels * 2,
|
||||
self.filters,
|
||||
kernel_size=self.kernel_size,
|
||||
padding=self.get_2d_padding(self.kernel_size),
|
||||
)
|
||||
)
|
||||
|
||||
in_chs = min(self.filters_scale * self.filters, self.max_filters)
|
||||
for i, dilation in enumerate(self.dilations):
|
||||
out_chs = min((self.filters_scale ** (i + 1)) * self.filters, self.max_filters)
|
||||
self.convs.append(
|
||||
weight_norm(
|
||||
nn.Conv2d(
|
||||
in_chs,
|
||||
out_chs,
|
||||
kernel_size=self.kernel_size,
|
||||
stride=self.stride,
|
||||
dilation=(dilation, 1),
|
||||
padding=self.get_2d_padding(self.kernel_size, (dilation, 1)),
|
||||
)
|
||||
)
|
||||
)
|
||||
in_chs = out_chs
|
||||
out_chs = min(
|
||||
(self.filters_scale ** (len(self.dilations) + 1)) * self.filters,
|
||||
self.max_filters,
|
||||
)
|
||||
self.convs.append(
|
||||
weight_norm(
|
||||
nn.Conv2d(
|
||||
in_chs,
|
||||
out_chs,
|
||||
kernel_size=(self.kernel_size[0], self.kernel_size[0]),
|
||||
padding=self.get_2d_padding((self.kernel_size[0], self.kernel_size[0])),
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
self.conv_post = weight_norm(
|
||||
nn.Conv2d(
|
||||
out_chs,
|
||||
self.out_channels,
|
||||
kernel_size=(self.kernel_size[0], self.kernel_size[0]),
|
||||
padding=self.get_2d_padding((self.kernel_size[0], self.kernel_size[0])),
|
||||
)
|
||||
)
|
||||
|
||||
self.activation = torch.nn.LeakyReLU(negative_slope=0.1)
|
||||
self.resample = Resample(orig_freq=self.fs, new_freq=self.fs * 2)
|
||||
|
||||
self.cqtd_normalize_volume = self.cfg.get("cqtd_normalize_volume", False)
|
||||
if self.cqtd_normalize_volume:
|
||||
print(
|
||||
"[INFO] cqtd_normalize_volume set to True. Will apply DC offset removal & peak volume normalization in CQTD!"
|
||||
)
|
||||
|
||||
def get_2d_padding(
|
||||
self,
|
||||
kernel_size: typing.Tuple[int, int],
|
||||
dilation: typing.Tuple[int, int] = (1, 1),
|
||||
):
|
||||
return (
|
||||
((kernel_size[0] - 1) * dilation[0]) // 2,
|
||||
((kernel_size[1] - 1) * dilation[1]) // 2,
|
||||
)
|
||||
|
||||
def forward(self, x: torch.tensor) -> Tuple[torch.Tensor, List[torch.Tensor]]:
|
||||
fmap = []
|
||||
|
||||
if self.cqtd_normalize_volume:
|
||||
# Remove DC offset
|
||||
x = x - x.mean(dim=-1, keepdims=True)
|
||||
# Peak normalize the volume of input audio
|
||||
x = 0.8 * x / (x.abs().max(dim=-1, keepdim=True)[0] + 1e-9)
|
||||
|
||||
x = self.resample(x)
|
||||
|
||||
z = self.cqt_transform(x)
|
||||
|
||||
z_amplitude = z[:, :, :, 0].unsqueeze(1)
|
||||
z_phase = z[:, :, :, 1].unsqueeze(1)
|
||||
|
||||
z = torch.cat([z_amplitude, z_phase], dim=1)
|
||||
z = torch.permute(z, (0, 1, 3, 2)) # [B, C, W, T] -> [B, C, T, W]
|
||||
|
||||
latent_z = []
|
||||
for i in range(self.n_octaves):
|
||||
latent_z.append(
|
||||
self.conv_pres[i](
|
||||
z[
|
||||
:,
|
||||
:,
|
||||
:,
|
||||
i * self.bins_per_octave : (i + 1) * self.bins_per_octave,
|
||||
]
|
||||
)
|
||||
)
|
||||
latent_z = torch.cat(latent_z, dim=-1)
|
||||
|
||||
for i, l in enumerate(self.convs):
|
||||
latent_z = l(latent_z)
|
||||
|
||||
latent_z = self.activation(latent_z)
|
||||
fmap.append(latent_z)
|
||||
|
||||
latent_z = self.conv_post(latent_z)
|
||||
|
||||
return latent_z, fmap
|
||||
|
||||
|
||||
class MultiScaleSubbandCQTDiscriminator(nn.Module):
|
||||
def __init__(self, cfg: AttrDict):
|
||||
super().__init__()
|
||||
|
||||
self.cfg = cfg
|
||||
# Using get with defaults
|
||||
self.cfg["cqtd_filters"] = self.cfg.get("cqtd_filters", 32)
|
||||
self.cfg["cqtd_max_filters"] = self.cfg.get("cqtd_max_filters", 1024)
|
||||
self.cfg["cqtd_filters_scale"] = self.cfg.get("cqtd_filters_scale", 1)
|
||||
self.cfg["cqtd_dilations"] = self.cfg.get("cqtd_dilations", [1, 2, 4])
|
||||
self.cfg["cqtd_in_channels"] = self.cfg.get("cqtd_in_channels", 1)
|
||||
self.cfg["cqtd_out_channels"] = self.cfg.get("cqtd_out_channels", 1)
|
||||
# Multi-scale params to loop over
|
||||
self.cfg["cqtd_hop_lengths"] = self.cfg.get("cqtd_hop_lengths", [512, 256, 256])
|
||||
self.cfg["cqtd_n_octaves"] = self.cfg.get("cqtd_n_octaves", [9, 9, 9])
|
||||
self.cfg["cqtd_bins_per_octaves"] = self.cfg.get("cqtd_bins_per_octaves", [24, 36, 48])
|
||||
|
||||
self.discriminators = nn.ModuleList(
|
||||
[
|
||||
DiscriminatorCQT(
|
||||
self.cfg,
|
||||
hop_length=self.cfg["cqtd_hop_lengths"][i],
|
||||
n_octaves=self.cfg["cqtd_n_octaves"][i],
|
||||
bins_per_octave=self.cfg["cqtd_bins_per_octaves"][i],
|
||||
)
|
||||
for i in range(len(self.cfg["cqtd_hop_lengths"]))
|
||||
]
|
||||
)
|
||||
|
||||
def forward(
|
||||
self, y: torch.Tensor, y_hat: torch.Tensor
|
||||
) -> Tuple[
|
||||
List[torch.Tensor],
|
||||
List[torch.Tensor],
|
||||
List[List[torch.Tensor]],
|
||||
List[List[torch.Tensor]],
|
||||
]:
|
||||
y_d_rs = []
|
||||
y_d_gs = []
|
||||
fmap_rs = []
|
||||
fmap_gs = []
|
||||
|
||||
for disc in self.discriminators:
|
||||
y_d_r, fmap_r = disc(y)
|
||||
y_d_g, fmap_g = disc(y_hat)
|
||||
y_d_rs.append(y_d_r)
|
||||
fmap_rs.append(fmap_r)
|
||||
y_d_gs.append(y_d_g)
|
||||
fmap_gs.append(fmap_g)
|
||||
|
||||
return y_d_rs, y_d_gs, fmap_rs, fmap_gs
|
||||
|
||||
|
||||
class CombinedDiscriminator(nn.Module):
|
||||
"""
|
||||
Wrapper of chaining multiple discrimiantor architectures.
|
||||
Example: combine mbd and cqtd as a single class
|
||||
"""
|
||||
|
||||
def __init__(self, list_discriminator: List[nn.Module]):
|
||||
super().__init__()
|
||||
self.discrimiantor = nn.ModuleList(list_discriminator)
|
||||
|
||||
def forward(
|
||||
self, y: torch.Tensor, y_hat: torch.Tensor
|
||||
) -> Tuple[
|
||||
List[torch.Tensor],
|
||||
List[torch.Tensor],
|
||||
List[List[torch.Tensor]],
|
||||
List[List[torch.Tensor]],
|
||||
]:
|
||||
y_d_rs = []
|
||||
y_d_gs = []
|
||||
fmap_rs = []
|
||||
fmap_gs = []
|
||||
|
||||
for disc in self.discrimiantor:
|
||||
y_d_r, y_d_g, fmap_r, fmap_g = disc(y, y_hat)
|
||||
y_d_rs.extend(y_d_r)
|
||||
fmap_rs.extend(fmap_r)
|
||||
y_d_gs.extend(y_d_g)
|
||||
fmap_gs.extend(fmap_g)
|
||||
|
||||
return y_d_rs, y_d_gs, fmap_rs, fmap_gs
|
||||
@ -1,18 +1,8 @@
|
||||
# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
|
||||
import os
|
||||
import shutil
|
||||
# LICENSE is in the incl_licenses directory.
|
||||
|
||||
|
||||
class AttrDict(dict):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super(AttrDict, self).__init__(*args, **kwargs)
|
||||
self.__dict__ = self
|
||||
|
||||
|
||||
def build_env(config, config_name, path):
|
||||
t_path = os.path.join(path, config_name)
|
||||
if config != t_path:
|
||||
os.makedirs(path, exist_ok=True)
|
||||
shutil.copyfile(config, os.path.join(path, config_name))
|
||||
|
||||
@ -1,85 +0,0 @@
|
||||
# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
|
||||
from __future__ import absolute_import, division, print_function, unicode_literals
|
||||
|
||||
import os
|
||||
import argparse
|
||||
import json
|
||||
import torch
|
||||
import librosa
|
||||
from utils import load_checkpoint
|
||||
from meldataset import get_mel_spectrogram
|
||||
from scipy.io.wavfile import write
|
||||
from env import AttrDict
|
||||
from meldataset import MAX_WAV_VALUE
|
||||
from bigvgan import BigVGAN as Generator
|
||||
|
||||
h = None
|
||||
device = None
|
||||
torch.backends.cudnn.benchmark = False
|
||||
|
||||
|
||||
def inference(a, h):
|
||||
generator = Generator(h, use_cuda_kernel=a.use_cuda_kernel).to(device)
|
||||
|
||||
state_dict_g = load_checkpoint(a.checkpoint_file, device)
|
||||
generator.load_state_dict(state_dict_g["generator"])
|
||||
|
||||
filelist = os.listdir(a.input_wavs_dir)
|
||||
|
||||
os.makedirs(a.output_dir, exist_ok=True)
|
||||
|
||||
generator.eval()
|
||||
generator.remove_weight_norm()
|
||||
with torch.no_grad():
|
||||
for i, filname in enumerate(filelist):
|
||||
# Load the ground truth audio and resample if necessary
|
||||
wav, sr = librosa.load(os.path.join(a.input_wavs_dir, filname), sr=h.sampling_rate, mono=True)
|
||||
wav = torch.FloatTensor(wav).to(device)
|
||||
# Compute mel spectrogram from the ground truth audio
|
||||
x = get_mel_spectrogram(wav.unsqueeze(0), generator.h)
|
||||
|
||||
y_g_hat = generator(x)
|
||||
|
||||
audio = y_g_hat.squeeze()
|
||||
audio = audio * MAX_WAV_VALUE
|
||||
audio = audio.cpu().numpy().astype("int16")
|
||||
|
||||
output_file = os.path.join(a.output_dir, os.path.splitext(filname)[0] + "_generated.wav")
|
||||
write(output_file, h.sampling_rate, audio)
|
||||
print(output_file)
|
||||
|
||||
|
||||
def main():
|
||||
print("Initializing Inference Process..")
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--input_wavs_dir", default="test_files")
|
||||
parser.add_argument("--output_dir", default="generated_files")
|
||||
parser.add_argument("--checkpoint_file", required=True)
|
||||
parser.add_argument("--use_cuda_kernel", action="store_true", default=False)
|
||||
|
||||
a = parser.parse_args()
|
||||
|
||||
config_file = os.path.join(os.path.split(a.checkpoint_file)[0], "config.json")
|
||||
with open(config_file) as f:
|
||||
data = f.read()
|
||||
|
||||
global h
|
||||
json_config = json.loads(data)
|
||||
h = AttrDict(json_config)
|
||||
|
||||
torch.manual_seed(h.seed)
|
||||
global device
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed(h.seed)
|
||||
device = torch.device("cuda")
|
||||
else:
|
||||
device = torch.device("cpu")
|
||||
|
||||
inference(a, h)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -1,100 +0,0 @@
|
||||
# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
|
||||
from __future__ import absolute_import, division, print_function, unicode_literals
|
||||
|
||||
import glob
|
||||
import os
|
||||
import numpy as np
|
||||
import argparse
|
||||
import json
|
||||
import torch
|
||||
from scipy.io.wavfile import write
|
||||
from env import AttrDict
|
||||
from meldataset import MAX_WAV_VALUE
|
||||
from bigvgan import BigVGAN as Generator
|
||||
|
||||
h = None
|
||||
device = None
|
||||
torch.backends.cudnn.benchmark = False
|
||||
|
||||
|
||||
def load_checkpoint(filepath, device):
|
||||
assert os.path.isfile(filepath)
|
||||
print(f"Loading '{filepath}'")
|
||||
checkpoint_dict = torch.load(filepath, map_location=device)
|
||||
print("Complete.")
|
||||
return checkpoint_dict
|
||||
|
||||
|
||||
def scan_checkpoint(cp_dir, prefix):
|
||||
pattern = os.path.join(cp_dir, prefix + "*")
|
||||
cp_list = glob.glob(pattern)
|
||||
if len(cp_list) == 0:
|
||||
return ""
|
||||
return sorted(cp_list)[-1]
|
||||
|
||||
|
||||
def inference(a, h):
|
||||
generator = Generator(h, use_cuda_kernel=a.use_cuda_kernel).to(device)
|
||||
|
||||
state_dict_g = load_checkpoint(a.checkpoint_file, device)
|
||||
generator.load_state_dict(state_dict_g["generator"])
|
||||
|
||||
filelist = os.listdir(a.input_mels_dir)
|
||||
|
||||
os.makedirs(a.output_dir, exist_ok=True)
|
||||
|
||||
generator.eval()
|
||||
generator.remove_weight_norm()
|
||||
with torch.no_grad():
|
||||
for i, filname in enumerate(filelist):
|
||||
# Load the mel spectrogram in .npy format
|
||||
x = np.load(os.path.join(a.input_mels_dir, filname))
|
||||
x = torch.FloatTensor(x).to(device)
|
||||
if len(x.shape) == 2:
|
||||
x = x.unsqueeze(0)
|
||||
|
||||
y_g_hat = generator(x)
|
||||
|
||||
audio = y_g_hat.squeeze()
|
||||
audio = audio * MAX_WAV_VALUE
|
||||
audio = audio.cpu().numpy().astype("int16")
|
||||
|
||||
output_file = os.path.join(a.output_dir, os.path.splitext(filname)[0] + "_generated_e2e.wav")
|
||||
write(output_file, h.sampling_rate, audio)
|
||||
print(output_file)
|
||||
|
||||
|
||||
def main():
|
||||
print("Initializing Inference Process..")
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--input_mels_dir", default="test_mel_files")
|
||||
parser.add_argument("--output_dir", default="generated_files_from_mel")
|
||||
parser.add_argument("--checkpoint_file", required=True)
|
||||
parser.add_argument("--use_cuda_kernel", action="store_true", default=False)
|
||||
|
||||
a = parser.parse_args()
|
||||
|
||||
config_file = os.path.join(os.path.split(a.checkpoint_file)[0], "config.json")
|
||||
with open(config_file) as f:
|
||||
data = f.read()
|
||||
|
||||
global h
|
||||
json_config = json.loads(data)
|
||||
h = AttrDict(json_config)
|
||||
|
||||
torch.manual_seed(h.seed)
|
||||
global device
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed(h.seed)
|
||||
device = torch.device("cuda")
|
||||
else:
|
||||
device = torch.device("cpu")
|
||||
|
||||
inference(a, h)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -1,238 +0,0 @@
|
||||
# Copyright (c) 2024 NVIDIA CORPORATION.
|
||||
# Licensed under the MIT license.
|
||||
|
||||
# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from librosa.filters import mel as librosa_mel_fn
|
||||
from scipy import signal
|
||||
|
||||
import typing
|
||||
from typing import List, Tuple
|
||||
from collections import namedtuple
|
||||
import math
|
||||
import functools
|
||||
|
||||
|
||||
# Adapted from https://github.com/descriptinc/descript-audio-codec/blob/main/dac/nn/loss.py under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
class MultiScaleMelSpectrogramLoss(nn.Module):
|
||||
"""Compute distance between mel spectrograms. Can be used
|
||||
in a multi-scale way.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
n_mels : List[int]
|
||||
Number of mels per STFT, by default [5, 10, 20, 40, 80, 160, 320],
|
||||
window_lengths : List[int], optional
|
||||
Length of each window of each STFT, by default [32, 64, 128, 256, 512, 1024, 2048]
|
||||
loss_fn : typing.Callable, optional
|
||||
How to compare each loss, by default nn.L1Loss()
|
||||
clamp_eps : float, optional
|
||||
Clamp on the log magnitude, below, by default 1e-5
|
||||
mag_weight : float, optional
|
||||
Weight of raw magnitude portion of loss, by default 0.0 (no ampliciation on mag part)
|
||||
log_weight : float, optional
|
||||
Weight of log magnitude portion of loss, by default 1.0
|
||||
pow : float, optional
|
||||
Power to raise magnitude to before taking log, by default 1.0
|
||||
weight : float, optional
|
||||
Weight of this loss, by default 1.0
|
||||
match_stride : bool, optional
|
||||
Whether to match the stride of convolutional layers, by default False
|
||||
|
||||
Implementation copied from: https://github.com/descriptinc/lyrebird-audiotools/blob/961786aa1a9d628cca0c0486e5885a457fe70c1a/audiotools/metrics/spectral.py
|
||||
Additional code copied and modified from https://github.com/descriptinc/audiotools/blob/master/audiotools/core/audio_signal.py
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
sampling_rate: int,
|
||||
n_mels: List[int] = [5, 10, 20, 40, 80, 160, 320],
|
||||
window_lengths: List[int] = [32, 64, 128, 256, 512, 1024, 2048],
|
||||
loss_fn: typing.Callable = nn.L1Loss(),
|
||||
clamp_eps: float = 1e-5,
|
||||
mag_weight: float = 0.0,
|
||||
log_weight: float = 1.0,
|
||||
pow: float = 1.0,
|
||||
weight: float = 1.0,
|
||||
match_stride: bool = False,
|
||||
mel_fmin: List[float] = [0, 0, 0, 0, 0, 0, 0],
|
||||
mel_fmax: List[float] = [None, None, None, None, None, None, None],
|
||||
window_type: str = "hann",
|
||||
):
|
||||
super().__init__()
|
||||
self.sampling_rate = sampling_rate
|
||||
|
||||
STFTParams = namedtuple(
|
||||
"STFTParams",
|
||||
["window_length", "hop_length", "window_type", "match_stride"],
|
||||
)
|
||||
|
||||
self.stft_params = [
|
||||
STFTParams(
|
||||
window_length=w,
|
||||
hop_length=w // 4,
|
||||
match_stride=match_stride,
|
||||
window_type=window_type,
|
||||
)
|
||||
for w in window_lengths
|
||||
]
|
||||
self.n_mels = n_mels
|
||||
self.loss_fn = loss_fn
|
||||
self.clamp_eps = clamp_eps
|
||||
self.log_weight = log_weight
|
||||
self.mag_weight = mag_weight
|
||||
self.weight = weight
|
||||
self.mel_fmin = mel_fmin
|
||||
self.mel_fmax = mel_fmax
|
||||
self.pow = pow
|
||||
|
||||
@staticmethod
|
||||
@functools.lru_cache(None)
|
||||
def get_window(
|
||||
window_type,
|
||||
window_length,
|
||||
):
|
||||
return signal.get_window(window_type, window_length)
|
||||
|
||||
@staticmethod
|
||||
@functools.lru_cache(None)
|
||||
def get_mel_filters(sr, n_fft, n_mels, fmin, fmax):
|
||||
return librosa_mel_fn(sr=sr, n_fft=n_fft, n_mels=n_mels, fmin=fmin, fmax=fmax)
|
||||
|
||||
def mel_spectrogram(
|
||||
self,
|
||||
wav,
|
||||
n_mels,
|
||||
fmin,
|
||||
fmax,
|
||||
window_length,
|
||||
hop_length,
|
||||
match_stride,
|
||||
window_type,
|
||||
):
|
||||
"""
|
||||
Mirrors AudioSignal.mel_spectrogram used by BigVGAN-v2 training from:
|
||||
https://github.com/descriptinc/audiotools/blob/master/audiotools/core/audio_signal.py
|
||||
"""
|
||||
B, C, T = wav.shape
|
||||
|
||||
if match_stride:
|
||||
assert hop_length == window_length // 4, "For match_stride, hop must equal n_fft // 4"
|
||||
right_pad = math.ceil(T / hop_length) * hop_length - T
|
||||
pad = (window_length - hop_length) // 2
|
||||
else:
|
||||
right_pad = 0
|
||||
pad = 0
|
||||
|
||||
wav = torch.nn.functional.pad(wav, (pad, pad + right_pad), mode="reflect")
|
||||
|
||||
window = self.get_window(window_type, window_length)
|
||||
window = torch.from_numpy(window).to(wav.device).float()
|
||||
|
||||
stft = torch.stft(
|
||||
wav.reshape(-1, T),
|
||||
n_fft=window_length,
|
||||
hop_length=hop_length,
|
||||
window=window,
|
||||
return_complex=True,
|
||||
center=True,
|
||||
)
|
||||
_, nf, nt = stft.shape
|
||||
stft = stft.reshape(B, C, nf, nt)
|
||||
if match_stride:
|
||||
"""
|
||||
Drop first two and last two frames, which are added, because of padding. Now num_frames * hop_length = num_samples.
|
||||
"""
|
||||
stft = stft[..., 2:-2]
|
||||
magnitude = torch.abs(stft)
|
||||
|
||||
nf = magnitude.shape[2]
|
||||
mel_basis = self.get_mel_filters(self.sampling_rate, 2 * (nf - 1), n_mels, fmin, fmax)
|
||||
mel_basis = torch.from_numpy(mel_basis).to(wav.device)
|
||||
mel_spectrogram = magnitude.transpose(2, -1) @ mel_basis.T
|
||||
mel_spectrogram = mel_spectrogram.transpose(-1, 2)
|
||||
|
||||
return mel_spectrogram
|
||||
|
||||
def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
||||
"""Computes mel loss between an estimate and a reference
|
||||
signal.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
x : torch.Tensor
|
||||
Estimate signal
|
||||
y : torch.Tensor
|
||||
Reference signal
|
||||
|
||||
Returns
|
||||
-------
|
||||
torch.Tensor
|
||||
Mel loss.
|
||||
"""
|
||||
|
||||
loss = 0.0
|
||||
for n_mels, fmin, fmax, s in zip(self.n_mels, self.mel_fmin, self.mel_fmax, self.stft_params):
|
||||
kwargs = {
|
||||
"n_mels": n_mels,
|
||||
"fmin": fmin,
|
||||
"fmax": fmax,
|
||||
"window_length": s.window_length,
|
||||
"hop_length": s.hop_length,
|
||||
"match_stride": s.match_stride,
|
||||
"window_type": s.window_type,
|
||||
}
|
||||
|
||||
x_mels = self.mel_spectrogram(x, **kwargs)
|
||||
y_mels = self.mel_spectrogram(y, **kwargs)
|
||||
x_logmels = torch.log(x_mels.clamp(min=self.clamp_eps).pow(self.pow)) / torch.log(torch.tensor(10.0))
|
||||
y_logmels = torch.log(y_mels.clamp(min=self.clamp_eps).pow(self.pow)) / torch.log(torch.tensor(10.0))
|
||||
|
||||
loss += self.log_weight * self.loss_fn(x_logmels, y_logmels)
|
||||
loss += self.mag_weight * self.loss_fn(x_logmels, y_logmels)
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
# Loss functions
|
||||
def feature_loss(fmap_r: List[List[torch.Tensor]], fmap_g: List[List[torch.Tensor]]) -> torch.Tensor:
|
||||
loss = 0
|
||||
for dr, dg in zip(fmap_r, fmap_g):
|
||||
for rl, gl in zip(dr, dg):
|
||||
loss += torch.mean(torch.abs(rl - gl))
|
||||
|
||||
return loss * 2 # This equates to lambda=2.0 for the feature matching loss
|
||||
|
||||
|
||||
def discriminator_loss(
|
||||
disc_real_outputs: List[torch.Tensor], disc_generated_outputs: List[torch.Tensor]
|
||||
) -> Tuple[torch.Tensor, List[torch.Tensor], List[torch.Tensor]]:
|
||||
loss = 0
|
||||
r_losses = []
|
||||
g_losses = []
|
||||
for dr, dg in zip(disc_real_outputs, disc_generated_outputs):
|
||||
r_loss = torch.mean((1 - dr) ** 2)
|
||||
g_loss = torch.mean(dg**2)
|
||||
loss += r_loss + g_loss
|
||||
r_losses.append(r_loss.item())
|
||||
g_losses.append(g_loss.item())
|
||||
|
||||
return loss, r_losses, g_losses
|
||||
|
||||
|
||||
def generator_loss(
|
||||
disc_outputs: List[torch.Tensor],
|
||||
) -> Tuple[torch.Tensor, List[torch.Tensor]]:
|
||||
loss = 0
|
||||
gen_losses = []
|
||||
for dg in disc_outputs:
|
||||
l = torch.mean((1 - dg) ** 2)
|
||||
gen_losses.append(l)
|
||||
loss += l
|
||||
|
||||
return loss, gen_losses
|
||||
@ -1,370 +0,0 @@
|
||||
# Copyright (c) 2024 NVIDIA CORPORATION.
|
||||
# Licensed under the MIT license.
|
||||
|
||||
# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
|
||||
import math
|
||||
import os
|
||||
import random
|
||||
import torch
|
||||
import torch.utils.data
|
||||
import numpy as np
|
||||
import librosa
|
||||
from librosa.filters import mel as librosa_mel_fn
|
||||
import pathlib
|
||||
from tqdm import tqdm
|
||||
from typing import List, Tuple, Optional
|
||||
from .env import AttrDict
|
||||
|
||||
MAX_WAV_VALUE = 32767.0 # NOTE: 32768.0 -1 to prevent int16 overflow (results in popping sound in corner cases)
|
||||
|
||||
|
||||
def dynamic_range_compression(x, C=1, clip_val=1e-5):
|
||||
return np.log(np.clip(x, a_min=clip_val, a_max=None) * C)
|
||||
|
||||
|
||||
def dynamic_range_decompression(x, C=1):
|
||||
return np.exp(x) / C
|
||||
|
||||
|
||||
def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
|
||||
return torch.log(torch.clamp(x, min=clip_val) * C)
|
||||
|
||||
|
||||
def dynamic_range_decompression_torch(x, C=1):
|
||||
return torch.exp(x) / C
|
||||
|
||||
|
||||
def spectral_normalize_torch(magnitudes):
|
||||
return dynamic_range_compression_torch(magnitudes)
|
||||
|
||||
|
||||
def spectral_de_normalize_torch(magnitudes):
|
||||
return dynamic_range_decompression_torch(magnitudes)
|
||||
|
||||
|
||||
mel_basis_cache = {}
|
||||
hann_window_cache = {}
|
||||
|
||||
|
||||
def mel_spectrogram(
|
||||
y: torch.Tensor,
|
||||
n_fft: int,
|
||||
num_mels: int,
|
||||
sampling_rate: int,
|
||||
hop_size: int,
|
||||
win_size: int,
|
||||
fmin: int,
|
||||
fmax: int = None,
|
||||
center: bool = False,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Calculate the mel spectrogram of an input signal.
|
||||
This function uses slaney norm for the librosa mel filterbank (using librosa.filters.mel) and uses Hann window for STFT (using torch.stft).
|
||||
|
||||
Args:
|
||||
y (torch.Tensor): Input signal.
|
||||
n_fft (int): FFT size.
|
||||
num_mels (int): Number of mel bins.
|
||||
sampling_rate (int): Sampling rate of the input signal.
|
||||
hop_size (int): Hop size for STFT.
|
||||
win_size (int): Window size for STFT.
|
||||
fmin (int): Minimum frequency for mel filterbank.
|
||||
fmax (int): Maximum frequency for mel filterbank. If None, defaults to half the sampling rate (fmax = sr / 2.0) inside librosa_mel_fn
|
||||
center (bool): Whether to pad the input to center the frames. Default is False.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Mel spectrogram.
|
||||
"""
|
||||
if torch.min(y) < -1.0:
|
||||
print(f"[WARNING] Min value of input waveform signal is {torch.min(y)}")
|
||||
if torch.max(y) > 1.0:
|
||||
print(f"[WARNING] Max value of input waveform signal is {torch.max(y)}")
|
||||
|
||||
device = y.device
|
||||
key = f"{n_fft}_{num_mels}_{sampling_rate}_{hop_size}_{win_size}_{fmin}_{fmax}_{device}"
|
||||
|
||||
if key not in mel_basis_cache:
|
||||
mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax)
|
||||
mel_basis_cache[key] = torch.from_numpy(mel).float().to(device)
|
||||
hann_window_cache[key] = torch.hann_window(win_size).to(device)
|
||||
|
||||
mel_basis = mel_basis_cache[key]
|
||||
hann_window = hann_window_cache[key]
|
||||
|
||||
padding = (n_fft - hop_size) // 2
|
||||
y = torch.nn.functional.pad(y.unsqueeze(1), (padding, padding), mode="reflect").squeeze(1)
|
||||
|
||||
spec = torch.stft(
|
||||
y,
|
||||
n_fft,
|
||||
hop_length=hop_size,
|
||||
win_length=win_size,
|
||||
window=hann_window,
|
||||
center=center,
|
||||
pad_mode="reflect",
|
||||
normalized=False,
|
||||
onesided=True,
|
||||
return_complex=True,
|
||||
)
|
||||
spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9)
|
||||
|
||||
mel_spec = torch.matmul(mel_basis, spec)
|
||||
mel_spec = spectral_normalize_torch(mel_spec)
|
||||
|
||||
return mel_spec
|
||||
|
||||
|
||||
def get_mel_spectrogram(wav, h):
|
||||
"""
|
||||
Generate mel spectrogram from a waveform using given hyperparameters.
|
||||
|
||||
Args:
|
||||
wav (torch.Tensor): Input waveform.
|
||||
h: Hyperparameters object with attributes n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Mel spectrogram.
|
||||
"""
|
||||
return mel_spectrogram(
|
||||
wav,
|
||||
h.n_fft,
|
||||
h.num_mels,
|
||||
h.sampling_rate,
|
||||
h.hop_size,
|
||||
h.win_size,
|
||||
h.fmin,
|
||||
h.fmax,
|
||||
)
|
||||
|
||||
|
||||
def get_dataset_filelist(a):
|
||||
training_files = []
|
||||
validation_files = []
|
||||
list_unseen_validation_files = []
|
||||
|
||||
with open(a.input_training_file, "r", encoding="utf-8") as fi:
|
||||
training_files = [
|
||||
os.path.join(a.input_wavs_dir, x.split("|")[0] + ".wav") for x in fi.read().split("\n") if len(x) > 0
|
||||
]
|
||||
print(f"first training file: {training_files[0]}")
|
||||
|
||||
with open(a.input_validation_file, "r", encoding="utf-8") as fi:
|
||||
validation_files = [
|
||||
os.path.join(a.input_wavs_dir, x.split("|")[0] + ".wav") for x in fi.read().split("\n") if len(x) > 0
|
||||
]
|
||||
print(f"first validation file: {validation_files[0]}")
|
||||
|
||||
for i in range(len(a.list_input_unseen_validation_file)):
|
||||
with open(a.list_input_unseen_validation_file[i], "r", encoding="utf-8") as fi:
|
||||
unseen_validation_files = [
|
||||
os.path.join(a.list_input_unseen_wavs_dir[i], x.split("|")[0] + ".wav")
|
||||
for x in fi.read().split("\n")
|
||||
if len(x) > 0
|
||||
]
|
||||
print(f"first unseen {i}th validation fileset: {unseen_validation_files[0]}")
|
||||
list_unseen_validation_files.append(unseen_validation_files)
|
||||
|
||||
return training_files, validation_files, list_unseen_validation_files
|
||||
|
||||
|
||||
class MelDataset(torch.utils.data.Dataset):
|
||||
def __init__(
|
||||
self,
|
||||
training_files: List[str],
|
||||
hparams: AttrDict,
|
||||
segment_size: int,
|
||||
n_fft: int,
|
||||
num_mels: int,
|
||||
hop_size: int,
|
||||
win_size: int,
|
||||
sampling_rate: int,
|
||||
fmin: int,
|
||||
fmax: Optional[int],
|
||||
split: bool = True,
|
||||
shuffle: bool = True,
|
||||
device: str = None,
|
||||
fmax_loss: Optional[int] = None,
|
||||
fine_tuning: bool = False,
|
||||
base_mels_path: str = None,
|
||||
is_seen: bool = True,
|
||||
):
|
||||
self.audio_files = training_files
|
||||
random.seed(1234)
|
||||
if shuffle:
|
||||
random.shuffle(self.audio_files)
|
||||
self.hparams = hparams
|
||||
self.is_seen = is_seen
|
||||
if self.is_seen:
|
||||
self.name = pathlib.Path(self.audio_files[0]).parts[0]
|
||||
else:
|
||||
self.name = "-".join(pathlib.Path(self.audio_files[0]).parts[:2]).strip("/")
|
||||
|
||||
self.segment_size = segment_size
|
||||
self.sampling_rate = sampling_rate
|
||||
self.split = split
|
||||
self.n_fft = n_fft
|
||||
self.num_mels = num_mels
|
||||
self.hop_size = hop_size
|
||||
self.win_size = win_size
|
||||
self.fmin = fmin
|
||||
self.fmax = fmax
|
||||
self.fmax_loss = fmax_loss
|
||||
self.device = device
|
||||
self.fine_tuning = fine_tuning
|
||||
self.base_mels_path = base_mels_path
|
||||
|
||||
print("[INFO] checking dataset integrity...")
|
||||
for i in tqdm(range(len(self.audio_files))):
|
||||
assert os.path.exists(self.audio_files[i]), f"{self.audio_files[i]} not found"
|
||||
|
||||
def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor, str, torch.Tensor]:
|
||||
try:
|
||||
filename = self.audio_files[index]
|
||||
|
||||
# Use librosa.load that ensures loading waveform into mono with [-1, 1] float values
|
||||
# Audio is ndarray with shape [T_time]. Disable auto-resampling here to minimize overhead
|
||||
# The on-the-fly resampling during training will be done only for the obtained random chunk
|
||||
audio, source_sampling_rate = librosa.load(filename, sr=None, mono=True)
|
||||
|
||||
# Main logic that uses <mel, audio> pair for training BigVGAN
|
||||
if not self.fine_tuning:
|
||||
if self.split: # Training step
|
||||
# Obtain randomized audio chunk
|
||||
if source_sampling_rate != self.sampling_rate:
|
||||
# Adjust segment size to crop if the source sr is different
|
||||
target_segment_size = math.ceil(self.segment_size * (source_sampling_rate / self.sampling_rate))
|
||||
else:
|
||||
target_segment_size = self.segment_size
|
||||
|
||||
# Compute upper bound index for the random chunk
|
||||
random_chunk_upper_bound = max(0, audio.shape[0] - target_segment_size)
|
||||
|
||||
# Crop or pad audio to obtain random chunk with target_segment_size
|
||||
if audio.shape[0] >= target_segment_size:
|
||||
audio_start = random.randint(0, random_chunk_upper_bound)
|
||||
audio = audio[audio_start : audio_start + target_segment_size]
|
||||
else:
|
||||
audio = np.pad(
|
||||
audio,
|
||||
(0, target_segment_size - audio.shape[0]),
|
||||
mode="constant",
|
||||
)
|
||||
|
||||
# Resample audio chunk to self.sampling rate
|
||||
if source_sampling_rate != self.sampling_rate:
|
||||
audio = librosa.resample(
|
||||
audio,
|
||||
orig_sr=source_sampling_rate,
|
||||
target_sr=self.sampling_rate,
|
||||
)
|
||||
if audio.shape[0] > self.segment_size:
|
||||
# trim last elements to match self.segment_size (e.g., 16385 for 44khz downsampled to 24khz -> 16384)
|
||||
audio = audio[: self.segment_size]
|
||||
|
||||
else: # Validation step
|
||||
# Resample full audio clip to target sampling rate
|
||||
if source_sampling_rate != self.sampling_rate:
|
||||
audio = librosa.resample(
|
||||
audio,
|
||||
orig_sr=source_sampling_rate,
|
||||
target_sr=self.sampling_rate,
|
||||
)
|
||||
# Trim last elements to match audio length to self.hop_size * n for evaluation
|
||||
if (audio.shape[0] % self.hop_size) != 0:
|
||||
audio = audio[: -(audio.shape[0] % self.hop_size)]
|
||||
|
||||
# BigVGAN is trained using volume-normalized waveform
|
||||
audio = librosa.util.normalize(audio) * 0.95
|
||||
|
||||
# Cast ndarray to torch tensor
|
||||
audio = torch.FloatTensor(audio)
|
||||
audio = audio.unsqueeze(0) # [B(1), self.segment_size]
|
||||
|
||||
# Compute mel spectrogram corresponding to audio
|
||||
mel = mel_spectrogram(
|
||||
audio,
|
||||
self.n_fft,
|
||||
self.num_mels,
|
||||
self.sampling_rate,
|
||||
self.hop_size,
|
||||
self.win_size,
|
||||
self.fmin,
|
||||
self.fmax,
|
||||
center=False,
|
||||
) # [B(1), self.num_mels, self.segment_size // self.hop_size]
|
||||
|
||||
# Fine-tuning logic that uses pre-computed mel. Example: Using TTS model-generated mel as input
|
||||
else:
|
||||
# For fine-tuning, assert that the waveform is in the defined sampling_rate
|
||||
# Fine-tuning won't support on-the-fly resampling to be fool-proof (the dataset should have been prepared properly)
|
||||
assert source_sampling_rate == self.sampling_rate, (
|
||||
f"For fine_tuning, waveform must be in the spcified sampling rate {self.sampling_rate}, got {source_sampling_rate}"
|
||||
)
|
||||
|
||||
# Cast ndarray to torch tensor
|
||||
audio = torch.FloatTensor(audio)
|
||||
audio = audio.unsqueeze(0) # [B(1), T_time]
|
||||
|
||||
# Load pre-computed mel from disk
|
||||
mel = np.load(
|
||||
os.path.join(
|
||||
self.base_mels_path,
|
||||
os.path.splitext(os.path.split(filename)[-1])[0] + ".npy",
|
||||
)
|
||||
)
|
||||
mel = torch.from_numpy(mel)
|
||||
|
||||
if len(mel.shape) < 3:
|
||||
mel = mel.unsqueeze(0) # ensure [B, C, T]
|
||||
|
||||
if self.split:
|
||||
frames_per_seg = math.ceil(self.segment_size / self.hop_size)
|
||||
|
||||
if audio.size(1) >= self.segment_size:
|
||||
mel_start = random.randint(0, mel.size(2) - frames_per_seg - 1)
|
||||
mel = mel[:, :, mel_start : mel_start + frames_per_seg]
|
||||
audio = audio[
|
||||
:,
|
||||
mel_start * self.hop_size : (mel_start + frames_per_seg) * self.hop_size,
|
||||
]
|
||||
|
||||
# Pad pre-computed mel and audio to match length to ensuring fine-tuning without error.
|
||||
# NOTE: this may introduce a single-frame misalignment of the <pre-computed mel, audio>
|
||||
# To remove possible misalignment, it is recommended to prepare the <pre-computed mel, audio> pair where the audio length is the integer multiple of self.hop_size
|
||||
mel = torch.nn.functional.pad(mel, (0, frames_per_seg - mel.size(2)), "constant")
|
||||
audio = torch.nn.functional.pad(audio, (0, self.segment_size - audio.size(1)), "constant")
|
||||
|
||||
# Compute mel_loss used by spectral regression objective. Uses self.fmax_loss instead (usually None)
|
||||
mel_loss = mel_spectrogram(
|
||||
audio,
|
||||
self.n_fft,
|
||||
self.num_mels,
|
||||
self.sampling_rate,
|
||||
self.hop_size,
|
||||
self.win_size,
|
||||
self.fmin,
|
||||
self.fmax_loss,
|
||||
center=False,
|
||||
) # [B(1), self.num_mels, self.segment_size // self.hop_size]
|
||||
|
||||
# Shape sanity checks
|
||||
assert (
|
||||
audio.shape[1] == mel.shape[2] * self.hop_size and audio.shape[1] == mel_loss.shape[2] * self.hop_size
|
||||
), (
|
||||
f"Audio length must be mel frame length * hop_size. Got audio shape {audio.shape} mel shape {mel.shape} mel_loss shape {mel_loss.shape}"
|
||||
)
|
||||
|
||||
return (mel.squeeze(), audio.squeeze(0), filename, mel_loss.squeeze())
|
||||
|
||||
# If it encounters error during loading the data, skip this sample and load random other sample to the batch
|
||||
except Exception as e:
|
||||
if self.fine_tuning:
|
||||
raise e # Terminate training if it is fine-tuning. The dataset should have been prepared properly.
|
||||
else:
|
||||
print(f"[WARNING] Failed to load waveform, skipping! filename: {filename} Error: {e}")
|
||||
return self[random.randrange(len(self))]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.audio_files)
|
||||
@ -1 +0,0 @@
|
||||
|
||||
@ -1,4 +0,0 @@
|
||||
| Field | Response |
|
||||
| :--------------------------------------------------------------------------------------------------------- | :--------------------------------------------------- |
|
||||
| Participation considerations from adversely impacted groups protected classes in model design and testing: | None |
|
||||
| Measures taken to mitigate against unwanted bias: | No measures taken to mitigate against unwanted bias. |
|
||||
@ -1,13 +0,0 @@
|
||||
| Field | Response |
|
||||
| :---------------------------------------------------------------------------------------------------- | :--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| Intended Application & Domain: | Generating waveform from mel spectrogram. |
|
||||
| Model Type: | Convolutional Neural Network (CNN) |
|
||||
| Intended Users: | This model is intended for developers to synthesize and generate waveforms from the AI-generated mel spectrograms. |
|
||||
| Output: | Audio Waveform |
|
||||
| Describe how the model works: | Model generates audio waveform corresponding to the input mel spectrogram. |
|
||||
| Name the adversely impacted groups this has been tested to deliver comparable outcomes regardless of: | Not Applicable |
|
||||
| Technical Limitations: | This may not perform well on synthetically-generated mel spectrograms that deviate significantly from the profile of mel spectrograms on which this was trained. |
|
||||
| Verified to have met prescribed NVIDIA quality standards: | Yes |
|
||||
| Performance Metrics: | Perceptual Evaluation of Speech Quality (PESQ), Virtual Speech Quality Objective Listener (VISQOL), Multi-resolution STFT (MRSTFT), Mel cepstral distortion (MCD), Periodicity RMSE, Voice/Unvoiced F1 Score (V/UV F1) |
|
||||
| Potential Known Risks: | This model may generate low-quality or distorted soundwaves. |
|
||||
| Licensing: | https://github.com/NVIDIA/BigVGAN/blob/main/LICENSE |
|
||||
@ -1,126 +0,0 @@
|
||||
# Model Overview
|
||||
|
||||
## Description:
|
||||
|
||||
BigVGAN is a generative AI model specialized in synthesizing audio waveforms using Mel spectrogram as inputs.
|
||||
|
||||
<center><img src="https://user-images.githubusercontent.com/15963413/218609148-881e39df-33af-4af9-ab95-1427c4ebf062.png" width="800"></center>
|
||||
|
||||
BigVGAN is a fully convolutional architecture with several upsampling blocks using transposed convolution followed by multiple residual dilated convolution layers.
|
||||
|
||||
BigVGAN consists of a novel module, called anti-aliased multi-periodicity composition (AMP), which is specifically designed for generating waveforms. AMP is specialized in synthesizing high-frequency and periodic soundwaves drawing inspiration from audio signal processing principles.
|
||||
|
||||
It applies a periodic activation function, called Snake, which provides an inductive bias to the architecture in generating periodic soundwaves. It also applies anti-aliasing filters to reduce undesired artifacts in the generated waveforms. <br>
|
||||
|
||||
This model is ready for commercial use.<br>
|
||||
|
||||
## References(s):
|
||||
|
||||
- [BigVGAN: A Universal Neural Vocoder with Large-Scale Training](https://arxiv.org/abs/2206.04658) <br>
|
||||
- [Project Page](https://research.nvidia.com/labs/adlr/projects/bigvgan/) <br>
|
||||
- [Audio Demo](https://bigvgan-demo.github.io/) <br>
|
||||
|
||||
## Model Architecture:
|
||||
|
||||
**Architecture Type:** Convolution Neural Network (CNN) <br>
|
||||
**Network Architecture:** You can see the details of this model on this link: https://github.com/NVIDIA/BigVGAN and the related paper can be found here: https://arxiv.org/abs/2206.04658<br>
|
||||
**Model Version:** 2.0 <br>
|
||||
|
||||
## Input:
|
||||
|
||||
**Input Type:** Audio <br>
|
||||
**Input Format:** Mel Spectrogram <br>
|
||||
**Input Parameters:** None <br>
|
||||
**Other Properties Related to Input:** The input mel spectrogram has shape `[batch, channels, frames]`, where `channels` refers to the number of mel bands defined by the model and `frames` refers to the temporal length. The model supports arbitrary long `frames` that fits into the GPU memory.
|
||||
|
||||
## Output:
|
||||
|
||||
**Input Type:** Audio <br>
|
||||
**Output Format:** Audio Waveform <br>
|
||||
**Output Parameters:** None <br>
|
||||
**Other Properties Related to Output:** The output audio waveform has shape `[batch, 1, time]`, where `1` refers to the mono audio channels and `time` refers to the temporal length. `time` is defined as a fixed integer multiple of input `frames`, which is an upsampling ratio of the model (`time = upsampling ratio * frames`). The output audio waveform consitutes float values with a range of `[-1, 1]`.
|
||||
|
||||
## Software Integration:
|
||||
|
||||
**Runtime Engine(s):** PyTorch
|
||||
|
||||
**Supported Hardware Microarchitecture Compatibility:** NVIDIA Ampere, NVIDIA Hopper, NVIDIA Lovelace, NVIDIA Turing, NVIDIA Volta <br>
|
||||
|
||||
## Preferred/Supported Operating System(s):
|
||||
|
||||
Linux
|
||||
|
||||
## Model Version(s):
|
||||
|
||||
v2.0
|
||||
|
||||
## Training, Testing, and Evaluation Datasets:
|
||||
|
||||
### Training Dataset:
|
||||
|
||||
The dataset contains diverse audio types, including speech in multiple languages, environmental sounds, and instruments.
|
||||
|
||||
**Links:**
|
||||
|
||||
- [AAM: Artificial Audio Multitracks Dataset](https://zenodo.org/records/5794629)
|
||||
- [AudioCaps](https://audiocaps.github.io/)
|
||||
- [AudioSet](https://research.google.com/audioset/index.html)
|
||||
- [common-accent](https://huggingface.co/datasets/DTU54DL/common-accent)
|
||||
- [Crowd Sourced Emotional Multimodal Actors Dataset (CREMA-D)](https://ieeexplore.ieee.org/document/6849440)
|
||||
- [DCASE2017 Challenge, Task 4: Large-scale weakly supervised sound event detection for smart cars](https://dcase.community/challenge2017/task-large-scale-sound-event-detection)
|
||||
- [FSDnoisy18k](https://zenodo.org/records/2529934)
|
||||
- [Free Universal Sound Separation Dataset](https://zenodo.org/records/3694384)
|
||||
- [Greatest Hits dataset](https://andrewowens.com/vis/)
|
||||
- [GTZAN](https://ieeexplore.ieee.org/document/1021072)
|
||||
- [JL corpus](https://www.kaggle.com/datasets/tli725/jl-corpus)
|
||||
- [Medley-solos-DB: a cross-collection dataset for musical instrument recognition](https://zenodo.org/records/3464194)
|
||||
- [MUSAN: A Music, Speech, and Noise Corpus](https://www.openslr.org/17/)
|
||||
- [MusicBench](https://huggingface.co/datasets/amaai-lab/MusicBench)
|
||||
- [MusicCaps](https://www.kaggle.com/datasets/googleai/musiccaps)
|
||||
- [MusicNet](https://www.kaggle.com/datasets/imsparsh/musicnet-dataset)
|
||||
- [NSynth](https://magenta.tensorflow.org/datasets/nsynth)
|
||||
- [OnAir-Music-Dataset](https://github.com/sevagh/OnAir-Music-Dataset)
|
||||
- [Audio Piano Triads Dataset](https://zenodo.org/records/4740877)
|
||||
- [Pitch Audio Dataset (Surge synthesizer)](https://zenodo.org/records/4677097)
|
||||
- [SONYC Urban Sound Tagging (SONYC-UST): a multilabel dataset from an urban acoustic sensor network](https://zenodo.org/records/3966543)
|
||||
- [VocalSound: A Dataset for Improving Human Vocal Sounds Recognition](https://arxiv.org/abs/2205.03433)
|
||||
- [WavText5K](https://github.com/microsoft/WavText5K)
|
||||
- [CSS10: A Collection of Single Speaker Speech Datasets for 10 Languages](https://github.com/Kyubyong/css10)
|
||||
- [Hi-Fi Multi-Speaker English TTS Dataset (Hi-Fi TTS)](https://www.openslr.org/109/)
|
||||
- [IIIT-H Indic Speech Databases](http://festvox.org/databases/iiit_voices/)
|
||||
- [Libri-Light: A Benchmark for ASR with Limited or No Supervision](https://arxiv.org/abs/1912.07875)
|
||||
- [LibriTTS: A Corpus Derived from LibriSpeech for Text-to-Speech](https://www.openslr.org/60)
|
||||
- [LibriTTS-R: A Restored Multi-Speaker Text-to-Speech Corpus](https://www.openslr.org/141/)
|
||||
- [The SIWIS French Speech Synthesis Database](https://datashare.ed.ac.uk/handle/10283/2353)
|
||||
- [Crowdsourced high-quality Colombian Spanish speech data set](https://openslr.org/72/)
|
||||
- [TTS-Portuguese Corpus](https://github.com/Edresson/TTS-Portuguese-Corpus)
|
||||
- [CSTR VCTK Corpus: English Multi-speaker Corpus for CSTR Voice Cloning Toolkit](https://datashare.ed.ac.uk/handle/10283/3443)
|
||||
|
||||
\*\* Data Collection Method by dataset <br>
|
||||
|
||||
- Human <br>
|
||||
|
||||
\*\* Labeling Method by dataset (for those with labels) <br>
|
||||
|
||||
- Hybrid: Automated, Human, Unknown <br>
|
||||
|
||||
### Evaluating Dataset:
|
||||
|
||||
Properties: The audio generation quality of BigVGAN is evaluated using `dev` splits of the [LibriTTS dataset](https://www.openslr.org/60/) and [Hi-Fi TTS dataset](https://www.openslr.org/109/). The datasets include speech in English language with equal balance of genders.
|
||||
|
||||
\*\* Data Collection Method by dataset <br>
|
||||
|
||||
- Human <br>
|
||||
|
||||
\*\* Labeling Method by dataset <br>
|
||||
|
||||
- Automated <br>
|
||||
|
||||
## Inference:
|
||||
|
||||
**Engine:** PyTorch <br>
|
||||
**Test Hardware:** NVIDIA A100 GPU <br>
|
||||
|
||||
## Ethical Considerations:
|
||||
|
||||
NVIDIA believes Trustworthy AI is a shared responsibility and we have established policies and practices to enable development for a wide array of AI applications. When downloaded or used in accordance with our terms of service, developers should work with their internal model team to ensure this model meets requirements for the relevant industry and use case and addresses unforeseen product misuse. For more detailed information on ethical considerations for this model, please see the Model Card++ Explainability, Bias, Safety & Security, and Privacy Subcards. Please report security vulnerabilities or NVIDIA AI Concerns [here](https://www.nvidia.com/en-us/support/submit-security-vulnerability/).
|
||||
@ -1,14 +0,0 @@
|
||||
| Field | Response |
|
||||
| :------------------------------------------------------------------------------------------------------------------------------------- | :--------------------------------------------- |
|
||||
| Generatable or reverse engineerable personal information? | None |
|
||||
| Protected class data used to create this model? | None |
|
||||
| Was consent obtained for any personal data used? | Not Applicable (No Personal Data) |
|
||||
| How often is dataset reviewed? | Before Release |
|
||||
| Is a mechanism in place to honor data subject right of access or deletion of personal data? | Not Applicable |
|
||||
| If personal collected for the development of the model, was it collected directly by NVIDIA? | Not Applicable |
|
||||
| If personal collected for the development of the model by NVIDIA, do you maintain or have access to disclosures made to data subjects? | Not Applicable |
|
||||
| If personal collected for the development of this AI model, was it minimized to only what was required? | Not Applicable |
|
||||
| Is data in dataset traceable? | Yes |
|
||||
| Is there provenance for all datasets used in training? | Yes |
|
||||
| Does data labeling (annotation, metadata) comply with privacy laws? | Yes |
|
||||
| Is data compliant with data subject requests for data correction or removal, if such a request was made? | No, not possible with externally-sourced data. |
|
||||
@ -1,6 +0,0 @@
|
||||
| Field | Response |
|
||||
| :---------------------------------------------- | :---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| Model Application(s): | Synethic Audio Generation |
|
||||
| Describe the life critical impact (if present). | Not Applicable |
|
||||
| Use Case Restrictions: | None |
|
||||
| Model and dataset restrictions: | The Principle of least privilege (PoLP) is applied limiting access for dataset generation and model development. Restrictions enforce dataset access during training, and dataset license constraints adhered to. |
|
||||
@ -1,13 +0,0 @@
|
||||
torch
|
||||
numpy
|
||||
librosa>=0.8.1
|
||||
scipy
|
||||
tensorboard
|
||||
soundfile
|
||||
matplotlib
|
||||
pesq
|
||||
auraloss
|
||||
tqdm
|
||||
nnAudio
|
||||
ninja
|
||||
huggingface_hub>=0.23.4
|
||||
@ -1,62 +0,0 @@
|
||||
# Copyright (c) 2024 NVIDIA CORPORATION.
|
||||
# Licensed under the MIT license.
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
# to import modules from parent_dir
|
||||
parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
||||
sys.path.append(parent_dir)
|
||||
|
||||
import torch
|
||||
from alias_free_activation.cuda import activation1d
|
||||
from activations import Snake
|
||||
|
||||
|
||||
def test_load_fused_kernels():
|
||||
try:
|
||||
print("[Success] load_fused_kernels")
|
||||
except ImportError as e:
|
||||
print("[Fail] load_fused_kernels")
|
||||
raise e
|
||||
|
||||
|
||||
def test_anti_alias_activation():
|
||||
data = torch.rand((10, 10, 200), device="cuda")
|
||||
|
||||
# Check activations.Snake cuda vs. torch
|
||||
fused_anti_alias_activation = activation1d.Activation1d(activation=Snake(10), fused=True).cuda()
|
||||
fused_activation_output = fused_anti_alias_activation(data)
|
||||
|
||||
torch_anti_alias_activation = activation1d.Activation1d(activation=Snake(10), fused=False).cuda()
|
||||
torch_activation_output = torch_anti_alias_activation(data)
|
||||
|
||||
test_result = (fused_activation_output - torch_activation_output).abs()
|
||||
|
||||
while test_result.dim() != 1:
|
||||
test_result = test_result.mean(dim=-1)
|
||||
|
||||
diff = test_result.mean(dim=-1)
|
||||
|
||||
if diff <= 1e-3:
|
||||
print(
|
||||
f"\n[Success] test_fused_anti_alias_activation"
|
||||
f"\n > mean_difference={diff}"
|
||||
f"\n > fused_values={fused_activation_output[-1][-1][:].tolist()}"
|
||||
f"\n > torch_values={torch_activation_output[-1][-1][:].tolist()}"
|
||||
)
|
||||
else:
|
||||
print(
|
||||
f"\n[Fail] test_fused_anti_alias_activation"
|
||||
f"\n > mean_difference={diff}, "
|
||||
f"\n > fused_values={fused_activation_output[-1][-1][:].tolist()}, "
|
||||
f"\n > torch_values={torch_activation_output[-1][-1][:].tolist()}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from alias_free_activation.cuda import load
|
||||
|
||||
load.load()
|
||||
test_load_fused_kernels()
|
||||
test_anti_alias_activation()
|
||||
@ -1,62 +0,0 @@
|
||||
# Copyright (c) 2024 NVIDIA CORPORATION.
|
||||
# Licensed under the MIT license.
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
# to import modules from parent_dir
|
||||
parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
||||
sys.path.append(parent_dir)
|
||||
|
||||
import torch
|
||||
from alias_free_activation.cuda import activation1d
|
||||
from activations import SnakeBeta
|
||||
|
||||
|
||||
def test_load_fused_kernels():
|
||||
try:
|
||||
print("[Success] load_fused_kernels")
|
||||
except ImportError as e:
|
||||
print("[Fail] load_fused_kernels")
|
||||
raise e
|
||||
|
||||
|
||||
def test_anti_alias_activation():
|
||||
data = torch.rand((10, 10, 200), device="cuda")
|
||||
|
||||
# Check activations, Snake CUDA vs. Torch
|
||||
fused_anti_alias_activation = activation1d.Activation1d(activation=SnakeBeta(10), fused=True).cuda()
|
||||
fused_activation_output = fused_anti_alias_activation(data)
|
||||
|
||||
torch_anti_alias_activation = activation1d.Activation1d(activation=SnakeBeta(10), fused=False).cuda()
|
||||
torch_activation_output = torch_anti_alias_activation(data)
|
||||
|
||||
test_result = (fused_activation_output - torch_activation_output).abs()
|
||||
|
||||
while test_result.dim() != 1:
|
||||
test_result = test_result.mean(dim=-1)
|
||||
|
||||
diff = test_result.mean(dim=-1)
|
||||
|
||||
if diff <= 1e-3:
|
||||
print(
|
||||
f"\n[Success] test_fused_anti_alias_activation"
|
||||
f"\n > mean_difference={diff}"
|
||||
f"\n > fused_values={fused_activation_output[-1][-1][:].tolist()}"
|
||||
f"\n > torch_values={torch_activation_output[-1][-1][:].tolist()}"
|
||||
)
|
||||
else:
|
||||
print(
|
||||
f"\n[Fail] test_fused_anti_alias_activation"
|
||||
f"\n > mean_difference={diff}, "
|
||||
f"\n > fused_values={fused_activation_output[-1][-1][:].tolist()}, "
|
||||
f"\n > torch_values={torch_activation_output[-1][-1][:].tolist()}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from alias_free_activation.cuda import load
|
||||
|
||||
load.load()
|
||||
test_load_fused_kernels()
|
||||
test_anti_alias_activation()
|
||||
@ -1,215 +0,0 @@
|
||||
# Copyright (c) 2024 NVIDIA CORPORATION.
|
||||
# Licensed under the MIT license.
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
# to import modules from parent_dir
|
||||
parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
|
||||
sys.path.append(parent_dir)
|
||||
|
||||
import torch
|
||||
import json
|
||||
from env import AttrDict
|
||||
from bigvgan import BigVGAN
|
||||
from time import time
|
||||
from tqdm import tqdm
|
||||
from meldataset import mel_spectrogram, MAX_WAV_VALUE
|
||||
from scipy.io.wavfile import write
|
||||
import numpy as np
|
||||
|
||||
import argparse
|
||||
|
||||
torch.backends.cudnn.benchmark = True
|
||||
|
||||
# For easier debugging
|
||||
torch.set_printoptions(linewidth=200, threshold=10_000)
|
||||
|
||||
|
||||
def generate_soundwave(duration=5.0, sr=24000):
|
||||
t = np.linspace(0, duration, int(sr * duration), False, dtype=np.float32)
|
||||
|
||||
modulation = np.sin(2 * np.pi * t / duration)
|
||||
|
||||
min_freq = 220
|
||||
max_freq = 1760
|
||||
frequencies = min_freq + (max_freq - min_freq) * (modulation + 1) / 2
|
||||
soundwave = np.sin(2 * np.pi * frequencies * t)
|
||||
|
||||
soundwave = soundwave / np.max(np.abs(soundwave)) * 0.95
|
||||
|
||||
return soundwave, sr
|
||||
|
||||
|
||||
def get_mel(x, h):
|
||||
return mel_spectrogram(x, h.n_fft, h.num_mels, h.sampling_rate, h.hop_size, h.win_size, h.fmin, h.fmax)
|
||||
|
||||
|
||||
def load_checkpoint(filepath, device):
|
||||
assert os.path.isfile(filepath)
|
||||
print(f"Loading '{filepath}'")
|
||||
checkpoint_dict = torch.load(filepath, map_location=device)
|
||||
print("Complete.")
|
||||
return checkpoint_dict
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Test script to check CUDA kernel correctness.")
|
||||
parser.add_argument(
|
||||
"--checkpoint_file",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path to the checkpoint file. Assumes config.json exists in the directory.",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
config_file = os.path.join(os.path.split(args.checkpoint_file)[0], "config.json")
|
||||
with open(config_file) as f:
|
||||
config = f.read()
|
||||
json_config = json.loads(config)
|
||||
h = AttrDict({**json_config})
|
||||
|
||||
print("loading plain Pytorch BigVGAN")
|
||||
generator_original = BigVGAN(h).to("cuda")
|
||||
print("loading CUDA kernel BigVGAN with auto-build")
|
||||
generator_cuda_kernel = BigVGAN(h, use_cuda_kernel=True).to("cuda")
|
||||
|
||||
state_dict_g = load_checkpoint(args.checkpoint_file, "cuda")
|
||||
generator_original.load_state_dict(state_dict_g["generator"])
|
||||
generator_cuda_kernel.load_state_dict(state_dict_g["generator"])
|
||||
|
||||
generator_original.remove_weight_norm()
|
||||
generator_original.eval()
|
||||
generator_cuda_kernel.remove_weight_norm()
|
||||
generator_cuda_kernel.eval()
|
||||
|
||||
# define number of samples and length of mel frame to benchmark
|
||||
num_sample = 10
|
||||
num_mel_frame = 16384
|
||||
|
||||
# CUDA kernel correctness check
|
||||
diff = 0.0
|
||||
for i in tqdm(range(num_sample)):
|
||||
# Random mel
|
||||
data = torch.rand((1, h.num_mels, num_mel_frame), device="cuda")
|
||||
|
||||
with torch.inference_mode():
|
||||
audio_original = generator_original(data)
|
||||
|
||||
with torch.inference_mode():
|
||||
audio_cuda_kernel = generator_cuda_kernel(data)
|
||||
|
||||
# Both outputs should be (almost) the same
|
||||
test_result = (audio_original - audio_cuda_kernel).abs()
|
||||
diff += test_result.mean(dim=-1).item()
|
||||
|
||||
diff /= num_sample
|
||||
if diff <= 2e-3: # We can expect a small difference (~1e-3) which does not affect perceptual quality
|
||||
print(
|
||||
f"\n[Success] test CUDA fused vs. plain torch BigVGAN inference"
|
||||
f"\n > mean_difference={diff}"
|
||||
f"\n > fused_values={audio_cuda_kernel[-1][-1][-30:].tolist()}"
|
||||
f"\n > torch_values={audio_original[-1][-1][-30:].tolist()}"
|
||||
)
|
||||
else:
|
||||
print(
|
||||
f"\n[Fail] test CUDA fused vs. plain torch BigVGAN inference"
|
||||
f"\n > mean_difference={diff}"
|
||||
f"\n > fused_values={audio_cuda_kernel[-1][-1][-30:].tolist()}, "
|
||||
f"\n > torch_values={audio_original[-1][-1][-30:].tolist()}"
|
||||
)
|
||||
|
||||
del data, audio_original, audio_cuda_kernel
|
||||
|
||||
# Variables for tracking total time and VRAM usage
|
||||
toc_total_original = 0
|
||||
toc_total_cuda_kernel = 0
|
||||
vram_used_original_total = 0
|
||||
vram_used_cuda_kernel_total = 0
|
||||
audio_length_total = 0
|
||||
|
||||
# Measure Original inference in isolation
|
||||
for i in tqdm(range(num_sample)):
|
||||
torch.cuda.reset_peak_memory_stats(device="cuda")
|
||||
data = torch.rand((1, h.num_mels, num_mel_frame), device="cuda")
|
||||
torch.cuda.synchronize()
|
||||
tic = time()
|
||||
with torch.inference_mode():
|
||||
audio_original = generator_original(data)
|
||||
torch.cuda.synchronize()
|
||||
toc = time() - tic
|
||||
toc_total_original += toc
|
||||
|
||||
vram_used_original_total += torch.cuda.max_memory_allocated(device="cuda")
|
||||
|
||||
del data, audio_original
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# Measure CUDA kernel inference in isolation
|
||||
for i in tqdm(range(num_sample)):
|
||||
torch.cuda.reset_peak_memory_stats(device="cuda")
|
||||
data = torch.rand((1, h.num_mels, num_mel_frame), device="cuda")
|
||||
torch.cuda.synchronize()
|
||||
tic = time()
|
||||
with torch.inference_mode():
|
||||
audio_cuda_kernel = generator_cuda_kernel(data)
|
||||
torch.cuda.synchronize()
|
||||
toc = time() - tic
|
||||
toc_total_cuda_kernel += toc
|
||||
|
||||
audio_length_total += audio_cuda_kernel.shape[-1]
|
||||
|
||||
vram_used_cuda_kernel_total += torch.cuda.max_memory_allocated(device="cuda")
|
||||
|
||||
del data, audio_cuda_kernel
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# Calculate metrics
|
||||
audio_second = audio_length_total / h.sampling_rate
|
||||
khz_original = audio_length_total / toc_total_original / 1000
|
||||
khz_cuda_kernel = audio_length_total / toc_total_cuda_kernel / 1000
|
||||
vram_used_original_gb = vram_used_original_total / num_sample / (1024**3)
|
||||
vram_used_cuda_kernel_gb = vram_used_cuda_kernel_total / num_sample / (1024**3)
|
||||
|
||||
# Print results
|
||||
print(
|
||||
f"Original BigVGAN: took {toc_total_original:.2f} seconds to generate {audio_second:.2f} seconds of audio, {khz_original:.1f}kHz, {audio_second / toc_total_original:.1f} faster than realtime, VRAM used {vram_used_original_gb:.1f} GB"
|
||||
)
|
||||
print(
|
||||
f"CUDA kernel BigVGAN: took {toc_total_cuda_kernel:.2f} seconds to generate {audio_second:.2f} seconds of audio, {khz_cuda_kernel:.1f}kHz, {audio_second / toc_total_cuda_kernel:.1f} faster than realtime, VRAM used {vram_used_cuda_kernel_gb:.1f} GB"
|
||||
)
|
||||
print(f"speedup of CUDA kernel: {khz_cuda_kernel / khz_original}")
|
||||
print(f"VRAM saving of CUDA kernel: {vram_used_original_gb / vram_used_cuda_kernel_gb}")
|
||||
|
||||
# Use artificial sine waves for inference test
|
||||
audio_real, sr = generate_soundwave(duration=5.0, sr=h.sampling_rate)
|
||||
audio_real = torch.tensor(audio_real).to("cuda")
|
||||
# Compute mel spectrogram from the ground truth audio
|
||||
x = get_mel(audio_real.unsqueeze(0), h)
|
||||
|
||||
with torch.inference_mode():
|
||||
y_g_hat_original = generator_original(x)
|
||||
y_g_hat_cuda_kernel = generator_cuda_kernel(x)
|
||||
|
||||
audio_real = audio_real.squeeze()
|
||||
audio_real = audio_real * MAX_WAV_VALUE
|
||||
audio_real = audio_real.cpu().numpy().astype("int16")
|
||||
|
||||
audio_original = y_g_hat_original.squeeze()
|
||||
audio_original = audio_original * MAX_WAV_VALUE
|
||||
audio_original = audio_original.cpu().numpy().astype("int16")
|
||||
|
||||
audio_cuda_kernel = y_g_hat_cuda_kernel.squeeze()
|
||||
audio_cuda_kernel = audio_cuda_kernel * MAX_WAV_VALUE
|
||||
audio_cuda_kernel = audio_cuda_kernel.cpu().numpy().astype("int16")
|
||||
|
||||
os.makedirs("tmp", exist_ok=True)
|
||||
output_file_real = os.path.join("tmp", "audio_real.wav")
|
||||
output_file_original = os.path.join("tmp", "audio_generated_original.wav")
|
||||
output_file_cuda_kernel = os.path.join("tmp", "audio_generated_cuda_kernel.wav")
|
||||
write(output_file_real, h.sampling_rate, audio_real)
|
||||
write(output_file_original, h.sampling_rate, audio_original)
|
||||
write(output_file_cuda_kernel, h.sampling_rate, audio_cuda_kernel)
|
||||
print("Example generated audios of original vs. fused CUDA kernel written to tmp!")
|
||||
print("Done")
|
||||
@ -1,716 +0,0 @@
|
||||
# Copyright (c) 2024 NVIDIA CORPORATION.
|
||||
# Licensed under the MIT license.
|
||||
|
||||
# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
|
||||
|
||||
import warnings
|
||||
|
||||
warnings.simplefilter(action="ignore", category=FutureWarning)
|
||||
import itertools
|
||||
import os
|
||||
import time
|
||||
import argparse
|
||||
import json
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
from torch.utils.data import DistributedSampler, DataLoader
|
||||
import torch.multiprocessing as mp
|
||||
from torch.distributed import init_process_group
|
||||
from torch.nn.parallel import DistributedDataParallel
|
||||
from env import AttrDict, build_env
|
||||
from meldataset import MelDataset, mel_spectrogram, get_dataset_filelist, MAX_WAV_VALUE
|
||||
|
||||
from bigvgan import BigVGAN
|
||||
from discriminators import (
|
||||
MultiPeriodDiscriminator,
|
||||
MultiResolutionDiscriminator,
|
||||
MultiBandDiscriminator,
|
||||
MultiScaleSubbandCQTDiscriminator,
|
||||
)
|
||||
from loss import (
|
||||
feature_loss,
|
||||
generator_loss,
|
||||
discriminator_loss,
|
||||
MultiScaleMelSpectrogramLoss,
|
||||
)
|
||||
|
||||
from utils import (
|
||||
plot_spectrogram,
|
||||
plot_spectrogram_clipped,
|
||||
scan_checkpoint,
|
||||
load_checkpoint,
|
||||
save_checkpoint,
|
||||
save_audio,
|
||||
)
|
||||
import torchaudio as ta
|
||||
from pesq import pesq
|
||||
from tqdm import tqdm
|
||||
import auraloss
|
||||
|
||||
torch.backends.cudnn.benchmark = False
|
||||
|
||||
|
||||
def train(rank, a, h):
|
||||
if h.num_gpus > 1:
|
||||
# initialize distributed
|
||||
init_process_group(
|
||||
backend=h.dist_config["dist_backend"],
|
||||
init_method=h.dist_config["dist_url"],
|
||||
world_size=h.dist_config["world_size"] * h.num_gpus,
|
||||
rank=rank,
|
||||
)
|
||||
|
||||
# Set seed and device
|
||||
torch.cuda.manual_seed(h.seed)
|
||||
torch.cuda.set_device(rank)
|
||||
device = torch.device(f"cuda:{rank:d}")
|
||||
|
||||
# Define BigVGAN generator
|
||||
generator = BigVGAN(h).to(device)
|
||||
|
||||
# Define discriminators. MPD is used by default
|
||||
mpd = MultiPeriodDiscriminator(h).to(device)
|
||||
|
||||
# Define additional discriminators. BigVGAN-v1 uses UnivNet's MRD as default
|
||||
# New in BigVGAN-v2: option to switch to new discriminators: MultiBandDiscriminator / MultiScaleSubbandCQTDiscriminator
|
||||
if h.get("use_mbd_instead_of_mrd", False): # Switch to MBD
|
||||
print("[INFO] using MultiBandDiscriminator of BigVGAN-v2 instead of MultiResolutionDiscriminator")
|
||||
# Variable name is kept as "mrd" for backward compatibility & minimal code change
|
||||
mrd = MultiBandDiscriminator(h).to(device)
|
||||
elif h.get("use_cqtd_instead_of_mrd", False): # Switch to CQTD
|
||||
print("[INFO] using MultiScaleSubbandCQTDiscriminator of BigVGAN-v2 instead of MultiResolutionDiscriminator")
|
||||
mrd = MultiScaleSubbandCQTDiscriminator(h).to(device)
|
||||
else: # Fallback to original MRD in BigVGAN-v1
|
||||
mrd = MultiResolutionDiscriminator(h).to(device)
|
||||
|
||||
# New in BigVGAN-v2: option to switch to multi-scale L1 mel loss
|
||||
if h.get("use_multiscale_melloss", False):
|
||||
print("[INFO] using multi-scale Mel l1 loss of BigVGAN-v2 instead of the original single-scale loss")
|
||||
fn_mel_loss_multiscale = MultiScaleMelSpectrogramLoss(
|
||||
sampling_rate=h.sampling_rate
|
||||
) # NOTE: accepts waveform as input
|
||||
else:
|
||||
fn_mel_loss_singlescale = F.l1_loss
|
||||
|
||||
# Print the model & number of parameters, and create or scan the latest checkpoint from checkpoints directory
|
||||
if rank == 0:
|
||||
print(generator)
|
||||
print(mpd)
|
||||
print(mrd)
|
||||
print(f"Generator params: {sum(p.numel() for p in generator.parameters())}")
|
||||
print(f"Discriminator mpd params: {sum(p.numel() for p in mpd.parameters())}")
|
||||
print(f"Discriminator mrd params: {sum(p.numel() for p in mrd.parameters())}")
|
||||
os.makedirs(a.checkpoint_path, exist_ok=True)
|
||||
print(f"Checkpoints directory: {a.checkpoint_path}")
|
||||
|
||||
if os.path.isdir(a.checkpoint_path):
|
||||
# New in v2.1: If the step prefix pattern-based checkpoints are not found, also check for renamed files in Hugging Face Hub to resume training
|
||||
cp_g = scan_checkpoint(a.checkpoint_path, prefix="g_", renamed_file="bigvgan_generator.pt")
|
||||
cp_do = scan_checkpoint(
|
||||
a.checkpoint_path,
|
||||
prefix="do_",
|
||||
renamed_file="bigvgan_discriminator_optimizer.pt",
|
||||
)
|
||||
|
||||
# Load the latest checkpoint if exists
|
||||
steps = 0
|
||||
if cp_g is None or cp_do is None:
|
||||
state_dict_do = None
|
||||
last_epoch = -1
|
||||
else:
|
||||
state_dict_g = load_checkpoint(cp_g, device)
|
||||
state_dict_do = load_checkpoint(cp_do, device)
|
||||
generator.load_state_dict(state_dict_g["generator"])
|
||||
mpd.load_state_dict(state_dict_do["mpd"])
|
||||
mrd.load_state_dict(state_dict_do["mrd"])
|
||||
steps = state_dict_do["steps"] + 1
|
||||
last_epoch = state_dict_do["epoch"]
|
||||
|
||||
# Initialize DDP, optimizers, and schedulers
|
||||
if h.num_gpus > 1:
|
||||
generator = DistributedDataParallel(generator, device_ids=[rank]).to(device)
|
||||
mpd = DistributedDataParallel(mpd, device_ids=[rank]).to(device)
|
||||
mrd = DistributedDataParallel(mrd, device_ids=[rank]).to(device)
|
||||
|
||||
optim_g = torch.optim.AdamW(generator.parameters(), h.learning_rate, betas=[h.adam_b1, h.adam_b2])
|
||||
optim_d = torch.optim.AdamW(
|
||||
itertools.chain(mrd.parameters(), mpd.parameters()),
|
||||
h.learning_rate,
|
||||
betas=[h.adam_b1, h.adam_b2],
|
||||
)
|
||||
|
||||
if state_dict_do is not None:
|
||||
optim_g.load_state_dict(state_dict_do["optim_g"])
|
||||
optim_d.load_state_dict(state_dict_do["optim_d"])
|
||||
|
||||
scheduler_g = torch.optim.lr_scheduler.ExponentialLR(optim_g, gamma=h.lr_decay, last_epoch=last_epoch)
|
||||
scheduler_d = torch.optim.lr_scheduler.ExponentialLR(optim_d, gamma=h.lr_decay, last_epoch=last_epoch)
|
||||
|
||||
# Define training and validation datasets
|
||||
|
||||
"""
|
||||
unseen_validation_filelist will contain sample filepaths outside the seen training & validation dataset
|
||||
Example: trained on LibriTTS, validate on VCTK
|
||||
"""
|
||||
training_filelist, validation_filelist, list_unseen_validation_filelist = get_dataset_filelist(a)
|
||||
|
||||
trainset = MelDataset(
|
||||
training_filelist,
|
||||
h,
|
||||
h.segment_size,
|
||||
h.n_fft,
|
||||
h.num_mels,
|
||||
h.hop_size,
|
||||
h.win_size,
|
||||
h.sampling_rate,
|
||||
h.fmin,
|
||||
h.fmax,
|
||||
shuffle=False if h.num_gpus > 1 else True,
|
||||
fmax_loss=h.fmax_for_loss,
|
||||
device=device,
|
||||
fine_tuning=a.fine_tuning,
|
||||
base_mels_path=a.input_mels_dir,
|
||||
is_seen=True,
|
||||
)
|
||||
|
||||
train_sampler = DistributedSampler(trainset) if h.num_gpus > 1 else None
|
||||
|
||||
train_loader = DataLoader(
|
||||
trainset,
|
||||
num_workers=h.num_workers,
|
||||
shuffle=False,
|
||||
sampler=train_sampler,
|
||||
batch_size=h.batch_size,
|
||||
pin_memory=True,
|
||||
drop_last=True,
|
||||
)
|
||||
|
||||
if rank == 0:
|
||||
validset = MelDataset(
|
||||
validation_filelist,
|
||||
h,
|
||||
h.segment_size,
|
||||
h.n_fft,
|
||||
h.num_mels,
|
||||
h.hop_size,
|
||||
h.win_size,
|
||||
h.sampling_rate,
|
||||
h.fmin,
|
||||
h.fmax,
|
||||
False,
|
||||
False,
|
||||
fmax_loss=h.fmax_for_loss,
|
||||
device=device,
|
||||
fine_tuning=a.fine_tuning,
|
||||
base_mels_path=a.input_mels_dir,
|
||||
is_seen=True,
|
||||
)
|
||||
validation_loader = DataLoader(
|
||||
validset,
|
||||
num_workers=1,
|
||||
shuffle=False,
|
||||
sampler=None,
|
||||
batch_size=1,
|
||||
pin_memory=True,
|
||||
drop_last=True,
|
||||
)
|
||||
|
||||
list_unseen_validset = []
|
||||
list_unseen_validation_loader = []
|
||||
for i in range(len(list_unseen_validation_filelist)):
|
||||
unseen_validset = MelDataset(
|
||||
list_unseen_validation_filelist[i],
|
||||
h,
|
||||
h.segment_size,
|
||||
h.n_fft,
|
||||
h.num_mels,
|
||||
h.hop_size,
|
||||
h.win_size,
|
||||
h.sampling_rate,
|
||||
h.fmin,
|
||||
h.fmax,
|
||||
False,
|
||||
False,
|
||||
fmax_loss=h.fmax_for_loss,
|
||||
device=device,
|
||||
fine_tuning=a.fine_tuning,
|
||||
base_mels_path=a.input_mels_dir,
|
||||
is_seen=False,
|
||||
)
|
||||
unseen_validation_loader = DataLoader(
|
||||
unseen_validset,
|
||||
num_workers=1,
|
||||
shuffle=False,
|
||||
sampler=None,
|
||||
batch_size=1,
|
||||
pin_memory=True,
|
||||
drop_last=True,
|
||||
)
|
||||
list_unseen_validset.append(unseen_validset)
|
||||
list_unseen_validation_loader.append(unseen_validation_loader)
|
||||
|
||||
# Tensorboard logger
|
||||
sw = SummaryWriter(os.path.join(a.checkpoint_path, "logs"))
|
||||
if a.save_audio: # Also save audio to disk if --save_audio is set to True
|
||||
os.makedirs(os.path.join(a.checkpoint_path, "samples"), exist_ok=True)
|
||||
|
||||
"""
|
||||
Validation loop, "mode" parameter is automatically defined as (seen or unseen)_(name of the dataset).
|
||||
If the name of the dataset contains "nonspeech", it skips PESQ calculation to prevent errors
|
||||
"""
|
||||
|
||||
def validate(rank, a, h, loader, mode="seen"):
|
||||
assert rank == 0, "validate should only run on rank=0"
|
||||
generator.eval()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
val_err_tot = 0
|
||||
val_pesq_tot = 0
|
||||
val_mrstft_tot = 0
|
||||
|
||||
# Modules for evaluation metrics
|
||||
pesq_resampler = ta.transforms.Resample(h.sampling_rate, 16000).cuda()
|
||||
loss_mrstft = auraloss.freq.MultiResolutionSTFTLoss(device="cuda")
|
||||
|
||||
if a.save_audio: # Also save audio to disk if --save_audio is set to True
|
||||
os.makedirs(
|
||||
os.path.join(a.checkpoint_path, "samples", f"gt_{mode}"),
|
||||
exist_ok=True,
|
||||
)
|
||||
os.makedirs(
|
||||
os.path.join(a.checkpoint_path, "samples", f"{mode}_{steps:08d}"),
|
||||
exist_ok=True,
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
print(f"step {steps} {mode} speaker validation...")
|
||||
|
||||
# Loop over validation set and compute metrics
|
||||
for j, batch in enumerate(tqdm(loader)):
|
||||
x, y, _, y_mel = batch
|
||||
y = y.to(device)
|
||||
if hasattr(generator, "module"):
|
||||
y_g_hat = generator.module(x.to(device))
|
||||
else:
|
||||
y_g_hat = generator(x.to(device))
|
||||
y_mel = y_mel.to(device, non_blocking=True)
|
||||
y_g_hat_mel = mel_spectrogram(
|
||||
y_g_hat.squeeze(1),
|
||||
h.n_fft,
|
||||
h.num_mels,
|
||||
h.sampling_rate,
|
||||
h.hop_size,
|
||||
h.win_size,
|
||||
h.fmin,
|
||||
h.fmax_for_loss,
|
||||
)
|
||||
min_t = min(y_mel.size(-1), y_g_hat_mel.size(-1))
|
||||
val_err_tot += F.l1_loss(y_mel[..., :min_t], y_g_hat_mel[..., :min_t]).item()
|
||||
|
||||
# PESQ calculation. only evaluate PESQ if it's speech signal (nonspeech PESQ will error out)
|
||||
if "nonspeech" not in mode: # Skips if the name of dataset (in mode string) contains "nonspeech"
|
||||
# Resample to 16000 for pesq
|
||||
y_16k = pesq_resampler(y)
|
||||
y_g_hat_16k = pesq_resampler(y_g_hat.squeeze(1))
|
||||
y_int_16k = (y_16k[0] * MAX_WAV_VALUE).short().cpu().numpy()
|
||||
y_g_hat_int_16k = (y_g_hat_16k[0] * MAX_WAV_VALUE).short().cpu().numpy()
|
||||
val_pesq_tot += pesq(16000, y_int_16k, y_g_hat_int_16k, "wb")
|
||||
|
||||
# MRSTFT calculation
|
||||
min_t = min(y.size(-1), y_g_hat.size(-1))
|
||||
val_mrstft_tot += loss_mrstft(y_g_hat[..., :min_t], y[..., :min_t]).item()
|
||||
|
||||
# Log audio and figures to Tensorboard
|
||||
if j % a.eval_subsample == 0: # Subsample every nth from validation set
|
||||
if steps >= 0:
|
||||
sw.add_audio(f"gt_{mode}/y_{j}", y[0], steps, h.sampling_rate)
|
||||
if a.save_audio: # Also save audio to disk if --save_audio is set to True
|
||||
save_audio(
|
||||
y[0],
|
||||
os.path.join(
|
||||
a.checkpoint_path,
|
||||
"samples",
|
||||
f"gt_{mode}",
|
||||
f"{j:04d}.wav",
|
||||
),
|
||||
h.sampling_rate,
|
||||
)
|
||||
sw.add_figure(
|
||||
f"gt_{mode}/y_spec_{j}",
|
||||
plot_spectrogram(x[0]),
|
||||
steps,
|
||||
)
|
||||
|
||||
sw.add_audio(
|
||||
f"generated_{mode}/y_hat_{j}",
|
||||
y_g_hat[0],
|
||||
steps,
|
||||
h.sampling_rate,
|
||||
)
|
||||
if a.save_audio: # Also save audio to disk if --save_audio is set to True
|
||||
save_audio(
|
||||
y_g_hat[0, 0],
|
||||
os.path.join(
|
||||
a.checkpoint_path,
|
||||
"samples",
|
||||
f"{mode}_{steps:08d}",
|
||||
f"{j:04d}.wav",
|
||||
),
|
||||
h.sampling_rate,
|
||||
)
|
||||
# Spectrogram of synthesized audio
|
||||
y_hat_spec = mel_spectrogram(
|
||||
y_g_hat.squeeze(1),
|
||||
h.n_fft,
|
||||
h.num_mels,
|
||||
h.sampling_rate,
|
||||
h.hop_size,
|
||||
h.win_size,
|
||||
h.fmin,
|
||||
h.fmax,
|
||||
)
|
||||
sw.add_figure(
|
||||
f"generated_{mode}/y_hat_spec_{j}",
|
||||
plot_spectrogram(y_hat_spec.squeeze(0).cpu().numpy()),
|
||||
steps,
|
||||
)
|
||||
|
||||
"""
|
||||
Visualization of spectrogram difference between GT and synthesized audio, difference higher than 1 is clipped for better visualization.
|
||||
"""
|
||||
spec_delta = torch.clamp(
|
||||
torch.abs(x[0] - y_hat_spec.squeeze(0).cpu()),
|
||||
min=1e-6,
|
||||
max=1.0,
|
||||
)
|
||||
sw.add_figure(
|
||||
f"delta_dclip1_{mode}/spec_{j}",
|
||||
plot_spectrogram_clipped(spec_delta.numpy(), clip_max=1.0),
|
||||
steps,
|
||||
)
|
||||
|
||||
val_err = val_err_tot / (j + 1)
|
||||
val_pesq = val_pesq_tot / (j + 1)
|
||||
val_mrstft = val_mrstft_tot / (j + 1)
|
||||
# Log evaluation metrics to Tensorboard
|
||||
sw.add_scalar(f"validation_{mode}/mel_spec_error", val_err, steps)
|
||||
sw.add_scalar(f"validation_{mode}/pesq", val_pesq, steps)
|
||||
sw.add_scalar(f"validation_{mode}/mrstft", val_mrstft, steps)
|
||||
|
||||
generator.train()
|
||||
|
||||
# If the checkpoint is loaded, start with validation loop
|
||||
if steps != 0 and rank == 0 and not a.debug:
|
||||
if not a.skip_seen:
|
||||
validate(
|
||||
rank,
|
||||
a,
|
||||
h,
|
||||
validation_loader,
|
||||
mode=f"seen_{train_loader.dataset.name}",
|
||||
)
|
||||
for i in range(len(list_unseen_validation_loader)):
|
||||
validate(
|
||||
rank,
|
||||
a,
|
||||
h,
|
||||
list_unseen_validation_loader[i],
|
||||
mode=f"unseen_{list_unseen_validation_loader[i].dataset.name}",
|
||||
)
|
||||
# Exit the script if --evaluate is set to True
|
||||
if a.evaluate:
|
||||
exit()
|
||||
|
||||
# Main training loop
|
||||
generator.train()
|
||||
mpd.train()
|
||||
mrd.train()
|
||||
for epoch in range(max(0, last_epoch), a.training_epochs):
|
||||
if rank == 0:
|
||||
start = time.time()
|
||||
print(f"Epoch: {epoch + 1}")
|
||||
|
||||
if h.num_gpus > 1:
|
||||
train_sampler.set_epoch(epoch)
|
||||
|
||||
for i, batch in enumerate(train_loader):
|
||||
if rank == 0:
|
||||
start_b = time.time()
|
||||
x, y, _, y_mel = batch
|
||||
|
||||
x = x.to(device, non_blocking=True)
|
||||
y = y.to(device, non_blocking=True)
|
||||
y_mel = y_mel.to(device, non_blocking=True)
|
||||
y = y.unsqueeze(1)
|
||||
|
||||
y_g_hat = generator(x)
|
||||
y_g_hat_mel = mel_spectrogram(
|
||||
y_g_hat.squeeze(1),
|
||||
h.n_fft,
|
||||
h.num_mels,
|
||||
h.sampling_rate,
|
||||
h.hop_size,
|
||||
h.win_size,
|
||||
h.fmin,
|
||||
h.fmax_for_loss,
|
||||
)
|
||||
|
||||
optim_d.zero_grad()
|
||||
|
||||
# MPD
|
||||
y_df_hat_r, y_df_hat_g, _, _ = mpd(y, y_g_hat.detach())
|
||||
loss_disc_f, losses_disc_f_r, losses_disc_f_g = discriminator_loss(y_df_hat_r, y_df_hat_g)
|
||||
|
||||
# MRD
|
||||
y_ds_hat_r, y_ds_hat_g, _, _ = mrd(y, y_g_hat.detach())
|
||||
loss_disc_s, losses_disc_s_r, losses_disc_s_g = discriminator_loss(y_ds_hat_r, y_ds_hat_g)
|
||||
|
||||
loss_disc_all = loss_disc_s + loss_disc_f
|
||||
|
||||
# Set clip_grad_norm value
|
||||
clip_grad_norm = h.get("clip_grad_norm", 1000.0) # Default to 1000
|
||||
|
||||
# Whether to freeze D for initial training steps
|
||||
if steps >= a.freeze_step:
|
||||
loss_disc_all.backward()
|
||||
grad_norm_mpd = torch.nn.utils.clip_grad_norm_(mpd.parameters(), clip_grad_norm)
|
||||
grad_norm_mrd = torch.nn.utils.clip_grad_norm_(mrd.parameters(), clip_grad_norm)
|
||||
optim_d.step()
|
||||
else:
|
||||
print(f"[WARNING] skipping D training for the first {a.freeze_step} steps")
|
||||
grad_norm_mpd = 0.0
|
||||
grad_norm_mrd = 0.0
|
||||
|
||||
# Generator
|
||||
optim_g.zero_grad()
|
||||
|
||||
# L1 Mel-Spectrogram Loss
|
||||
lambda_melloss = h.get("lambda_melloss", 45.0) # Defaults to 45 in BigVGAN-v1 if not set
|
||||
if h.get("use_multiscale_melloss", False): # uses wav <y, y_g_hat> for loss
|
||||
loss_mel = fn_mel_loss_multiscale(y, y_g_hat) * lambda_melloss
|
||||
else: # Uses mel <y_mel, y_g_hat_mel> for loss
|
||||
loss_mel = fn_mel_loss_singlescale(y_mel, y_g_hat_mel) * lambda_melloss
|
||||
|
||||
# MPD loss
|
||||
y_df_hat_r, y_df_hat_g, fmap_f_r, fmap_f_g = mpd(y, y_g_hat)
|
||||
loss_fm_f = feature_loss(fmap_f_r, fmap_f_g)
|
||||
loss_gen_f, losses_gen_f = generator_loss(y_df_hat_g)
|
||||
|
||||
# MRD loss
|
||||
y_ds_hat_r, y_ds_hat_g, fmap_s_r, fmap_s_g = mrd(y, y_g_hat)
|
||||
loss_fm_s = feature_loss(fmap_s_r, fmap_s_g)
|
||||
loss_gen_s, losses_gen_s = generator_loss(y_ds_hat_g)
|
||||
|
||||
if steps >= a.freeze_step:
|
||||
loss_gen_all = loss_gen_s + loss_gen_f + loss_fm_s + loss_fm_f + loss_mel
|
||||
else:
|
||||
print(f"[WARNING] using regression loss only for G for the first {a.freeze_step} steps")
|
||||
loss_gen_all = loss_mel
|
||||
|
||||
loss_gen_all.backward()
|
||||
grad_norm_g = torch.nn.utils.clip_grad_norm_(generator.parameters(), clip_grad_norm)
|
||||
optim_g.step()
|
||||
|
||||
if rank == 0:
|
||||
# STDOUT logging
|
||||
if steps % a.stdout_interval == 0:
|
||||
mel_error = loss_mel.item() / lambda_melloss # Log training mel regression loss to stdout
|
||||
print(
|
||||
f"Steps: {steps:d}, "
|
||||
f"Gen Loss Total: {loss_gen_all:4.3f}, "
|
||||
f"Mel Error: {mel_error:4.3f}, "
|
||||
f"s/b: {time.time() - start_b:4.3f} "
|
||||
f"lr: {optim_g.param_groups[0]['lr']:4.7f} "
|
||||
f"grad_norm_g: {grad_norm_g:4.3f}"
|
||||
)
|
||||
|
||||
# Checkpointing
|
||||
if steps % a.checkpoint_interval == 0 and steps != 0:
|
||||
checkpoint_path = f"{a.checkpoint_path}/g_{steps:08d}"
|
||||
save_checkpoint(
|
||||
checkpoint_path,
|
||||
{"generator": (generator.module if h.num_gpus > 1 else generator).state_dict()},
|
||||
)
|
||||
checkpoint_path = f"{a.checkpoint_path}/do_{steps:08d}"
|
||||
save_checkpoint(
|
||||
checkpoint_path,
|
||||
{
|
||||
"mpd": (mpd.module if h.num_gpus > 1 else mpd).state_dict(),
|
||||
"mrd": (mrd.module if h.num_gpus > 1 else mrd).state_dict(),
|
||||
"optim_g": optim_g.state_dict(),
|
||||
"optim_d": optim_d.state_dict(),
|
||||
"steps": steps,
|
||||
"epoch": epoch,
|
||||
},
|
||||
)
|
||||
|
||||
# Tensorboard summary logging
|
||||
if steps % a.summary_interval == 0:
|
||||
mel_error = loss_mel.item() / lambda_melloss # Log training mel regression loss to tensorboard
|
||||
sw.add_scalar("training/gen_loss_total", loss_gen_all.item(), steps)
|
||||
sw.add_scalar("training/mel_spec_error", mel_error, steps)
|
||||
sw.add_scalar("training/fm_loss_mpd", loss_fm_f.item(), steps)
|
||||
sw.add_scalar("training/gen_loss_mpd", loss_gen_f.item(), steps)
|
||||
sw.add_scalar("training/disc_loss_mpd", loss_disc_f.item(), steps)
|
||||
sw.add_scalar("training/grad_norm_mpd", grad_norm_mpd, steps)
|
||||
sw.add_scalar("training/fm_loss_mrd", loss_fm_s.item(), steps)
|
||||
sw.add_scalar("training/gen_loss_mrd", loss_gen_s.item(), steps)
|
||||
sw.add_scalar("training/disc_loss_mrd", loss_disc_s.item(), steps)
|
||||
sw.add_scalar("training/grad_norm_mrd", grad_norm_mrd, steps)
|
||||
sw.add_scalar("training/grad_norm_g", grad_norm_g, steps)
|
||||
sw.add_scalar("training/learning_rate_d", scheduler_d.get_last_lr()[0], steps)
|
||||
sw.add_scalar("training/learning_rate_g", scheduler_g.get_last_lr()[0], steps)
|
||||
sw.add_scalar("training/epoch", epoch + 1, steps)
|
||||
|
||||
# Validation
|
||||
if steps % a.validation_interval == 0:
|
||||
# Plot training input x so far used
|
||||
for i_x in range(x.shape[0]):
|
||||
sw.add_figure(
|
||||
f"training_input/x_{i_x}",
|
||||
plot_spectrogram(x[i_x].cpu()),
|
||||
steps,
|
||||
)
|
||||
sw.add_audio(
|
||||
f"training_input/y_{i_x}",
|
||||
y[i_x][0],
|
||||
steps,
|
||||
h.sampling_rate,
|
||||
)
|
||||
|
||||
# Seen and unseen speakers validation loops
|
||||
if not a.debug and steps != 0:
|
||||
validate(
|
||||
rank,
|
||||
a,
|
||||
h,
|
||||
validation_loader,
|
||||
mode=f"seen_{train_loader.dataset.name}",
|
||||
)
|
||||
for i in range(len(list_unseen_validation_loader)):
|
||||
validate(
|
||||
rank,
|
||||
a,
|
||||
h,
|
||||
list_unseen_validation_loader[i],
|
||||
mode=f"unseen_{list_unseen_validation_loader[i].dataset.name}",
|
||||
)
|
||||
steps += 1
|
||||
|
||||
# BigVGAN-v2 learning rate scheduler is changed from epoch-level to step-level
|
||||
scheduler_g.step()
|
||||
scheduler_d.step()
|
||||
|
||||
if rank == 0:
|
||||
print(f"Time taken for epoch {epoch + 1} is {int(time.time() - start)} sec\n")
|
||||
|
||||
|
||||
def main():
|
||||
print("Initializing Training Process..")
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument("--group_name", default=None)
|
||||
|
||||
parser.add_argument("--input_wavs_dir", default="LibriTTS")
|
||||
parser.add_argument("--input_mels_dir", default="ft_dataset")
|
||||
parser.add_argument("--input_training_file", default="tests/LibriTTS/train-full.txt")
|
||||
parser.add_argument("--input_validation_file", default="tests/LibriTTS/val-full.txt")
|
||||
|
||||
parser.add_argument(
|
||||
"--list_input_unseen_wavs_dir",
|
||||
nargs="+",
|
||||
default=["tests/LibriTTS", "tests/LibriTTS"],
|
||||
)
|
||||
parser.add_argument(
|
||||
"--list_input_unseen_validation_file",
|
||||
nargs="+",
|
||||
default=["tests/LibriTTS/dev-clean.txt", "tests/LibriTTS/dev-other.txt"],
|
||||
)
|
||||
|
||||
parser.add_argument("--checkpoint_path", default="exp/bigvgan")
|
||||
parser.add_argument("--config", default="")
|
||||
|
||||
parser.add_argument("--training_epochs", default=100000, type=int)
|
||||
parser.add_argument("--stdout_interval", default=5, type=int)
|
||||
parser.add_argument("--checkpoint_interval", default=50000, type=int)
|
||||
parser.add_argument("--summary_interval", default=100, type=int)
|
||||
parser.add_argument("--validation_interval", default=50000, type=int)
|
||||
|
||||
parser.add_argument(
|
||||
"--freeze_step",
|
||||
default=0,
|
||||
type=int,
|
||||
help="freeze D for the first specified steps. G only uses regression loss for these steps.",
|
||||
)
|
||||
|
||||
parser.add_argument("--fine_tuning", default=False, type=bool)
|
||||
|
||||
parser.add_argument(
|
||||
"--debug",
|
||||
default=False,
|
||||
type=bool,
|
||||
help="debug mode. skips validation loop throughout training",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--evaluate",
|
||||
default=False,
|
||||
type=bool,
|
||||
help="only run evaluation from checkpoint and exit",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--eval_subsample",
|
||||
default=5,
|
||||
type=int,
|
||||
help="subsampling during evaluation loop",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--skip_seen",
|
||||
default=False,
|
||||
type=bool,
|
||||
help="skip seen dataset. useful for test set inference",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--save_audio",
|
||||
default=False,
|
||||
type=bool,
|
||||
help="save audio of test set inference to disk",
|
||||
)
|
||||
|
||||
a = parser.parse_args()
|
||||
|
||||
with open(a.config) as f:
|
||||
data = f.read()
|
||||
|
||||
json_config = json.loads(data)
|
||||
h = AttrDict(json_config)
|
||||
|
||||
build_env(a.config, "config.json", a.checkpoint_path)
|
||||
|
||||
torch.manual_seed(h.seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed(h.seed)
|
||||
h.num_gpus = torch.cuda.device_count()
|
||||
h.batch_size = int(h.batch_size / h.num_gpus)
|
||||
print(f"Batch size per GPU: {h.batch_size}")
|
||||
else:
|
||||
pass
|
||||
|
||||
if h.num_gpus > 1:
|
||||
mp.spawn(
|
||||
train,
|
||||
nprocs=h.num_gpus,
|
||||
args=(
|
||||
a,
|
||||
h,
|
||||
),
|
||||
)
|
||||
else:
|
||||
train(0, a, h)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -1,45 +1,5 @@
|
||||
# Adapted from https://github.com/jik876/hifi-gan under the MIT license.
|
||||
# LICENSE is in incl_licenses directory.
|
||||
|
||||
import glob
|
||||
import os
|
||||
import matplotlib
|
||||
import torch
|
||||
from torch.nn.utils import weight_norm
|
||||
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pylab as plt
|
||||
from .meldataset import MAX_WAV_VALUE
|
||||
from scipy.io.wavfile import write
|
||||
|
||||
|
||||
def plot_spectrogram(spectrogram):
|
||||
fig, ax = plt.subplots(figsize=(10, 2))
|
||||
im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none")
|
||||
plt.colorbar(im, ax=ax)
|
||||
|
||||
fig.canvas.draw()
|
||||
plt.close()
|
||||
|
||||
return fig
|
||||
|
||||
|
||||
def plot_spectrogram_clipped(spectrogram, clip_max=2.0):
|
||||
fig, ax = plt.subplots(figsize=(10, 2))
|
||||
im = ax.imshow(
|
||||
spectrogram,
|
||||
aspect="auto",
|
||||
origin="lower",
|
||||
interpolation="none",
|
||||
vmin=1e-6,
|
||||
vmax=clip_max,
|
||||
)
|
||||
plt.colorbar(im, ax=ax)
|
||||
|
||||
fig.canvas.draw()
|
||||
plt.close()
|
||||
|
||||
return fig
|
||||
# LICENSE is in the incl_licenses directory.
|
||||
|
||||
|
||||
def init_weights(m, mean=0.0, std=0.01):
|
||||
@ -48,52 +8,5 @@ def init_weights(m, mean=0.0, std=0.01):
|
||||
m.weight.data.normal_(mean, std)
|
||||
|
||||
|
||||
def apply_weight_norm(m):
|
||||
classname = m.__class__.__name__
|
||||
if classname.find("Conv") != -1:
|
||||
weight_norm(m)
|
||||
|
||||
|
||||
def get_padding(kernel_size, dilation=1):
|
||||
return int((kernel_size * dilation - dilation) / 2)
|
||||
|
||||
|
||||
def load_checkpoint(filepath, device):
|
||||
assert os.path.isfile(filepath)
|
||||
print(f"Loading '{filepath}'")
|
||||
checkpoint_dict = torch.load(filepath, map_location=device)
|
||||
print("Complete.")
|
||||
return checkpoint_dict
|
||||
|
||||
|
||||
def save_checkpoint(filepath, obj):
|
||||
print(f"Saving checkpoint to {filepath}")
|
||||
torch.save(obj, filepath)
|
||||
print("Complete.")
|
||||
|
||||
|
||||
def scan_checkpoint(cp_dir, prefix, renamed_file=None):
|
||||
# Fallback to original scanning logic first
|
||||
pattern = os.path.join(cp_dir, prefix + "????????")
|
||||
cp_list = glob.glob(pattern)
|
||||
|
||||
if len(cp_list) > 0:
|
||||
last_checkpoint_path = sorted(cp_list)[-1]
|
||||
print(f"[INFO] Resuming from checkpoint: '{last_checkpoint_path}'")
|
||||
return last_checkpoint_path
|
||||
|
||||
# If no pattern-based checkpoints are found, check for renamed file
|
||||
if renamed_file:
|
||||
renamed_path = os.path.join(cp_dir, renamed_file)
|
||||
if os.path.isfile(renamed_path):
|
||||
print(f"[INFO] Resuming from renamed checkpoint: '{renamed_file}'")
|
||||
return renamed_path
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def save_audio(audio, path, sr):
|
||||
# wav: torch with 1d shape
|
||||
audio = audio * MAX_WAV_VALUE
|
||||
audio = audio.cpu().numpy().astype("int16")
|
||||
write(path, sr, audio)
|
||||
|
||||
@ -623,8 +623,7 @@ class TTS:
|
||||
|
||||
self.vocoder = BigVGAN.from_pretrained(
|
||||
"%s/GPT_SoVITS/pretrained_models/models--nvidia--bigvgan_v2_24khz_100band_256x" % (now_dir,),
|
||||
use_cuda_kernel=False,
|
||||
) # if True, RuntimeError: Ninja is required to load C++ extensions
|
||||
)
|
||||
# remove weight norm in the model and set to eval mode
|
||||
self.vocoder.remove_weight_norm()
|
||||
|
||||
|
||||
@ -1,264 +0,0 @@
|
||||
# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
|
||||
|
||||
"""
|
||||
Res2Net implementation is adapted from https://github.com/wenet-e2e/wespeaker.
|
||||
ERes2Net incorporates both local and global feature fusion techniques to improve the performance.
|
||||
The local feature fusion (LFF) fuses the features within one single residual block to extract the local signal.
|
||||
The global feature fusion (GFF) takes acoustic features of different scales as input to aggregate global signal.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import math
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import pooling_layers as pooling_layers
|
||||
from fusion import AFF
|
||||
|
||||
|
||||
class ReLU(nn.Hardtanh):
|
||||
def __init__(self, inplace=False):
|
||||
super(ReLU, self).__init__(0, 20, inplace)
|
||||
|
||||
def __repr__(self):
|
||||
inplace_str = "inplace" if self.inplace else ""
|
||||
return self.__class__.__name__ + " (" + inplace_str + ")"
|
||||
|
||||
|
||||
class BasicBlockERes2Net(nn.Module):
|
||||
expansion = 2
|
||||
|
||||
def __init__(self, in_planes, planes, stride=1, baseWidth=32, scale=2):
|
||||
super(BasicBlockERes2Net, self).__init__()
|
||||
width = int(math.floor(planes * (baseWidth / 64.0)))
|
||||
self.conv1 = nn.Conv2d(in_planes, width * scale, kernel_size=1, stride=stride, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(width * scale)
|
||||
self.nums = scale
|
||||
|
||||
convs = []
|
||||
bns = []
|
||||
for i in range(self.nums):
|
||||
convs.append(nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False))
|
||||
bns.append(nn.BatchNorm2d(width))
|
||||
self.convs = nn.ModuleList(convs)
|
||||
self.bns = nn.ModuleList(bns)
|
||||
self.relu = ReLU(inplace=True)
|
||||
|
||||
self.conv3 = nn.Conv2d(width * scale, planes * self.expansion, kernel_size=1, bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
|
||||
self.shortcut = nn.Sequential()
|
||||
if stride != 1 or in_planes != self.expansion * planes:
|
||||
self.shortcut = nn.Sequential(
|
||||
nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
|
||||
nn.BatchNorm2d(self.expansion * planes),
|
||||
)
|
||||
self.stride = stride
|
||||
self.width = width
|
||||
self.scale = scale
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
spx = torch.split(out, self.width, 1)
|
||||
for i in range(self.nums):
|
||||
if i == 0:
|
||||
sp = spx[i]
|
||||
else:
|
||||
sp = sp + spx[i]
|
||||
sp = self.convs[i](sp)
|
||||
sp = self.relu(self.bns[i](sp))
|
||||
if i == 0:
|
||||
out = sp
|
||||
else:
|
||||
out = torch.cat((out, sp), 1)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
residual = self.shortcut(x)
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class BasicBlockERes2Net_diff_AFF(nn.Module):
|
||||
expansion = 2
|
||||
|
||||
def __init__(self, in_planes, planes, stride=1, baseWidth=32, scale=2):
|
||||
super(BasicBlockERes2Net_diff_AFF, self).__init__()
|
||||
width = int(math.floor(planes * (baseWidth / 64.0)))
|
||||
self.conv1 = nn.Conv2d(in_planes, width * scale, kernel_size=1, stride=stride, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(width * scale)
|
||||
self.nums = scale
|
||||
|
||||
convs = []
|
||||
fuse_models = []
|
||||
bns = []
|
||||
for i in range(self.nums):
|
||||
convs.append(nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False))
|
||||
bns.append(nn.BatchNorm2d(width))
|
||||
for j in range(self.nums - 1):
|
||||
fuse_models.append(AFF(channels=width))
|
||||
|
||||
self.convs = nn.ModuleList(convs)
|
||||
self.bns = nn.ModuleList(bns)
|
||||
self.fuse_models = nn.ModuleList(fuse_models)
|
||||
self.relu = ReLU(inplace=True)
|
||||
|
||||
self.conv3 = nn.Conv2d(width * scale, planes * self.expansion, kernel_size=1, bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
|
||||
self.shortcut = nn.Sequential()
|
||||
if stride != 1 or in_planes != self.expansion * planes:
|
||||
self.shortcut = nn.Sequential(
|
||||
nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
|
||||
nn.BatchNorm2d(self.expansion * planes),
|
||||
)
|
||||
self.stride = stride
|
||||
self.width = width
|
||||
self.scale = scale
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
spx = torch.split(out, self.width, 1)
|
||||
for i in range(self.nums):
|
||||
if i == 0:
|
||||
sp = spx[i]
|
||||
else:
|
||||
sp = self.fuse_models[i - 1](sp, spx[i])
|
||||
|
||||
sp = self.convs[i](sp)
|
||||
sp = self.relu(self.bns[i](sp))
|
||||
if i == 0:
|
||||
out = sp
|
||||
else:
|
||||
out = torch.cat((out, sp), 1)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
residual = self.shortcut(x)
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class ERes2Net(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
block=BasicBlockERes2Net,
|
||||
block_fuse=BasicBlockERes2Net_diff_AFF,
|
||||
num_blocks=[3, 4, 6, 3],
|
||||
m_channels=32,
|
||||
feat_dim=80,
|
||||
embedding_size=192,
|
||||
pooling_func="TSTP",
|
||||
two_emb_layer=False,
|
||||
):
|
||||
super(ERes2Net, self).__init__()
|
||||
self.in_planes = m_channels
|
||||
self.feat_dim = feat_dim
|
||||
self.embedding_size = embedding_size
|
||||
self.stats_dim = int(feat_dim / 8) * m_channels * 8
|
||||
self.two_emb_layer = two_emb_layer
|
||||
|
||||
self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(m_channels)
|
||||
self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=1)
|
||||
self.layer2 = self._make_layer(block, m_channels * 2, num_blocks[1], stride=2)
|
||||
self.layer3 = self._make_layer(block_fuse, m_channels * 4, num_blocks[2], stride=2)
|
||||
self.layer4 = self._make_layer(block_fuse, m_channels * 8, num_blocks[3], stride=2)
|
||||
|
||||
# Downsampling module for each layer
|
||||
self.layer1_downsample = nn.Conv2d(
|
||||
m_channels * 2, m_channels * 4, kernel_size=3, stride=2, padding=1, bias=False
|
||||
)
|
||||
self.layer2_downsample = nn.Conv2d(
|
||||
m_channels * 4, m_channels * 8, kernel_size=3, padding=1, stride=2, bias=False
|
||||
)
|
||||
self.layer3_downsample = nn.Conv2d(
|
||||
m_channels * 8, m_channels * 16, kernel_size=3, padding=1, stride=2, bias=False
|
||||
)
|
||||
|
||||
# Bottom-up fusion module
|
||||
self.fuse_mode12 = AFF(channels=m_channels * 4)
|
||||
self.fuse_mode123 = AFF(channels=m_channels * 8)
|
||||
self.fuse_mode1234 = AFF(channels=m_channels * 16)
|
||||
|
||||
self.n_stats = 1 if pooling_func == "TAP" or pooling_func == "TSDP" else 2
|
||||
self.pool = getattr(pooling_layers, pooling_func)(in_dim=self.stats_dim * block.expansion)
|
||||
self.seg_1 = nn.Linear(self.stats_dim * block.expansion * self.n_stats, embedding_size)
|
||||
if self.two_emb_layer:
|
||||
self.seg_bn_1 = nn.BatchNorm1d(embedding_size, affine=False)
|
||||
self.seg_2 = nn.Linear(embedding_size, embedding_size)
|
||||
else:
|
||||
self.seg_bn_1 = nn.Identity()
|
||||
self.seg_2 = nn.Identity()
|
||||
|
||||
def _make_layer(self, block, planes, num_blocks, stride):
|
||||
strides = [stride] + [1] * (num_blocks - 1)
|
||||
layers = []
|
||||
for stride in strides:
|
||||
layers.append(block(self.in_planes, planes, stride))
|
||||
self.in_planes = planes * block.expansion
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T)
|
||||
x = x.unsqueeze_(1)
|
||||
out = F.relu(self.bn1(self.conv1(x)))
|
||||
out1 = self.layer1(out)
|
||||
out2 = self.layer2(out1)
|
||||
out1_downsample = self.layer1_downsample(out1)
|
||||
fuse_out12 = self.fuse_mode12(out2, out1_downsample)
|
||||
out3 = self.layer3(out2)
|
||||
fuse_out12_downsample = self.layer2_downsample(fuse_out12)
|
||||
fuse_out123 = self.fuse_mode123(out3, fuse_out12_downsample)
|
||||
out4 = self.layer4(out3)
|
||||
fuse_out123_downsample = self.layer3_downsample(fuse_out123)
|
||||
fuse_out1234 = self.fuse_mode1234(out4, fuse_out123_downsample)
|
||||
stats = self.pool(fuse_out1234)
|
||||
|
||||
embed_a = self.seg_1(stats)
|
||||
if self.two_emb_layer:
|
||||
out = F.relu(embed_a)
|
||||
out = self.seg_bn_1(out)
|
||||
embed_b = self.seg_2(out)
|
||||
return embed_b
|
||||
else:
|
||||
return embed_a
|
||||
|
||||
def forward3(self, x):
|
||||
x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T)
|
||||
x = x.unsqueeze_(1)
|
||||
out = F.relu(self.bn1(self.conv1(x)))
|
||||
out1 = self.layer1(out)
|
||||
out2 = self.layer2(out1)
|
||||
out1_downsample = self.layer1_downsample(out1)
|
||||
fuse_out12 = self.fuse_mode12(out2, out1_downsample)
|
||||
out3 = self.layer3(out2)
|
||||
fuse_out12_downsample = self.layer2_downsample(fuse_out12)
|
||||
fuse_out123 = self.fuse_mode123(out3, fuse_out12_downsample)
|
||||
out4 = self.layer4(out3)
|
||||
fuse_out123_downsample = self.layer3_downsample(fuse_out123)
|
||||
fuse_out1234 = self.fuse_mode1234(out4, fuse_out123_downsample).flatten(start_dim=1, end_dim=2).mean(-1)
|
||||
return fuse_out1234
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
x = torch.zeros(10, 300, 80)
|
||||
model = ERes2Net(feat_dim=80, embedding_size=192, pooling_func="TSTP")
|
||||
model.eval()
|
||||
out = model(x)
|
||||
print(out.shape) # torch.Size([10, 192])
|
||||
|
||||
num_params = sum(param.numel() for param in model.parameters())
|
||||
print("{} M".format(num_params / 1e6)) # 6.61M
|
||||
@ -259,14 +259,3 @@ class ERes2NetV2(nn.Module):
|
||||
# return embed_b
|
||||
# else:
|
||||
# return embed_a
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
x = torch.randn(1, 300, 80)
|
||||
model = ERes2NetV2(feat_dim=80, embedding_size=192, m_channels=64, baseWidth=26, scale=2, expansion=2)
|
||||
model.eval()
|
||||
y = model(x)
|
||||
print(y.size())
|
||||
macs, num_params = profile(model, inputs=(x,))
|
||||
print("Params: {} M".format(num_params / 1e6)) # 17.86 M
|
||||
print("MACs: {} G".format(macs / 1e9)) # 12.69 G
|
||||
|
||||
@ -1,289 +0,0 @@
|
||||
# Copyright 3D-Speaker (https://github.com/alibaba-damo-academy/3D-Speaker). All Rights Reserved.
|
||||
# Licensed under the Apache License, Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
|
||||
|
||||
"""Res2Net implementation is adapted from https://github.com/wenet-e2e/wespeaker.
|
||||
ERes2Net incorporates both local and global feature fusion techniques to improve the performance.
|
||||
The local feature fusion (LFF) fuses the features within one single residual block to extract the local signal.
|
||||
The global feature fusion (GFF) takes acoustic features of different scales as input to aggregate global signal.
|
||||
ERes2Net-huge is an upgraded version of ERes2Net that uses a larger number of parameters to achieve better
|
||||
recognition performance. Parameters expansion, baseWidth, and scale can be modified to obtain optimal performance.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import math
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import pooling_layers as pooling_layers
|
||||
from fusion import AFF
|
||||
|
||||
|
||||
class ReLU(nn.Hardtanh):
|
||||
def __init__(self, inplace=False):
|
||||
super(ReLU, self).__init__(0, 20, inplace)
|
||||
|
||||
def __repr__(self):
|
||||
inplace_str = "inplace" if self.inplace else ""
|
||||
return self.__class__.__name__ + " (" + inplace_str + ")"
|
||||
|
||||
|
||||
class BasicBlockERes2Net(nn.Module):
|
||||
expansion = 4
|
||||
|
||||
def __init__(self, in_planes, planes, stride=1, baseWidth=24, scale=3):
|
||||
super(BasicBlockERes2Net, self).__init__()
|
||||
width = int(math.floor(planes * (baseWidth / 64.0)))
|
||||
self.conv1 = nn.Conv2d(in_planes, width * scale, kernel_size=1, stride=stride, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(width * scale)
|
||||
self.nums = scale
|
||||
|
||||
convs = []
|
||||
bns = []
|
||||
for i in range(self.nums):
|
||||
convs.append(nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False))
|
||||
bns.append(nn.BatchNorm2d(width))
|
||||
self.convs = nn.ModuleList(convs)
|
||||
self.bns = nn.ModuleList(bns)
|
||||
self.relu = ReLU(inplace=True)
|
||||
|
||||
self.conv3 = nn.Conv2d(width * scale, planes * self.expansion, kernel_size=1, bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
|
||||
self.shortcut = nn.Sequential()
|
||||
if stride != 1 or in_planes != self.expansion * planes:
|
||||
self.shortcut = nn.Sequential(
|
||||
nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
|
||||
nn.BatchNorm2d(self.expansion * planes),
|
||||
)
|
||||
self.stride = stride
|
||||
self.width = width
|
||||
self.scale = scale
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
spx = torch.split(out, self.width, 1)
|
||||
for i in range(self.nums):
|
||||
if i == 0:
|
||||
sp = spx[i]
|
||||
else:
|
||||
sp = sp + spx[i]
|
||||
sp = self.convs[i](sp)
|
||||
sp = self.relu(self.bns[i](sp))
|
||||
if i == 0:
|
||||
out = sp
|
||||
else:
|
||||
out = torch.cat((out, sp), 1)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
residual = self.shortcut(x)
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class BasicBlockERes2Net_diff_AFF(nn.Module):
|
||||
expansion = 4
|
||||
|
||||
def __init__(self, in_planes, planes, stride=1, baseWidth=24, scale=3):
|
||||
super(BasicBlockERes2Net_diff_AFF, self).__init__()
|
||||
width = int(math.floor(planes * (baseWidth / 64.0)))
|
||||
self.conv1 = nn.Conv2d(in_planes, width * scale, kernel_size=1, stride=stride, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(width * scale)
|
||||
self.nums = scale
|
||||
|
||||
convs = []
|
||||
fuse_models = []
|
||||
bns = []
|
||||
for i in range(self.nums):
|
||||
convs.append(nn.Conv2d(width, width, kernel_size=3, padding=1, bias=False))
|
||||
bns.append(nn.BatchNorm2d(width))
|
||||
for j in range(self.nums - 1):
|
||||
fuse_models.append(AFF(channels=width))
|
||||
|
||||
self.convs = nn.ModuleList(convs)
|
||||
self.bns = nn.ModuleList(bns)
|
||||
self.fuse_models = nn.ModuleList(fuse_models)
|
||||
self.relu = ReLU(inplace=True)
|
||||
|
||||
self.conv3 = nn.Conv2d(width * scale, planes * self.expansion, kernel_size=1, bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes * self.expansion)
|
||||
self.shortcut = nn.Sequential()
|
||||
if stride != 1 or in_planes != self.expansion * planes:
|
||||
self.shortcut = nn.Sequential(
|
||||
nn.Conv2d(in_planes, self.expansion * planes, kernel_size=1, stride=stride, bias=False),
|
||||
nn.BatchNorm2d(self.expansion * planes),
|
||||
)
|
||||
self.stride = stride
|
||||
self.width = width
|
||||
self.scale = scale
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
spx = torch.split(out, self.width, 1)
|
||||
for i in range(self.nums):
|
||||
if i == 0:
|
||||
sp = spx[i]
|
||||
else:
|
||||
sp = self.fuse_models[i - 1](sp, spx[i])
|
||||
|
||||
sp = self.convs[i](sp)
|
||||
sp = self.relu(self.bns[i](sp))
|
||||
if i == 0:
|
||||
out = sp
|
||||
else:
|
||||
out = torch.cat((out, sp), 1)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
residual = self.shortcut(x)
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class ERes2Net(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
block=BasicBlockERes2Net,
|
||||
block_fuse=BasicBlockERes2Net_diff_AFF,
|
||||
num_blocks=[3, 4, 6, 3],
|
||||
m_channels=64,
|
||||
feat_dim=80,
|
||||
embedding_size=192,
|
||||
pooling_func="TSTP",
|
||||
two_emb_layer=False,
|
||||
):
|
||||
super(ERes2Net, self).__init__()
|
||||
self.in_planes = m_channels
|
||||
self.feat_dim = feat_dim
|
||||
self.embedding_size = embedding_size
|
||||
self.stats_dim = int(feat_dim / 8) * m_channels * 8
|
||||
self.two_emb_layer = two_emb_layer
|
||||
|
||||
self.conv1 = nn.Conv2d(1, m_channels, kernel_size=3, stride=1, padding=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(m_channels)
|
||||
|
||||
self.layer1 = self._make_layer(block, m_channels, num_blocks[0], stride=1)
|
||||
self.layer2 = self._make_layer(block, m_channels * 2, num_blocks[1], stride=2)
|
||||
self.layer3 = self._make_layer(block_fuse, m_channels * 4, num_blocks[2], stride=2)
|
||||
self.layer4 = self._make_layer(block_fuse, m_channels * 8, num_blocks[3], stride=2)
|
||||
|
||||
self.layer1_downsample = nn.Conv2d(
|
||||
m_channels * 4, m_channels * 8, kernel_size=3, padding=1, stride=2, bias=False
|
||||
)
|
||||
self.layer2_downsample = nn.Conv2d(
|
||||
m_channels * 8, m_channels * 16, kernel_size=3, padding=1, stride=2, bias=False
|
||||
)
|
||||
self.layer3_downsample = nn.Conv2d(
|
||||
m_channels * 16, m_channels * 32, kernel_size=3, padding=1, stride=2, bias=False
|
||||
)
|
||||
|
||||
self.fuse_mode12 = AFF(channels=m_channels * 8)
|
||||
self.fuse_mode123 = AFF(channels=m_channels * 16)
|
||||
self.fuse_mode1234 = AFF(channels=m_channels * 32)
|
||||
|
||||
self.n_stats = 1 if pooling_func == "TAP" or pooling_func == "TSDP" else 2
|
||||
self.pool = getattr(pooling_layers, pooling_func)(in_dim=self.stats_dim * block.expansion)
|
||||
self.seg_1 = nn.Linear(self.stats_dim * block.expansion * self.n_stats, embedding_size)
|
||||
if self.two_emb_layer:
|
||||
self.seg_bn_1 = nn.BatchNorm1d(embedding_size, affine=False)
|
||||
self.seg_2 = nn.Linear(embedding_size, embedding_size)
|
||||
else:
|
||||
self.seg_bn_1 = nn.Identity()
|
||||
self.seg_2 = nn.Identity()
|
||||
|
||||
def _make_layer(self, block, planes, num_blocks, stride):
|
||||
strides = [stride] + [1] * (num_blocks - 1)
|
||||
layers = []
|
||||
for stride in strides:
|
||||
layers.append(block(self.in_planes, planes, stride))
|
||||
self.in_planes = planes * block.expansion
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T)
|
||||
|
||||
x = x.unsqueeze_(1)
|
||||
out = F.relu(self.bn1(self.conv1(x)))
|
||||
out1 = self.layer1(out)
|
||||
out2 = self.layer2(out1)
|
||||
out1_downsample = self.layer1_downsample(out1)
|
||||
fuse_out12 = self.fuse_mode12(out2, out1_downsample)
|
||||
out3 = self.layer3(out2)
|
||||
fuse_out12_downsample = self.layer2_downsample(fuse_out12)
|
||||
fuse_out123 = self.fuse_mode123(out3, fuse_out12_downsample)
|
||||
out4 = self.layer4(out3)
|
||||
fuse_out123_downsample = self.layer3_downsample(fuse_out123)
|
||||
fuse_out1234 = self.fuse_mode1234(out4, fuse_out123_downsample)
|
||||
stats = self.pool(fuse_out1234)
|
||||
|
||||
embed_a = self.seg_1(stats)
|
||||
if self.two_emb_layer:
|
||||
out = F.relu(embed_a)
|
||||
out = self.seg_bn_1(out)
|
||||
embed_b = self.seg_2(out)
|
||||
return embed_b
|
||||
else:
|
||||
return embed_a
|
||||
|
||||
def forward2(self, x, if_mean):
|
||||
x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T)
|
||||
|
||||
x = x.unsqueeze_(1)
|
||||
out = F.relu(self.bn1(self.conv1(x)))
|
||||
out1 = self.layer1(out)
|
||||
out2 = self.layer2(out1)
|
||||
out1_downsample = self.layer1_downsample(out1)
|
||||
fuse_out12 = self.fuse_mode12(out2, out1_downsample)
|
||||
out3 = self.layer3(out2)
|
||||
fuse_out12_downsample = self.layer2_downsample(fuse_out12)
|
||||
fuse_out123 = self.fuse_mode123(out3, fuse_out12_downsample)
|
||||
out4 = self.layer4(out3)
|
||||
fuse_out123_downsample = self.layer3_downsample(fuse_out123)
|
||||
fuse_out1234 = self.fuse_mode1234(out4, fuse_out123_downsample).flatten(start_dim=1, end_dim=2) # bs,20480,T
|
||||
if if_mean == False:
|
||||
mean = fuse_out1234[0].transpose(1, 0) # (T,20480),bs=T
|
||||
else:
|
||||
mean = fuse_out1234.mean(2) # bs,20480
|
||||
mean_std = torch.cat([mean, torch.zeros_like(mean)], 1)
|
||||
return self.seg_1(mean_std) # (T,192)
|
||||
|
||||
# stats = self.pool(fuse_out1234)
|
||||
# if self.two_emb_layer:
|
||||
# out = F.relu(embed_a)
|
||||
# out = self.seg_bn_1(out)
|
||||
# embed_b = self.seg_2(out)
|
||||
# return embed_b
|
||||
# else:
|
||||
# return embed_a
|
||||
|
||||
def forward3(self, x):
|
||||
x = x.permute(0, 2, 1) # (B,T,F) => (B,F,T)
|
||||
|
||||
x = x.unsqueeze_(1)
|
||||
out = F.relu(self.bn1(self.conv1(x)))
|
||||
out1 = self.layer1(out)
|
||||
out2 = self.layer2(out1)
|
||||
out1_downsample = self.layer1_downsample(out1)
|
||||
fuse_out12 = self.fuse_mode12(out2, out1_downsample)
|
||||
out3 = self.layer3(out2)
|
||||
fuse_out12_downsample = self.layer2_downsample(fuse_out12)
|
||||
fuse_out123 = self.fuse_mode123(out3, fuse_out12_downsample)
|
||||
out4 = self.layer4(out3)
|
||||
fuse_out123_downsample = self.layer3_downsample(fuse_out123)
|
||||
fuse_out1234 = self.fuse_mode1234(out4, fuse_out123_downsample).flatten(start_dim=1, end_dim=2).mean(-1)
|
||||
return fuse_out1234
|
||||
# print(fuse_out1234.shape)
|
||||
# print(fuse_out1234.flatten(start_dim=1,end_dim=2).shape)
|
||||
# pdb.set_trace()
|
||||
@ -2,18 +2,14 @@ import math
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
import torchaudio
|
||||
from torch import Tensor
|
||||
|
||||
__all__ = [
|
||||
"get_mel_banks",
|
||||
"inverse_mel_scale",
|
||||
"inverse_mel_scale_scalar",
|
||||
"mel_scale",
|
||||
"mel_scale_scalar",
|
||||
"spectrogram",
|
||||
"fbank",
|
||||
"mfcc",
|
||||
"vtln_warp_freq",
|
||||
"vtln_warp_mel_freq",
|
||||
]
|
||||
@ -226,99 +222,6 @@ def _subtract_column_mean(tensor: Tensor, subtract_mean: bool) -> Tensor:
|
||||
return tensor
|
||||
|
||||
|
||||
def spectrogram(
|
||||
waveform: Tensor,
|
||||
blackman_coeff: float = 0.42,
|
||||
channel: int = -1,
|
||||
dither: float = 0.0,
|
||||
energy_floor: float = 1.0,
|
||||
frame_length: float = 25.0,
|
||||
frame_shift: float = 10.0,
|
||||
min_duration: float = 0.0,
|
||||
preemphasis_coefficient: float = 0.97,
|
||||
raw_energy: bool = True,
|
||||
remove_dc_offset: bool = True,
|
||||
round_to_power_of_two: bool = True,
|
||||
sample_frequency: float = 16000.0,
|
||||
snip_edges: bool = True,
|
||||
subtract_mean: bool = False,
|
||||
window_type: str = POVEY,
|
||||
) -> Tensor:
|
||||
r"""Create a spectrogram from a raw audio signal. This matches the input/output of Kaldi's
|
||||
compute-spectrogram-feats.
|
||||
|
||||
Args:
|
||||
waveform (Tensor): Tensor of audio of size (c, n) where c is in the range [0,2)
|
||||
blackman_coeff (float, optional): Constant coefficient for generalized Blackman window. (Default: ``0.42``)
|
||||
channel (int, optional): Channel to extract (-1 -> expect mono, 0 -> left, 1 -> right) (Default: ``-1``)
|
||||
dither (float, optional): Dithering constant (0.0 means no dither). If you turn this off, you should set
|
||||
the energy_floor option, e.g. to 1.0 or 0.1 (Default: ``0.0``)
|
||||
energy_floor (float, optional): Floor on energy (absolute, not relative) in Spectrogram computation. Caution:
|
||||
this floor is applied to the zeroth component, representing the total signal energy. The floor on the
|
||||
individual spectrogram elements is fixed at std::numeric_limits<float>::epsilon(). (Default: ``1.0``)
|
||||
frame_length (float, optional): Frame length in milliseconds (Default: ``25.0``)
|
||||
frame_shift (float, optional): Frame shift in milliseconds (Default: ``10.0``)
|
||||
min_duration (float, optional): Minimum duration of segments to process (in seconds). (Default: ``0.0``)
|
||||
preemphasis_coefficient (float, optional): Coefficient for use in signal preemphasis (Default: ``0.97``)
|
||||
raw_energy (bool, optional): If True, compute energy before preemphasis and windowing (Default: ``True``)
|
||||
remove_dc_offset (bool, optional): Subtract mean from waveform on each frame (Default: ``True``)
|
||||
round_to_power_of_two (bool, optional): If True, round window size to power of two by zero-padding input
|
||||
to FFT. (Default: ``True``)
|
||||
sample_frequency (float, optional): Waveform data sample frequency (must match the waveform file, if
|
||||
specified there) (Default: ``16000.0``)
|
||||
snip_edges (bool, optional): If True, end effects will be handled by outputting only frames that completely fit
|
||||
in the file, and the number of frames depends on the frame_length. If False, the number of frames
|
||||
depends only on the frame_shift, and we reflect the data at the ends. (Default: ``True``)
|
||||
subtract_mean (bool, optional): Subtract mean of each feature file [CMS]; not recommended to do
|
||||
it this way. (Default: ``False``)
|
||||
window_type (str, optional): Type of window ('hamming'|'hanning'|'povey'|'rectangular'|'blackman')
|
||||
(Default: ``'povey'``)
|
||||
|
||||
Returns:
|
||||
Tensor: A spectrogram identical to what Kaldi would output. The shape is
|
||||
(m, ``padded_window_size // 2 + 1``) where m is calculated in _get_strided
|
||||
"""
|
||||
device, dtype = waveform.device, waveform.dtype
|
||||
epsilon = _get_epsilon(device, dtype)
|
||||
|
||||
waveform, window_shift, window_size, padded_window_size = _get_waveform_and_window_properties(
|
||||
waveform, channel, sample_frequency, frame_shift, frame_length, round_to_power_of_two, preemphasis_coefficient
|
||||
)
|
||||
|
||||
if len(waveform) < min_duration * sample_frequency:
|
||||
# signal is too short
|
||||
return torch.empty(0)
|
||||
|
||||
strided_input, signal_log_energy = _get_window(
|
||||
waveform,
|
||||
padded_window_size,
|
||||
window_size,
|
||||
window_shift,
|
||||
window_type,
|
||||
blackman_coeff,
|
||||
snip_edges,
|
||||
raw_energy,
|
||||
energy_floor,
|
||||
dither,
|
||||
remove_dc_offset,
|
||||
preemphasis_coefficient,
|
||||
)
|
||||
|
||||
# size (m, padded_window_size // 2 + 1, 2)
|
||||
fft = torch.fft.rfft(strided_input)
|
||||
|
||||
# Convert the FFT into a power spectrum
|
||||
power_spectrum = torch.max(fft.abs().pow(2.0), epsilon).log() # size (m, padded_window_size // 2 + 1)
|
||||
power_spectrum[:, 0] = signal_log_energy
|
||||
|
||||
power_spectrum = _subtract_column_mean(power_spectrum, subtract_mean)
|
||||
return power_spectrum
|
||||
|
||||
|
||||
def inverse_mel_scale_scalar(mel_freq: float) -> float:
|
||||
return 700.0 * (math.exp(mel_freq / 1127.0) - 1.0)
|
||||
|
||||
|
||||
def inverse_mel_scale(mel_freq: Tensor) -> Tensor:
|
||||
return 700.0 * ((mel_freq / 1127.0).exp() - 1.0)
|
||||
|
||||
@ -674,171 +577,3 @@ def fbank(
|
||||
|
||||
mel_energies = _subtract_column_mean(mel_energies, subtract_mean)
|
||||
return mel_energies
|
||||
|
||||
|
||||
def _get_dct_matrix(num_ceps: int, num_mel_bins: int) -> Tensor:
|
||||
# returns a dct matrix of size (num_mel_bins, num_ceps)
|
||||
# size (num_mel_bins, num_mel_bins)
|
||||
dct_matrix = torchaudio.functional.create_dct(num_mel_bins, num_mel_bins, "ortho")
|
||||
# kaldi expects the first cepstral to be weighted sum of factor sqrt(1/num_mel_bins)
|
||||
# this would be the first column in the dct_matrix for torchaudio as it expects a
|
||||
# right multiply (which would be the first column of the kaldi's dct_matrix as kaldi
|
||||
# expects a left multiply e.g. dct_matrix * vector).
|
||||
dct_matrix[:, 0] = math.sqrt(1 / float(num_mel_bins))
|
||||
dct_matrix = dct_matrix[:, :num_ceps]
|
||||
return dct_matrix
|
||||
|
||||
|
||||
def _get_lifter_coeffs(num_ceps: int, cepstral_lifter: float) -> Tensor:
|
||||
# returns size (num_ceps)
|
||||
# Compute liftering coefficients (scaling on cepstral coeffs)
|
||||
# coeffs are numbered slightly differently from HTK: the zeroth index is C0, which is not affected.
|
||||
i = torch.arange(num_ceps)
|
||||
return 1.0 + 0.5 * cepstral_lifter * torch.sin(math.pi * i / cepstral_lifter)
|
||||
|
||||
|
||||
def mfcc(
|
||||
waveform: Tensor,
|
||||
blackman_coeff: float = 0.42,
|
||||
cepstral_lifter: float = 22.0,
|
||||
channel: int = -1,
|
||||
dither: float = 0.0,
|
||||
energy_floor: float = 1.0,
|
||||
frame_length: float = 25.0,
|
||||
frame_shift: float = 10.0,
|
||||
high_freq: float = 0.0,
|
||||
htk_compat: bool = False,
|
||||
low_freq: float = 20.0,
|
||||
num_ceps: int = 13,
|
||||
min_duration: float = 0.0,
|
||||
num_mel_bins: int = 23,
|
||||
preemphasis_coefficient: float = 0.97,
|
||||
raw_energy: bool = True,
|
||||
remove_dc_offset: bool = True,
|
||||
round_to_power_of_two: bool = True,
|
||||
sample_frequency: float = 16000.0,
|
||||
snip_edges: bool = True,
|
||||
subtract_mean: bool = False,
|
||||
use_energy: bool = False,
|
||||
vtln_high: float = -500.0,
|
||||
vtln_low: float = 100.0,
|
||||
vtln_warp: float = 1.0,
|
||||
window_type: str = POVEY,
|
||||
) -> Tensor:
|
||||
r"""Create a mfcc from a raw audio signal. This matches the input/output of Kaldi's
|
||||
compute-mfcc-feats.
|
||||
|
||||
Args:
|
||||
waveform (Tensor): Tensor of audio of size (c, n) where c is in the range [0,2)
|
||||
blackman_coeff (float, optional): Constant coefficient for generalized Blackman window. (Default: ``0.42``)
|
||||
cepstral_lifter (float, optional): Constant that controls scaling of MFCCs (Default: ``22.0``)
|
||||
channel (int, optional): Channel to extract (-1 -> expect mono, 0 -> left, 1 -> right) (Default: ``-1``)
|
||||
dither (float, optional): Dithering constant (0.0 means no dither). If you turn this off, you should set
|
||||
the energy_floor option, e.g. to 1.0 or 0.1 (Default: ``0.0``)
|
||||
energy_floor (float, optional): Floor on energy (absolute, not relative) in Spectrogram computation. Caution:
|
||||
this floor is applied to the zeroth component, representing the total signal energy. The floor on the
|
||||
individual spectrogram elements is fixed at std::numeric_limits<float>::epsilon(). (Default: ``1.0``)
|
||||
frame_length (float, optional): Frame length in milliseconds (Default: ``25.0``)
|
||||
frame_shift (float, optional): Frame shift in milliseconds (Default: ``10.0``)
|
||||
high_freq (float, optional): High cutoff frequency for mel bins (if <= 0, offset from Nyquist)
|
||||
(Default: ``0.0``)
|
||||
htk_compat (bool, optional): If true, put energy last. Warning: not sufficient to get HTK compatible
|
||||
features (need to change other parameters). (Default: ``False``)
|
||||
low_freq (float, optional): Low cutoff frequency for mel bins (Default: ``20.0``)
|
||||
num_ceps (int, optional): Number of cepstra in MFCC computation (including C0) (Default: ``13``)
|
||||
min_duration (float, optional): Minimum duration of segments to process (in seconds). (Default: ``0.0``)
|
||||
num_mel_bins (int, optional): Number of triangular mel-frequency bins (Default: ``23``)
|
||||
preemphasis_coefficient (float, optional): Coefficient for use in signal preemphasis (Default: ``0.97``)
|
||||
raw_energy (bool, optional): If True, compute energy before preemphasis and windowing (Default: ``True``)
|
||||
remove_dc_offset (bool, optional): Subtract mean from waveform on each frame (Default: ``True``)
|
||||
round_to_power_of_two (bool, optional): If True, round window size to power of two by zero-padding input
|
||||
to FFT. (Default: ``True``)
|
||||
sample_frequency (float, optional): Waveform data sample frequency (must match the waveform file, if
|
||||
specified there) (Default: ``16000.0``)
|
||||
snip_edges (bool, optional): If True, end effects will be handled by outputting only frames that completely fit
|
||||
in the file, and the number of frames depends on the frame_length. If False, the number of frames
|
||||
depends only on the frame_shift, and we reflect the data at the ends. (Default: ``True``)
|
||||
subtract_mean (bool, optional): Subtract mean of each feature file [CMS]; not recommended to do
|
||||
it this way. (Default: ``False``)
|
||||
use_energy (bool, optional): Add an extra dimension with energy to the FBANK output. (Default: ``False``)
|
||||
vtln_high (float, optional): High inflection point in piecewise linear VTLN warping function (if
|
||||
negative, offset from high-mel-freq (Default: ``-500.0``)
|
||||
vtln_low (float, optional): Low inflection point in piecewise linear VTLN warping function (Default: ``100.0``)
|
||||
vtln_warp (float, optional): Vtln warp factor (only applicable if vtln_map not specified) (Default: ``1.0``)
|
||||
window_type (str, optional): Type of window ('hamming'|'hanning'|'povey'|'rectangular'|'blackman')
|
||||
(Default: ``"povey"``)
|
||||
|
||||
Returns:
|
||||
Tensor: A mfcc identical to what Kaldi would output. The shape is (m, ``num_ceps``)
|
||||
where m is calculated in _get_strided
|
||||
"""
|
||||
assert num_ceps <= num_mel_bins, "num_ceps cannot be larger than num_mel_bins: %d vs %d" % (num_ceps, num_mel_bins)
|
||||
|
||||
device, dtype = waveform.device, waveform.dtype
|
||||
|
||||
# The mel_energies should not be squared (use_power=True), not have mean subtracted
|
||||
# (subtract_mean=False), and use log (use_log_fbank=True).
|
||||
# size (m, num_mel_bins + use_energy)
|
||||
feature = fbank(
|
||||
waveform=waveform,
|
||||
blackman_coeff=blackman_coeff,
|
||||
channel=channel,
|
||||
dither=dither,
|
||||
energy_floor=energy_floor,
|
||||
frame_length=frame_length,
|
||||
frame_shift=frame_shift,
|
||||
high_freq=high_freq,
|
||||
htk_compat=htk_compat,
|
||||
low_freq=low_freq,
|
||||
min_duration=min_duration,
|
||||
num_mel_bins=num_mel_bins,
|
||||
preemphasis_coefficient=preemphasis_coefficient,
|
||||
raw_energy=raw_energy,
|
||||
remove_dc_offset=remove_dc_offset,
|
||||
round_to_power_of_two=round_to_power_of_two,
|
||||
sample_frequency=sample_frequency,
|
||||
snip_edges=snip_edges,
|
||||
subtract_mean=False,
|
||||
use_energy=use_energy,
|
||||
use_log_fbank=True,
|
||||
use_power=True,
|
||||
vtln_high=vtln_high,
|
||||
vtln_low=vtln_low,
|
||||
vtln_warp=vtln_warp,
|
||||
window_type=window_type,
|
||||
)
|
||||
|
||||
if use_energy:
|
||||
# size (m)
|
||||
signal_log_energy = feature[:, num_mel_bins if htk_compat else 0]
|
||||
# offset is 0 if htk_compat==True else 1
|
||||
mel_offset = int(not htk_compat)
|
||||
feature = feature[:, mel_offset : (num_mel_bins + mel_offset)]
|
||||
|
||||
# size (num_mel_bins, num_ceps)
|
||||
dct_matrix = _get_dct_matrix(num_ceps, num_mel_bins).to(dtype=dtype, device=device)
|
||||
|
||||
# size (m, num_ceps)
|
||||
feature = feature.matmul(dct_matrix)
|
||||
|
||||
if cepstral_lifter != 0.0:
|
||||
# size (1, num_ceps)
|
||||
lifter_coeffs = _get_lifter_coeffs(num_ceps, cepstral_lifter).unsqueeze(0)
|
||||
feature *= lifter_coeffs.to(device=device, dtype=dtype)
|
||||
|
||||
# if use_energy then replace the last column for htk_compat == true else first column
|
||||
if use_energy:
|
||||
feature[:, 0] = signal_log_energy
|
||||
|
||||
if htk_compat:
|
||||
energy = feature[:, 0].unsqueeze(1) # size (m, 1)
|
||||
feature = feature[:, 1:] # size (m, num_ceps - 1)
|
||||
if not use_energy:
|
||||
# scale on C0 (actually removing a scale we previously added that's
|
||||
# part of one common definition of the cosine transform.)
|
||||
energy *= math.sqrt(2)
|
||||
|
||||
feature = torch.cat((feature, energy), dim=1)
|
||||
|
||||
feature = _subtract_column_mean(feature, subtract_mean)
|
||||
return feature
|
||||
|
||||
@ -506,8 +506,7 @@ def init_bigvgan():
|
||||
|
||||
bigvgan_model = bigvgan.BigVGAN.from_pretrained(
|
||||
"%s/GPT_SoVITS/pretrained_models/models--nvidia--bigvgan_v2_24khz_100band_256x" % (now_dir,),
|
||||
use_cuda_kernel=False,
|
||||
) # if True, RuntimeError: Ninja is required to load C++ extensions
|
||||
)
|
||||
# remove weight norm in the model and set to eval mode
|
||||
bigvgan_model.remove_weight_norm()
|
||||
bigvgan_model = bigvgan_model.eval()
|
||||
|
||||
@ -1,13 +1 @@
|
||||
# from f5_tts.model.cfm import CFM
|
||||
#
|
||||
# from f5_tts.model.backbones.unett import UNetT
|
||||
from GPT_SoVITS.f5_tts.model.backbones.dit import DiT
|
||||
# from f5_tts.model.backbones.dit import DiTNoCond
|
||||
# from f5_tts.model.backbones.dit import DiTNoCondNoT
|
||||
# from f5_tts.model.backbones.mmdit import MMDiT
|
||||
|
||||
# from f5_tts.model.trainer import Trainer
|
||||
|
||||
|
||||
# __all__ = ["CFM", "UNetT", "DiT", "MMDiT", "Trainer"]
|
||||
# __all__ = ["CFM", "UNetT", "DiTNoCond","DiT", "MMDiT"]
|
||||
|
||||
@ -1,146 +0,0 @@
|
||||
"""
|
||||
ein notation:
|
||||
b - batch
|
||||
n - sequence
|
||||
nt - text sequence
|
||||
nw - raw wave length
|
||||
d - dimension
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from x_transformers.x_transformers import RotaryEmbedding
|
||||
|
||||
from f5_tts.model.modules import (
|
||||
TimestepEmbedding,
|
||||
ConvPositionEmbedding,
|
||||
MMDiTBlock,
|
||||
AdaLayerNormZero_Final,
|
||||
precompute_freqs_cis,
|
||||
get_pos_embed_indices,
|
||||
)
|
||||
|
||||
|
||||
# text embedding
|
||||
|
||||
|
||||
class TextEmbedding(nn.Module):
|
||||
def __init__(self, out_dim, text_num_embeds):
|
||||
super().__init__()
|
||||
self.text_embed = nn.Embedding(text_num_embeds + 1, out_dim) # will use 0 as filler token
|
||||
|
||||
self.precompute_max_pos = 1024
|
||||
self.register_buffer("freqs_cis", precompute_freqs_cis(out_dim, self.precompute_max_pos), persistent=False)
|
||||
|
||||
def forward(self, text: int["b nt"], drop_text=False) -> int["b nt d"]: # noqa: F722
|
||||
text = text + 1
|
||||
if drop_text:
|
||||
text = torch.zeros_like(text)
|
||||
text = self.text_embed(text)
|
||||
|
||||
# sinus pos emb
|
||||
batch_start = torch.zeros((text.shape[0],), dtype=torch.long)
|
||||
batch_text_len = text.shape[1]
|
||||
pos_idx = get_pos_embed_indices(batch_start, batch_text_len, max_pos=self.precompute_max_pos)
|
||||
text_pos_embed = self.freqs_cis[pos_idx]
|
||||
|
||||
text = text + text_pos_embed
|
||||
|
||||
return text
|
||||
|
||||
|
||||
# noised input & masked cond audio embedding
|
||||
|
||||
|
||||
class AudioEmbedding(nn.Module):
|
||||
def __init__(self, in_dim, out_dim):
|
||||
super().__init__()
|
||||
self.linear = nn.Linear(2 * in_dim, out_dim)
|
||||
self.conv_pos_embed = ConvPositionEmbedding(out_dim)
|
||||
|
||||
def forward(self, x: float["b n d"], cond: float["b n d"], drop_audio_cond=False): # noqa: F722
|
||||
if drop_audio_cond:
|
||||
cond = torch.zeros_like(cond)
|
||||
x = torch.cat((x, cond), dim=-1)
|
||||
x = self.linear(x)
|
||||
x = self.conv_pos_embed(x) + x
|
||||
return x
|
||||
|
||||
|
||||
# Transformer backbone using MM-DiT blocks
|
||||
|
||||
|
||||
class MMDiT(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
dim,
|
||||
depth=8,
|
||||
heads=8,
|
||||
dim_head=64,
|
||||
dropout=0.1,
|
||||
ff_mult=4,
|
||||
text_num_embeds=256,
|
||||
mel_dim=100,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.time_embed = TimestepEmbedding(dim)
|
||||
self.text_embed = TextEmbedding(dim, text_num_embeds)
|
||||
self.audio_embed = AudioEmbedding(mel_dim, dim)
|
||||
|
||||
self.rotary_embed = RotaryEmbedding(dim_head)
|
||||
|
||||
self.dim = dim
|
||||
self.depth = depth
|
||||
|
||||
self.transformer_blocks = nn.ModuleList(
|
||||
[
|
||||
MMDiTBlock(
|
||||
dim=dim,
|
||||
heads=heads,
|
||||
dim_head=dim_head,
|
||||
dropout=dropout,
|
||||
ff_mult=ff_mult,
|
||||
context_pre_only=i == depth - 1,
|
||||
)
|
||||
for i in range(depth)
|
||||
]
|
||||
)
|
||||
self.norm_out = AdaLayerNormZero_Final(dim) # final modulation
|
||||
self.proj_out = nn.Linear(dim, mel_dim)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: float["b n d"], # nosied input audio # noqa: F722
|
||||
cond: float["b n d"], # masked cond audio # noqa: F722
|
||||
text: int["b nt"], # text # noqa: F722
|
||||
time: float["b"] | float[""], # time step # noqa: F821 F722
|
||||
drop_audio_cond, # cfg for cond audio
|
||||
drop_text, # cfg for text
|
||||
mask: bool["b n"] | None = None, # noqa: F722
|
||||
):
|
||||
batch = x.shape[0]
|
||||
if time.ndim == 0:
|
||||
time = time.repeat(batch)
|
||||
|
||||
# t: conditioning (time), c: context (text + masked cond audio), x: noised input audio
|
||||
t = self.time_embed(time)
|
||||
c = self.text_embed(text, drop_text=drop_text)
|
||||
x = self.audio_embed(x, cond, drop_audio_cond=drop_audio_cond)
|
||||
|
||||
seq_len = x.shape[1]
|
||||
text_len = text.shape[1]
|
||||
rope_audio = self.rotary_embed.forward_from_seq_len(seq_len)
|
||||
rope_text = self.rotary_embed.forward_from_seq_len(text_len)
|
||||
|
||||
for block in self.transformer_blocks:
|
||||
c, x = block(x, c, t, mask=mask, rope=rope_audio, c_rope=rope_text)
|
||||
|
||||
x = self.norm_out(x, t)
|
||||
output = self.proj_out(x)
|
||||
|
||||
return output
|
||||
@ -1,219 +0,0 @@
|
||||
"""
|
||||
ein notation:
|
||||
b - batch
|
||||
n - sequence
|
||||
nt - text sequence
|
||||
nw - raw wave length
|
||||
d - dimension
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
from typing import Literal
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from x_transformers import RMSNorm
|
||||
from x_transformers.x_transformers import RotaryEmbedding
|
||||
|
||||
from f5_tts.model.modules import (
|
||||
TimestepEmbedding,
|
||||
ConvNeXtV2Block,
|
||||
ConvPositionEmbedding,
|
||||
Attention,
|
||||
AttnProcessor,
|
||||
FeedForward,
|
||||
precompute_freqs_cis,
|
||||
get_pos_embed_indices,
|
||||
)
|
||||
|
||||
|
||||
# Text embedding
|
||||
|
||||
|
||||
class TextEmbedding(nn.Module):
|
||||
def __init__(self, text_num_embeds, text_dim, conv_layers=0, conv_mult=2):
|
||||
super().__init__()
|
||||
self.text_embed = nn.Embedding(text_num_embeds + 1, text_dim) # use 0 as filler token
|
||||
|
||||
if conv_layers > 0:
|
||||
self.extra_modeling = True
|
||||
self.precompute_max_pos = 4096 # ~44s of 24khz audio
|
||||
self.register_buffer("freqs_cis", precompute_freqs_cis(text_dim, self.precompute_max_pos), persistent=False)
|
||||
self.text_blocks = nn.Sequential(
|
||||
*[ConvNeXtV2Block(text_dim, text_dim * conv_mult) for _ in range(conv_layers)]
|
||||
)
|
||||
else:
|
||||
self.extra_modeling = False
|
||||
|
||||
def forward(self, text: int["b nt"], seq_len, drop_text=False): # noqa: F722
|
||||
text = text + 1 # use 0 as filler token. preprocess of batch pad -1, see list_str_to_idx()
|
||||
text = text[:, :seq_len] # curtail if character tokens are more than the mel spec tokens
|
||||
batch, text_len = text.shape[0], text.shape[1]
|
||||
text = F.pad(text, (0, seq_len - text_len), value=0)
|
||||
|
||||
if drop_text: # cfg for text
|
||||
text = torch.zeros_like(text)
|
||||
|
||||
text = self.text_embed(text) # b n -> b n d
|
||||
|
||||
# possible extra modeling
|
||||
if self.extra_modeling:
|
||||
# sinus pos emb
|
||||
batch_start = torch.zeros((batch,), dtype=torch.long)
|
||||
pos_idx = get_pos_embed_indices(batch_start, seq_len, max_pos=self.precompute_max_pos)
|
||||
text_pos_embed = self.freqs_cis[pos_idx]
|
||||
text = text + text_pos_embed
|
||||
|
||||
# convnextv2 blocks
|
||||
text = self.text_blocks(text)
|
||||
|
||||
return text
|
||||
|
||||
|
||||
# noised input audio and context mixing embedding
|
||||
|
||||
|
||||
class InputEmbedding(nn.Module):
|
||||
def __init__(self, mel_dim, text_dim, out_dim):
|
||||
super().__init__()
|
||||
self.proj = nn.Linear(mel_dim * 2 + text_dim, out_dim)
|
||||
self.conv_pos_embed = ConvPositionEmbedding(dim=out_dim)
|
||||
|
||||
def forward(self, x: float["b n d"], cond: float["b n d"], text_embed: float["b n d"], drop_audio_cond=False): # noqa: F722
|
||||
if drop_audio_cond: # cfg for cond audio
|
||||
cond = torch.zeros_like(cond)
|
||||
|
||||
x = self.proj(torch.cat((x, cond, text_embed), dim=-1))
|
||||
x = self.conv_pos_embed(x) + x
|
||||
return x
|
||||
|
||||
|
||||
# Flat UNet Transformer backbone
|
||||
|
||||
|
||||
class UNetT(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
dim,
|
||||
depth=8,
|
||||
heads=8,
|
||||
dim_head=64,
|
||||
dropout=0.1,
|
||||
ff_mult=4,
|
||||
mel_dim=100,
|
||||
text_num_embeds=256,
|
||||
text_dim=None,
|
||||
conv_layers=0,
|
||||
skip_connect_type: Literal["add", "concat", "none"] = "concat",
|
||||
):
|
||||
super().__init__()
|
||||
assert depth % 2 == 0, "UNet-Transformer's depth should be even."
|
||||
|
||||
self.time_embed = TimestepEmbedding(dim)
|
||||
if text_dim is None:
|
||||
text_dim = mel_dim
|
||||
self.text_embed = TextEmbedding(text_num_embeds, text_dim, conv_layers=conv_layers)
|
||||
self.input_embed = InputEmbedding(mel_dim, text_dim, dim)
|
||||
|
||||
self.rotary_embed = RotaryEmbedding(dim_head)
|
||||
|
||||
# transformer layers & skip connections
|
||||
|
||||
self.dim = dim
|
||||
self.skip_connect_type = skip_connect_type
|
||||
needs_skip_proj = skip_connect_type == "concat"
|
||||
|
||||
self.depth = depth
|
||||
self.layers = nn.ModuleList([])
|
||||
|
||||
for idx in range(depth):
|
||||
is_later_half = idx >= (depth // 2)
|
||||
|
||||
attn_norm = RMSNorm(dim)
|
||||
attn = Attention(
|
||||
processor=AttnProcessor(),
|
||||
dim=dim,
|
||||
heads=heads,
|
||||
dim_head=dim_head,
|
||||
dropout=dropout,
|
||||
)
|
||||
|
||||
ff_norm = RMSNorm(dim)
|
||||
ff = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
|
||||
|
||||
skip_proj = nn.Linear(dim * 2, dim, bias=False) if needs_skip_proj and is_later_half else None
|
||||
|
||||
self.layers.append(
|
||||
nn.ModuleList(
|
||||
[
|
||||
skip_proj,
|
||||
attn_norm,
|
||||
attn,
|
||||
ff_norm,
|
||||
ff,
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
self.norm_out = RMSNorm(dim)
|
||||
self.proj_out = nn.Linear(dim, mel_dim)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: float["b n d"], # nosied input audio # noqa: F722
|
||||
cond: float["b n d"], # masked cond audio # noqa: F722
|
||||
text: int["b nt"], # text # noqa: F722
|
||||
time: float["b"] | float[""], # time step # noqa: F821 F722
|
||||
drop_audio_cond, # cfg for cond audio
|
||||
drop_text, # cfg for text
|
||||
mask: bool["b n"] | None = None, # noqa: F722
|
||||
):
|
||||
batch, seq_len = x.shape[0], x.shape[1]
|
||||
if time.ndim == 0:
|
||||
time = time.repeat(batch)
|
||||
|
||||
# t: conditioning time, c: context (text + masked cond audio), x: noised input audio
|
||||
t = self.time_embed(time)
|
||||
text_embed = self.text_embed(text, seq_len, drop_text=drop_text)
|
||||
x = self.input_embed(x, cond, text_embed, drop_audio_cond=drop_audio_cond)
|
||||
|
||||
# postfix time t to input x, [b n d] -> [b n+1 d]
|
||||
x = torch.cat([t.unsqueeze(1), x], dim=1) # pack t to x
|
||||
if mask is not None:
|
||||
mask = F.pad(mask, (1, 0), value=1)
|
||||
|
||||
rope = self.rotary_embed.forward_from_seq_len(seq_len + 1)
|
||||
|
||||
# flat unet transformer
|
||||
skip_connect_type = self.skip_connect_type
|
||||
skips = []
|
||||
for idx, (maybe_skip_proj, attn_norm, attn, ff_norm, ff) in enumerate(self.layers):
|
||||
layer = idx + 1
|
||||
|
||||
# skip connection logic
|
||||
is_first_half = layer <= (self.depth // 2)
|
||||
is_later_half = not is_first_half
|
||||
|
||||
if is_first_half:
|
||||
skips.append(x)
|
||||
|
||||
if is_later_half:
|
||||
skip = skips.pop()
|
||||
if skip_connect_type == "concat":
|
||||
x = torch.cat((x, skip), dim=-1)
|
||||
x = maybe_skip_proj(x)
|
||||
elif skip_connect_type == "add":
|
||||
x = x + skip
|
||||
|
||||
# attention and feedforward blocks
|
||||
x = attn(attn_norm(x), rope=rope, mask=mask) + x
|
||||
x = ff(ff_norm(x)) + x
|
||||
|
||||
assert len(skips) == 0
|
||||
|
||||
x = self.norm_out(x)[:, 1:, :] # unpack t from x
|
||||
|
||||
return self.proj_out(x)
|
||||
@ -10,139 +10,13 @@ d - dimension
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torchaudio
|
||||
from librosa.filters import mel as librosa_mel_fn
|
||||
from torch import nn
|
||||
from x_transformers.x_transformers import apply_rotary_pos_emb
|
||||
|
||||
|
||||
# raw wav to mel spec
|
||||
|
||||
|
||||
mel_basis_cache = {}
|
||||
hann_window_cache = {}
|
||||
|
||||
|
||||
def get_bigvgan_mel_spectrogram(
|
||||
waveform,
|
||||
n_fft=1024,
|
||||
n_mel_channels=100,
|
||||
target_sample_rate=24000,
|
||||
hop_length=256,
|
||||
win_length=1024,
|
||||
fmin=0,
|
||||
fmax=None,
|
||||
center=False,
|
||||
): # Copy from https://github.com/NVIDIA/BigVGAN/tree/main
|
||||
device = waveform.device
|
||||
key = f"{n_fft}_{n_mel_channels}_{target_sample_rate}_{hop_length}_{win_length}_{fmin}_{fmax}_{device}"
|
||||
|
||||
if key not in mel_basis_cache:
|
||||
mel = librosa_mel_fn(sr=target_sample_rate, n_fft=n_fft, n_mels=n_mel_channels, fmin=fmin, fmax=fmax)
|
||||
mel_basis_cache[key] = torch.from_numpy(mel).float().to(device) # TODO: why they need .float()?
|
||||
hann_window_cache[key] = torch.hann_window(win_length).to(device)
|
||||
|
||||
mel_basis = mel_basis_cache[key]
|
||||
hann_window = hann_window_cache[key]
|
||||
|
||||
padding = (n_fft - hop_length) // 2
|
||||
waveform = torch.nn.functional.pad(waveform.unsqueeze(1), (padding, padding), mode="reflect").squeeze(1)
|
||||
|
||||
spec = torch.stft(
|
||||
waveform,
|
||||
n_fft,
|
||||
hop_length=hop_length,
|
||||
win_length=win_length,
|
||||
window=hann_window,
|
||||
center=center,
|
||||
pad_mode="reflect",
|
||||
normalized=False,
|
||||
onesided=True,
|
||||
return_complex=True,
|
||||
)
|
||||
spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9)
|
||||
|
||||
mel_spec = torch.matmul(mel_basis, spec)
|
||||
mel_spec = torch.log(torch.clamp(mel_spec, min=1e-5))
|
||||
|
||||
return mel_spec
|
||||
|
||||
|
||||
def get_vocos_mel_spectrogram(
|
||||
waveform,
|
||||
n_fft=1024,
|
||||
n_mel_channels=100,
|
||||
target_sample_rate=24000,
|
||||
hop_length=256,
|
||||
win_length=1024,
|
||||
):
|
||||
mel_stft = torchaudio.transforms.MelSpectrogram(
|
||||
sample_rate=target_sample_rate,
|
||||
n_fft=n_fft,
|
||||
win_length=win_length,
|
||||
hop_length=hop_length,
|
||||
n_mels=n_mel_channels,
|
||||
power=1,
|
||||
center=True,
|
||||
normalized=False,
|
||||
norm=None,
|
||||
).to(waveform.device)
|
||||
if len(waveform.shape) == 3:
|
||||
waveform = waveform.squeeze(1) # 'b 1 nw -> b nw'
|
||||
|
||||
assert len(waveform.shape) == 2
|
||||
|
||||
mel = mel_stft(waveform)
|
||||
mel = mel.clamp(min=1e-5).log()
|
||||
return mel
|
||||
|
||||
|
||||
class MelSpec(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
n_fft=1024,
|
||||
hop_length=256,
|
||||
win_length=1024,
|
||||
n_mel_channels=100,
|
||||
target_sample_rate=24_000,
|
||||
mel_spec_type="vocos",
|
||||
):
|
||||
super().__init__()
|
||||
assert mel_spec_type in ["vocos", "bigvgan"], print("We only support two extract mel backend: vocos or bigvgan")
|
||||
|
||||
self.n_fft = n_fft
|
||||
self.hop_length = hop_length
|
||||
self.win_length = win_length
|
||||
self.n_mel_channels = n_mel_channels
|
||||
self.target_sample_rate = target_sample_rate
|
||||
|
||||
if mel_spec_type == "vocos":
|
||||
self.extractor = get_vocos_mel_spectrogram
|
||||
elif mel_spec_type == "bigvgan":
|
||||
self.extractor = get_bigvgan_mel_spectrogram
|
||||
|
||||
self.register_buffer("dummy", torch.tensor(0), persistent=False)
|
||||
|
||||
def forward(self, wav):
|
||||
if self.dummy.device != wav.device:
|
||||
self.to(wav.device)
|
||||
|
||||
mel = self.extractor(
|
||||
waveform=wav,
|
||||
n_fft=self.n_fft,
|
||||
n_mel_channels=self.n_mel_channels,
|
||||
target_sample_rate=self.target_sample_rate,
|
||||
hop_length=self.hop_length,
|
||||
win_length=self.win_length,
|
||||
)
|
||||
|
||||
return mel
|
||||
|
||||
|
||||
# sinusoidal position embedding
|
||||
|
||||
|
||||
@ -328,20 +202,18 @@ class FeedForward(nn.Module):
|
||||
return self.ff(x)
|
||||
|
||||
|
||||
# Attention with possible joint part
|
||||
# Attention
|
||||
# modified from diffusers/src/diffusers/models/attention_processor.py
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
processor: JointAttnProcessor | AttnProcessor,
|
||||
processor: AttnProcessor,
|
||||
dim: int,
|
||||
heads: int = 8,
|
||||
dim_head: int = 64,
|
||||
dropout: float = 0.0,
|
||||
context_dim: Optional[int] = None, # if not None -> joint attention
|
||||
context_pre_only=None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
@ -355,38 +227,21 @@ class Attention(nn.Module):
|
||||
self.inner_dim = dim_head * heads
|
||||
self.dropout = dropout
|
||||
|
||||
self.context_dim = context_dim
|
||||
self.context_pre_only = context_pre_only
|
||||
|
||||
self.to_q = nn.Linear(dim, self.inner_dim)
|
||||
self.to_k = nn.Linear(dim, self.inner_dim)
|
||||
self.to_v = nn.Linear(dim, self.inner_dim)
|
||||
|
||||
if self.context_dim is not None:
|
||||
self.to_k_c = nn.Linear(context_dim, self.inner_dim)
|
||||
self.to_v_c = nn.Linear(context_dim, self.inner_dim)
|
||||
if self.context_pre_only is not None:
|
||||
self.to_q_c = nn.Linear(context_dim, self.inner_dim)
|
||||
|
||||
self.to_out = nn.ModuleList([])
|
||||
self.to_out.append(nn.Linear(self.inner_dim, dim))
|
||||
self.to_out.append(nn.Dropout(dropout))
|
||||
|
||||
if self.context_pre_only is not None and not self.context_pre_only:
|
||||
self.to_out_c = nn.Linear(self.inner_dim, dim)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: float["b n d"], # noised input x # noqa: F722
|
||||
c: float["b n d"] = None, # context c # noqa: F722
|
||||
mask: bool["b n"] | None = None, # noqa: F722
|
||||
rope=None, # rotary position embedding for x
|
||||
c_rope=None, # rotary position embedding for c
|
||||
) -> torch.Tensor:
|
||||
if c is not None:
|
||||
return self.processor(self, x, c=c, mask=mask, rope=rope, c_rope=c_rope)
|
||||
else:
|
||||
return self.processor(self, x, mask=mask, rope=rope)
|
||||
return self.processor(self, x, mask=mask, rope=rope)
|
||||
|
||||
|
||||
# Attention processor
|
||||
@ -457,93 +312,6 @@ class AttnProcessor:
|
||||
return x
|
||||
|
||||
|
||||
# Joint Attention processor for MM-DiT
|
||||
# modified from diffusers/src/diffusers/models/attention_processor.py
|
||||
|
||||
|
||||
class JointAttnProcessor:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
attn: Attention,
|
||||
x: float["b n d"], # noised input x # noqa: F722
|
||||
c: float["b nt d"] = None, # context c, here text # noqa: F722
|
||||
mask: bool["b n"] | None = None, # noqa: F722
|
||||
rope=None, # rotary position embedding for x
|
||||
c_rope=None, # rotary position embedding for c
|
||||
) -> torch.FloatTensor:
|
||||
residual = x
|
||||
|
||||
batch_size = c.shape[0]
|
||||
|
||||
# `sample` projections.
|
||||
query = attn.to_q(x)
|
||||
key = attn.to_k(x)
|
||||
value = attn.to_v(x)
|
||||
|
||||
# `context` projections.
|
||||
c_query = attn.to_q_c(c)
|
||||
c_key = attn.to_k_c(c)
|
||||
c_value = attn.to_v_c(c)
|
||||
|
||||
# apply rope for context and noised input independently
|
||||
if rope is not None:
|
||||
freqs, xpos_scale = rope
|
||||
q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
|
||||
query = apply_rotary_pos_emb(query, freqs, q_xpos_scale)
|
||||
key = apply_rotary_pos_emb(key, freqs, k_xpos_scale)
|
||||
if c_rope is not None:
|
||||
freqs, xpos_scale = c_rope
|
||||
q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0)
|
||||
c_query = apply_rotary_pos_emb(c_query, freqs, q_xpos_scale)
|
||||
c_key = apply_rotary_pos_emb(c_key, freqs, k_xpos_scale)
|
||||
|
||||
# attention
|
||||
query = torch.cat([query, c_query], dim=1)
|
||||
key = torch.cat([key, c_key], dim=1)
|
||||
value = torch.cat([value, c_value], dim=1)
|
||||
|
||||
inner_dim = key.shape[-1]
|
||||
head_dim = inner_dim // attn.heads
|
||||
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
|
||||
|
||||
# mask. e.g. inference got a batch with different target durations, mask out the padding
|
||||
if mask is not None:
|
||||
attn_mask = F.pad(mask, (0, c.shape[1]), value=True) # no mask for c (text)
|
||||
attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) # 'b n -> b 1 1 n'
|
||||
attn_mask = attn_mask.expand(batch_size, attn.heads, query.shape[-2], key.shape[-2])
|
||||
else:
|
||||
attn_mask = None
|
||||
|
||||
x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False)
|
||||
x = x.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
|
||||
x = x.to(query.dtype)
|
||||
|
||||
# Split the attention outputs.
|
||||
x, c = (
|
||||
x[:, : residual.shape[1]],
|
||||
x[:, residual.shape[1] :],
|
||||
)
|
||||
|
||||
# linear proj
|
||||
x = attn.to_out[0](x)
|
||||
# dropout
|
||||
x = attn.to_out[1](x)
|
||||
if not attn.context_pre_only:
|
||||
c = attn.to_out_c(c)
|
||||
|
||||
if mask is not None:
|
||||
mask = mask.unsqueeze(-1)
|
||||
x = x.masked_fill(~mask, 0.0)
|
||||
# c = c.masked_fill(~mask, 0.) # no mask for c (text)
|
||||
|
||||
return x, c
|
||||
|
||||
|
||||
# DiT Block
|
||||
|
||||
|
||||
@ -580,76 +348,6 @@ class DiTBlock(nn.Module):
|
||||
return x
|
||||
|
||||
|
||||
# MMDiT Block https://arxiv.org/abs/2403.03206
|
||||
|
||||
|
||||
class MMDiTBlock(nn.Module):
|
||||
r"""
|
||||
modified from diffusers/src/diffusers/models/attention.py
|
||||
|
||||
notes.
|
||||
_c: context related. text, cond, etc. (left part in sd3 fig2.b)
|
||||
_x: noised input related. (right part)
|
||||
context_pre_only: last layer only do prenorm + modulation cuz no more ffn
|
||||
"""
|
||||
|
||||
def __init__(self, dim, heads, dim_head, ff_mult=4, dropout=0.1, context_pre_only=False):
|
||||
super().__init__()
|
||||
|
||||
self.context_pre_only = context_pre_only
|
||||
|
||||
self.attn_norm_c = AdaLayerNormZero_Final(dim) if context_pre_only else AdaLayerNormZero(dim)
|
||||
self.attn_norm_x = AdaLayerNormZero(dim)
|
||||
self.attn = Attention(
|
||||
processor=JointAttnProcessor(),
|
||||
dim=dim,
|
||||
heads=heads,
|
||||
dim_head=dim_head,
|
||||
dropout=dropout,
|
||||
context_dim=dim,
|
||||
context_pre_only=context_pre_only,
|
||||
)
|
||||
|
||||
if not context_pre_only:
|
||||
self.ff_norm_c = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
|
||||
self.ff_c = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
|
||||
else:
|
||||
self.ff_norm_c = None
|
||||
self.ff_c = None
|
||||
self.ff_norm_x = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
|
||||
self.ff_x = FeedForward(dim=dim, mult=ff_mult, dropout=dropout, approximate="tanh")
|
||||
|
||||
def forward(self, x, c, t, mask=None, rope=None, c_rope=None): # x: noised input, c: context, t: time embedding
|
||||
# pre-norm & modulation for attention input
|
||||
if self.context_pre_only:
|
||||
norm_c = self.attn_norm_c(c, t)
|
||||
else:
|
||||
norm_c, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = self.attn_norm_c(c, emb=t)
|
||||
norm_x, x_gate_msa, x_shift_mlp, x_scale_mlp, x_gate_mlp = self.attn_norm_x(x, emb=t)
|
||||
|
||||
# attention
|
||||
x_attn_output, c_attn_output = self.attn(x=norm_x, c=norm_c, mask=mask, rope=rope, c_rope=c_rope)
|
||||
|
||||
# process attention output for context c
|
||||
if self.context_pre_only:
|
||||
c = None
|
||||
else: # if not last layer
|
||||
c = c + c_gate_msa.unsqueeze(1) * c_attn_output
|
||||
|
||||
norm_c = self.ff_norm_c(c) * (1 + c_scale_mlp[:, None]) + c_shift_mlp[:, None]
|
||||
c_ff_output = self.ff_c(norm_c)
|
||||
c = c + c_gate_mlp.unsqueeze(1) * c_ff_output
|
||||
|
||||
# process attention output for input x
|
||||
x = x + x_gate_msa.unsqueeze(1) * x_attn_output
|
||||
|
||||
norm_x = self.ff_norm_x(x) * (1 + x_scale_mlp[:, None]) + x_shift_mlp[:, None]
|
||||
x_ff_output = self.ff_x(norm_x)
|
||||
x = x + x_gate_mlp.unsqueeze(1) * x_ff_output
|
||||
|
||||
return c, x
|
||||
|
||||
|
||||
# time step conditioning embedding
|
||||
|
||||
|
||||
|
||||
@ -482,8 +482,7 @@ def init_bigvgan():
|
||||
|
||||
bigvgan_model = bigvgan.BigVGAN.from_pretrained(
|
||||
"%s/GPT_SoVITS/pretrained_models/models--nvidia--bigvgan_v2_24khz_100band_256x" % (now_dir,),
|
||||
use_cuda_kernel=False,
|
||||
) # if True, RuntimeError: Ninja is required to load C++ extensions
|
||||
)
|
||||
# remove weight norm in the model and set to eval mode
|
||||
bigvgan_model.remove_weight_norm()
|
||||
bigvgan_model = bigvgan_model.eval()
|
||||
|
||||
39
README.md
39
README.md
@ -27,6 +27,44 @@ A Powerful Few-shot Voice Conversion and Text-to-Speech WebUI.<br><br>
|
||||
|
||||
---
|
||||
|
||||
## 简化接口 / 测试前端
|
||||
|
||||
本工作区新增了一个用于 MVP 调用的中间层接口和测试前端：
|
||||
|
||||
| 页面 | 地址 | 说明 |
|
||||
|------|------|------|
|
||||
| 测试前端 | http://127.0.0.1:9881/test/ | 带波形裁剪、视频转音频的测试 UI |
|
||||
| Swagger UI | http://127.0.0.1:9881/docs | 交互式 API 文档 |
|
||||
| ReDoc | http://127.0.0.1:9881/redoc | 可读式 API 文档 |
|
||||
|
||||
文档：
|
||||
|
||||
- 快速启动：[docs/simple_api_quickstart.md](./docs/simple_api_quickstart.md)
|
||||
- 完整教程：[docs/simple_api.md](./docs/simple_api.md)
|
||||
|
||||
核心功能：
|
||||
|
||||
- `/api/tts` — 上传参考音频/视频 + 文字，直接返回生成的音频
|
||||
- 前端支持视频上传，自动提取音频
|
||||
- 前端波形裁剪工具，可选择 3-10 秒片段
|
||||
- 5 种情绪预设（neutral/happy/calm/sad/angry）
|
||||
- 7 个契约测试，无需 GPU 即可运行
|
||||
|
||||
启动命令：
|
||||
|
||||
```
|
||||
cd D:\tts\GPT-SoVITS
|
||||
.\go-simple-api.ps1
|
||||
```
|
||||
|
||||
运行测试：
|
||||
|
||||
```
|
||||
python -m unittest tests.test_simple_api_contract -v
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Features:
|
||||
|
||||
1. **Zero-shot TTS:** Input a 5-second vocal sample and experience instant text-to-speech conversion.
|
||||
@ -480,3 +518,4 @@ Thankful to @Naozumi520 for providing the Cantonese training set and for the gui
|
||||
<a href="https://github.com/RVC-Boss/GPT-SoVITS/graphs/contributors" target="_blank">
|
||||
<img src="https://contrib.rocks/image?repo=RVC-Boss/GPT-SoVITS" />
|
||||
</a>
|
||||
|
||||
|
||||
3
api.py
3
api.py
@ -240,8 +240,7 @@ def init_bigvgan():
|
||||
|
||||
bigvgan_model = bigvgan.BigVGAN.from_pretrained(
|
||||
"%s/GPT_SoVITS/pretrained_models/models--nvidia--bigvgan_v2_24khz_100band_256x" % (now_dir,),
|
||||
use_cuda_kernel=False,
|
||||
) # if True, RuntimeError: Ninja is required to load C++ extensions
|
||||
)
|
||||
# remove weight norm in the model and set to eval mode
|
||||
bigvgan_model.remove_weight_norm()
|
||||
bigvgan_model = bigvgan_model.eval()
|
||||
|
||||
402
docs/simple_api.md
Normal file
402
docs/simple_api.md
Normal file
@ -0,0 +1,402 @@
|
||||
# GPT-SoVITS 简化接口文档
|
||||
|
||||
本项目新增 `simple_api.py` 作为中间层，封装 GPT-SoVITS 推理引擎，提供更简洁的调用方式。
|
||||
|
||||
## 快速开始
|
||||
|
||||
```bash
|
||||
# 安装依赖
|
||||
python -m pip install -r requirements.txt
|
||||
|
||||
# 启动
|
||||
python simple_api.py -c simple_api.yaml
|
||||
|
||||
# 访问
|
||||
Swagger UI: http://127.0.0.1:9881/docs
|
||||
ReDoc: http://127.0.0.1:9881/redoc
|
||||
测试前端: http://127.0.0.1:9881/test/
|
||||
```
|
||||
|
||||
## 接口总览
|
||||
|
||||
| 方法 | 路径 | 说明 | 标签 |
|
||||
|------|------|------|------|
|
||||
| GET | `/health` | 健康检查（含 GPU 信息） | System |
|
||||
| GET | `/voices` | 列出 voice profiles | System |
|
||||
| **POST** | **`/api/tts`** | **核心 TTS 接口（MVP）** | **MVP** |
|
||||
| GET | `/speak` | voice profile TTS (GET) | Profile |
|
||||
| POST | `/speak` | voice profile TTS (POST) | Profile |
|
||||
| POST | `/v1/tts` | OpenAI 兼容格式 TTS | Profile |
|
||||
| POST | `/speak/base64` | 返回 Base64 音频 | Profile |
|
||||
| POST | `/admin/reload-config` | 热加载配置 | Admin |
|
||||
| POST | `/admin/weights` | 切换模型权重 | Admin |
|
||||
|
||||
---
|
||||
|
||||
## 1. POST /api/tts — 核心 TTS 接口
|
||||
|
||||
**推荐使用此接口**。上传参考音频和文字，直接返回生成的音频。
|
||||
|
||||
### 请求格式
|
||||
|
||||
```
|
||||
Content-Type: multipart/form-data
|
||||
```
|
||||
|
||||
### 字段说明
|
||||
|
||||
| 字段 | 类型 | 必填 | 默认值 | 说明 |
|
||||
|------|------|------|--------|------|
|
||||
| `text` | string | **是** | — | 需要生成的文字 |
|
||||
| `ref_audio` | file | **是** | — | 主参考音频，3-10 秒（支持 wav/flac/ogg/mp3/m4a/aac） |
|
||||
| `aux_ref_audio` | file[] | 否 | — | 辅助参考音频，可上传多个 |
|
||||
| `prompt_text` | string | 否 | `""` | 主参考音频对应文字（v2 可留空；v3/v4 必填） |
|
||||
| `text_lang` | string | 否 | `zh` | 生成文字语言：zh/en/ja/ko/yue/auto |
|
||||
| `prompt_lang` | string | 否 | `zh` | 参考音频语言：zh/en/ja/ko/yue/auto |
|
||||
| `format` | string | 否 | `wav` | 返回格式：wav/ogg/aac/raw |
|
||||
| `emotion` | string | 否 | `neutral` | 情绪预设：neutral/happy/calm/sad/angry |
|
||||
| `speed` | float | 否 | — | 语速（0.5-2.0），覆盖情绪预设中的语速 |
|
||||
| `seed` | int | 否 | `-1` | 随机种子，-1 为随机 |
|
||||
|
||||
### 情绪预设参数映射
|
||||
|
||||
| 情绪 | temperature | top_p | top_k | speed_factor | repetition_penalty |
|
||||
|------|-------------|-------|-------|--------------|-------------------|
|
||||
| neutral | — | — | — | — | — |
|
||||
| happy | 1.1 | 0.95 | — | — | — |
|
||||
| calm | 0.8 | 0.85 | — | 0.92 | — |
|
||||
| sad | 0.75 | 0.85 | — | 0.9 | — |
|
||||
| angry | 1.2 | — | 20 | — | 1.25 |
|
||||
|
||||
> 显式传入 `speed` 会覆盖情绪预设中的 `speed_factor`。
|
||||
|
||||
### curl 示例
|
||||
|
||||
**基础调用：**
|
||||
|
||||
```powershell
|
||||
curl.exe -X POST http://127.0.0.1:9881/api/tts `
|
||||
-F "text=你好，欢迎使用这个声音。" `
|
||||
-F "ref_audio=@D:\audio\ref.wav" `
|
||||
--output output.wav
|
||||
```
|
||||
|
||||
**带辅助参考音频和情绪：**
|
||||
|
||||
```powershell
|
||||
curl.exe -X POST http://127.0.0.1:9881/api/tts `
|
||||
-F "text=你好，欢迎使用这个声音。" `
|
||||
-F "ref_audio=@D:\audio\ref.wav" `
|
||||
-F "aux_ref_audio=@D:\audio\aux1.wav" `
|
||||
-F "emotion=happy" `
|
||||
-F "speed=1.1" `
|
||||
--output output.wav
|
||||
```
|
||||
|
||||
**Linux/macOS：**
|
||||
|
||||
```bash
|
||||
curl -X POST http://127.0.0.1:9881/api/tts \
|
||||
-F "text=你好，欢迎使用这个声音。" \
|
||||
-F "ref_audio=@/path/to/ref.wav" \
|
||||
-F "emotion=calm" \
|
||||
--output output.wav
|
||||
```
|
||||
|
||||
### 返回
|
||||
|
||||
- 成功：音频二进制流（Content-Type: `audio/wav` 等）
|
||||
- 失败：JSON 错误信息
|
||||
|
||||
```json
|
||||
{"message": "tts failed", "exception": "..."}
|
||||
```
|
||||
|
||||
### 常见错误
|
||||
|
||||
| HTTP 状态码 | 原因 |
|
||||
|------------|------|
|
||||
| 400 | text 为空 / ref_audio 缺失 / 音频时长不在 3-10 秒 / 不支持的 format / v3/v4 时 prompt_text 为空 |
|
||||
| 404 | voice profile 不存在（仅 /speak 接口） |
|
||||
| 503 | TTS pipeline 未就绪（模型未加载） |
|
||||
|
||||
---
|
||||
|
||||
## 2. GET /health — 健康检查
|
||||
|
||||
```bash
|
||||
curl http://127.0.0.1:9881/health
|
||||
```
|
||||
|
||||
返回示例：
|
||||
|
||||
```json
|
||||
{
|
||||
"status": "ok",
|
||||
"tts_config": "GPT_SoVITS/configs/tts_infer.yaml",
|
||||
"version": "v2",
|
||||
"languages": ["auto", "en", "zh"],
|
||||
"pid": 12345,
|
||||
"memory_mb": 2048.5,
|
||||
"gpu": {
|
||||
"name": "NVIDIA GeForce RTX 3080",
|
||||
"memory_used_mb": 4096.2,
|
||||
"memory_total_mb": 10240.0
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 3. GET /voices — 列出 voice profiles
|
||||
|
||||
```bash
|
||||
curl http://127.0.0.1:9881/voices
|
||||
```
|
||||
|
||||
返回示例：
|
||||
|
||||
```json
|
||||
{
|
||||
"default_voice": "default",
|
||||
"voices": [
|
||||
{
|
||||
"name": "default",
|
||||
"description": "Replace this profile with your reference voice.",
|
||||
"text_lang": "zh",
|
||||
"prompt_lang": "zh",
|
||||
"ref_audio_path": "reference.wav",
|
||||
"ready": true
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 4. POST /speak — voice profile TTS
|
||||
|
||||
基于 `simple_api.yaml` 中配置的 voice profile 调用 TTS。
|
||||
|
||||
### 请求体（JSON）
|
||||
|
||||
```json
|
||||
{
|
||||
"text": "hello world",
|
||||
"voice": "default",
|
||||
"text_lang": "zh",
|
||||
"format": "wav",
|
||||
"speed": 1.0
|
||||
}
|
||||
```
|
||||
|
||||
| 字段 | 类型 | 必填 | 说明 |
|
||||
|------|------|------|------|
|
||||
| `text` | string | **是** | 需要生成的文字 |
|
||||
| `voice` | string | 否 | voice profile 名称，不传则使用 default |
|
||||
| `text_lang` | string | 否 | 生成文字语言 |
|
||||
| `format` | string | 否 | 返回格式 |
|
||||
| `stream` | bool | 否 | 是否流式返回 |
|
||||
| `speed` | float | 否 | 语速 |
|
||||
|
||||
### curl 示例
|
||||
|
||||
```bash
|
||||
curl -X POST http://127.0.0.1:9881/speak \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"text":"你好世界","voice":"default"}' \
|
||||
--output output.wav
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 5. GET /speak — voice profile TTS (GET)
|
||||
|
||||
与 POST /speak 相同，但通过 URL 参数传递。
|
||||
|
||||
```
|
||||
GET /speak?text=hello&voice=default&format=wav
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 6. POST /speak/base64 — 返回 Base64 音频
|
||||
|
||||
返回 Base64 编码的音频，适合 Web 前端直接使用。
|
||||
|
||||
```bash
|
||||
curl -X POST http://127.0.0.1:9881/speak/base64 \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"text":"hello","voice":"default"}'
|
||||
```
|
||||
|
||||
返回：
|
||||
|
||||
```json
|
||||
{
|
||||
"media_type": "audio/wav",
|
||||
"audio_base64": "UklGRi..."
|
||||
}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 7. POST /v1/tts — OpenAI 兼容格式
|
||||
|
||||
请求格式与 POST /speak 相同，路径兼容 OpenAI TTS API 风格。
|
||||
|
||||
---
|
||||
|
||||
## 8. POST /admin/reload-config — 热加载配置
|
||||
|
||||
重新加载 `simple_api.yaml`，无需重启服务。
|
||||
|
||||
```bash
|
||||
curl -X POST http://127.0.0.1:9881/admin/reload-config
|
||||
```
|
||||
|
||||
返回：`{"message": "success", "default_voice": "default"}`
|
||||
|
||||
---
|
||||
|
||||
## 9. POST /admin/weights — 切换模型权重
|
||||
|
||||
运行时切换 GPT-SoVITS 模型权重文件。
|
||||
|
||||
```bash
|
||||
curl -X POST http://127.0.0.1:9881/admin/weights \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"gpt_weights_path":"path/to/gpt.pt","sovits_weights_path":"path/to/sovits.pt"}'
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 配置文件
|
||||
|
||||
`simple_api.yaml`：
|
||||
|
||||
```yaml
|
||||
server:
|
||||
host: 127.0.0.1
|
||||
port: 9881
|
||||
tts_config: GPT_SoVITS/configs/tts_infer.yaml
|
||||
|
||||
cors_allow_origins:
|
||||
- "*"
|
||||
|
||||
upload:
|
||||
dir: runtime/uploads
|
||||
min_ref_seconds: 3
|
||||
max_ref_seconds: 10
|
||||
max_upload_mb: 80
|
||||
|
||||
defaults:
|
||||
text_lang: zh
|
||||
prompt_lang: zh
|
||||
media_type: wav
|
||||
text_split_method: cut5
|
||||
batch_size: 1
|
||||
speed_factor: 1.0
|
||||
seed: -1
|
||||
|
||||
emotion_presets:
|
||||
neutral: {}
|
||||
happy:
|
||||
temperature: 1.1
|
||||
top_p: 0.95
|
||||
calm:
|
||||
temperature: 0.8
|
||||
top_p: 0.85
|
||||
speed_factor: 0.92
|
||||
sad:
|
||||
temperature: 0.75
|
||||
top_p: 0.85
|
||||
speed_factor: 0.9
|
||||
angry:
|
||||
temperature: 1.2
|
||||
top_k: 20
|
||||
repetition_penalty: 1.25
|
||||
|
||||
voices:
|
||||
default:
|
||||
description: Replace this profile with your reference voice.
|
||||
ref_audio_path: reference.wav
|
||||
prompt_text: Replace this with the exact text spoken in reference.wav.
|
||||
prompt_lang: zh
|
||||
text_lang: zh
|
||||
```
|
||||
|
||||
### 配置说明
|
||||
|
||||
| 配置项 | 说明 |
|
||||
|--------|------|
|
||||
| `server.host` | 监听地址 |
|
||||
| `server.port` | 监听端口 |
|
||||
| `server.tts_config` | GPT-SoVITS 推理配置文件路径 |
|
||||
| `upload.dir` | 临时上传目录 |
|
||||
| `upload.min_ref_seconds` | 主参考音频最短秒数 |
|
||||
| `upload.max_ref_seconds` | 主参考音频最长秒数 |
|
||||
| `upload.max_upload_mb` | 单个上传文件最大体积 (MB) |
|
||||
| `defaults.*` | 所有接口的默认参数 |
|
||||
| `emotion_presets.*` | 情绪预设参数映射 |
|
||||
| `voices.*` | 固定音色 profile |
|
||||
|
||||
---
|
||||
|
||||
## 添加自定义音色
|
||||
|
||||
编辑 `simple_api.yaml`，在 `voices` 下添加：
|
||||
|
||||
```yaml
|
||||
voices:
|
||||
narrator:
|
||||
description: "男声旁白"
|
||||
ref_audio_path: voices/narrator.wav
|
||||
prompt_text: "旁白参考音频的逐字稿"
|
||||
prompt_lang: zh
|
||||
text_lang: zh
|
||||
```
|
||||
|
||||
然后热加载：
|
||||
|
||||
```bash
|
||||
curl -X POST http://127.0.0.1:9881/admin/reload-config
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 测试
|
||||
|
||||
### 契约测试（无需 GPU）
|
||||
|
||||
```bash
|
||||
python -m unittest tests.test_simple_api_contract -v
|
||||
```
|
||||
|
||||
覆盖：
|
||||
|
||||
- `/api/tts` 路由注册
|
||||
- 上传接口参数构造
|
||||
- 主参考音频 3-10 秒校验
|
||||
- v2 空 prompt_text 允许 / v3/v4 空 prompt_text 拒绝
|
||||
- 临时上传目录清理
|
||||
- 情绪预设应用与 speed 覆盖
|
||||
|
||||
### 前端测试
|
||||
|
||||
1. 启动后端
|
||||
2. 访问 `http://127.0.0.1:9881/test/`
|
||||
3. 上传音频或视频（视频会自动提取音频）
|
||||
4. 使用波形裁剪工具选择 3-10 秒片段
|
||||
5. 填写文字，选择情绪和语速
|
||||
6. 点击生成
|
||||
|
||||
---
|
||||
|
||||
## 启动脚本
|
||||
|
||||
| 脚本 | 平台 | 说明 |
|
||||
|------|------|------|
|
||||
| `go-simple-api.ps1` | Windows PowerShell | 自动检测 runtime\python.exe |
|
||||
| `go-simple-api.bat` | Windows CMD | 同上 |
|
||||
| `open-test-frontend.ps1` | Windows PowerShell | 直接打开测试前端 HTML |
|
||||
85
docs/simple_api_quickstart.md
Normal file
85
docs/simple_api_quickstart.md
Normal file
@ -0,0 +1,85 @@
|
||||
# 简化 TTS 接口快速启动
|
||||
|
||||
完整教程见：`docs/simple_api.md`
|
||||
|
||||
## 一句话流程
|
||||
|
||||
启动后端，打开测试页，上传 3-10 秒参考音频或视频（视频自动提取音频），填写生成文字，点击生成。
|
||||
|
||||
## 1. 启动后端
|
||||
|
||||
```powershell
|
||||
cd D:\tts\GPT-SoVITS
|
||||
python -m pip install -r requirements.txt
|
||||
.\go-simple-api.ps1
|
||||
```
|
||||
|
||||
也可以运行：
|
||||
|
||||
```bat
|
||||
go-simple-api.bat
|
||||
```
|
||||
|
||||
默认后端地址：
|
||||
|
||||
```text
|
||||
http://127.0.0.1:9881
|
||||
```
|
||||
|
||||
## 2. 打开页面
|
||||
|
||||
| 页面 | 地址 | 说明 |
|
||||
|------|------|------|
|
||||
| 测试前端 | http://127.0.0.1:9881/test/ | 带波形裁剪的测试 UI |
|
||||
| Swagger UI | http://127.0.0.1:9881/docs | 交互式 API 文档 |
|
||||
| ReDoc | http://127.0.0.1:9881/redoc | 可读式 API 文档 |
|
||||
|
||||
## 3. 调用接口
|
||||
|
||||
接口：
|
||||
|
||||
```http
|
||||
POST /api/tts
|
||||
Content-Type: multipart/form-data
|
||||
```
|
||||
|
||||
最小字段：
|
||||
|
||||
```text
|
||||
text 要生成的文字
|
||||
ref_audio 3-10 秒主参考音频（或视频，前端自动提取）
|
||||
```
|
||||
|
||||
常用可选字段：
|
||||
|
||||
```text
|
||||
aux_ref_audio 辅助参考音频，可多个
|
||||
prompt_text 参考音频文本，可留空
|
||||
text_lang 默认 zh
|
||||
prompt_lang 默认 zh
|
||||
emotion neutral / happy / calm / sad / angry
|
||||
speed 语速，默认 1
|
||||
seed 默认 -1
|
||||
```
|
||||
|
||||
## 4. PowerShell 示例
|
||||
|
||||
```powershell
|
||||
curl.exe -X POST http://127.0.0.1:9881/api/tts `
|
||||
-F "text=你好，欢迎使用这个声音。" `
|
||||
-F "ref_audio=@D:\audio\ref.wav" `
|
||||
-F "prompt_text=" `
|
||||
-F "text_lang=zh" `
|
||||
-F "prompt_lang=zh" `
|
||||
--output output.wav
|
||||
```
|
||||
|
||||
## 5. 注意事项
|
||||
|
||||
- 主参考音频必须是 3-10 秒。
|
||||
- 前端支持上传视频文件，会自动提取音频。
|
||||
- 前端提供波形裁剪工具，可直接选择 3-10 秒片段。
|
||||
- `prompt_text` 在当前 v2 配置下可以为空。
|
||||
- 如果切换到 v3/v4，空 `prompt_text` 会返回 400。
|
||||
- 生成文字固定按标点符号切句。
|
||||
- 更详细的配置、profile、base64 接口见 `docs/simple_api.md`。
|
||||
12
go-simple-api.bat
Normal file
12
go-simple-api.bat
Normal file
@ -0,0 +1,12 @@
|
||||
@echo off
|
||||
set "SCRIPT_DIR=%~dp0"
|
||||
set "SCRIPT_DIR=%SCRIPT_DIR:~0,-1%"
|
||||
cd /d "%SCRIPT_DIR%"
|
||||
|
||||
if exist "%SCRIPT_DIR%\runtime\python.exe" (
|
||||
"%SCRIPT_DIR%\runtime\python.exe" "%SCRIPT_DIR%\simple_api.py" -c "%SCRIPT_DIR%\simple_api.yaml" %*
|
||||
) else (
|
||||
python "%SCRIPT_DIR%\simple_api.py" -c "%SCRIPT_DIR%\simple_api.yaml" %*
|
||||
)
|
||||
|
||||
pause
|
||||
13
go-simple-api.ps1
Normal file
13
go-simple-api.ps1
Normal file
@ -0,0 +1,13 @@
|
||||
$ErrorActionPreference = "Stop"
|
||||
chcp 65001
|
||||
Set-Location $PSScriptRoot
|
||||
|
||||
$runtimePython = Join-Path $PSScriptRoot "runtime\python.exe"
|
||||
$apiScript = Join-Path $PSScriptRoot "simple_api.py"
|
||||
$apiConfig = Join-Path $PSScriptRoot "simple_api.yaml"
|
||||
|
||||
if (Test-Path $runtimePython) {
|
||||
& $runtimePython $apiScript -c $apiConfig @args
|
||||
} else {
|
||||
python $apiScript -c $apiConfig @args
|
||||
}
|
||||
3
open-test-frontend.ps1
Normal file
3
open-test-frontend.ps1
Normal file
@ -0,0 +1,3 @@
|
||||
$ErrorActionPreference = "Stop"
|
||||
Set-Location $PSScriptRoot
|
||||
Start-Process (Join-Path $PSScriptRoot "test_frontend\index.html")
|
||||
@ -3,6 +3,7 @@ numpy<2.0
|
||||
scipy
|
||||
tensorboard
|
||||
librosa==0.10.2
|
||||
soundfile
|
||||
numba
|
||||
pytorch-lightning>=2.4
|
||||
gradio<5
|
||||
@ -18,10 +19,11 @@ g2p_en
|
||||
torchaudio
|
||||
modelscope
|
||||
sentencepiece
|
||||
transformers>=4.43,<=4.50
|
||||
transformers>=4.51,<5
|
||||
peft<0.18.0
|
||||
chardet
|
||||
PyYAML
|
||||
python-multipart
|
||||
psutil
|
||||
jieba_fast
|
||||
jieba
|
||||
|
||||
787
simple_api.py
Normal file
787
simple_api.py
Normal file
@ -0,0 +1,787 @@
|
||||
"""
|
||||
Small profile-based API layer for GPT-SoVITS.
|
||||
|
||||
Run:
|
||||
python simple_api.py -c simple_api.yaml
|
||||
|
||||
Then call:
|
||||
POST http://127.0.0.1:9881/speak
|
||||
{"text": "hello", "voice": "default"}
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import traceback
|
||||
import uuid
|
||||
import wave
|
||||
from copy import deepcopy
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, Generator, List, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import yaml
|
||||
from fastapi import FastAPI, File, Form, HTTPException, Response, UploadFile
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import JSONResponse, StreamingResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parent
|
||||
os.chdir(PROJECT_ROOT)
|
||||
sys.path.append(str(PROJECT_ROOT))
|
||||
sys.path.append(str(PROJECT_ROOT / "GPT_SoVITS"))
|
||||
|
||||
from GPT_SoVITS.TTS_infer_pack.TTS import TTS, TTS_Config # noqa: E402
|
||||
from GPT_SoVITS.TTS_infer_pack.text_segmentation_method import ( # noqa: E402
|
||||
get_method_names as get_cut_method_names,
|
||||
)
|
||||
|
||||
|
||||
DEFAULT_TTS_PARAMS: Dict[str, Any] = {
|
||||
"text": "",
|
||||
"text_lang": "zh",
|
||||
"ref_audio_path": "",
|
||||
"aux_ref_audio_paths": [],
|
||||
"prompt_text": "",
|
||||
"prompt_lang": "zh",
|
||||
"top_k": 15,
|
||||
"top_p": 1.0,
|
||||
"temperature": 1.0,
|
||||
"text_split_method": "cut5",
|
||||
"batch_size": 1,
|
||||
"batch_threshold": 0.75,
|
||||
"split_bucket": True,
|
||||
"speed_factor": 1.0,
|
||||
"fragment_interval": 0.3,
|
||||
"seed": -1,
|
||||
"media_type": "wav",
|
||||
"streaming_mode": False,
|
||||
"return_fragment": False,
|
||||
"fixed_length_chunk": False,
|
||||
"parallel_infer": True,
|
||||
"repetition_penalty": 1.35,
|
||||
"sample_steps": 32,
|
||||
"super_sampling": False,
|
||||
"overlap_length": 2,
|
||||
"min_chunk_length": 16,
|
||||
}
|
||||
|
||||
SUPPORTED_MEDIA_TYPES = {"wav", "raw", "ogg", "aac"}
|
||||
SUPPORTED_UPLOAD_EXTENSIONS = {".wav", ".flac", ".ogg", ".mp3", ".m4a", ".aac"}
|
||||
VOICE_META_KEYS = {"description"}
|
||||
|
||||
|
||||
class SpeakRequest(BaseModel):
|
||||
text: str
|
||||
voice: Optional[str] = None
|
||||
text_lang: Optional[str] = None
|
||||
ref_audio_path: Optional[str] = None
|
||||
aux_ref_audio_paths: Optional[List[str]] = None
|
||||
prompt_text: Optional[str] = None
|
||||
prompt_lang: Optional[str] = None
|
||||
format: Optional[str] = Field(default=None, description="wav, raw, ogg, or aac")
|
||||
stream: Optional[bool] = Field(default=None, description="true maps to streaming mode 2")
|
||||
streaming_mode: Optional[Union[bool, int]] = Field(default=None, description="0, 1, 2, 3, true, or false")
|
||||
speed: Optional[float] = Field(default=None, description="Alias of speed_factor")
|
||||
speed_factor: Optional[float] = None
|
||||
top_k: Optional[int] = None
|
||||
top_p: Optional[float] = None
|
||||
temperature: Optional[float] = None
|
||||
text_split_method: Optional[str] = None
|
||||
batch_size: Optional[int] = None
|
||||
batch_threshold: Optional[float] = None
|
||||
split_bucket: Optional[bool] = None
|
||||
fragment_interval: Optional[float] = None
|
||||
seed: Optional[int] = None
|
||||
parallel_infer: Optional[bool] = None
|
||||
repetition_penalty: Optional[float] = None
|
||||
sample_steps: Optional[int] = None
|
||||
super_sampling: Optional[bool] = None
|
||||
overlap_length: Optional[int] = None
|
||||
min_chunk_length: Optional[int] = None
|
||||
|
||||
|
||||
class WeightsRequest(BaseModel):
|
||||
gpt_weights_path: Optional[str] = None
|
||||
sovits_weights_path: Optional[str] = None
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="GPT-SoVITS simple API")
|
||||
parser.add_argument("-c", "--config", default="simple_api.yaml", help="simple API config path")
|
||||
parser.add_argument("--tts-config", default=None, help="GPT-SoVITS tts_infer.yaml path")
|
||||
parser.add_argument("-a", "--bind-addr", default=None, help="bind address")
|
||||
parser.add_argument("-p", "--port", type=int, default=None, help="bind port")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def load_yaml_config(config_path: Union[str, Path]) -> Dict[str, Any]:
|
||||
path = Path(config_path)
|
||||
if not path.is_absolute():
|
||||
path = PROJECT_ROOT / path
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(f"simple API config not found: {path}")
|
||||
with path.open("r", encoding="utf-8") as f:
|
||||
data = yaml.safe_load(f) or {}
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("simple API config must be a YAML object")
|
||||
return data
|
||||
|
||||
|
||||
args = parse_args()
|
||||
simple_config = load_yaml_config(args.config)
|
||||
server_config = simple_config.get("server", {}) or {}
|
||||
host = args.bind_addr or server_config.get("host", "127.0.0.1")
|
||||
port = args.port or int(server_config.get("port", 9881))
|
||||
tts_config_path = args.tts_config or server_config.get("tts_config", "GPT_SoVITS/configs/tts_infer.yaml")
|
||||
|
||||
cut_method_names = get_cut_method_names()
|
||||
tts_config: Optional[TTS_Config] = None
|
||||
tts_pipeline: Optional[TTS] = None
|
||||
infer_lock = threading.Lock()
|
||||
|
||||
APP = FastAPI(
|
||||
title="GPT-SoVITS Simple API",
|
||||
description=(
|
||||
"简化接口层，封装 GPT-SoVITS 推理引擎。\n\n"
|
||||
"## 核心流程\n"
|
||||
"1. 上传 3-10 秒参考音频（或视频，前端自动提取音频）\n"
|
||||
"2. 填写需要生成的文字\n"
|
||||
"3. 调用 `/api/tts` 获取生成的音频\n\n"
|
||||
"## 其他接口\n"
|
||||
"- `/speak` — 基于 voice profile 的调用方式\n"
|
||||
"- `/v1/tts` — OpenAI 兼容格式\n"
|
||||
"- `/admin/*` — 管理接口（热加载配置、切换模型）"
|
||||
),
|
||||
version="1.1.0",
|
||||
docs_url="/docs",
|
||||
redoc_url="/redoc",
|
||||
openapi_tags=[
|
||||
{"name": "MVP", "description": "核心 TTS 接口，上传参考音频直接生成"},
|
||||
{"name": "Profile", "description": "基于 voice profile 的调用方式"},
|
||||
{"name": "Admin", "description": "管理接口：热加载配置、切换模型权重"},
|
||||
{"name": "System", "description": "健康检查与信息查询"},
|
||||
],
|
||||
)
|
||||
APP.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=simple_config.get("cors_allow_origins", ["*"]),
|
||||
allow_credentials=False,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
test_frontend_dir = PROJECT_ROOT / "test_frontend"
|
||||
if test_frontend_dir.exists():
|
||||
APP.mount("/test", StaticFiles(directory=str(test_frontend_dir), html=True), name="test_frontend")
|
||||
|
||||
|
||||
def pack_ogg(io_buffer: BytesIO, data: np.ndarray, rate: int) -> BytesIO:
|
||||
def handle_pack_ogg() -> None:
|
||||
with sf.SoundFile(io_buffer, mode="w", samplerate=rate, channels=1, format="ogg") as audio_file:
|
||||
audio_file.write(data)
|
||||
|
||||
try:
|
||||
threading.stack_size(4096 * 4096)
|
||||
pack_thread = threading.Thread(target=handle_pack_ogg)
|
||||
pack_thread.start()
|
||||
pack_thread.join()
|
||||
except (RuntimeError, ValueError):
|
||||
handle_pack_ogg()
|
||||
return io_buffer
|
||||
|
||||
|
||||
def pack_raw(io_buffer: BytesIO, data: np.ndarray, rate: int) -> BytesIO:
|
||||
del rate
|
||||
io_buffer.write(data.tobytes())
|
||||
return io_buffer
|
||||
|
||||
|
||||
def pack_wav(io_buffer: BytesIO, data: np.ndarray, rate: int) -> BytesIO:
|
||||
del io_buffer
|
||||
wav_buffer = BytesIO()
|
||||
sf.write(wav_buffer, data, rate, format="wav")
|
||||
return wav_buffer
|
||||
|
||||
|
||||
def pack_aac(io_buffer: BytesIO, data: np.ndarray, rate: int) -> BytesIO:
|
||||
process = subprocess.Popen(
|
||||
[
|
||||
"ffmpeg",
|
||||
"-f",
|
||||
"s16le",
|
||||
"-ar",
|
||||
str(rate),
|
||||
"-ac",
|
||||
"1",
|
||||
"-i",
|
||||
"pipe:0",
|
||||
"-c:a",
|
||||
"aac",
|
||||
"-b:a",
|
||||
"192k",
|
||||
"-vn",
|
||||
"-f",
|
||||
"adts",
|
||||
"pipe:1",
|
||||
],
|
||||
stdin=subprocess.PIPE,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
)
|
||||
out, _ = process.communicate(input=data.tobytes())
|
||||
io_buffer.write(out)
|
||||
return io_buffer
|
||||
|
||||
|
||||
def pack_audio(io_buffer: BytesIO, data: np.ndarray, rate: int, media_type: str) -> BytesIO:
|
||||
if media_type == "ogg":
|
||||
io_buffer = pack_ogg(io_buffer, data, rate)
|
||||
elif media_type == "aac":
|
||||
io_buffer = pack_aac(io_buffer, data, rate)
|
||||
elif media_type == "wav":
|
||||
io_buffer = pack_wav(io_buffer, data, rate)
|
||||
else:
|
||||
io_buffer = pack_raw(io_buffer, data, rate)
|
||||
io_buffer.seek(0)
|
||||
return io_buffer
|
||||
|
||||
|
||||
def wave_header_chunk(frame_input: bytes = b"", channels: int = 1, sample_width: int = 2, sample_rate: int = 32000) -> bytes:
|
||||
wav_buf = BytesIO()
|
||||
with wave.open(wav_buf, "wb") as vfout:
|
||||
vfout.setnchannels(channels)
|
||||
vfout.setsampwidth(sample_width)
|
||||
vfout.setframerate(sample_rate)
|
||||
vfout.writeframes(frame_input)
|
||||
wav_buf.seek(0)
|
||||
return wav_buf.read()
|
||||
|
||||
|
||||
def request_to_dict(request: BaseModel) -> Dict[str, Any]:
|
||||
if hasattr(request, "model_dump"):
|
||||
return request.model_dump(exclude_none=True)
|
||||
return request.dict(exclude_none=True)
|
||||
|
||||
|
||||
def get_default_voice_name() -> Optional[str]:
|
||||
default_voice = simple_config.get("default_voice")
|
||||
if default_voice:
|
||||
return str(default_voice)
|
||||
voices = simple_config.get("voices", {}) or {}
|
||||
if "default" in voices:
|
||||
return "default"
|
||||
return next(iter(voices.keys()), None)
|
||||
|
||||
|
||||
def resolve_project_path(path_value: Optional[str]) -> Optional[str]:
|
||||
if path_value in [None, ""]:
|
||||
return None
|
||||
path = Path(str(path_value))
|
||||
if not path.is_absolute():
|
||||
path = PROJECT_ROOT / path
|
||||
return str(path)
|
||||
|
||||
|
||||
def normalize_streaming(streaming_mode: Optional[Union[bool, int]], stream: Optional[bool]) -> Tuple[bool, bool, bool, bool]:
|
||||
if streaming_mode is None:
|
||||
streaming_mode = 2 if stream else 0
|
||||
elif isinstance(streaming_mode, bool):
|
||||
streaming_mode = 2 if streaming_mode else 0
|
||||
|
||||
if streaming_mode == 0:
|
||||
return False, False, False, False
|
||||
if streaming_mode == 1:
|
||||
return False, True, False, True
|
||||
if streaming_mode == 2:
|
||||
return True, False, False, True
|
||||
if streaming_mode == 3:
|
||||
return True, False, True, True
|
||||
raise HTTPException(status_code=400, detail="streaming_mode must be 0, 1, 2, 3, true, or false")
|
||||
|
||||
|
||||
def get_upload_config() -> Dict[str, Any]:
|
||||
return simple_config.get("upload", {}) or {}
|
||||
|
||||
|
||||
def get_upload_root() -> Path:
|
||||
upload_dir = str(get_upload_config().get("dir", "runtime/uploads"))
|
||||
root = Path(upload_dir)
|
||||
if not root.is_absolute():
|
||||
root = PROJECT_ROOT / root
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
return root
|
||||
|
||||
|
||||
def get_request_upload_dir() -> Path:
|
||||
request_dir = get_upload_root() / uuid.uuid4().hex
|
||||
request_dir.mkdir(parents=True, exist_ok=True)
|
||||
return request_dir
|
||||
|
||||
|
||||
def get_upload_limits() -> Tuple[float, float, int]:
|
||||
upload_config = get_upload_config()
|
||||
min_seconds = float(upload_config.get("min_ref_seconds", 3.0))
|
||||
max_seconds = float(upload_config.get("max_ref_seconds", 10.0))
|
||||
max_upload_mb = int(upload_config.get("max_upload_mb", 80))
|
||||
return min_seconds, max_seconds, max_upload_mb
|
||||
|
||||
|
||||
def get_upload_suffix(upload: UploadFile) -> str:
|
||||
suffix = Path(upload.filename or "").suffix.lower()
|
||||
if suffix:
|
||||
return suffix
|
||||
content_type = (upload.content_type or "").lower()
|
||||
if content_type in {"audio/wav", "audio/x-wav", "audio/wave"}:
|
||||
return ".wav"
|
||||
if content_type == "audio/mpeg":
|
||||
return ".mp3"
|
||||
if content_type == "audio/ogg":
|
||||
return ".ogg"
|
||||
if content_type in {"audio/aac", "audio/aacp"}:
|
||||
return ".aac"
|
||||
return ".wav"
|
||||
|
||||
|
||||
def validate_audio_upload(upload: UploadFile) -> None:
|
||||
suffix = get_upload_suffix(upload)
|
||||
content_type = (upload.content_type or "").lower()
|
||||
is_audio_type = content_type.startswith("audio/") or content_type in {"application/octet-stream", ""}
|
||||
if suffix not in SUPPORTED_UPLOAD_EXTENSIONS or not is_audio_type:
|
||||
raise HTTPException(status_code=400, detail=f"unsupported audio upload: {upload.filename or content_type}")
|
||||
|
||||
|
||||
async def save_upload_file(upload: UploadFile, target_dir: Path, name_prefix: str) -> str:
|
||||
validate_audio_upload(upload)
|
||||
_, _, max_upload_mb = get_upload_limits()
|
||||
max_bytes = max_upload_mb * 1024 * 1024
|
||||
suffix = get_upload_suffix(upload)
|
||||
target_path = target_dir / f"{name_prefix}{suffix}"
|
||||
total_size = 0
|
||||
|
||||
try:
|
||||
with target_path.open("wb") as f:
|
||||
while True:
|
||||
chunk = await upload.read(1024 * 1024)
|
||||
if not chunk:
|
||||
break
|
||||
total_size += len(chunk)
|
||||
if total_size > max_bytes:
|
||||
raise HTTPException(status_code=400, detail=f"audio upload exceeds {max_upload_mb} MB")
|
||||
f.write(chunk)
|
||||
if total_size == 0:
|
||||
raise HTTPException(status_code=400, detail=f"audio upload is empty: {upload.filename or name_prefix}")
|
||||
finally:
|
||||
await upload.close()
|
||||
|
||||
return str(target_path)
|
||||
|
||||
|
||||
def get_audio_duration_with_ffprobe(audio_path: str) -> float:
|
||||
result = subprocess.run(
|
||||
[
|
||||
"ffprobe",
|
||||
"-v",
|
||||
"error",
|
||||
"-show_entries",
|
||||
"format=duration",
|
||||
"-of",
|
||||
"default=noprint_wrappers=1:nokey=1",
|
||||
audio_path,
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=15,
|
||||
)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError(result.stderr.strip() or "ffprobe failed")
|
||||
return float(result.stdout.strip())
|
||||
|
||||
|
||||
def get_audio_duration_seconds(audio_path: str) -> float:
|
||||
try:
|
||||
info = sf.info(audio_path)
|
||||
if info.samplerate <= 0:
|
||||
raise ValueError("invalid sample rate")
|
||||
return float(info.frames) / float(info.samplerate)
|
||||
except Exception as sf_exc:
|
||||
try:
|
||||
return get_audio_duration_with_ffprobe(audio_path)
|
||||
except Exception as ffprobe_exc:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"unable to read audio duration: {audio_path}. soundfile={sf_exc}; ffprobe={ffprobe_exc}",
|
||||
) from ffprobe_exc
|
||||
|
||||
|
||||
def validate_main_ref_duration(audio_path: str) -> None:
|
||||
min_seconds, max_seconds, _ = get_upload_limits()
|
||||
duration = get_audio_duration_seconds(audio_path)
|
||||
if duration < min_seconds or duration > max_seconds:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"ref_audio duration must be between {min_seconds:g}s and {max_seconds:g}s, got {duration:.2f}s",
|
||||
)
|
||||
|
||||
|
||||
def apply_emotion_preset(payload: Dict[str, Any], emotion: Optional[str]) -> None:
|
||||
if emotion in [None, ""]:
|
||||
return
|
||||
emotion_name = str(emotion).strip().lower()
|
||||
emotion_presets = simple_config.get("emotion_presets", {}) or {}
|
||||
if emotion_name not in emotion_presets:
|
||||
supported = ", ".join(sorted(emotion_presets.keys())) or "none"
|
||||
raise HTTPException(status_code=400, detail=f"emotion is not configured: {emotion_name}. supported: {supported}")
|
||||
payload.update(emotion_presets[emotion_name] or {})
|
||||
|
||||
|
||||
def prompt_text_required_for_current_model() -> bool:
|
||||
if tts_config is None:
|
||||
return False
|
||||
use_vocoder = getattr(tts_config, "use_vocoder", None)
|
||||
if use_vocoder is not None:
|
||||
return bool(use_vocoder)
|
||||
return getattr(tts_config, "version", None) in {"v3", "v4"}
|
||||
|
||||
|
||||
def voice_public_info(name: str, profile: Dict[str, Any]) -> Dict[str, Any]:
|
||||
ref_audio_path = resolve_project_path(profile.get("ref_audio_path"))
|
||||
return {
|
||||
"name": name,
|
||||
"description": profile.get("description", ""),
|
||||
"text_lang": profile.get("text_lang"),
|
||||
"prompt_lang": profile.get("prompt_lang"),
|
||||
"ref_audio_path": profile.get("ref_audio_path"),
|
||||
"ready": bool(ref_audio_path and Path(ref_audio_path).exists() and profile.get("prompt_lang")),
|
||||
}
|
||||
|
||||
|
||||
def build_tts_request(payload: Dict[str, Any]) -> Tuple[Dict[str, Any], str, bool]:
|
||||
voices = simple_config.get("voices", {}) or {}
|
||||
explicit_ref_audio = bool(payload.get("ref_audio_path"))
|
||||
voice_name = payload.pop("voice", None)
|
||||
if voice_name is None and not explicit_ref_audio:
|
||||
voice_name = get_default_voice_name()
|
||||
voice_profile: Dict[str, Any] = {}
|
||||
|
||||
if voice_name:
|
||||
if voice_name not in voices:
|
||||
raise HTTPException(status_code=404, detail=f"voice not found: {voice_name}")
|
||||
voice_profile = deepcopy(voices[voice_name] or {})
|
||||
|
||||
tts_req = deepcopy(DEFAULT_TTS_PARAMS)
|
||||
tts_req.update(simple_config.get("defaults", {}) or {})
|
||||
tts_req.update({k: v for k, v in voice_profile.items() if k not in VOICE_META_KEYS})
|
||||
|
||||
stream = payload.pop("stream", None)
|
||||
media_type = payload.pop("format", None)
|
||||
speed = payload.pop("speed", None)
|
||||
if media_type is not None:
|
||||
tts_req["media_type"] = media_type
|
||||
|
||||
for key, value in payload.items():
|
||||
if key in DEFAULT_TTS_PARAMS:
|
||||
tts_req[key] = value
|
||||
if speed is not None:
|
||||
tts_req["speed_factor"] = speed
|
||||
|
||||
ref_audio_path = resolve_project_path(tts_req.get("ref_audio_path"))
|
||||
if ref_audio_path:
|
||||
tts_req["ref_audio_path"] = ref_audio_path
|
||||
|
||||
aux_paths = tts_req.get("aux_ref_audio_paths") or []
|
||||
tts_req["aux_ref_audio_paths"] = [resolve_project_path(item) for item in aux_paths if item]
|
||||
|
||||
media_type = str(tts_req.get("media_type", "wav")).lower()
|
||||
tts_req["media_type"] = media_type
|
||||
if media_type not in SUPPORTED_MEDIA_TYPES:
|
||||
raise HTTPException(status_code=400, detail=f"format is not supported: {media_type}")
|
||||
|
||||
text = str(tts_req.get("text") or "").strip()
|
||||
if not text:
|
||||
raise HTTPException(status_code=400, detail="text is required")
|
||||
tts_req["text"] = text
|
||||
|
||||
if not tts_req.get("ref_audio_path"):
|
||||
raise HTTPException(status_code=400, detail="ref_audio_path is required in voice profile or request")
|
||||
if not Path(tts_req["ref_audio_path"]).exists():
|
||||
raise HTTPException(status_code=400, detail=f"ref_audio_path does not exist: {tts_req['ref_audio_path']}")
|
||||
|
||||
for aux_path in tts_req["aux_ref_audio_paths"]:
|
||||
if aux_path and not Path(aux_path).exists():
|
||||
raise HTTPException(status_code=400, detail=f"aux_ref_audio_path does not exist: {aux_path}")
|
||||
|
||||
if tts_config is None:
|
||||
raise HTTPException(status_code=503, detail="TTS pipeline is not ready")
|
||||
|
||||
text_lang = str(tts_req.get("text_lang") or "").lower()
|
||||
prompt_lang = str(tts_req.get("prompt_lang") or "").lower()
|
||||
tts_req["text_lang"] = text_lang
|
||||
tts_req["prompt_lang"] = prompt_lang
|
||||
|
||||
if text_lang not in tts_config.languages:
|
||||
raise HTTPException(status_code=400, detail=f"text_lang is not supported: {text_lang}")
|
||||
if prompt_lang not in tts_config.languages:
|
||||
raise HTTPException(status_code=400, detail=f"prompt_lang is not supported: {prompt_lang}")
|
||||
if not str(tts_req.get("prompt_text") or "").strip() and prompt_text_required_for_current_model():
|
||||
version = getattr(tts_config, "version", "current")
|
||||
raise HTTPException(status_code=400, detail=f"prompt_text is required when using GPT-SoVITS {version}")
|
||||
|
||||
split_method = str(tts_req.get("text_split_method") or "cut5")
|
||||
if split_method not in cut_method_names:
|
||||
raise HTTPException(status_code=400, detail=f"text_split_method is not supported: {split_method}")
|
||||
|
||||
streaming_mode = tts_req.pop("streaming_mode", None)
|
||||
streaming_enabled, return_fragment, fixed_length_chunk, response_stream = normalize_streaming(streaming_mode, stream)
|
||||
tts_req["streaming_mode"] = streaming_enabled
|
||||
tts_req["return_fragment"] = return_fragment
|
||||
tts_req["fixed_length_chunk"] = fixed_length_chunk
|
||||
|
||||
return tts_req, media_type, response_stream
|
||||
|
||||
|
||||
def synthesize_once(payload: Dict[str, Any]) -> Tuple[bytes, str]:
|
||||
tts_req, media_type, response_stream = build_tts_request(payload)
|
||||
if response_stream:
|
||||
raise HTTPException(status_code=400, detail="base64 output does not support streaming")
|
||||
|
||||
if tts_pipeline is None:
|
||||
raise HTTPException(status_code=503, detail="TTS pipeline is not ready")
|
||||
|
||||
try:
|
||||
with infer_lock:
|
||||
sr, audio_data = next(tts_pipeline.run(tts_req))
|
||||
audio_bytes = pack_audio(BytesIO(), audio_data, sr, media_type).getvalue()
|
||||
return audio_bytes, media_type
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=400, detail={"message": "tts failed", "exception": str(exc)}) from exc
|
||||
|
||||
|
||||
def synthesize_response(payload: Dict[str, Any]) -> Response:
|
||||
tts_req, media_type, response_stream = build_tts_request(payload)
|
||||
if tts_pipeline is None:
|
||||
raise HTTPException(status_code=503, detail="TTS pipeline is not ready")
|
||||
|
||||
if response_stream:
|
||||
|
||||
def streaming_generator() -> Generator[bytes, None, None]:
|
||||
first_chunk = True
|
||||
chunk_media_type = media_type
|
||||
with infer_lock:
|
||||
for sr, chunk in tts_pipeline.run(tts_req):
|
||||
if first_chunk and chunk_media_type == "wav":
|
||||
yield wave_header_chunk(sample_rate=sr)
|
||||
chunk_media_type = "raw"
|
||||
first_chunk = False
|
||||
yield pack_audio(BytesIO(), chunk, sr, chunk_media_type).getvalue()
|
||||
|
||||
return StreamingResponse(streaming_generator(), media_type=f"audio/{media_type}")
|
||||
|
||||
try:
|
||||
with infer_lock:
|
||||
sr, audio_data = next(tts_pipeline.run(tts_req))
|
||||
audio_bytes = pack_audio(BytesIO(), audio_data, sr, media_type).getvalue()
|
||||
return Response(audio_bytes, media_type=f"audio/{media_type}")
|
||||
except Exception as exc:
|
||||
return JSONResponse(status_code=400, content={"message": "tts failed", "exception": str(exc)})
|
||||
|
||||
|
||||
@APP.on_event("startup")
|
||||
def startup() -> None:
|
||||
global tts_config, tts_pipeline
|
||||
tts_config = TTS_Config(tts_config_path)
|
||||
print(tts_config)
|
||||
tts_pipeline = TTS(tts_config)
|
||||
|
||||
|
||||
@APP.get("/", tags=["System"])
|
||||
def index() -> Dict[str, Any]:
|
||||
return {
|
||||
"name": "GPT-SoVITS Simple API",
|
||||
"version": "1.1.0",
|
||||
"docs": "/docs",
|
||||
"endpoints": {
|
||||
"system": ["/health", "/voices"],
|
||||
"mvp": ["/api/tts"],
|
||||
"profile": ["/speak", "/speak/base64", "/v1/tts"],
|
||||
"admin": ["/admin/reload-config", "/admin/weights"],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@APP.get("/health", tags=["System"], summary="健康检查")
|
||||
def health() -> Dict[str, Any]:
|
||||
import os
|
||||
|
||||
result: Dict[str, Any] = {
|
||||
"status": "ok" if tts_pipeline is not None else "starting",
|
||||
"tts_config": tts_config_path,
|
||||
"version": getattr(tts_config, "version", None),
|
||||
"languages": getattr(tts_config, "languages", []),
|
||||
"pid": os.getpid(),
|
||||
}
|
||||
try:
|
||||
import psutil
|
||||
result["memory_mb"] = round(psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024, 1)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
import torch
|
||||
if torch.cuda.is_available():
|
||||
result["gpu"] = {
|
||||
"name": torch.cuda.get_device_name(0),
|
||||
"memory_used_mb": round(torch.cuda.memory_allocated(0) / 1024 / 1024, 1),
|
||||
"memory_total_mb": round(torch.cuda.get_device_properties(0).total_mem / 1024 / 1024, 1),
|
||||
}
|
||||
except Exception:
|
||||
pass
|
||||
return result
|
||||
|
||||
|
||||
@APP.get("/voices", tags=["System"], summary="列出可用 voice profiles")
|
||||
def list_voices() -> Dict[str, Any]:
|
||||
voices = simple_config.get("voices", {}) or {}
|
||||
return {
|
||||
"default_voice": get_default_voice_name(),
|
||||
"voices": [voice_public_info(name, profile or {}) for name, profile in voices.items()],
|
||||
}
|
||||
|
||||
|
||||
@APP.post(
|
||||
"/api/tts",
|
||||
tags=["MVP"],
|
||||
summary="核心 TTS 接口",
|
||||
description=(
|
||||
"上传参考音频和需要生成的文字，返回生成的音频。\n\n"
|
||||
"**主参考音频要求**：3-10 秒，支持 wav/flac/ogg/mp3/m4a/aac 格式。\n\n"
|
||||
"**文字切句**：固定使用 `cut5`（按标点符号切句）。\n\n"
|
||||
"**情绪预设**：neutral / happy / calm / sad / angry，本质是映射到采样和语速参数。"
|
||||
),
|
||||
)
|
||||
async def mvp_tts(
|
||||
text: str = Form(...),
|
||||
ref_audio: UploadFile = File(...),
|
||||
aux_ref_audio: Optional[List[UploadFile]] = File(default=None),
|
||||
prompt_text: str = Form(default=""),
|
||||
text_lang: str = Form(default="zh"),
|
||||
prompt_lang: str = Form(default="zh"),
|
||||
format: str = Form(default="wav"),
|
||||
emotion: Optional[str] = Form(default=None),
|
||||
speed: Optional[float] = Form(default=None),
|
||||
seed: int = Form(default=-1),
|
||||
) -> Response:
|
||||
request_dir = get_request_upload_dir()
|
||||
try:
|
||||
ref_audio_path = await save_upload_file(ref_audio, request_dir, "ref")
|
||||
validate_main_ref_duration(ref_audio_path)
|
||||
|
||||
aux_ref_audio_paths = []
|
||||
for index, upload in enumerate(aux_ref_audio or []):
|
||||
aux_path = await save_upload_file(upload, request_dir, f"aux_{index}")
|
||||
get_audio_duration_seconds(aux_path)
|
||||
aux_ref_audio_paths.append(aux_path)
|
||||
|
||||
payload: Dict[str, Any] = {
|
||||
"text": text,
|
||||
"text_lang": text_lang,
|
||||
"ref_audio_path": ref_audio_path,
|
||||
"aux_ref_audio_paths": aux_ref_audio_paths,
|
||||
"prompt_text": prompt_text or "",
|
||||
"prompt_lang": prompt_lang,
|
||||
"format": format,
|
||||
"text_split_method": "cut5",
|
||||
"streaming_mode": 0,
|
||||
"seed": seed,
|
||||
}
|
||||
apply_emotion_preset(payload, emotion)
|
||||
payload["text_split_method"] = "cut5"
|
||||
payload["streaming_mode"] = 0
|
||||
if speed is not None:
|
||||
payload["speed"] = speed
|
||||
|
||||
return synthesize_response(payload)
|
||||
finally:
|
||||
shutil.rmtree(request_dir, ignore_errors=True)
|
||||
|
||||
|
||||
@APP.get("/speak", tags=["Profile"], summary="GET 方式调用 voice profile TTS")
|
||||
def speak_get(
|
||||
text: str,
|
||||
voice: Optional[str] = None,
|
||||
text_lang: Optional[str] = None,
|
||||
format: Optional[str] = None,
|
||||
stream: Optional[bool] = None,
|
||||
speed: Optional[float] = None,
|
||||
) -> Response:
|
||||
payload = {
|
||||
"text": text,
|
||||
"voice": voice,
|
||||
"text_lang": text_lang,
|
||||
"format": format,
|
||||
"stream": stream,
|
||||
"speed": speed,
|
||||
}
|
||||
return synthesize_response({k: v for k, v in payload.items() if v is not None})
|
||||
|
||||
|
||||
@APP.post("/speak", tags=["Profile"], summary="POST 方式调用 voice profile TTS")
|
||||
def speak_post(request: SpeakRequest) -> Response:
|
||||
return synthesize_response(request_to_dict(request))
|
||||
|
||||
|
||||
@APP.post("/v1/tts", tags=["Profile"], summary="OpenAI 兼容格式 TTS")
|
||||
def openai_style_tts(request: SpeakRequest) -> Response:
|
||||
return synthesize_response(request_to_dict(request))
|
||||
|
||||
|
||||
@APP.post("/speak/base64", tags=["Profile"], summary="返回 Base64 编码的音频")
|
||||
def speak_base64(request: SpeakRequest) -> Dict[str, Any]:
|
||||
audio_bytes, media_type = synthesize_once(request_to_dict(request))
|
||||
return {
|
||||
"media_type": f"audio/{media_type}",
|
||||
"audio_base64": base64.b64encode(audio_bytes).decode("ascii"),
|
||||
}
|
||||
|
||||
|
||||
@APP.post("/admin/reload-config", tags=["Admin"], summary="热加载 simple_api.yaml 配置")
|
||||
def reload_config() -> Dict[str, Any]:
|
||||
global simple_config
|
||||
simple_config = load_yaml_config(args.config)
|
||||
return {"message": "success", "default_voice": get_default_voice_name()}
|
||||
|
||||
|
||||
@APP.post("/admin/weights", tags=["Admin"], summary="运行时切换 GPT-SoVITS 模型权重")
|
||||
def set_weights(request: WeightsRequest) -> Dict[str, Any]:
|
||||
if tts_pipeline is None:
|
||||
raise HTTPException(status_code=503, detail="TTS pipeline is not ready")
|
||||
if not request.gpt_weights_path and not request.sovits_weights_path:
|
||||
raise HTTPException(status_code=400, detail="gpt_weights_path or sovits_weights_path is required")
|
||||
|
||||
try:
|
||||
with infer_lock:
|
||||
if request.gpt_weights_path:
|
||||
tts_pipeline.init_t2s_weights(request.gpt_weights_path)
|
||||
if request.sovits_weights_path:
|
||||
tts_pipeline.init_vits_weights(request.sovits_weights_path)
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=400, detail={"message": "change weights failed", "exception": str(exc)}) from exc
|
||||
|
||||
return {"message": "success"}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
|
||||
try:
|
||||
uvicorn.run(app=APP, host=None if host == "None" else host, port=port, workers=1)
|
||||
except Exception:
|
||||
traceback.print_exc()
|
||||
sys.exit(1)
|
||||
59
simple_api.yaml
Normal file
59
simple_api.yaml
Normal file
@ -0,0 +1,59 @@
|
||||
server:
|
||||
host: 127.0.0.1
|
||||
port: 9881
|
||||
tts_config: GPT_SoVITS/configs/tts_infer.yaml
|
||||
|
||||
cors_allow_origins:
|
||||
- "*"
|
||||
|
||||
upload:
|
||||
dir: runtime/uploads
|
||||
min_ref_seconds: 3
|
||||
max_ref_seconds: 10
|
||||
max_upload_mb: 80
|
||||
|
||||
default_voice: default
|
||||
|
||||
defaults:
|
||||
text_lang: zh
|
||||
prompt_lang: zh
|
||||
media_type: wav
|
||||
text_split_method: cut5
|
||||
batch_size: 1
|
||||
batch_threshold: 0.75
|
||||
split_bucket: true
|
||||
speed_factor: 1.0
|
||||
fragment_interval: 0.3
|
||||
seed: -1
|
||||
parallel_infer: true
|
||||
repetition_penalty: 1.35
|
||||
sample_steps: 32
|
||||
super_sampling: false
|
||||
overlap_length: 2
|
||||
min_chunk_length: 16
|
||||
|
||||
emotion_presets:
|
||||
neutral: {}
|
||||
happy:
|
||||
temperature: 1.1
|
||||
top_p: 0.95
|
||||
calm:
|
||||
temperature: 0.8
|
||||
top_p: 0.85
|
||||
speed_factor: 0.92
|
||||
sad:
|
||||
temperature: 0.75
|
||||
top_p: 0.85
|
||||
speed_factor: 0.9
|
||||
angry:
|
||||
temperature: 1.2
|
||||
top_k: 20
|
||||
repetition_penalty: 1.25
|
||||
|
||||
voices:
|
||||
default:
|
||||
description: Replace this profile with your reference voice.
|
||||
ref_audio_path: reference.wav
|
||||
prompt_text: Replace this with the exact text spoken in reference.wav.
|
||||
prompt_lang: zh
|
||||
text_lang: zh
|
||||
1006
test_frontend/index.html
Normal file
1006
test_frontend/index.html
Normal file
File diff suppressed because it is too large
Load Diff
@ -1,71 +0,0 @@
|
||||
import importlib.util
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
MODULE_PATH = ROOT / "tools" / "asr" / "funasr_asr.py"
|
||||
|
||||
|
||||
def load_funasr_asr(monkeypatch, auto_model):
|
||||
funasr = types.ModuleType("funasr")
|
||||
funasr.AutoModel = auto_model
|
||||
monkeypatch.setitem(sys.modules, "funasr", funasr)
|
||||
|
||||
modelscope = types.ModuleType("modelscope")
|
||||
modelscope.snapshot_download = lambda *args, **kwargs: None
|
||||
monkeypatch.setitem(sys.modules, "modelscope", modelscope)
|
||||
|
||||
tqdm = types.ModuleType("tqdm")
|
||||
tqdm.tqdm = lambda iterable: iterable
|
||||
monkeypatch.setitem(sys.modules, "tqdm", tqdm)
|
||||
|
||||
torch = types.ModuleType("torch")
|
||||
torch.cuda = types.SimpleNamespace(is_available=lambda: False)
|
||||
monkeypatch.setitem(sys.modules, "torch", torch)
|
||||
|
||||
module_name = "funasr_asr_under_test"
|
||||
sys.modules.pop(module_name, None)
|
||||
spec = importlib.util.spec_from_file_location(module_name, MODULE_PATH)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def test_fun_asr_nano_falls_back_to_modelscope_when_hf_config_is_not_resolved(monkeypatch):
|
||||
calls = []
|
||||
|
||||
def fake_auto_model(**kwargs):
|
||||
calls.append(kwargs.copy())
|
||||
if kwargs["hub"] == "hf":
|
||||
raise RuntimeError("model 'FunAudioLLM/Fun-ASR-Nano-2512' is not registered.")
|
||||
if kwargs["hub"] == "ms":
|
||||
return types.SimpleNamespace(hub=kwargs["hub"], model_name=kwargs["model"])
|
||||
raise AssertionError(f"unexpected hub {kwargs['hub']}")
|
||||
|
||||
module = load_funasr_asr(monkeypatch, fake_auto_model)
|
||||
|
||||
model = module.create_model("zh", backend="fun-asr-nano")
|
||||
|
||||
assert model.hub == "ms"
|
||||
assert model.model_name == "FunAudioLLM/Fun-ASR-Nano-2512"
|
||||
assert [(call["model"], call["hub"], call["trust_remote_code"]) for call in calls] == [
|
||||
("FunAudioLLM/Fun-ASR-Nano-2512", "hf", True),
|
||||
("FunAudioLLM/Fun-ASR-Nano-2512", "ms", False),
|
||||
]
|
||||
assert all(call["vad_model"] == "fsmn-vad" for call in calls)
|
||||
assert all(call["device"] == "cpu" for call in calls)
|
||||
assert all(call["disable_update"] is True for call in calls)
|
||||
|
||||
|
||||
def test_fun_asr_nano_keeps_non_registration_errors_visible(monkeypatch):
|
||||
def fake_auto_model(**kwargs):
|
||||
raise RuntimeError("network unavailable")
|
||||
|
||||
module = load_funasr_asr(monkeypatch, fake_auto_model)
|
||||
|
||||
with pytest.raises(RuntimeError, match="network unavailable"):
|
||||
module.create_model("zh", backend="fun-asr-nano")
|
||||
303
tests/test_simple_api_contract.py
Normal file
303
tests/test_simple_api_contract.py
Normal file
@ -0,0 +1,303 @@
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import sys
|
||||
import tempfile
|
||||
import types
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
PROJECT_ROOT = Path(__file__).resolve().parents[1]
|
||||
SIMPLE_API_PATH = PROJECT_ROOT / "simple_api.py"
|
||||
|
||||
|
||||
class DummyHTTPException(Exception):
|
||||
def __init__(self, status_code=500, detail=None):
|
||||
super().__init__(detail)
|
||||
self.status_code = status_code
|
||||
self.detail = detail
|
||||
|
||||
|
||||
class DummyFastAPI:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.routes = []
|
||||
|
||||
def add_middleware(self, *args, **kwargs):
|
||||
pass
|
||||
|
||||
def mount(self, path, app, name=None):
|
||||
self.routes.append(types.SimpleNamespace(path=path, endpoint=app, name=name))
|
||||
|
||||
def get(self, path, **kwargs):
|
||||
return self._route(path)
|
||||
|
||||
def post(self, path, **kwargs):
|
||||
return self._route(path)
|
||||
|
||||
def on_event(self, event):
|
||||
return lambda func: func
|
||||
|
||||
def _route(self, path):
|
||||
def decorator(func):
|
||||
self.routes.append(types.SimpleNamespace(path=path, endpoint=func))
|
||||
return func
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
class DummyBaseModel:
|
||||
def dict(self, exclude_none=False):
|
||||
data = dict(self.__dict__)
|
||||
if exclude_none:
|
||||
data = {key: value for key, value in data.items() if value is not None}
|
||||
return data
|
||||
|
||||
|
||||
class DummyTTSConfig:
|
||||
version = "v2"
|
||||
languages = ["zh", "en", "auto"]
|
||||
use_vocoder = False
|
||||
|
||||
|
||||
class FakeUpload:
|
||||
def __init__(self, filename, chunks, content_type="audio/wav"):
|
||||
self.filename = filename
|
||||
self.content_type = content_type
|
||||
self._chunks = list(chunks)
|
||||
self.closed = False
|
||||
|
||||
async def read(self, size):
|
||||
del size
|
||||
if self._chunks:
|
||||
return self._chunks.pop(0)
|
||||
return b""
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
|
||||
def install_dummy_modules():
|
||||
fastapi = types.ModuleType("fastapi")
|
||||
fastapi.FastAPI = DummyFastAPI
|
||||
fastapi.File = lambda default=None, **kwargs: default
|
||||
fastapi.Form = lambda default=None, **kwargs: default
|
||||
fastapi.HTTPException = DummyHTTPException
|
||||
fastapi.Response = type("Response", (), {})
|
||||
fastapi.UploadFile = type("UploadFile", (), {})
|
||||
sys.modules["fastapi"] = fastapi
|
||||
|
||||
middleware = types.ModuleType("fastapi.middleware")
|
||||
cors = types.ModuleType("fastapi.middleware.cors")
|
||||
cors.CORSMiddleware = type("CORSMiddleware", (), {})
|
||||
sys.modules["fastapi.middleware"] = middleware
|
||||
sys.modules["fastapi.middleware.cors"] = cors
|
||||
|
||||
responses = types.ModuleType("fastapi.responses")
|
||||
responses.JSONResponse = type("JSONResponse", (), {})
|
||||
responses.StreamingResponse = type("StreamingResponse", (), {})
|
||||
sys.modules["fastapi.responses"] = responses
|
||||
|
||||
staticfiles = types.ModuleType("fastapi.staticfiles")
|
||||
staticfiles.StaticFiles = type("StaticFiles", (), {"__init__": lambda self, *args, **kwargs: None})
|
||||
sys.modules["fastapi.staticfiles"] = staticfiles
|
||||
|
||||
pydantic = types.ModuleType("pydantic")
|
||||
pydantic.BaseModel = DummyBaseModel
|
||||
pydantic.Field = lambda default=None, **kwargs: default
|
||||
sys.modules["pydantic"] = pydantic
|
||||
|
||||
numpy = types.ModuleType("numpy")
|
||||
numpy.ndarray = object
|
||||
sys.modules["numpy"] = numpy
|
||||
|
||||
soundfile = types.ModuleType("soundfile")
|
||||
soundfile.info = lambda path: types.SimpleNamespace(frames=16000 * 5, samplerate=16000)
|
||||
sys.modules["soundfile"] = soundfile
|
||||
|
||||
yaml = types.ModuleType("yaml")
|
||||
yaml.safe_load = lambda stream: {
|
||||
"server": {"host": "127.0.0.1", "port": 9881, "tts_config": "dummy.yaml"},
|
||||
"upload": {"dir": "runtime/test_uploads", "min_ref_seconds": 3, "max_ref_seconds": 10, "max_upload_mb": 1},
|
||||
"defaults": {"text_lang": "zh", "prompt_lang": "zh", "media_type": "wav", "text_split_method": "cut5"},
|
||||
"emotion_presets": {"calm": {"temperature": 0.8, "speed_factor": 0.9}},
|
||||
"voices": {},
|
||||
}
|
||||
sys.modules["yaml"] = yaml
|
||||
|
||||
sys.modules["GPT_SoVITS"] = types.ModuleType("GPT_SoVITS")
|
||||
sys.modules["GPT_SoVITS.TTS_infer_pack"] = types.ModuleType("GPT_SoVITS.TTS_infer_pack")
|
||||
|
||||
tts_module = types.ModuleType("GPT_SoVITS.TTS_infer_pack.TTS")
|
||||
tts_module.TTS = type("DummyTTS", (), {})
|
||||
tts_module.TTS_Config = DummyTTSConfig
|
||||
sys.modules["GPT_SoVITS.TTS_infer_pack.TTS"] = tts_module
|
||||
|
||||
segmentation = types.ModuleType("GPT_SoVITS.TTS_infer_pack.text_segmentation_method")
|
||||
segmentation.get_method_names = lambda: ["cut0", "cut5"]
|
||||
sys.modules["GPT_SoVITS.TTS_infer_pack.text_segmentation_method"] = segmentation
|
||||
|
||||
|
||||
def load_simple_api():
|
||||
install_dummy_modules()
|
||||
sys.argv = ["simple_api.py"]
|
||||
spec = importlib.util.spec_from_file_location("simple_api_under_test", SIMPLE_API_PATH)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
module.tts_config = DummyTTSConfig()
|
||||
return module
|
||||
|
||||
|
||||
class SimpleApiContractTest(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.simple_api = load_simple_api()
|
||||
|
||||
def test_api_tts_route_is_registered(self):
|
||||
routes = {route.path for route in self.simple_api.APP.routes}
|
||||
self.assertIn("/api/tts", routes)
|
||||
|
||||
def test_build_request_accepts_explicit_ref_without_voice_profile(self):
|
||||
with tempfile.NamedTemporaryFile(suffix=".wav") as ref:
|
||||
req, media_type, response_stream = self.simple_api.build_tts_request(
|
||||
{
|
||||
"text": "Hello, test.",
|
||||
"ref_audio_path": ref.name,
|
||||
"prompt_text": "",
|
||||
"prompt_lang": "zh",
|
||||
"text_lang": "zh",
|
||||
"format": "wav",
|
||||
"text_split_method": "cut5",
|
||||
"streaming_mode": 0,
|
||||
}
|
||||
)
|
||||
|
||||
self.assertEqual(req["prompt_text"], "")
|
||||
self.assertEqual(req["text_split_method"], "cut5")
|
||||
self.assertEqual(req["ref_audio_path"], ref.name)
|
||||
self.assertEqual(media_type, "wav")
|
||||
self.assertFalse(response_stream)
|
||||
|
||||
def test_explicit_speed_overrides_emotion_speed_preset(self):
|
||||
with tempfile.NamedTemporaryFile(suffix=".wav") as ref:
|
||||
payload = {
|
||||
"text": "Hello.",
|
||||
"ref_audio_path": ref.name,
|
||||
"prompt_lang": "zh",
|
||||
"text_lang": "zh",
|
||||
"format": "wav",
|
||||
"streaming_mode": 0,
|
||||
"speed": 1.05,
|
||||
}
|
||||
self.simple_api.apply_emotion_preset(payload, "calm")
|
||||
payload["speed"] = 1.05
|
||||
req, _, _ = self.simple_api.build_tts_request(payload)
|
||||
|
||||
self.assertEqual(req["temperature"], 0.8)
|
||||
self.assertEqual(req["speed_factor"], 1.05)
|
||||
|
||||
def test_empty_prompt_text_is_rejected_for_vocoder_models(self):
|
||||
original_use_vocoder = self.simple_api.tts_config.use_vocoder
|
||||
original_version = self.simple_api.tts_config.version
|
||||
self.simple_api.tts_config.use_vocoder = True
|
||||
self.simple_api.tts_config.version = "v3"
|
||||
|
||||
try:
|
||||
with tempfile.NamedTemporaryFile(suffix=".wav") as ref:
|
||||
with self.assertRaises(DummyHTTPException) as exc:
|
||||
self.simple_api.build_tts_request(
|
||||
{
|
||||
"text": "Hello.",
|
||||
"ref_audio_path": ref.name,
|
||||
"prompt_text": "",
|
||||
"prompt_lang": "zh",
|
||||
"text_lang": "zh",
|
||||
"format": "wav",
|
||||
"text_split_method": "cut5",
|
||||
"streaming_mode": 0,
|
||||
}
|
||||
)
|
||||
self.assertEqual(exc.exception.status_code, 400)
|
||||
self.assertIn("prompt_text is required", str(exc.exception.detail))
|
||||
finally:
|
||||
self.simple_api.tts_config.use_vocoder = original_use_vocoder
|
||||
self.simple_api.tts_config.version = original_version
|
||||
|
||||
def test_ref_audio_duration_limits(self):
|
||||
self.simple_api.sf.info = lambda path: types.SimpleNamespace(frames=16000 * 5, samplerate=16000)
|
||||
self.simple_api.validate_main_ref_duration("ok.wav")
|
||||
|
||||
self.simple_api.sf.info = lambda path: types.SimpleNamespace(frames=16000 * 2, samplerate=16000)
|
||||
with self.assertRaises(DummyHTTPException) as too_short:
|
||||
self.simple_api.validate_main_ref_duration("short.wav")
|
||||
self.assertEqual(too_short.exception.status_code, 400)
|
||||
|
||||
self.simple_api.sf.info = lambda path: types.SimpleNamespace(frames=16000 * 11, samplerate=16000)
|
||||
with self.assertRaises(DummyHTTPException) as too_long:
|
||||
self.simple_api.validate_main_ref_duration("long.wav")
|
||||
self.assertEqual(too_long.exception.status_code, 400)
|
||||
|
||||
def test_save_upload_file_writes_and_closes(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
upload = FakeUpload("ref.wav", [b"abc", b"def"])
|
||||
saved_path = asyncio.run(self.simple_api.save_upload_file(upload, Path(tmp), "ref"))
|
||||
|
||||
self.assertTrue(upload.closed)
|
||||
self.assertEqual(Path(saved_path).read_bytes(), b"abcdef")
|
||||
|
||||
def test_mvp_tts_builds_payload_from_uploads_and_cleans_tmp_dir(self):
|
||||
captured = {}
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
request_dir = Path(tmp) / "request"
|
||||
request_dir.mkdir()
|
||||
|
||||
original_get_request_upload_dir = self.simple_api.get_request_upload_dir
|
||||
original_synthesize_response = self.simple_api.synthesize_response
|
||||
original_sf_info = self.simple_api.sf.info
|
||||
|
||||
def fake_synthesize_response(payload):
|
||||
captured.update(payload)
|
||||
return "audio-response"
|
||||
|
||||
try:
|
||||
self.simple_api.get_request_upload_dir = lambda: request_dir
|
||||
self.simple_api.sf.info = lambda path: types.SimpleNamespace(frames=16000 * 5, samplerate=16000)
|
||||
self.simple_api.synthesize_response = fake_synthesize_response
|
||||
|
||||
response = asyncio.run(
|
||||
self.simple_api.mvp_tts(
|
||||
text="Hello, test.",
|
||||
ref_audio=FakeUpload("ref.wav", [b"main"]),
|
||||
aux_ref_audio=[FakeUpload("aux.wav", [b"aux"])],
|
||||
prompt_text="",
|
||||
text_lang="zh",
|
||||
prompt_lang="zh",
|
||||
format="wav",
|
||||
emotion="calm",
|
||||
speed=1.05,
|
||||
seed=123,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
self.simple_api.get_request_upload_dir = original_get_request_upload_dir
|
||||
self.simple_api.synthesize_response = original_synthesize_response
|
||||
self.simple_api.sf.info = original_sf_info
|
||||
|
||||
self.assertEqual(response, "audio-response")
|
||||
self.assertFalse(request_dir.exists())
|
||||
|
||||
self.assertEqual(captured["text"], "Hello, test.")
|
||||
self.assertEqual(captured["prompt_text"], "")
|
||||
self.assertEqual(captured["text_lang"], "zh")
|
||||
self.assertEqual(captured["prompt_lang"], "zh")
|
||||
self.assertEqual(captured["format"], "wav")
|
||||
self.assertEqual(captured["text_split_method"], "cut5")
|
||||
self.assertEqual(captured["streaming_mode"], 0)
|
||||
self.assertEqual(captured["seed"], 123)
|
||||
self.assertEqual(captured["temperature"], 0.8)
|
||||
self.assertEqual(captured["speed"], 1.05)
|
||||
self.assertEqual(len(captured["aux_ref_audio_paths"]), 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@ -1,167 +0,0 @@
|
||||
import os
|
||||
import random
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.utils.data
|
||||
from tqdm import tqdm
|
||||
|
||||
from . import spec_utils
|
||||
|
||||
|
||||
class VocalRemoverValidationSet(torch.utils.data.Dataset):
|
||||
def __init__(self, patch_list):
|
||||
self.patch_list = patch_list
|
||||
|
||||
def __len__(self):
|
||||
return len(self.patch_list)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
path = self.patch_list[idx]
|
||||
data = np.load(path)
|
||||
|
||||
X, y = data["X"], data["y"]
|
||||
|
||||
X_mag = np.abs(X)
|
||||
y_mag = np.abs(y)
|
||||
|
||||
return X_mag, y_mag
|
||||
|
||||
|
||||
def make_pair(mix_dir, inst_dir):
|
||||
input_exts = [".wav", ".m4a", ".mp3", ".mp4", ".flac"]
|
||||
|
||||
X_list = sorted(
|
||||
[os.path.join(mix_dir, fname) for fname in os.listdir(mix_dir) if os.path.splitext(fname)[1] in input_exts]
|
||||
)
|
||||
y_list = sorted(
|
||||
[os.path.join(inst_dir, fname) for fname in os.listdir(inst_dir) if os.path.splitext(fname)[1] in input_exts]
|
||||
)
|
||||
|
||||
filelist = list(zip(X_list, y_list))
|
||||
|
||||
return filelist
|
||||
|
||||
|
||||
def train_val_split(dataset_dir, split_mode, val_rate, val_filelist):
|
||||
if split_mode == "random":
|
||||
filelist = make_pair(
|
||||
os.path.join(dataset_dir, "mixtures"),
|
||||
os.path.join(dataset_dir, "instruments"),
|
||||
)
|
||||
|
||||
random.shuffle(filelist)
|
||||
|
||||
if len(val_filelist) == 0:
|
||||
val_size = int(len(filelist) * val_rate)
|
||||
train_filelist = filelist[:-val_size]
|
||||
val_filelist = filelist[-val_size:]
|
||||
else:
|
||||
train_filelist = [pair for pair in filelist if list(pair) not in val_filelist]
|
||||
elif split_mode == "subdirs":
|
||||
if len(val_filelist) != 0:
|
||||
raise ValueError("The `val_filelist` option is not available in `subdirs` mode")
|
||||
|
||||
train_filelist = make_pair(
|
||||
os.path.join(dataset_dir, "training/mixtures"),
|
||||
os.path.join(dataset_dir, "training/instruments"),
|
||||
)
|
||||
|
||||
val_filelist = make_pair(
|
||||
os.path.join(dataset_dir, "validation/mixtures"),
|
||||
os.path.join(dataset_dir, "validation/instruments"),
|
||||
)
|
||||
|
||||
return train_filelist, val_filelist
|
||||
|
||||
|
||||
def augment(X, y, reduction_rate, reduction_mask, mixup_rate, mixup_alpha):
|
||||
perm = np.random.permutation(len(X))
|
||||
for i, idx in enumerate(tqdm(perm)):
|
||||
if np.random.uniform() < reduction_rate:
|
||||
y[idx] = spec_utils.reduce_vocal_aggressively(X[idx], y[idx], reduction_mask)
|
||||
|
||||
if np.random.uniform() < 0.5:
|
||||
# swap channel
|
||||
X[idx] = X[idx, ::-1]
|
||||
y[idx] = y[idx, ::-1]
|
||||
if np.random.uniform() < 0.02:
|
||||
# mono
|
||||
X[idx] = X[idx].mean(axis=0, keepdims=True)
|
||||
y[idx] = y[idx].mean(axis=0, keepdims=True)
|
||||
if np.random.uniform() < 0.02:
|
||||
# inst
|
||||
X[idx] = y[idx]
|
||||
|
||||
if np.random.uniform() < mixup_rate and i < len(perm) - 1:
|
||||
lam = np.random.beta(mixup_alpha, mixup_alpha)
|
||||
X[idx] = lam * X[idx] + (1 - lam) * X[perm[i + 1]]
|
||||
y[idx] = lam * y[idx] + (1 - lam) * y[perm[i + 1]]
|
||||
|
||||
return X, y
|
||||
|
||||
|
||||
def make_padding(width, cropsize, offset):
|
||||
left = offset
|
||||
roi_size = cropsize - left * 2
|
||||
if roi_size == 0:
|
||||
roi_size = cropsize
|
||||
right = roi_size - (width % roi_size) + left
|
||||
|
||||
return left, right, roi_size
|
||||
|
||||
|
||||
def make_training_set(filelist, cropsize, patches, sr, hop_length, n_fft, offset):
|
||||
len_dataset = patches * len(filelist)
|
||||
|
||||
X_dataset = np.zeros((len_dataset, 2, n_fft // 2 + 1, cropsize), dtype=np.complex64)
|
||||
y_dataset = np.zeros((len_dataset, 2, n_fft // 2 + 1, cropsize), dtype=np.complex64)
|
||||
|
||||
for i, (X_path, y_path) in enumerate(tqdm(filelist)):
|
||||
X, y = spec_utils.cache_or_load(X_path, y_path, sr, hop_length, n_fft)
|
||||
coef = np.max([np.abs(X).max(), np.abs(y).max()])
|
||||
X, y = X / coef, y / coef
|
||||
|
||||
l, r, roi_size = make_padding(X.shape[2], cropsize, offset)
|
||||
X_pad = np.pad(X, ((0, 0), (0, 0), (l, r)), mode="constant")
|
||||
y_pad = np.pad(y, ((0, 0), (0, 0), (l, r)), mode="constant")
|
||||
|
||||
starts = np.random.randint(0, X_pad.shape[2] - cropsize, patches)
|
||||
ends = starts + cropsize
|
||||
for j in range(patches):
|
||||
idx = i * patches + j
|
||||
X_dataset[idx] = X_pad[:, :, starts[j] : ends[j]]
|
||||
y_dataset[idx] = y_pad[:, :, starts[j] : ends[j]]
|
||||
|
||||
return X_dataset, y_dataset
|
||||
|
||||
|
||||
def make_validation_set(filelist, cropsize, sr, hop_length, n_fft, offset):
|
||||
patch_list = []
|
||||
patch_dir = "cs{}_sr{}_hl{}_nf{}_of{}".format(cropsize, sr, hop_length, n_fft, offset)
|
||||
os.makedirs(patch_dir, exist_ok=True)
|
||||
|
||||
for i, (X_path, y_path) in enumerate(tqdm(filelist)):
|
||||
basename = os.path.splitext(os.path.basename(X_path))[0]
|
||||
|
||||
X, y = spec_utils.cache_or_load(X_path, y_path, sr, hop_length, n_fft)
|
||||
coef = np.max([np.abs(X).max(), np.abs(y).max()])
|
||||
X, y = X / coef, y / coef
|
||||
|
||||
l, r, roi_size = make_padding(X.shape[2], cropsize, offset)
|
||||
X_pad = np.pad(X, ((0, 0), (0, 0), (l, r)), mode="constant")
|
||||
y_pad = np.pad(y, ((0, 0), (0, 0), (l, r)), mode="constant")
|
||||
|
||||
len_dataset = int(np.ceil(X.shape[2] / roi_size))
|
||||
for j in range(len_dataset):
|
||||
outpath = os.path.join(patch_dir, "{}_p{}.npz".format(basename, j))
|
||||
start = j * roi_size
|
||||
if not os.path.exists(outpath):
|
||||
np.savez(
|
||||
outpath,
|
||||
X=X_pad[:, :, start : start + cropsize],
|
||||
y=y_pad[:, :, start : start + cropsize],
|
||||
)
|
||||
patch_list.append(outpath)
|
||||
|
||||
return VocalRemoverValidationSet(patch_list)
|
||||
@ -1,106 +0,0 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from . import spec_utils
|
||||
|
||||
|
||||
class Conv2DBNActiv(nn.Module):
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU):
|
||||
super(Conv2DBNActiv, self).__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
nin,
|
||||
nout,
|
||||
kernel_size=ksize,
|
||||
stride=stride,
|
||||
padding=pad,
|
||||
dilation=dilation,
|
||||
bias=False,
|
||||
),
|
||||
nn.BatchNorm2d(nout),
|
||||
activ(),
|
||||
)
|
||||
|
||||
def __call__(self, x):
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class SeperableConv2DBNActiv(nn.Module):
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU):
|
||||
super(SeperableConv2DBNActiv, self).__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
nin,
|
||||
nin,
|
||||
kernel_size=ksize,
|
||||
stride=stride,
|
||||
padding=pad,
|
||||
dilation=dilation,
|
||||
groups=nin,
|
||||
bias=False,
|
||||
),
|
||||
nn.Conv2d(nin, nout, kernel_size=1, bias=False),
|
||||
nn.BatchNorm2d(nout),
|
||||
activ(),
|
||||
)
|
||||
|
||||
def __call__(self, x):
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.LeakyReLU):
|
||||
super(Encoder, self).__init__()
|
||||
self.conv1 = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ)
|
||||
self.conv2 = Conv2DBNActiv(nout, nout, ksize, stride, pad, activ=activ)
|
||||
|
||||
def __call__(self, x):
|
||||
skip = self.conv1(x)
|
||||
h = self.conv2(skip)
|
||||
|
||||
return h, skip
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.ReLU, dropout=False):
|
||||
super(Decoder, self).__init__()
|
||||
self.conv = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ)
|
||||
self.dropout = nn.Dropout2d(0.1) if dropout else None
|
||||
|
||||
def __call__(self, x, skip=None):
|
||||
x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=True)
|
||||
if skip is not None:
|
||||
skip = spec_utils.crop_center(skip, x)
|
||||
x = torch.cat([x, skip], dim=1)
|
||||
h = self.conv(x)
|
||||
|
||||
if self.dropout is not None:
|
||||
h = self.dropout(h)
|
||||
|
||||
return h
|
||||
|
||||
|
||||
class ASPPModule(nn.Module):
|
||||
def __init__(self, nin, nout, dilations=(4, 8, 16), activ=nn.ReLU):
|
||||
super(ASPPModule, self).__init__()
|
||||
self.conv1 = nn.Sequential(
|
||||
nn.AdaptiveAvgPool2d((1, None)),
|
||||
Conv2DBNActiv(nin, nin, 1, 1, 0, activ=activ),
|
||||
)
|
||||
self.conv2 = Conv2DBNActiv(nin, nin, 1, 1, 0, activ=activ)
|
||||
self.conv3 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[0], dilations[0], activ=activ)
|
||||
self.conv4 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[1], dilations[1], activ=activ)
|
||||
self.conv5 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[2], dilations[2], activ=activ)
|
||||
self.bottleneck = nn.Sequential(Conv2DBNActiv(nin * 5, nout, 1, 1, 0, activ=activ), nn.Dropout2d(0.1))
|
||||
|
||||
def forward(self, x):
|
||||
_, _, h, w = x.size()
|
||||
feat1 = F.interpolate(self.conv1(x), size=(h, w), mode="bilinear", align_corners=True)
|
||||
feat2 = self.conv2(x)
|
||||
feat3 = self.conv3(x)
|
||||
feat4 = self.conv4(x)
|
||||
feat5 = self.conv5(x)
|
||||
out = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1)
|
||||
bottle = self.bottleneck(out)
|
||||
return bottle
|
||||
@ -1,106 +0,0 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from . import spec_utils
|
||||
|
||||
|
||||
class Conv2DBNActiv(nn.Module):
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU):
|
||||
super(Conv2DBNActiv, self).__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
nin,
|
||||
nout,
|
||||
kernel_size=ksize,
|
||||
stride=stride,
|
||||
padding=pad,
|
||||
dilation=dilation,
|
||||
bias=False,
|
||||
),
|
||||
nn.BatchNorm2d(nout),
|
||||
activ(),
|
||||
)
|
||||
|
||||
def __call__(self, x):
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class SeperableConv2DBNActiv(nn.Module):
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU):
|
||||
super(SeperableConv2DBNActiv, self).__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
nin,
|
||||
nin,
|
||||
kernel_size=ksize,
|
||||
stride=stride,
|
||||
padding=pad,
|
||||
dilation=dilation,
|
||||
groups=nin,
|
||||
bias=False,
|
||||
),
|
||||
nn.Conv2d(nin, nout, kernel_size=1, bias=False),
|
||||
nn.BatchNorm2d(nout),
|
||||
activ(),
|
||||
)
|
||||
|
||||
def __call__(self, x):
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.LeakyReLU):
|
||||
super(Encoder, self).__init__()
|
||||
self.conv1 = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ)
|
||||
self.conv2 = Conv2DBNActiv(nout, nout, ksize, stride, pad, activ=activ)
|
||||
|
||||
def __call__(self, x):
|
||||
skip = self.conv1(x)
|
||||
h = self.conv2(skip)
|
||||
|
||||
return h, skip
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.ReLU, dropout=False):
|
||||
super(Decoder, self).__init__()
|
||||
self.conv = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ)
|
||||
self.dropout = nn.Dropout2d(0.1) if dropout else None
|
||||
|
||||
def __call__(self, x, skip=None):
|
||||
x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=True)
|
||||
if skip is not None:
|
||||
skip = spec_utils.crop_center(skip, x)
|
||||
x = torch.cat([x, skip], dim=1)
|
||||
h = self.conv(x)
|
||||
|
||||
if self.dropout is not None:
|
||||
h = self.dropout(h)
|
||||
|
||||
return h
|
||||
|
||||
|
||||
class ASPPModule(nn.Module):
|
||||
def __init__(self, nin, nout, dilations=(4, 8, 16), activ=nn.ReLU):
|
||||
super(ASPPModule, self).__init__()
|
||||
self.conv1 = nn.Sequential(
|
||||
nn.AdaptiveAvgPool2d((1, None)),
|
||||
Conv2DBNActiv(nin, nin, 1, 1, 0, activ=activ),
|
||||
)
|
||||
self.conv2 = Conv2DBNActiv(nin, nin, 1, 1, 0, activ=activ)
|
||||
self.conv3 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[0], dilations[0], activ=activ)
|
||||
self.conv4 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[1], dilations[1], activ=activ)
|
||||
self.conv5 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[2], dilations[2], activ=activ)
|
||||
self.bottleneck = nn.Sequential(Conv2DBNActiv(nin * 5, nout, 1, 1, 0, activ=activ), nn.Dropout2d(0.1))
|
||||
|
||||
def forward(self, x):
|
||||
_, _, h, w = x.size()
|
||||
feat1 = F.interpolate(self.conv1(x), size=(h, w), mode="bilinear", align_corners=True)
|
||||
feat2 = self.conv2(x)
|
||||
feat3 = self.conv3(x)
|
||||
feat4 = self.conv4(x)
|
||||
feat5 = self.conv5(x)
|
||||
out = torch.cat((feat1, feat2, feat3, feat4, feat5), dim=1)
|
||||
bottle = self.bottleneck(out)
|
||||
return bottle
|
||||
@ -1,110 +0,0 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from . import spec_utils
|
||||
|
||||
|
||||
class Conv2DBNActiv(nn.Module):
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU):
|
||||
super(Conv2DBNActiv, self).__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
nin,
|
||||
nout,
|
||||
kernel_size=ksize,
|
||||
stride=stride,
|
||||
padding=pad,
|
||||
dilation=dilation,
|
||||
bias=False,
|
||||
),
|
||||
nn.BatchNorm2d(nout),
|
||||
activ(),
|
||||
)
|
||||
|
||||
def __call__(self, x):
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class SeperableConv2DBNActiv(nn.Module):
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU):
|
||||
super(SeperableConv2DBNActiv, self).__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
nin,
|
||||
nin,
|
||||
kernel_size=ksize,
|
||||
stride=stride,
|
||||
padding=pad,
|
||||
dilation=dilation,
|
||||
groups=nin,
|
||||
bias=False,
|
||||
),
|
||||
nn.Conv2d(nin, nout, kernel_size=1, bias=False),
|
||||
nn.BatchNorm2d(nout),
|
||||
activ(),
|
||||
)
|
||||
|
||||
def __call__(self, x):
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.LeakyReLU):
|
||||
super(Encoder, self).__init__()
|
||||
self.conv1 = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ)
|
||||
self.conv2 = Conv2DBNActiv(nout, nout, ksize, stride, pad, activ=activ)
|
||||
|
||||
def __call__(self, x):
|
||||
skip = self.conv1(x)
|
||||
h = self.conv2(skip)
|
||||
|
||||
return h, skip
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.ReLU, dropout=False):
|
||||
super(Decoder, self).__init__()
|
||||
self.conv = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ)
|
||||
self.dropout = nn.Dropout2d(0.1) if dropout else None
|
||||
|
||||
def __call__(self, x, skip=None):
|
||||
x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=True)
|
||||
if skip is not None:
|
||||
skip = spec_utils.crop_center(skip, x)
|
||||
x = torch.cat([x, skip], dim=1)
|
||||
h = self.conv(x)
|
||||
|
||||
if self.dropout is not None:
|
||||
h = self.dropout(h)
|
||||
|
||||
return h
|
||||
|
||||
|
||||
class ASPPModule(nn.Module):
|
||||
def __init__(self, nin, nout, dilations=(4, 8, 16, 32, 64), activ=nn.ReLU):
|
||||
super(ASPPModule, self).__init__()
|
||||
self.conv1 = nn.Sequential(
|
||||
nn.AdaptiveAvgPool2d((1, None)),
|
||||
Conv2DBNActiv(nin, nin, 1, 1, 0, activ=activ),
|
||||
)
|
||||
self.conv2 = Conv2DBNActiv(nin, nin, 1, 1, 0, activ=activ)
|
||||
self.conv3 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[0], dilations[0], activ=activ)
|
||||
self.conv4 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[1], dilations[1], activ=activ)
|
||||
self.conv5 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[2], dilations[2], activ=activ)
|
||||
self.conv6 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[2], dilations[2], activ=activ)
|
||||
self.conv7 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[2], dilations[2], activ=activ)
|
||||
self.bottleneck = nn.Sequential(Conv2DBNActiv(nin * 7, nout, 1, 1, 0, activ=activ), nn.Dropout2d(0.1))
|
||||
|
||||
def forward(self, x):
|
||||
_, _, h, w = x.size()
|
||||
feat1 = F.interpolate(self.conv1(x), size=(h, w), mode="bilinear", align_corners=True)
|
||||
feat2 = self.conv2(x)
|
||||
feat3 = self.conv3(x)
|
||||
feat4 = self.conv4(x)
|
||||
feat5 = self.conv5(x)
|
||||
feat6 = self.conv6(x)
|
||||
feat7 = self.conv7(x)
|
||||
out = torch.cat((feat1, feat2, feat3, feat4, feat5, feat6, feat7), dim=1)
|
||||
bottle = self.bottleneck(out)
|
||||
return bottle
|
||||
@ -1,110 +0,0 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from . import spec_utils
|
||||
|
||||
|
||||
class Conv2DBNActiv(nn.Module):
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU):
|
||||
super(Conv2DBNActiv, self).__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
nin,
|
||||
nout,
|
||||
kernel_size=ksize,
|
||||
stride=stride,
|
||||
padding=pad,
|
||||
dilation=dilation,
|
||||
bias=False,
|
||||
),
|
||||
nn.BatchNorm2d(nout),
|
||||
activ(),
|
||||
)
|
||||
|
||||
def __call__(self, x):
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class SeperableConv2DBNActiv(nn.Module):
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU):
|
||||
super(SeperableConv2DBNActiv, self).__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
nin,
|
||||
nin,
|
||||
kernel_size=ksize,
|
||||
stride=stride,
|
||||
padding=pad,
|
||||
dilation=dilation,
|
||||
groups=nin,
|
||||
bias=False,
|
||||
),
|
||||
nn.Conv2d(nin, nout, kernel_size=1, bias=False),
|
||||
nn.BatchNorm2d(nout),
|
||||
activ(),
|
||||
)
|
||||
|
||||
def __call__(self, x):
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.LeakyReLU):
|
||||
super(Encoder, self).__init__()
|
||||
self.conv1 = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ)
|
||||
self.conv2 = Conv2DBNActiv(nout, nout, ksize, stride, pad, activ=activ)
|
||||
|
||||
def __call__(self, x):
|
||||
skip = self.conv1(x)
|
||||
h = self.conv2(skip)
|
||||
|
||||
return h, skip
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.ReLU, dropout=False):
|
||||
super(Decoder, self).__init__()
|
||||
self.conv = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ)
|
||||
self.dropout = nn.Dropout2d(0.1) if dropout else None
|
||||
|
||||
def __call__(self, x, skip=None):
|
||||
x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=True)
|
||||
if skip is not None:
|
||||
skip = spec_utils.crop_center(skip, x)
|
||||
x = torch.cat([x, skip], dim=1)
|
||||
h = self.conv(x)
|
||||
|
||||
if self.dropout is not None:
|
||||
h = self.dropout(h)
|
||||
|
||||
return h
|
||||
|
||||
|
||||
class ASPPModule(nn.Module):
|
||||
def __init__(self, nin, nout, dilations=(4, 8, 16, 32, 64), activ=nn.ReLU):
|
||||
super(ASPPModule, self).__init__()
|
||||
self.conv1 = nn.Sequential(
|
||||
nn.AdaptiveAvgPool2d((1, None)),
|
||||
Conv2DBNActiv(nin, nin, 1, 1, 0, activ=activ),
|
||||
)
|
||||
self.conv2 = Conv2DBNActiv(nin, nin, 1, 1, 0, activ=activ)
|
||||
self.conv3 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[0], dilations[0], activ=activ)
|
||||
self.conv4 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[1], dilations[1], activ=activ)
|
||||
self.conv5 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[2], dilations[2], activ=activ)
|
||||
self.conv6 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[2], dilations[2], activ=activ)
|
||||
self.conv7 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[2], dilations[2], activ=activ)
|
||||
self.bottleneck = nn.Sequential(Conv2DBNActiv(nin * 7, nout, 1, 1, 0, activ=activ), nn.Dropout2d(0.1))
|
||||
|
||||
def forward(self, x):
|
||||
_, _, h, w = x.size()
|
||||
feat1 = F.interpolate(self.conv1(x), size=(h, w), mode="bilinear", align_corners=True)
|
||||
feat2 = self.conv2(x)
|
||||
feat3 = self.conv3(x)
|
||||
feat4 = self.conv4(x)
|
||||
feat5 = self.conv5(x)
|
||||
feat6 = self.conv6(x)
|
||||
feat7 = self.conv7(x)
|
||||
out = torch.cat((feat1, feat2, feat3, feat4, feat5, feat6, feat7), dim=1)
|
||||
bottle = self.bottleneck(out)
|
||||
return bottle
|
||||
@ -1,110 +0,0 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from . import spec_utils
|
||||
|
||||
|
||||
class Conv2DBNActiv(nn.Module):
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU):
|
||||
super(Conv2DBNActiv, self).__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
nin,
|
||||
nout,
|
||||
kernel_size=ksize,
|
||||
stride=stride,
|
||||
padding=pad,
|
||||
dilation=dilation,
|
||||
bias=False,
|
||||
),
|
||||
nn.BatchNorm2d(nout),
|
||||
activ(),
|
||||
)
|
||||
|
||||
def __call__(self, x):
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class SeperableConv2DBNActiv(nn.Module):
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, dilation=1, activ=nn.ReLU):
|
||||
super(SeperableConv2DBNActiv, self).__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
nin,
|
||||
nin,
|
||||
kernel_size=ksize,
|
||||
stride=stride,
|
||||
padding=pad,
|
||||
dilation=dilation,
|
||||
groups=nin,
|
||||
bias=False,
|
||||
),
|
||||
nn.Conv2d(nin, nout, kernel_size=1, bias=False),
|
||||
nn.BatchNorm2d(nout),
|
||||
activ(),
|
||||
)
|
||||
|
||||
def __call__(self, x):
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
class Encoder(nn.Module):
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.LeakyReLU):
|
||||
super(Encoder, self).__init__()
|
||||
self.conv1 = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ)
|
||||
self.conv2 = Conv2DBNActiv(nout, nout, ksize, stride, pad, activ=activ)
|
||||
|
||||
def __call__(self, x):
|
||||
skip = self.conv1(x)
|
||||
h = self.conv2(skip)
|
||||
|
||||
return h, skip
|
||||
|
||||
|
||||
class Decoder(nn.Module):
|
||||
def __init__(self, nin, nout, ksize=3, stride=1, pad=1, activ=nn.ReLU, dropout=False):
|
||||
super(Decoder, self).__init__()
|
||||
self.conv = Conv2DBNActiv(nin, nout, ksize, 1, pad, activ=activ)
|
||||
self.dropout = nn.Dropout2d(0.1) if dropout else None
|
||||
|
||||
def __call__(self, x, skip=None):
|
||||
x = F.interpolate(x, scale_factor=2, mode="bilinear", align_corners=True)
|
||||
if skip is not None:
|
||||
skip = spec_utils.crop_center(skip, x)
|
||||
x = torch.cat([x, skip], dim=1)
|
||||
h = self.conv(x)
|
||||
|
||||
if self.dropout is not None:
|
||||
h = self.dropout(h)
|
||||
|
||||
return h
|
||||
|
||||
|
||||
class ASPPModule(nn.Module):
|
||||
def __init__(self, nin, nout, dilations=(4, 8, 16, 32, 64), activ=nn.ReLU):
|
||||
super(ASPPModule, self).__init__()
|
||||
self.conv1 = nn.Sequential(
|
||||
nn.AdaptiveAvgPool2d((1, None)),
|
||||
Conv2DBNActiv(nin, nin, 1, 1, 0, activ=activ),
|
||||
)
|
||||
self.conv2 = Conv2DBNActiv(nin, nin, 1, 1, 0, activ=activ)
|
||||
self.conv3 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[0], dilations[0], activ=activ)
|
||||
self.conv4 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[1], dilations[1], activ=activ)
|
||||
self.conv5 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[2], dilations[2], activ=activ)
|
||||
self.conv6 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[2], dilations[2], activ=activ)
|
||||
self.conv7 = SeperableConv2DBNActiv(nin, nin, 3, 1, dilations[2], dilations[2], activ=activ)
|
||||
self.bottleneck = nn.Sequential(Conv2DBNActiv(nin * 7, nout, 1, 1, 0, activ=activ), nn.Dropout2d(0.1))
|
||||
|
||||
def forward(self, x):
|
||||
_, _, h, w = x.size()
|
||||
feat1 = F.interpolate(self.conv1(x), size=(h, w), mode="bilinear", align_corners=True)
|
||||
feat2 = self.conv2(x)
|
||||
feat3 = self.conv3(x)
|
||||
feat4 = self.conv4(x)
|
||||
feat5 = self.conv5(x)
|
||||
feat6 = self.conv6(x)
|
||||
feat7 = self.conv7(x)
|
||||
out = torch.cat((feat1, feat2, feat3, feat4, feat5, feat6, feat7), dim=1)
|
||||
bottle = self.bottleneck(out)
|
||||
return bottle
|
||||
@ -1,19 +0,0 @@
|
||||
{
|
||||
"bins": 1024,
|
||||
"unstable_bins": 0,
|
||||
"reduction_bins": 0,
|
||||
"band": {
|
||||
"1": {
|
||||
"sr": 16000,
|
||||
"hl": 512,
|
||||
"n_fft": 2048,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 1024,
|
||||
"hpf_start": -1,
|
||||
"res_type": "sinc_best"
|
||||
}
|
||||
},
|
||||
"sr": 16000,
|
||||
"pre_filter_start": 1023,
|
||||
"pre_filter_stop": 1024
|
||||
}
|
||||
@ -1,19 +0,0 @@
|
||||
{
|
||||
"bins": 1024,
|
||||
"unstable_bins": 0,
|
||||
"reduction_bins": 0,
|
||||
"band": {
|
||||
"1": {
|
||||
"sr": 32000,
|
||||
"hl": 512,
|
||||
"n_fft": 2048,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 1024,
|
||||
"hpf_start": -1,
|
||||
"res_type": "kaiser_fast"
|
||||
}
|
||||
},
|
||||
"sr": 32000,
|
||||
"pre_filter_start": 1000,
|
||||
"pre_filter_stop": 1021
|
||||
}
|
||||
@ -1,19 +0,0 @@
|
||||
{
|
||||
"bins": 1024,
|
||||
"unstable_bins": 0,
|
||||
"reduction_bins": 0,
|
||||
"band": {
|
||||
"1": {
|
||||
"sr": 33075,
|
||||
"hl": 384,
|
||||
"n_fft": 2048,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 1024,
|
||||
"hpf_start": -1,
|
||||
"res_type": "sinc_best"
|
||||
}
|
||||
},
|
||||
"sr": 33075,
|
||||
"pre_filter_start": 1000,
|
||||
"pre_filter_stop": 1021
|
||||
}
|
||||
@ -1,19 +0,0 @@
|
||||
{
|
||||
"bins": 1024,
|
||||
"unstable_bins": 0,
|
||||
"reduction_bins": 0,
|
||||
"band": {
|
||||
"1": {
|
||||
"sr": 44100,
|
||||
"hl": 1024,
|
||||
"n_fft": 2048,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 1024,
|
||||
"hpf_start": -1,
|
||||
"res_type": "sinc_best"
|
||||
}
|
||||
},
|
||||
"sr": 44100,
|
||||
"pre_filter_start": 1023,
|
||||
"pre_filter_stop": 1024
|
||||
}
|
||||
@ -1,19 +0,0 @@
|
||||
{
|
||||
"bins": 256,
|
||||
"unstable_bins": 0,
|
||||
"reduction_bins": 0,
|
||||
"band": {
|
||||
"1": {
|
||||
"sr": 44100,
|
||||
"hl": 256,
|
||||
"n_fft": 512,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 256,
|
||||
"hpf_start": -1,
|
||||
"res_type": "sinc_best"
|
||||
}
|
||||
},
|
||||
"sr": 44100,
|
||||
"pre_filter_start": 256,
|
||||
"pre_filter_stop": 256
|
||||
}
|
||||
@ -1,19 +0,0 @@
|
||||
{
|
||||
"bins": 1024,
|
||||
"unstable_bins": 0,
|
||||
"reduction_bins": 0,
|
||||
"band": {
|
||||
"1": {
|
||||
"sr": 44100,
|
||||
"hl": 512,
|
||||
"n_fft": 2048,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 1024,
|
||||
"hpf_start": -1,
|
||||
"res_type": "sinc_best"
|
||||
}
|
||||
},
|
||||
"sr": 44100,
|
||||
"pre_filter_start": 1023,
|
||||
"pre_filter_stop": 1024
|
||||
}
|
||||
@ -1,19 +0,0 @@
|
||||
{
|
||||
"bins": 1024,
|
||||
"unstable_bins": 0,
|
||||
"reduction_bins": 0,
|
||||
"band": {
|
||||
"1": {
|
||||
"sr": 44100,
|
||||
"hl": 512,
|
||||
"n_fft": 2048,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 700,
|
||||
"hpf_start": -1,
|
||||
"res_type": "sinc_best"
|
||||
}
|
||||
},
|
||||
"sr": 44100,
|
||||
"pre_filter_start": 1023,
|
||||
"pre_filter_stop": 700
|
||||
}
|
||||
@ -1,30 +0,0 @@
|
||||
{
|
||||
"bins": 768,
|
||||
"unstable_bins": 7,
|
||||
"reduction_bins": 705,
|
||||
"band": {
|
||||
"1": {
|
||||
"sr": 6000,
|
||||
"hl": 66,
|
||||
"n_fft": 512,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 240,
|
||||
"lpf_start": 60,
|
||||
"lpf_stop": 118,
|
||||
"res_type": "sinc_fastest"
|
||||
},
|
||||
"2": {
|
||||
"sr": 32000,
|
||||
"hl": 352,
|
||||
"n_fft": 1024,
|
||||
"crop_start": 22,
|
||||
"crop_stop": 505,
|
||||
"hpf_start": 44,
|
||||
"hpf_stop": 23,
|
||||
"res_type": "sinc_medium"
|
||||
}
|
||||
},
|
||||
"sr": 32000,
|
||||
"pre_filter_start": 710,
|
||||
"pre_filter_stop": 731
|
||||
}
|
||||
@ -1,30 +0,0 @@
|
||||
{
|
||||
"bins": 512,
|
||||
"unstable_bins": 7,
|
||||
"reduction_bins": 510,
|
||||
"band": {
|
||||
"1": {
|
||||
"sr": 11025,
|
||||
"hl": 160,
|
||||
"n_fft": 768,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 192,
|
||||
"lpf_start": 41,
|
||||
"lpf_stop": 139,
|
||||
"res_type": "sinc_fastest"
|
||||
},
|
||||
"2": {
|
||||
"sr": 44100,
|
||||
"hl": 640,
|
||||
"n_fft": 1024,
|
||||
"crop_start": 10,
|
||||
"crop_stop": 320,
|
||||
"hpf_start": 47,
|
||||
"hpf_stop": 15,
|
||||
"res_type": "sinc_medium"
|
||||
}
|
||||
},
|
||||
"sr": 44100,
|
||||
"pre_filter_start": 510,
|
||||
"pre_filter_stop": 512
|
||||
}
|
||||
@ -1,30 +0,0 @@
|
||||
{
|
||||
"bins": 768,
|
||||
"unstable_bins": 7,
|
||||
"reduction_bins": 705,
|
||||
"band": {
|
||||
"1": {
|
||||
"sr": 6000,
|
||||
"hl": 66,
|
||||
"n_fft": 512,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 240,
|
||||
"lpf_start": 60,
|
||||
"lpf_stop": 240,
|
||||
"res_type": "sinc_fastest"
|
||||
},
|
||||
"2": {
|
||||
"sr": 48000,
|
||||
"hl": 528,
|
||||
"n_fft": 1536,
|
||||
"crop_start": 22,
|
||||
"crop_stop": 505,
|
||||
"hpf_start": 82,
|
||||
"hpf_stop": 22,
|
||||
"res_type": "sinc_medium"
|
||||
}
|
||||
},
|
||||
"sr": 48000,
|
||||
"pre_filter_start": 710,
|
||||
"pre_filter_stop": 731
|
||||
}
|
||||
@ -1,42 +0,0 @@
|
||||
{
|
||||
"bins": 768,
|
||||
"unstable_bins": 5,
|
||||
"reduction_bins": 733,
|
||||
"band": {
|
||||
"1": {
|
||||
"sr": 11025,
|
||||
"hl": 128,
|
||||
"n_fft": 768,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 278,
|
||||
"lpf_start": 28,
|
||||
"lpf_stop": 140,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"2": {
|
||||
"sr": 22050,
|
||||
"hl": 256,
|
||||
"n_fft": 768,
|
||||
"crop_start": 14,
|
||||
"crop_stop": 322,
|
||||
"hpf_start": 70,
|
||||
"hpf_stop": 14,
|
||||
"lpf_start": 283,
|
||||
"lpf_stop": 314,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"3": {
|
||||
"sr": 44100,
|
||||
"hl": 512,
|
||||
"n_fft": 768,
|
||||
"crop_start": 131,
|
||||
"crop_stop": 313,
|
||||
"hpf_start": 154,
|
||||
"hpf_stop": 141,
|
||||
"res_type": "sinc_medium"
|
||||
}
|
||||
},
|
||||
"sr": 44100,
|
||||
"pre_filter_start": 757,
|
||||
"pre_filter_stop": 768
|
||||
}
|
||||
@ -1,43 +0,0 @@
|
||||
{
|
||||
"mid_side": true,
|
||||
"bins": 768,
|
||||
"unstable_bins": 5,
|
||||
"reduction_bins": 733,
|
||||
"band": {
|
||||
"1": {
|
||||
"sr": 11025,
|
||||
"hl": 128,
|
||||
"n_fft": 768,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 278,
|
||||
"lpf_start": 28,
|
||||
"lpf_stop": 140,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"2": {
|
||||
"sr": 22050,
|
||||
"hl": 256,
|
||||
"n_fft": 768,
|
||||
"crop_start": 14,
|
||||
"crop_stop": 322,
|
||||
"hpf_start": 70,
|
||||
"hpf_stop": 14,
|
||||
"lpf_start": 283,
|
||||
"lpf_stop": 314,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"3": {
|
||||
"sr": 44100,
|
||||
"hl": 512,
|
||||
"n_fft": 768,
|
||||
"crop_start": 131,
|
||||
"crop_stop": 313,
|
||||
"hpf_start": 154,
|
||||
"hpf_stop": 141,
|
||||
"res_type": "sinc_medium"
|
||||
}
|
||||
},
|
||||
"sr": 44100,
|
||||
"pre_filter_start": 757,
|
||||
"pre_filter_stop": 768
|
||||
}
|
||||
@ -1,43 +0,0 @@
|
||||
{
|
||||
"mid_side_b2": true,
|
||||
"bins": 640,
|
||||
"unstable_bins": 7,
|
||||
"reduction_bins": 565,
|
||||
"band": {
|
||||
"1": {
|
||||
"sr": 11025,
|
||||
"hl": 108,
|
||||
"n_fft": 1024,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 187,
|
||||
"lpf_start": 92,
|
||||
"lpf_stop": 186,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"2": {
|
||||
"sr": 22050,
|
||||
"hl": 216,
|
||||
"n_fft": 768,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 212,
|
||||
"hpf_start": 68,
|
||||
"hpf_stop": 34,
|
||||
"lpf_start": 174,
|
||||
"lpf_stop": 209,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"3": {
|
||||
"sr": 44100,
|
||||
"hl": 432,
|
||||
"n_fft": 640,
|
||||
"crop_start": 66,
|
||||
"crop_stop": 307,
|
||||
"hpf_start": 86,
|
||||
"hpf_stop": 72,
|
||||
"res_type": "kaiser_fast"
|
||||
}
|
||||
},
|
||||
"sr": 44100,
|
||||
"pre_filter_start": 639,
|
||||
"pre_filter_stop": 640
|
||||
}
|
||||
@ -1,54 +0,0 @@
|
||||
{
|
||||
"bins": 768,
|
||||
"unstable_bins": 7,
|
||||
"reduction_bins": 668,
|
||||
"band": {
|
||||
"1": {
|
||||
"sr": 11025,
|
||||
"hl": 128,
|
||||
"n_fft": 1024,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 186,
|
||||
"lpf_start": 37,
|
||||
"lpf_stop": 73,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"2": {
|
||||
"sr": 11025,
|
||||
"hl": 128,
|
||||
"n_fft": 512,
|
||||
"crop_start": 4,
|
||||
"crop_stop": 185,
|
||||
"hpf_start": 36,
|
||||
"hpf_stop": 18,
|
||||
"lpf_start": 93,
|
||||
"lpf_stop": 185,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"3": {
|
||||
"sr": 22050,
|
||||
"hl": 256,
|
||||
"n_fft": 512,
|
||||
"crop_start": 46,
|
||||
"crop_stop": 186,
|
||||
"hpf_start": 93,
|
||||
"hpf_stop": 46,
|
||||
"lpf_start": 164,
|
||||
"lpf_stop": 186,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"4": {
|
||||
"sr": 44100,
|
||||
"hl": 512,
|
||||
"n_fft": 768,
|
||||
"crop_start": 121,
|
||||
"crop_stop": 382,
|
||||
"hpf_start": 138,
|
||||
"hpf_stop": 123,
|
||||
"res_type": "sinc_medium"
|
||||
}
|
||||
},
|
||||
"sr": 44100,
|
||||
"pre_filter_start": 740,
|
||||
"pre_filter_stop": 768
|
||||
}
|
||||
@ -1,55 +0,0 @@
|
||||
{
|
||||
"bins": 768,
|
||||
"unstable_bins": 7,
|
||||
"mid_side": true,
|
||||
"reduction_bins": 668,
|
||||
"band": {
|
||||
"1": {
|
||||
"sr": 11025,
|
||||
"hl": 128,
|
||||
"n_fft": 1024,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 186,
|
||||
"lpf_start": 37,
|
||||
"lpf_stop": 73,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"2": {
|
||||
"sr": 11025,
|
||||
"hl": 128,
|
||||
"n_fft": 512,
|
||||
"crop_start": 4,
|
||||
"crop_stop": 185,
|
||||
"hpf_start": 36,
|
||||
"hpf_stop": 18,
|
||||
"lpf_start": 93,
|
||||
"lpf_stop": 185,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"3": {
|
||||
"sr": 22050,
|
||||
"hl": 256,
|
||||
"n_fft": 512,
|
||||
"crop_start": 46,
|
||||
"crop_stop": 186,
|
||||
"hpf_start": 93,
|
||||
"hpf_stop": 46,
|
||||
"lpf_start": 164,
|
||||
"lpf_stop": 186,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"4": {
|
||||
"sr": 44100,
|
||||
"hl": 512,
|
||||
"n_fft": 768,
|
||||
"crop_start": 121,
|
||||
"crop_stop": 382,
|
||||
"hpf_start": 138,
|
||||
"hpf_stop": 123,
|
||||
"res_type": "sinc_medium"
|
||||
}
|
||||
},
|
||||
"sr": 44100,
|
||||
"pre_filter_start": 740,
|
||||
"pre_filter_stop": 768
|
||||
}
|
||||
@ -1,55 +0,0 @@
|
||||
{
|
||||
"mid_side_b": true,
|
||||
"bins": 768,
|
||||
"unstable_bins": 7,
|
||||
"reduction_bins": 668,
|
||||
"band": {
|
||||
"1": {
|
||||
"sr": 11025,
|
||||
"hl": 128,
|
||||
"n_fft": 1024,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 186,
|
||||
"lpf_start": 37,
|
||||
"lpf_stop": 73,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"2": {
|
||||
"sr": 11025,
|
||||
"hl": 128,
|
||||
"n_fft": 512,
|
||||
"crop_start": 4,
|
||||
"crop_stop": 185,
|
||||
"hpf_start": 36,
|
||||
"hpf_stop": 18,
|
||||
"lpf_start": 93,
|
||||
"lpf_stop": 185,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"3": {
|
||||
"sr": 22050,
|
||||
"hl": 256,
|
||||
"n_fft": 512,
|
||||
"crop_start": 46,
|
||||
"crop_stop": 186,
|
||||
"hpf_start": 93,
|
||||
"hpf_stop": 46,
|
||||
"lpf_start": 164,
|
||||
"lpf_stop": 186,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"4": {
|
||||
"sr": 44100,
|
||||
"hl": 512,
|
||||
"n_fft": 768,
|
||||
"crop_start": 121,
|
||||
"crop_stop": 382,
|
||||
"hpf_start": 138,
|
||||
"hpf_stop": 123,
|
||||
"res_type": "sinc_medium"
|
||||
}
|
||||
},
|
||||
"sr": 44100,
|
||||
"pre_filter_start": 740,
|
||||
"pre_filter_stop": 768
|
||||
}
|
||||
@ -1,55 +0,0 @@
|
||||
{
|
||||
"mid_side_b": true,
|
||||
"bins": 768,
|
||||
"unstable_bins": 7,
|
||||
"reduction_bins": 668,
|
||||
"band": {
|
||||
"1": {
|
||||
"sr": 11025,
|
||||
"hl": 128,
|
||||
"n_fft": 1024,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 186,
|
||||
"lpf_start": 37,
|
||||
"lpf_stop": 73,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"2": {
|
||||
"sr": 11025,
|
||||
"hl": 128,
|
||||
"n_fft": 512,
|
||||
"crop_start": 4,
|
||||
"crop_stop": 185,
|
||||
"hpf_start": 36,
|
||||
"hpf_stop": 18,
|
||||
"lpf_start": 93,
|
||||
"lpf_stop": 185,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"3": {
|
||||
"sr": 22050,
|
||||
"hl": 256,
|
||||
"n_fft": 512,
|
||||
"crop_start": 46,
|
||||
"crop_stop": 186,
|
||||
"hpf_start": 93,
|
||||
"hpf_stop": 46,
|
||||
"lpf_start": 164,
|
||||
"lpf_stop": 186,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"4": {
|
||||
"sr": 44100,
|
||||
"hl": 512,
|
||||
"n_fft": 768,
|
||||
"crop_start": 121,
|
||||
"crop_stop": 382,
|
||||
"hpf_start": 138,
|
||||
"hpf_stop": 123,
|
||||
"res_type": "sinc_medium"
|
||||
}
|
||||
},
|
||||
"sr": 44100,
|
||||
"pre_filter_start": 740,
|
||||
"pre_filter_stop": 768
|
||||
}
|
||||
@ -1,55 +0,0 @@
|
||||
{
|
||||
"reverse": true,
|
||||
"bins": 768,
|
||||
"unstable_bins": 7,
|
||||
"reduction_bins": 668,
|
||||
"band": {
|
||||
"1": {
|
||||
"sr": 11025,
|
||||
"hl": 128,
|
||||
"n_fft": 1024,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 186,
|
||||
"lpf_start": 37,
|
||||
"lpf_stop": 73,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"2": {
|
||||
"sr": 11025,
|
||||
"hl": 128,
|
||||
"n_fft": 512,
|
||||
"crop_start": 4,
|
||||
"crop_stop": 185,
|
||||
"hpf_start": 36,
|
||||
"hpf_stop": 18,
|
||||
"lpf_start": 93,
|
||||
"lpf_stop": 185,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"3": {
|
||||
"sr": 22050,
|
||||
"hl": 256,
|
||||
"n_fft": 512,
|
||||
"crop_start": 46,
|
||||
"crop_stop": 186,
|
||||
"hpf_start": 93,
|
||||
"hpf_stop": 46,
|
||||
"lpf_start": 164,
|
||||
"lpf_stop": 186,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"4": {
|
||||
"sr": 44100,
|
||||
"hl": 512,
|
||||
"n_fft": 768,
|
||||
"crop_start": 121,
|
||||
"crop_stop": 382,
|
||||
"hpf_start": 138,
|
||||
"hpf_stop": 123,
|
||||
"res_type": "sinc_medium"
|
||||
}
|
||||
},
|
||||
"sr": 44100,
|
||||
"pre_filter_start": 740,
|
||||
"pre_filter_stop": 768
|
||||
}
|
||||
@ -1,55 +0,0 @@
|
||||
{
|
||||
"stereo_w": true,
|
||||
"bins": 768,
|
||||
"unstable_bins": 7,
|
||||
"reduction_bins": 668,
|
||||
"band": {
|
||||
"1": {
|
||||
"sr": 11025,
|
||||
"hl": 128,
|
||||
"n_fft": 1024,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 186,
|
||||
"lpf_start": 37,
|
||||
"lpf_stop": 73,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"2": {
|
||||
"sr": 11025,
|
||||
"hl": 128,
|
||||
"n_fft": 512,
|
||||
"crop_start": 4,
|
||||
"crop_stop": 185,
|
||||
"hpf_start": 36,
|
||||
"hpf_stop": 18,
|
||||
"lpf_start": 93,
|
||||
"lpf_stop": 185,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"3": {
|
||||
"sr": 22050,
|
||||
"hl": 256,
|
||||
"n_fft": 512,
|
||||
"crop_start": 46,
|
||||
"crop_stop": 186,
|
||||
"hpf_start": 93,
|
||||
"hpf_stop": 46,
|
||||
"lpf_start": 164,
|
||||
"lpf_stop": 186,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"4": {
|
||||
"sr": 44100,
|
||||
"hl": 512,
|
||||
"n_fft": 768,
|
||||
"crop_start": 121,
|
||||
"crop_stop": 382,
|
||||
"hpf_start": 138,
|
||||
"hpf_stop": 123,
|
||||
"res_type": "sinc_medium"
|
||||
}
|
||||
},
|
||||
"sr": 44100,
|
||||
"pre_filter_start": 740,
|
||||
"pre_filter_stop": 768
|
||||
}
|
||||
@ -1,55 +0,0 @@
|
||||
{
|
||||
"bins": 672,
|
||||
"unstable_bins": 8,
|
||||
"reduction_bins": 637,
|
||||
"band": {
|
||||
"1": {
|
||||
"sr": 7350,
|
||||
"hl": 80,
|
||||
"n_fft": 640,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 85,
|
||||
"lpf_start": 25,
|
||||
"lpf_stop": 53,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"2": {
|
||||
"sr": 7350,
|
||||
"hl": 80,
|
||||
"n_fft": 320,
|
||||
"crop_start": 4,
|
||||
"crop_stop": 87,
|
||||
"hpf_start": 25,
|
||||
"hpf_stop": 12,
|
||||
"lpf_start": 31,
|
||||
"lpf_stop": 62,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"3": {
|
||||
"sr": 14700,
|
||||
"hl": 160,
|
||||
"n_fft": 512,
|
||||
"crop_start": 17,
|
||||
"crop_stop": 216,
|
||||
"hpf_start": 48,
|
||||
"hpf_stop": 24,
|
||||
"lpf_start": 139,
|
||||
"lpf_stop": 210,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"4": {
|
||||
"sr": 44100,
|
||||
"hl": 480,
|
||||
"n_fft": 960,
|
||||
"crop_start": 78,
|
||||
"crop_stop": 383,
|
||||
"hpf_start": 130,
|
||||
"hpf_stop": 86,
|
||||
"convert_channels": "stereo_n",
|
||||
"res_type": "kaiser_fast"
|
||||
}
|
||||
},
|
||||
"sr": 44100,
|
||||
"pre_filter_start": 668,
|
||||
"pre_filter_stop": 672
|
||||
}
|
||||
@ -1,43 +0,0 @@
|
||||
{
|
||||
"mid_side_b2": true,
|
||||
"bins": 1280,
|
||||
"unstable_bins": 7,
|
||||
"reduction_bins": 565,
|
||||
"band": {
|
||||
"1": {
|
||||
"sr": 11025,
|
||||
"hl": 108,
|
||||
"n_fft": 2048,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 374,
|
||||
"lpf_start": 92,
|
||||
"lpf_stop": 186,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"2": {
|
||||
"sr": 22050,
|
||||
"hl": 216,
|
||||
"n_fft": 1536,
|
||||
"crop_start": 0,
|
||||
"crop_stop": 424,
|
||||
"hpf_start": 68,
|
||||
"hpf_stop": 34,
|
||||
"lpf_start": 348,
|
||||
"lpf_stop": 418,
|
||||
"res_type": "polyphase"
|
||||
},
|
||||
"3": {
|
||||
"sr": 44100,
|
||||
"hl": 432,
|
||||
"n_fft": 1280,
|
||||
"crop_start": 132,
|
||||
"crop_stop": 614,
|
||||
"hpf_start": 172,
|
||||
"hpf_stop": 144,
|
||||
"res_type": "polyphase"
|
||||
}
|
||||
},
|
||||
"sr": 44100,
|
||||
"pre_filter_start": 1280,
|
||||
"pre_filter_stop": 1280
|
||||
}
|
||||
@ -1,121 +0,0 @@
|
||||
import layers
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
|
||||
class BaseASPPNet(nn.Module):
|
||||
def __init__(self, nin, ch, dilations=(4, 8, 16)):
|
||||
super(BaseASPPNet, self).__init__()
|
||||
self.enc1 = layers.Encoder(nin, ch, 3, 2, 1)
|
||||
self.enc2 = layers.Encoder(ch, ch * 2, 3, 2, 1)
|
||||
self.enc3 = layers.Encoder(ch * 2, ch * 4, 3, 2, 1)
|
||||
self.enc4 = layers.Encoder(ch * 4, ch * 8, 3, 2, 1)
|
||||
|
||||
self.aspp = layers.ASPPModule(ch * 8, ch * 16, dilations)
|
||||
|
||||
self.dec4 = layers.Decoder(ch * (8 + 16), ch * 8, 3, 1, 1)
|
||||
self.dec3 = layers.Decoder(ch * (4 + 8), ch * 4, 3, 1, 1)
|
||||
self.dec2 = layers.Decoder(ch * (2 + 4), ch * 2, 3, 1, 1)
|
||||
self.dec1 = layers.Decoder(ch * (1 + 2), ch, 3, 1, 1)
|
||||
|
||||
def __call__(self, x):
|
||||
h, e1 = self.enc1(x)
|
||||
h, e2 = self.enc2(h)
|
||||
h, e3 = self.enc3(h)
|
||||
h, e4 = self.enc4(h)
|
||||
|
||||
h = self.aspp(h)
|
||||
|
||||
h = self.dec4(h, e4)
|
||||
h = self.dec3(h, e3)
|
||||
h = self.dec2(h, e2)
|
||||
h = self.dec1(h, e1)
|
||||
|
||||
return h
|
||||
|
||||
|
||||
class CascadedASPPNet(nn.Module):
|
||||
def __init__(self, n_fft):
|
||||
super(CascadedASPPNet, self).__init__()
|
||||
self.stg1_low_band_net = BaseASPPNet(2, 16)
|
||||
self.stg1_high_band_net = BaseASPPNet(2, 16)
|
||||
|
||||
self.stg2_bridge = layers.Conv2DBNActiv(18, 8, 1, 1, 0)
|
||||
self.stg2_full_band_net = BaseASPPNet(8, 16)
|
||||
|
||||
self.stg3_bridge = layers.Conv2DBNActiv(34, 16, 1, 1, 0)
|
||||
self.stg3_full_band_net = BaseASPPNet(16, 32)
|
||||
|
||||
self.out = nn.Conv2d(32, 2, 1, bias=False)
|
||||
self.aux1_out = nn.Conv2d(16, 2, 1, bias=False)
|
||||
self.aux2_out = nn.Conv2d(16, 2, 1, bias=False)
|
||||
|
||||
self.max_bin = n_fft // 2
|
||||
self.output_bin = n_fft // 2 + 1
|
||||
|
||||
self.offset = 128
|
||||
|
||||
def forward(self, x, aggressiveness=None):
|
||||
mix = x.detach()
|
||||
x = x.clone()
|
||||
|
||||
x = x[:, :, : self.max_bin]
|
||||
|
||||
bandw = x.size()[2] // 2
|
||||
aux1 = torch.cat(
|
||||
[
|
||||
self.stg1_low_band_net(x[:, :, :bandw]),
|
||||
self.stg1_high_band_net(x[:, :, bandw:]),
|
||||
],
|
||||
dim=2,
|
||||
)
|
||||
|
||||
h = torch.cat([x, aux1], dim=1)
|
||||
aux2 = self.stg2_full_band_net(self.stg2_bridge(h))
|
||||
|
||||
h = torch.cat([x, aux1, aux2], dim=1)
|
||||
h = self.stg3_full_band_net(self.stg3_bridge(h))
|
||||
|
||||
mask = torch.sigmoid(self.out(h))
|
||||
mask = F.pad(
|
||||
input=mask,
|
||||
pad=(0, 0, 0, self.output_bin - mask.size()[2]),
|
||||
mode="replicate",
|
||||
)
|
||||
|
||||
if self.training:
|
||||
aux1 = torch.sigmoid(self.aux1_out(aux1))
|
||||
aux1 = F.pad(
|
||||
input=aux1,
|
||||
pad=(0, 0, 0, self.output_bin - aux1.size()[2]),
|
||||
mode="replicate",
|
||||
)
|
||||
aux2 = torch.sigmoid(self.aux2_out(aux2))
|
||||
aux2 = F.pad(
|
||||
input=aux2,
|
||||
pad=(0, 0, 0, self.output_bin - aux2.size()[2]),
|
||||
mode="replicate",
|
||||
)
|
||||
return mask * mix, aux1 * mix, aux2 * mix
|
||||
else:
|
||||
if aggressiveness:
|
||||
mask[:, :, : aggressiveness["split_bin"]] = torch.pow(
|
||||
mask[:, :, : aggressiveness["split_bin"]],
|
||||
1 + aggressiveness["value"] / 3,
|
||||
)
|
||||
mask[:, :, aggressiveness["split_bin"] :] = torch.pow(
|
||||
mask[:, :, aggressiveness["split_bin"] :],
|
||||
1 + aggressiveness["value"],
|
||||
)
|
||||
|
||||
return mask * mix
|
||||
|
||||
def predict(self, x_mag, aggressiveness=None):
|
||||
h = self.forward(x_mag, aggressiveness)
|
||||
|
||||
if self.offset > 0:
|
||||
h = h[:, :, :, self.offset : -self.offset]
|
||||
assert h.size()[3] > 0
|
||||
|
||||
return h
|
||||
@ -1,122 +0,0 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from . import layers_123821KB as layers
|
||||
|
||||
|
||||
class BaseASPPNet(nn.Module):
|
||||
def __init__(self, nin, ch, dilations=(4, 8, 16)):
|
||||
super(BaseASPPNet, self).__init__()
|
||||
self.enc1 = layers.Encoder(nin, ch, 3, 2, 1)
|
||||
self.enc2 = layers.Encoder(ch, ch * 2, 3, 2, 1)
|
||||
self.enc3 = layers.Encoder(ch * 2, ch * 4, 3, 2, 1)
|
||||
self.enc4 = layers.Encoder(ch * 4, ch * 8, 3, 2, 1)
|
||||
|
||||
self.aspp = layers.ASPPModule(ch * 8, ch * 16, dilations)
|
||||
|
||||
self.dec4 = layers.Decoder(ch * (8 + 16), ch * 8, 3, 1, 1)
|
||||
self.dec3 = layers.Decoder(ch * (4 + 8), ch * 4, 3, 1, 1)
|
||||
self.dec2 = layers.Decoder(ch * (2 + 4), ch * 2, 3, 1, 1)
|
||||
self.dec1 = layers.Decoder(ch * (1 + 2), ch, 3, 1, 1)
|
||||
|
||||
def __call__(self, x):
|
||||
h, e1 = self.enc1(x)
|
||||
h, e2 = self.enc2(h)
|
||||
h, e3 = self.enc3(h)
|
||||
h, e4 = self.enc4(h)
|
||||
|
||||
h = self.aspp(h)
|
||||
|
||||
h = self.dec4(h, e4)
|
||||
h = self.dec3(h, e3)
|
||||
h = self.dec2(h, e2)
|
||||
h = self.dec1(h, e1)
|
||||
|
||||
return h
|
||||
|
||||
|
||||
class CascadedASPPNet(nn.Module):
|
||||
def __init__(self, n_fft):
|
||||
super(CascadedASPPNet, self).__init__()
|
||||
self.stg1_low_band_net = BaseASPPNet(2, 32)
|
||||
self.stg1_high_band_net = BaseASPPNet(2, 32)
|
||||
|
||||
self.stg2_bridge = layers.Conv2DBNActiv(34, 16, 1, 1, 0)
|
||||
self.stg2_full_band_net = BaseASPPNet(16, 32)
|
||||
|
||||
self.stg3_bridge = layers.Conv2DBNActiv(66, 32, 1, 1, 0)
|
||||
self.stg3_full_band_net = BaseASPPNet(32, 64)
|
||||
|
||||
self.out = nn.Conv2d(64, 2, 1, bias=False)
|
||||
self.aux1_out = nn.Conv2d(32, 2, 1, bias=False)
|
||||
self.aux2_out = nn.Conv2d(32, 2, 1, bias=False)
|
||||
|
||||
self.max_bin = n_fft // 2
|
||||
self.output_bin = n_fft // 2 + 1
|
||||
|
||||
self.offset = 128
|
||||
|
||||
def forward(self, x, aggressiveness=None):
|
||||
mix = x.detach()
|
||||
x = x.clone()
|
||||
|
||||
x = x[:, :, : self.max_bin]
|
||||
|
||||
bandw = x.size()[2] // 2
|
||||
aux1 = torch.cat(
|
||||
[
|
||||
self.stg1_low_band_net(x[:, :, :bandw]),
|
||||
self.stg1_high_band_net(x[:, :, bandw:]),
|
||||
],
|
||||
dim=2,
|
||||
)
|
||||
|
||||
h = torch.cat([x, aux1], dim=1)
|
||||
aux2 = self.stg2_full_band_net(self.stg2_bridge(h))
|
||||
|
||||
h = torch.cat([x, aux1, aux2], dim=1)
|
||||
h = self.stg3_full_band_net(self.stg3_bridge(h))
|
||||
|
||||
mask = torch.sigmoid(self.out(h))
|
||||
mask = F.pad(
|
||||
input=mask,
|
||||
pad=(0, 0, 0, self.output_bin - mask.size()[2]),
|
||||
mode="replicate",
|
||||
)
|
||||
|
||||
if self.training:
|
||||
aux1 = torch.sigmoid(self.aux1_out(aux1))
|
||||
aux1 = F.pad(
|
||||
input=aux1,
|
||||
pad=(0, 0, 0, self.output_bin - aux1.size()[2]),
|
||||
mode="replicate",
|
||||
)
|
||||
aux2 = torch.sigmoid(self.aux2_out(aux2))
|
||||
aux2 = F.pad(
|
||||
input=aux2,
|
||||
pad=(0, 0, 0, self.output_bin - aux2.size()[2]),
|
||||
mode="replicate",
|
||||
)
|
||||
return mask * mix, aux1 * mix, aux2 * mix
|
||||
else:
|
||||
if aggressiveness:
|
||||
mask[:, :, : aggressiveness["split_bin"]] = torch.pow(
|
||||
mask[:, :, : aggressiveness["split_bin"]],
|
||||
1 + aggressiveness["value"] / 3,
|
||||
)
|
||||
mask[:, :, aggressiveness["split_bin"] :] = torch.pow(
|
||||
mask[:, :, aggressiveness["split_bin"] :],
|
||||
1 + aggressiveness["value"],
|
||||
)
|
||||
|
||||
return mask * mix
|
||||
|
||||
def predict(self, x_mag, aggressiveness=None):
|
||||
h = self.forward(x_mag, aggressiveness)
|
||||
|
||||
if self.offset > 0:
|
||||
h = h[:, :, :, self.offset : -self.offset]
|
||||
assert h.size()[3] > 0
|
||||
|
||||
return h
|
||||
@ -1,122 +0,0 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from . import layers_123821KB as layers
|
||||
|
||||
|
||||
class BaseASPPNet(nn.Module):
|
||||
def __init__(self, nin, ch, dilations=(4, 8, 16)):
|
||||
super(BaseASPPNet, self).__init__()
|
||||
self.enc1 = layers.Encoder(nin, ch, 3, 2, 1)
|
||||
self.enc2 = layers.Encoder(ch, ch * 2, 3, 2, 1)
|
||||
self.enc3 = layers.Encoder(ch * 2, ch * 4, 3, 2, 1)
|
||||
self.enc4 = layers.Encoder(ch * 4, ch * 8, 3, 2, 1)
|
||||
|
||||
self.aspp = layers.ASPPModule(ch * 8, ch * 16, dilations)
|
||||
|
||||
self.dec4 = layers.Decoder(ch * (8 + 16), ch * 8, 3, 1, 1)
|
||||
self.dec3 = layers.Decoder(ch * (4 + 8), ch * 4, 3, 1, 1)
|
||||
self.dec2 = layers.Decoder(ch * (2 + 4), ch * 2, 3, 1, 1)
|
||||
self.dec1 = layers.Decoder(ch * (1 + 2), ch, 3, 1, 1)
|
||||
|
||||
def __call__(self, x):
|
||||
h, e1 = self.enc1(x)
|
||||
h, e2 = self.enc2(h)
|
||||
h, e3 = self.enc3(h)
|
||||
h, e4 = self.enc4(h)
|
||||
|
||||
h = self.aspp(h)
|
||||
|
||||
h = self.dec4(h, e4)
|
||||
h = self.dec3(h, e3)
|
||||
h = self.dec2(h, e2)
|
||||
h = self.dec1(h, e1)
|
||||
|
||||
return h
|
||||
|
||||
|
||||
class CascadedASPPNet(nn.Module):
|
||||
def __init__(self, n_fft):
|
||||
super(CascadedASPPNet, self).__init__()
|
||||
self.stg1_low_band_net = BaseASPPNet(2, 32)
|
||||
self.stg1_high_band_net = BaseASPPNet(2, 32)
|
||||
|
||||
self.stg2_bridge = layers.Conv2DBNActiv(34, 16, 1, 1, 0)
|
||||
self.stg2_full_band_net = BaseASPPNet(16, 32)
|
||||
|
||||
self.stg3_bridge = layers.Conv2DBNActiv(66, 32, 1, 1, 0)
|
||||
self.stg3_full_band_net = BaseASPPNet(32, 64)
|
||||
|
||||
self.out = nn.Conv2d(64, 2, 1, bias=False)
|
||||
self.aux1_out = nn.Conv2d(32, 2, 1, bias=False)
|
||||
self.aux2_out = nn.Conv2d(32, 2, 1, bias=False)
|
||||
|
||||
self.max_bin = n_fft // 2
|
||||
self.output_bin = n_fft // 2 + 1
|
||||
|
||||
self.offset = 128
|
||||
|
||||
def forward(self, x, aggressiveness=None):
|
||||
mix = x.detach()
|
||||
x = x.clone()
|
||||
|
||||
x = x[:, :, : self.max_bin]
|
||||
|
||||
bandw = x.size()[2] // 2
|
||||
aux1 = torch.cat(
|
||||
[
|
||||
self.stg1_low_band_net(x[:, :, :bandw]),
|
||||
self.stg1_high_band_net(x[:, :, bandw:]),
|
||||
],
|
||||
dim=2,
|
||||
)
|
||||
|
||||
h = torch.cat([x, aux1], dim=1)
|
||||
aux2 = self.stg2_full_band_net(self.stg2_bridge(h))
|
||||
|
||||
h = torch.cat([x, aux1, aux2], dim=1)
|
||||
h = self.stg3_full_band_net(self.stg3_bridge(h))
|
||||
|
||||
mask = torch.sigmoid(self.out(h))
|
||||
mask = F.pad(
|
||||
input=mask,
|
||||
pad=(0, 0, 0, self.output_bin - mask.size()[2]),
|
||||
mode="replicate",
|
||||
)
|
||||
|
||||
if self.training:
|
||||
aux1 = torch.sigmoid(self.aux1_out(aux1))
|
||||
aux1 = F.pad(
|
||||
input=aux1,
|
||||
pad=(0, 0, 0, self.output_bin - aux1.size()[2]),
|
||||
mode="replicate",
|
||||
)
|
||||
aux2 = torch.sigmoid(self.aux2_out(aux2))
|
||||
aux2 = F.pad(
|
||||
input=aux2,
|
||||
pad=(0, 0, 0, self.output_bin - aux2.size()[2]),
|
||||
mode="replicate",
|
||||
)
|
||||
return mask * mix, aux1 * mix, aux2 * mix
|
||||
else:
|
||||
if aggressiveness:
|
||||
mask[:, :, : aggressiveness["split_bin"]] = torch.pow(
|
||||
mask[:, :, : aggressiveness["split_bin"]],
|
||||
1 + aggressiveness["value"] / 3,
|
||||
)
|
||||
mask[:, :, aggressiveness["split_bin"] :] = torch.pow(
|
||||
mask[:, :, aggressiveness["split_bin"] :],
|
||||
1 + aggressiveness["value"],
|
||||
)
|
||||
|
||||
return mask * mix
|
||||
|
||||
def predict(self, x_mag, aggressiveness=None):
|
||||
h = self.forward(x_mag, aggressiveness)
|
||||
|
||||
if self.offset > 0:
|
||||
h = h[:, :, :, self.offset : -self.offset]
|
||||
assert h.size()[3] > 0
|
||||
|
||||
return h
|
||||
@ -1,122 +0,0 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from . import layers_33966KB as layers
|
||||
|
||||
|
||||
class BaseASPPNet(nn.Module):
|
||||
def __init__(self, nin, ch, dilations=(4, 8, 16, 32)):
|
||||
super(BaseASPPNet, self).__init__()
|
||||
self.enc1 = layers.Encoder(nin, ch, 3, 2, 1)
|
||||
self.enc2 = layers.Encoder(ch, ch * 2, 3, 2, 1)
|
||||
self.enc3 = layers.Encoder(ch * 2, ch * 4, 3, 2, 1)
|
||||
self.enc4 = layers.Encoder(ch * 4, ch * 8, 3, 2, 1)
|
||||
|
||||
self.aspp = layers.ASPPModule(ch * 8, ch * 16, dilations)
|
||||
|
||||
self.dec4 = layers.Decoder(ch * (8 + 16), ch * 8, 3, 1, 1)
|
||||
self.dec3 = layers.Decoder(ch * (4 + 8), ch * 4, 3, 1, 1)
|
||||
self.dec2 = layers.Decoder(ch * (2 + 4), ch * 2, 3, 1, 1)
|
||||
self.dec1 = layers.Decoder(ch * (1 + 2), ch, 3, 1, 1)
|
||||
|
||||
def __call__(self, x):
|
||||
h, e1 = self.enc1(x)
|
||||
h, e2 = self.enc2(h)
|
||||
h, e3 = self.enc3(h)
|
||||
h, e4 = self.enc4(h)
|
||||
|
||||
h = self.aspp(h)
|
||||
|
||||
h = self.dec4(h, e4)
|
||||
h = self.dec3(h, e3)
|
||||
h = self.dec2(h, e2)
|
||||
h = self.dec1(h, e1)
|
||||
|
||||
return h
|
||||
|
||||
|
||||
class CascadedASPPNet(nn.Module):
|
||||
def __init__(self, n_fft):
|
||||
super(CascadedASPPNet, self).__init__()
|
||||
self.stg1_low_band_net = BaseASPPNet(2, 16)
|
||||
self.stg1_high_band_net = BaseASPPNet(2, 16)
|
||||
|
||||
self.stg2_bridge = layers.Conv2DBNActiv(18, 8, 1, 1, 0)
|
||||
self.stg2_full_band_net = BaseASPPNet(8, 16)
|
||||
|
||||
self.stg3_bridge = layers.Conv2DBNActiv(34, 16, 1, 1, 0)
|
||||
self.stg3_full_band_net = BaseASPPNet(16, 32)
|
||||
|
||||
self.out = nn.Conv2d(32, 2, 1, bias=False)
|
||||
self.aux1_out = nn.Conv2d(16, 2, 1, bias=False)
|
||||
self.aux2_out = nn.Conv2d(16, 2, 1, bias=False)
|
||||
|
||||
self.max_bin = n_fft // 2
|
||||
self.output_bin = n_fft // 2 + 1
|
||||
|
||||
self.offset = 128
|
||||
|
||||
def forward(self, x, aggressiveness=None):
|
||||
mix = x.detach()
|
||||
x = x.clone()
|
||||
|
||||
x = x[:, :, : self.max_bin]
|
||||
|
||||
bandw = x.size()[2] // 2
|
||||
aux1 = torch.cat(
|
||||
[
|
||||
self.stg1_low_band_net(x[:, :, :bandw]),
|
||||
self.stg1_high_band_net(x[:, :, bandw:]),
|
||||
],
|
||||
dim=2,
|
||||
)
|
||||
|
||||
h = torch.cat([x, aux1], dim=1)
|
||||
aux2 = self.stg2_full_band_net(self.stg2_bridge(h))
|
||||
|
||||
h = torch.cat([x, aux1, aux2], dim=1)
|
||||
h = self.stg3_full_band_net(self.stg3_bridge(h))
|
||||
|
||||
mask = torch.sigmoid(self.out(h))
|
||||
mask = F.pad(
|
||||
input=mask,
|
||||
pad=(0, 0, 0, self.output_bin - mask.size()[2]),
|
||||
mode="replicate",
|
||||
)
|
||||
|
||||
if self.training:
|
||||
aux1 = torch.sigmoid(self.aux1_out(aux1))
|
||||
aux1 = F.pad(
|
||||
input=aux1,
|
||||
pad=(0, 0, 0, self.output_bin - aux1.size()[2]),
|
||||
mode="replicate",
|
||||
)
|
||||
aux2 = torch.sigmoid(self.aux2_out(aux2))
|
||||
aux2 = F.pad(
|
||||
input=aux2,
|
||||
pad=(0, 0, 0, self.output_bin - aux2.size()[2]),
|
||||
mode="replicate",
|
||||
)
|
||||
return mask * mix, aux1 * mix, aux2 * mix
|
||||
else:
|
||||
if aggressiveness:
|
||||
mask[:, :, : aggressiveness["split_bin"]] = torch.pow(
|
||||
mask[:, :, : aggressiveness["split_bin"]],
|
||||
1 + aggressiveness["value"] / 3,
|
||||
)
|
||||
mask[:, :, aggressiveness["split_bin"] :] = torch.pow(
|
||||
mask[:, :, aggressiveness["split_bin"] :],
|
||||
1 + aggressiveness["value"],
|
||||
)
|
||||
|
||||
return mask * mix
|
||||
|
||||
def predict(self, x_mag, aggressiveness=None):
|
||||
h = self.forward(x_mag, aggressiveness)
|
||||
|
||||
if self.offset > 0:
|
||||
h = h[:, :, :, self.offset : -self.offset]
|
||||
assert h.size()[3] > 0
|
||||
|
||||
return h
|
||||
@ -1,122 +0,0 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from . import layers_537238KB as layers
|
||||
|
||||
|
||||
class BaseASPPNet(nn.Module):
|
||||
def __init__(self, nin, ch, dilations=(4, 8, 16)):
|
||||
super(BaseASPPNet, self).__init__()
|
||||
self.enc1 = layers.Encoder(nin, ch, 3, 2, 1)
|
||||
self.enc2 = layers.Encoder(ch, ch * 2, 3, 2, 1)
|
||||
self.enc3 = layers.Encoder(ch * 2, ch * 4, 3, 2, 1)
|
||||
self.enc4 = layers.Encoder(ch * 4, ch * 8, 3, 2, 1)
|
||||
|
||||
self.aspp = layers.ASPPModule(ch * 8, ch * 16, dilations)
|
||||
|
||||
self.dec4 = layers.Decoder(ch * (8 + 16), ch * 8, 3, 1, 1)
|
||||
self.dec3 = layers.Decoder(ch * (4 + 8), ch * 4, 3, 1, 1)
|
||||
self.dec2 = layers.Decoder(ch * (2 + 4), ch * 2, 3, 1, 1)
|
||||
self.dec1 = layers.Decoder(ch * (1 + 2), ch, 3, 1, 1)
|
||||
|
||||
def __call__(self, x):
|
||||
h, e1 = self.enc1(x)
|
||||
h, e2 = self.enc2(h)
|
||||
h, e3 = self.enc3(h)
|
||||
h, e4 = self.enc4(h)
|
||||
|
||||
h = self.aspp(h)
|
||||
|
||||
h = self.dec4(h, e4)
|
||||
h = self.dec3(h, e3)
|
||||
h = self.dec2(h, e2)
|
||||
h = self.dec1(h, e1)
|
||||
|
||||
return h
|
||||
|
||||
|
||||
class CascadedASPPNet(nn.Module):
|
||||
def __init__(self, n_fft):
|
||||
super(CascadedASPPNet, self).__init__()
|
||||
self.stg1_low_band_net = BaseASPPNet(2, 64)
|
||||
self.stg1_high_band_net = BaseASPPNet(2, 64)
|
||||
|
||||
self.stg2_bridge = layers.Conv2DBNActiv(66, 32, 1, 1, 0)
|
||||
self.stg2_full_band_net = BaseASPPNet(32, 64)
|
||||
|
||||
self.stg3_bridge = layers.Conv2DBNActiv(130, 64, 1, 1, 0)
|
||||
self.stg3_full_band_net = BaseASPPNet(64, 128)
|
||||
|
||||
self.out = nn.Conv2d(128, 2, 1, bias=False)
|
||||
self.aux1_out = nn.Conv2d(64, 2, 1, bias=False)
|
||||
self.aux2_out = nn.Conv2d(64, 2, 1, bias=False)
|
||||
|
||||
self.max_bin = n_fft // 2
|
||||
self.output_bin = n_fft // 2 + 1
|
||||
|
||||
self.offset = 128
|
||||
|
||||
def forward(self, x, aggressiveness=None):
|
||||
mix = x.detach()
|
||||
x = x.clone()
|
||||
|
||||
x = x[:, :, : self.max_bin]
|
||||
|
||||
bandw = x.size()[2] // 2
|
||||
aux1 = torch.cat(
|
||||
[
|
||||
self.stg1_low_band_net(x[:, :, :bandw]),
|
||||
self.stg1_high_band_net(x[:, :, bandw:]),
|
||||
],
|
||||
dim=2,
|
||||
)
|
||||
|
||||
h = torch.cat([x, aux1], dim=1)
|
||||
aux2 = self.stg2_full_band_net(self.stg2_bridge(h))
|
||||
|
||||
h = torch.cat([x, aux1, aux2], dim=1)
|
||||
h = self.stg3_full_band_net(self.stg3_bridge(h))
|
||||
|
||||
mask = torch.sigmoid(self.out(h))
|
||||
mask = F.pad(
|
||||
input=mask,
|
||||
pad=(0, 0, 0, self.output_bin - mask.size()[2]),
|
||||
mode="replicate",
|
||||
)
|
||||
|
||||
if self.training:
|
||||
aux1 = torch.sigmoid(self.aux1_out(aux1))
|
||||
aux1 = F.pad(
|
||||
input=aux1,
|
||||
pad=(0, 0, 0, self.output_bin - aux1.size()[2]),
|
||||
mode="replicate",
|
||||
)
|
||||
aux2 = torch.sigmoid(self.aux2_out(aux2))
|
||||
aux2 = F.pad(
|
||||
input=aux2,
|
||||
pad=(0, 0, 0, self.output_bin - aux2.size()[2]),
|
||||
mode="replicate",
|
||||
)
|
||||
return mask * mix, aux1 * mix, aux2 * mix
|
||||
else:
|
||||
if aggressiveness:
|
||||
mask[:, :, : aggressiveness["split_bin"]] = torch.pow(
|
||||
mask[:, :, : aggressiveness["split_bin"]],
|
||||
1 + aggressiveness["value"] / 3,
|
||||
)
|
||||
mask[:, :, aggressiveness["split_bin"] :] = torch.pow(
|
||||
mask[:, :, aggressiveness["split_bin"] :],
|
||||
1 + aggressiveness["value"],
|
||||
)
|
||||
|
||||
return mask * mix
|
||||
|
||||
def predict(self, x_mag, aggressiveness=None):
|
||||
h = self.forward(x_mag, aggressiveness)
|
||||
|
||||
if self.offset > 0:
|
||||
h = h[:, :, :, self.offset : -self.offset]
|
||||
assert h.size()[3] > 0
|
||||
|
||||
return h
|
||||
@ -1,122 +0,0 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
|
||||
from . import layers_537238KB as layers
|
||||
|
||||
|
||||
class BaseASPPNet(nn.Module):
|
||||
def __init__(self, nin, ch, dilations=(4, 8, 16)):
|
||||
super(BaseASPPNet, self).__init__()
|
||||
self.enc1 = layers.Encoder(nin, ch, 3, 2, 1)
|
||||
self.enc2 = layers.Encoder(ch, ch * 2, 3, 2, 1)
|
||||
self.enc3 = layers.Encoder(ch * 2, ch * 4, 3, 2, 1)
|
||||
self.enc4 = layers.Encoder(ch * 4, ch * 8, 3, 2, 1)
|
||||
|
||||
self.aspp = layers.ASPPModule(ch * 8, ch * 16, dilations)
|
||||
|
||||
self.dec4 = layers.Decoder(ch * (8 + 16), ch * 8, 3, 1, 1)
|
||||
self.dec3 = layers.Decoder(ch * (4 + 8), ch * 4, 3, 1, 1)
|
||||
self.dec2 = layers.Decoder(ch * (2 + 4), ch * 2, 3, 1, 1)
|
||||
self.dec1 = layers.Decoder(ch * (1 + 2), ch, 3, 1, 1)
|
||||
|
||||
def __call__(self, x):
|
||||
h, e1 = self.enc1(x)
|
||||
h, e2 = self.enc2(h)
|
||||
h, e3 = self.enc3(h)
|
||||
h, e4 = self.enc4(h)
|
||||
|
||||
h = self.aspp(h)
|
||||
|
||||
h = self.dec4(h, e4)
|
||||
h = self.dec3(h, e3)
|
||||
h = self.dec2(h, e2)
|
||||
h = self.dec1(h, e1)
|
||||
|
||||
return h
|
||||
|
||||
|
||||
class CascadedASPPNet(nn.Module):
|
||||
def __init__(self, n_fft):
|
||||
super(CascadedASPPNet, self).__init__()
|
||||
self.stg1_low_band_net = BaseASPPNet(2, 64)
|
||||
self.stg1_high_band_net = BaseASPPNet(2, 64)
|
||||
|
||||
self.stg2_bridge = layers.Conv2DBNActiv(66, 32, 1, 1, 0)
|
||||
self.stg2_full_band_net = BaseASPPNet(32, 64)
|
||||
|
||||
self.stg3_bridge = layers.Conv2DBNActiv(130, 64, 1, 1, 0)
|
||||
self.stg3_full_band_net = BaseASPPNet(64, 128)
|
||||
|
||||
self.out = nn.Conv2d(128, 2, 1, bias=False)
|
||||
self.aux1_out = nn.Conv2d(64, 2, 1, bias=False)
|
||||
self.aux2_out = nn.Conv2d(64, 2, 1, bias=False)
|
||||
|
||||
self.max_bin = n_fft // 2
|
||||
self.output_bin = n_fft // 2 + 1
|
||||
|
||||
self.offset = 128
|
||||
|
||||
def forward(self, x, aggressiveness=None):
|
||||
mix = x.detach()
|
||||
x = x.clone()
|
||||
|
||||
x = x[:, :, : self.max_bin]
|
||||
|
||||
bandw = x.size()[2] // 2
|
||||
aux1 = torch.cat(
|
||||
[
|
||||
self.stg1_low_band_net(x[:, :, :bandw]),
|
||||
self.stg1_high_band_net(x[:, :, bandw:]),
|
||||
],
|
||||
dim=2,
|
||||
)
|
||||
|
||||
h = torch.cat([x, aux1], dim=1)
|
||||
aux2 = self.stg2_full_band_net(self.stg2_bridge(h))
|
||||
|
||||
h = torch.cat([x, aux1, aux2], dim=1)
|
||||
h = self.stg3_full_band_net(self.stg3_bridge(h))
|
||||
|
||||
mask = torch.sigmoid(self.out(h))
|
||||
mask = F.pad(
|
||||
input=mask,
|
||||
pad=(0, 0, 0, self.output_bin - mask.size()[2]),
|
||||
mode="replicate",
|
||||
)
|
||||
|
||||
if self.training:
|
||||
aux1 = torch.sigmoid(self.aux1_out(aux1))
|
||||
aux1 = F.pad(
|
||||
input=aux1,
|
||||
pad=(0, 0, 0, self.output_bin - aux1.size()[2]),
|
||||
mode="replicate",
|
||||
)
|
||||
aux2 = torch.sigmoid(self.aux2_out(aux2))
|
||||
aux2 = F.pad(
|
||||
input=aux2,
|
||||
pad=(0, 0, 0, self.output_bin - aux2.size()[2]),
|
||||
mode="replicate",
|
||||
)
|
||||
return mask * mix, aux1 * mix, aux2 * mix
|
||||
else:
|
||||
if aggressiveness:
|
||||
mask[:, :, : aggressiveness["split_bin"]] = torch.pow(
|
||||
mask[:, :, : aggressiveness["split_bin"]],
|
||||
1 + aggressiveness["value"] / 3,
|
||||
)
|
||||
mask[:, :, aggressiveness["split_bin"] :] = torch.pow(
|
||||
mask[:, :, aggressiveness["split_bin"] :],
|
||||
1 + aggressiveness["value"],
|
||||
)
|
||||
|
||||
return mask * mix
|
||||
|
||||
def predict(self, x_mag, aggressiveness=None):
|
||||
h = self.forward(x_mag, aggressiveness)
|
||||
|
||||
if self.offset > 0:
|
||||
h = h[:, :, :, self.offset : -self.offset]
|
||||
assert h.size()[3] > 0
|
||||
|
||||
return h
|
||||
@ -1,12 +1,7 @@
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
|
||||
import librosa
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
def crop_center(h1, h2):
|
||||
@ -27,28 +22,6 @@ def crop_center(h1, h2):
|
||||
return h1
|
||||
|
||||
|
||||
def wave_to_spectrogram(wave, hop_length, n_fft, mid_side=False, mid_side_b2=False, reverse=False):
|
||||
if reverse:
|
||||
wave_left = np.flip(np.asfortranarray(wave[0]))
|
||||
wave_right = np.flip(np.asfortranarray(wave[1]))
|
||||
elif mid_side:
|
||||
wave_left = np.asfortranarray(np.add(wave[0], wave[1]) / 2)
|
||||
wave_right = np.asfortranarray(np.subtract(wave[0], wave[1]))
|
||||
elif mid_side_b2:
|
||||
wave_left = np.asfortranarray(np.add(wave[1], wave[0] * 0.5))
|
||||
wave_right = np.asfortranarray(np.subtract(wave[0], wave[1] * 0.5))
|
||||
else:
|
||||
wave_left = np.asfortranarray(wave[0])
|
||||
wave_right = np.asfortranarray(wave[1])
|
||||
|
||||
spec_left = librosa.stft(wave_left, n_fft=n_fft, hop_length=hop_length)
|
||||
spec_right = librosa.stft(wave_right, n_fft=n_fft, hop_length=hop_length)
|
||||
|
||||
spec = np.asfortranarray([spec_left, spec_right])
|
||||
|
||||
return spec
|
||||
|
||||
|
||||
def wave_to_spectrogram_mt(wave, hop_length, n_fft, mid_side=False, mid_side_b2=False, reverse=False):
|
||||
import threading
|
||||
|
||||
@ -112,41 +85,6 @@ def combine_spectrograms(specs, mp):
|
||||
return np.asfortranarray(spec_c)
|
||||
|
||||
|
||||
def spectrogram_to_image(spec, mode="magnitude"):
|
||||
if mode == "magnitude":
|
||||
if np.iscomplexobj(spec):
|
||||
y = np.abs(spec)
|
||||
else:
|
||||
y = spec
|
||||
y = np.log10(y**2 + 1e-8)
|
||||
elif mode == "phase":
|
||||
if np.iscomplexobj(spec):
|
||||
y = np.angle(spec)
|
||||
else:
|
||||
y = spec
|
||||
|
||||
y -= y.min()
|
||||
y *= 255 / y.max()
|
||||
img = np.uint8(y)
|
||||
|
||||
if y.ndim == 3:
|
||||
img = img.transpose(1, 2, 0)
|
||||
img = np.concatenate([np.max(img, axis=2, keepdims=True), img], axis=2)
|
||||
|
||||
return img
|
||||
|
||||
|
||||
def reduce_vocal_aggressively(X, y, softmask):
|
||||
v = X - y
|
||||
y_mag_tmp = np.abs(y)
|
||||
v_mag_tmp = np.abs(v)
|
||||
|
||||
v_mask = v_mag_tmp > y_mag_tmp
|
||||
y_mag = np.clip(y_mag_tmp - v_mag_tmp * v_mask * softmask, 0, np.inf)
|
||||
|
||||
return y_mag * np.exp(1.0j * np.angle(y))
|
||||
|
||||
|
||||
def mask_silence(mag, ref, thres=0.2, min_range=64, fade_size=32):
|
||||
if min_range < fade_size * 2:
|
||||
raise ValueError("min_range must be >= fade_area * 2")
|
||||
@ -183,95 +121,6 @@ def mask_silence(mag, ref, thres=0.2, min_range=64, fade_size=32):
|
||||
return mag
|
||||
|
||||
|
||||
def align_wave_head_and_tail(a, b):
|
||||
l = min([a[0].size, b[0].size])
|
||||
|
||||
return a[:l, :l], b[:l, :l]
|
||||
|
||||
|
||||
def cache_or_load(mix_path, inst_path, mp):
|
||||
mix_basename = os.path.splitext(os.path.basename(mix_path))[0]
|
||||
inst_basename = os.path.splitext(os.path.basename(inst_path))[0]
|
||||
|
||||
cache_dir = "mph{}".format(hashlib.sha1(json.dumps(mp.param, sort_keys=True).encode("utf-8")).hexdigest())
|
||||
mix_cache_dir = os.path.join("cache", cache_dir)
|
||||
inst_cache_dir = os.path.join("cache", cache_dir)
|
||||
|
||||
os.makedirs(mix_cache_dir, exist_ok=True)
|
||||
os.makedirs(inst_cache_dir, exist_ok=True)
|
||||
|
||||
mix_cache_path = os.path.join(mix_cache_dir, mix_basename + ".npy")
|
||||
inst_cache_path = os.path.join(inst_cache_dir, inst_basename + ".npy")
|
||||
|
||||
if os.path.exists(mix_cache_path) and os.path.exists(inst_cache_path):
|
||||
X_spec_m = np.load(mix_cache_path)
|
||||
y_spec_m = np.load(inst_cache_path)
|
||||
else:
|
||||
X_wave, y_wave, X_spec_s, y_spec_s = {}, {}, {}, {}
|
||||
|
||||
for d in range(len(mp.param["band"]), 0, -1):
|
||||
bp = mp.param["band"][d]
|
||||
|
||||
if d == len(mp.param["band"]): # high-end band
|
||||
X_wave[d], _ = librosa.load(
|
||||
mix_path, sr=bp["sr"], mono=False, dtype=np.float32, res_type=bp["res_type"]
|
||||
)
|
||||
y_wave[d], _ = librosa.load(
|
||||
inst_path,
|
||||
sr=bp["sr"],
|
||||
mono=False,
|
||||
dtype=np.float32,
|
||||
res_type=bp["res_type"],
|
||||
)
|
||||
else: # lower bands
|
||||
X_wave[d] = librosa.resample(
|
||||
X_wave[d + 1],
|
||||
orig_sr=mp.param["band"][d + 1]["sr"],
|
||||
target_sr=bp["sr"],
|
||||
res_type=bp["res_type"],
|
||||
)
|
||||
y_wave[d] = librosa.resample(
|
||||
y_wave[d + 1],
|
||||
orig_sr=mp.param["band"][d + 1]["sr"],
|
||||
target_sr=bp["sr"],
|
||||
res_type=bp["res_type"],
|
||||
)
|
||||
|
||||
X_wave[d], y_wave[d] = align_wave_head_and_tail(X_wave[d], y_wave[d])
|
||||
|
||||
X_spec_s[d] = wave_to_spectrogram(
|
||||
X_wave[d],
|
||||
bp["hl"],
|
||||
bp["n_fft"],
|
||||
mp.param["mid_side"],
|
||||
mp.param["mid_side_b2"],
|
||||
mp.param["reverse"],
|
||||
)
|
||||
y_spec_s[d] = wave_to_spectrogram(
|
||||
y_wave[d],
|
||||
bp["hl"],
|
||||
bp["n_fft"],
|
||||
mp.param["mid_side"],
|
||||
mp.param["mid_side_b2"],
|
||||
mp.param["reverse"],
|
||||
)
|
||||
|
||||
del X_wave, y_wave
|
||||
|
||||
X_spec_m = combine_spectrograms(X_spec_s, mp)
|
||||
y_spec_m = combine_spectrograms(y_spec_s, mp)
|
||||
|
||||
if X_spec_m.shape != y_spec_m.shape:
|
||||
raise ValueError("The combined spectrograms are different: " + mix_path)
|
||||
|
||||
_, ext = os.path.splitext(mix_path)
|
||||
|
||||
np.save(mix_cache_path, X_spec_m)
|
||||
np.save(inst_cache_path, y_spec_m)
|
||||
|
||||
return X_spec_m, y_spec_m
|
||||
|
||||
|
||||
def spectrogram_to_wave(spec, hop_length, mid_side, mid_side_b2, reverse):
|
||||
spec_left = np.asfortranarray(spec[0])
|
||||
spec_right = np.asfortranarray(spec[1])
|
||||
@ -294,36 +143,6 @@ def spectrogram_to_wave(spec, hop_length, mid_side, mid_side_b2, reverse):
|
||||
return np.asfortranarray([wave_left, wave_right])
|
||||
|
||||
|
||||
def spectrogram_to_wave_mt(spec, hop_length, mid_side, reverse, mid_side_b2):
|
||||
import threading
|
||||
|
||||
spec_left = np.asfortranarray(spec[0])
|
||||
spec_right = np.asfortranarray(spec[1])
|
||||
|
||||
def run_thread(**kwargs):
|
||||
global wave_left
|
||||
wave_left = librosa.istft(**kwargs)
|
||||
|
||||
thread = threading.Thread(target=run_thread, kwargs={"stft_matrix": spec_left, "hop_length": hop_length})
|
||||
thread.start()
|
||||
wave_right = librosa.istft(spec_right, hop_length=hop_length)
|
||||
thread.join()
|
||||
|
||||
if reverse:
|
||||
return np.asfortranarray([np.flip(wave_left), np.flip(wave_right)])
|
||||
elif mid_side:
|
||||
return np.asfortranarray([np.add(wave_left, wave_right / 2), np.subtract(wave_left, wave_right / 2)])
|
||||
elif mid_side_b2:
|
||||
return np.asfortranarray(
|
||||
[
|
||||
np.add(wave_right / 1.25, 0.4 * wave_left),
|
||||
np.subtract(wave_left / 1.25, 0.4 * wave_right),
|
||||
]
|
||||
)
|
||||
else:
|
||||
return np.asfortranarray([wave_left, wave_right])
|
||||
|
||||
|
||||
def cmb_spectrogram_to_wave(spec_m, mp, extra_bins_h=None, extra_bins=None):
|
||||
wave_band = {}
|
||||
bands_n = len(mp.param["band"])
|
||||
@ -448,190 +267,3 @@ def mirroring(a, spec_m, input_high_end, mp):
|
||||
mi = np.multiply(mirror, input_high_end * 1.7)
|
||||
|
||||
return np.where(np.abs(input_high_end) <= np.abs(mi), input_high_end, mi)
|
||||
|
||||
|
||||
def ensembling(a, specs):
|
||||
for i in range(1, len(specs)):
|
||||
if i == 1:
|
||||
spec = specs[0]
|
||||
|
||||
ln = min([spec.shape[2], specs[i].shape[2]])
|
||||
spec = spec[:, :, :ln]
|
||||
specs[i] = specs[i][:, :, :ln]
|
||||
|
||||
if "min_mag" == a:
|
||||
spec = np.where(np.abs(specs[i]) <= np.abs(spec), specs[i], spec)
|
||||
if "max_mag" == a:
|
||||
spec = np.where(np.abs(specs[i]) >= np.abs(spec), specs[i], spec)
|
||||
|
||||
return spec
|
||||
|
||||
|
||||
def stft(wave, nfft, hl):
|
||||
wave_left = np.asfortranarray(wave[0])
|
||||
wave_right = np.asfortranarray(wave[1])
|
||||
spec_left = librosa.stft(wave_left, n_fft=nfft, hop_length=hl)
|
||||
spec_right = librosa.stft(wave_right, n_fft=nfft, hop_length=hl)
|
||||
spec = np.asfortranarray([spec_left, spec_right])
|
||||
|
||||
return spec
|
||||
|
||||
|
||||
def istft(spec, hl):
|
||||
spec_left = np.asfortranarray(spec[0])
|
||||
spec_right = np.asfortranarray(spec[1])
|
||||
|
||||
wave_left = librosa.istft(spec_left, hop_length=hl)
|
||||
wave_right = librosa.istft(spec_right, hop_length=hl)
|
||||
wave = np.asfortranarray([wave_left, wave_right])
|
||||
|
||||
return wave
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
import time
|
||||
|
||||
import cv2
|
||||
from model_param_init import ModelParameters
|
||||
|
||||
p = argparse.ArgumentParser()
|
||||
p.add_argument(
|
||||
"--algorithm",
|
||||
"-a",
|
||||
type=str,
|
||||
choices=["invert", "invert_p", "min_mag", "max_mag", "deep", "align"],
|
||||
default="min_mag",
|
||||
)
|
||||
p.add_argument(
|
||||
"--model_params",
|
||||
"-m",
|
||||
type=str,
|
||||
default=os.path.join("modelparams", "1band_sr44100_hl512.json"),
|
||||
)
|
||||
p.add_argument("--output_name", "-o", type=str, default="output")
|
||||
p.add_argument("--vocals_only", "-v", action="store_true")
|
||||
p.add_argument("input", nargs="+")
|
||||
args = p.parse_args()
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
if args.algorithm.startswith("invert") and len(args.input) != 2:
|
||||
raise ValueError("There should be two input files.")
|
||||
|
||||
if not args.algorithm.startswith("invert") and len(args.input) < 2:
|
||||
raise ValueError("There must be at least two input files.")
|
||||
|
||||
wave, specs = {}, {}
|
||||
mp = ModelParameters(args.model_params)
|
||||
|
||||
for i in range(len(args.input)):
|
||||
spec = {}
|
||||
|
||||
for d in range(len(mp.param["band"]), 0, -1):
|
||||
bp = mp.param["band"][d]
|
||||
|
||||
if d == len(mp.param["band"]): # high-end band
|
||||
wave[d], _ = librosa.load(
|
||||
args.input[i],
|
||||
sr=bp["sr"],
|
||||
mono=False,
|
||||
dtype=np.float32,
|
||||
res_type=bp["res_type"],
|
||||
)
|
||||
|
||||
if len(wave[d].shape) == 1: # mono to stereo
|
||||
wave[d] = np.array([wave[d], wave[d]])
|
||||
else: # lower bands
|
||||
wave[d] = librosa.resample(
|
||||
wave[d + 1],
|
||||
orig_sr=mp.param["band"][d + 1]["sr"],
|
||||
target_sr=bp["sr"],
|
||||
res_type=bp["res_type"],
|
||||
)
|
||||
|
||||
spec[d] = wave_to_spectrogram(
|
||||
wave[d],
|
||||
bp["hl"],
|
||||
bp["n_fft"],
|
||||
mp.param["mid_side"],
|
||||
mp.param["mid_side_b2"],
|
||||
mp.param["reverse"],
|
||||
)
|
||||
|
||||
specs[i] = combine_spectrograms(spec, mp)
|
||||
|
||||
del wave
|
||||
|
||||
if args.algorithm == "deep":
|
||||
d_spec = np.where(np.abs(specs[0]) <= np.abs(spec[1]), specs[0], spec[1])
|
||||
v_spec = d_spec - specs[1]
|
||||
sf.write(
|
||||
os.path.join("{}.wav".format(args.output_name)),
|
||||
cmb_spectrogram_to_wave(v_spec, mp),
|
||||
mp.param["sr"],
|
||||
)
|
||||
|
||||
if args.algorithm.startswith("invert"):
|
||||
ln = min([specs[0].shape[2], specs[1].shape[2]])
|
||||
specs[0] = specs[0][:, :, :ln]
|
||||
specs[1] = specs[1][:, :, :ln]
|
||||
|
||||
if "invert_p" == args.algorithm:
|
||||
X_mag = np.abs(specs[0])
|
||||
y_mag = np.abs(specs[1])
|
||||
max_mag = np.where(X_mag >= y_mag, X_mag, y_mag)
|
||||
v_spec = specs[1] - max_mag * np.exp(1.0j * np.angle(specs[0]))
|
||||
else:
|
||||
specs[1] = reduce_vocal_aggressively(specs[0], specs[1], 0.2)
|
||||
v_spec = specs[0] - specs[1]
|
||||
|
||||
if not args.vocals_only:
|
||||
X_mag = np.abs(specs[0])
|
||||
y_mag = np.abs(specs[1])
|
||||
v_mag = np.abs(v_spec)
|
||||
|
||||
X_image = spectrogram_to_image(X_mag)
|
||||
y_image = spectrogram_to_image(y_mag)
|
||||
v_image = spectrogram_to_image(v_mag)
|
||||
|
||||
cv2.imwrite("{}_X.png".format(args.output_name), X_image)
|
||||
cv2.imwrite("{}_y.png".format(args.output_name), y_image)
|
||||
cv2.imwrite("{}_v.png".format(args.output_name), v_image)
|
||||
|
||||
sf.write(
|
||||
"{}_X.wav".format(args.output_name),
|
||||
cmb_spectrogram_to_wave(specs[0], mp),
|
||||
mp.param["sr"],
|
||||
)
|
||||
sf.write(
|
||||
"{}_y.wav".format(args.output_name),
|
||||
cmb_spectrogram_to_wave(specs[1], mp),
|
||||
mp.param["sr"],
|
||||
)
|
||||
|
||||
sf.write(
|
||||
"{}_v.wav".format(args.output_name),
|
||||
cmb_spectrogram_to_wave(v_spec, mp),
|
||||
mp.param["sr"],
|
||||
)
|
||||
else:
|
||||
if not args.algorithm == "deep":
|
||||
sf.write(
|
||||
os.path.join("ensembled", "{}.wav".format(args.output_name)),
|
||||
cmb_spectrogram_to_wave(ensembling(args.algorithm, specs), mp),
|
||||
mp.param["sr"],
|
||||
)
|
||||
|
||||
if args.algorithm == "align":
|
||||
trackalignment = [
|
||||
{
|
||||
"file1": '"{}"'.format(args.input[0]),
|
||||
"file2": '"{}"'.format(args.input[1]),
|
||||
}
|
||||
]
|
||||
|
||||
for i, e in tqdm(enumerate(trackalignment), desc="Performing Alignment..."):
|
||||
os.system(f"python lib/align_tracks.py {e['file1']} {e['file2']}")
|
||||
|
||||
# print('Total time: {0:.{1}f}s'.format(time.time() - start_time, 1))
|
||||
|
||||
@ -1,263 +0,0 @@
|
||||
{
|
||||
"equivalent" : [
|
||||
{
|
||||
"model_hash_name" : [
|
||||
{
|
||||
"hash_name": "47939caf0cfe52a0e81442b85b971dfd",
|
||||
"model_params": "lib/lib_v5/modelparams/4band_44100.json",
|
||||
"param_name": "4band_44100"
|
||||
},
|
||||
{
|
||||
"hash_name": "4e4ecb9764c50a8c414fee6e10395bbe",
|
||||
"model_params": "lib/lib_v5/modelparams/4band_v2.json",
|
||||
"param_name": "4band_v2"
|
||||
},
|
||||
{
|
||||
"hash_name": "ca106edd563e034bde0bdec4bb7a4b36",
|
||||
"model_params": "lib/lib_v5/modelparams/4band_v2.json",
|
||||
"param_name": "4band_v2"
|
||||
},
|
||||
{
|
||||
"hash_name": "e60a1e84803ce4efc0a6551206cc4b71",
|
||||
"model_params": "lib/lib_v5/modelparams/4band_44100.json",
|
||||
"param_name": "4band_44100"
|
||||
},
|
||||
{
|
||||
"hash_name": "a82f14e75892e55e994376edbf0c8435",
|
||||
"model_params": "lib/lib_v5/modelparams/4band_44100.json",
|
||||
"param_name": "4band_44100"
|
||||
},
|
||||
{
|
||||
"hash_name": "6dd9eaa6f0420af9f1d403aaafa4cc06",
|
||||
"model_params": "lib/lib_v5/modelparams/4band_v2_sn.json",
|
||||
"param_name": "4band_v2_sn"
|
||||
},
|
||||
{
|
||||
"hash_name": "08611fb99bd59eaa79ad27c58d137727",
|
||||
"model_params": "lib/lib_v5/modelparams/4band_v2_sn.json",
|
||||
"param_name": "4band_v2_sn"
|
||||
},
|
||||
{
|
||||
"hash_name": "5c7bbca45a187e81abbbd351606164e5",
|
||||
"model_params": "lib/lib_v5/modelparams/3band_44100_msb2.json",
|
||||
"param_name": "3band_44100_msb2"
|
||||
},
|
||||
{
|
||||
"hash_name": "d6b2cb685a058a091e5e7098192d3233",
|
||||
"model_params": "lib/lib_v5/modelparams/3band_44100_msb2.json",
|
||||
"param_name": "3band_44100_msb2"
|
||||
},
|
||||
{
|
||||
"hash_name": "c1b9f38170a7c90e96f027992eb7c62b",
|
||||
"model_params": "lib/lib_v5/modelparams/4band_44100.json",
|
||||
"param_name": "4band_44100"
|
||||
},
|
||||
{
|
||||
"hash_name": "c3448ec923fa0edf3d03a19e633faa53",
|
||||
"model_params": "lib/lib_v5/modelparams/4band_44100.json",
|
||||
"param_name": "4band_44100"
|
||||
},
|
||||
{
|
||||
"hash_name": "68aa2c8093d0080704b200d140f59e54",
|
||||
"model_params": "lib/lib_v5/modelparams/3band_44100.json",
|
||||
"param_name": "3band_44100"
|
||||
},
|
||||
{
|
||||
"hash_name": "fdc83be5b798e4bd29fe00fe6600e147",
|
||||
"model_params": "lib/lib_v5/modelparams/3band_44100_mid.json",
|
||||
"param_name": "3band_44100_mid.json"
|
||||
},
|
||||
{
|
||||
"hash_name": "2ce34bc92fd57f55db16b7a4def3d745",
|
||||
"model_params": "lib/lib_v5/modelparams/3band_44100_mid.json",
|
||||
"param_name": "3band_44100_mid.json"
|
||||
},
|
||||
{
|
||||
"hash_name": "52fdca89576f06cf4340b74a4730ee5f",
|
||||
"model_params": "lib/lib_v5/modelparams/4band_44100.json",
|
||||
"param_name": "4band_44100.json"
|
||||
},
|
||||
{
|
||||
"hash_name": "41191165b05d38fc77f072fa9e8e8a30",
|
||||
"model_params": "lib/lib_v5/modelparams/4band_44100.json",
|
||||
"param_name": "4band_44100.json"
|
||||
},
|
||||
{
|
||||
"hash_name": "89e83b511ad474592689e562d5b1f80e",
|
||||
"model_params": "lib/lib_v5/modelparams/2band_32000.json",
|
||||
"param_name": "2band_32000.json"
|
||||
},
|
||||
{
|
||||
"hash_name": "0b954da81d453b716b114d6d7c95177f",
|
||||
"model_params": "lib/lib_v5/modelparams/2band_32000.json",
|
||||
"param_name": "2band_32000.json"
|
||||
}
|
||||
|
||||
],
|
||||
"v4 Models": [
|
||||
{
|
||||
"hash_name": "6a00461c51c2920fd68937d4609ed6c8",
|
||||
"model_params": "lib/lib_v5/modelparams/1band_sr16000_hl512.json",
|
||||
"param_name": "1band_sr16000_hl512"
|
||||
},
|
||||
{
|
||||
"hash_name": "0ab504864d20f1bd378fe9c81ef37140",
|
||||
"model_params": "lib/lib_v5/modelparams/1band_sr32000_hl512.json",
|
||||
"param_name": "1band_sr32000_hl512"
|
||||
},
|
||||
{
|
||||
"hash_name": "7dd21065bf91c10f7fccb57d7d83b07f",
|
||||
"model_params": "lib/lib_v5/modelparams/1band_sr32000_hl512.json",
|
||||
"param_name": "1band_sr32000_hl512"
|
||||
},
|
||||
{
|
||||
"hash_name": "80ab74d65e515caa3622728d2de07d23",
|
||||
"model_params": "lib/lib_v5/modelparams/1band_sr32000_hl512.json",
|
||||
"param_name": "1band_sr32000_hl512"
|
||||
},
|
||||
{
|
||||
"hash_name": "edc115e7fc523245062200c00caa847f",
|
||||
"model_params": "lib/lib_v5/modelparams/1band_sr33075_hl384.json",
|
||||
"param_name": "1band_sr33075_hl384"
|
||||
},
|
||||
{
|
||||
"hash_name": "28063e9f6ab5b341c5f6d3c67f2045b7",
|
||||
"model_params": "lib/lib_v5/modelparams/1band_sr33075_hl384.json",
|
||||
"param_name": "1band_sr33075_hl384"
|
||||
},
|
||||
{
|
||||
"hash_name": "b58090534c52cbc3e9b5104bad666ef2",
|
||||
"model_params": "lib/lib_v5/modelparams/1band_sr44100_hl512.json",
|
||||
"param_name": "1band_sr44100_hl512"
|
||||
},
|
||||
{
|
||||
"hash_name": "0cdab9947f1b0928705f518f3c78ea8f",
|
||||
"model_params": "lib/lib_v5/modelparams/1band_sr44100_hl512.json",
|
||||
"param_name": "1band_sr44100_hl512"
|
||||
},
|
||||
{
|
||||
"hash_name": "ae702fed0238afb5346db8356fe25f13",
|
||||
"model_params": "lib/lib_v5/modelparams/1band_sr44100_hl1024.json",
|
||||
"param_name": "1band_sr44100_hl1024"
|
||||
}
|
||||
]
|
||||
}
|
||||
],
|
||||
"User Models" : [
|
||||
{
|
||||
"1 Band": [
|
||||
{
|
||||
"hash_name": "1band_sr16000_hl512",
|
||||
"model_params": "lib/lib_v5/modelparams/1band_sr16000_hl512.json",
|
||||
"param_name": "1band_sr16000_hl512"
|
||||
},
|
||||
{
|
||||
"hash_name": "1band_sr32000_hl512",
|
||||
"model_params": "lib/lib_v5/modelparams/1band_sr32000_hl512.json",
|
||||
"param_name": "1band_sr16000_hl512"
|
||||
},
|
||||
{
|
||||
"hash_name": "1band_sr33075_hl384",
|
||||
"model_params": "lib/lib_v5/modelparams/1band_sr33075_hl384.json",
|
||||
"param_name": "1band_sr33075_hl384"
|
||||
},
|
||||
{
|
||||
"hash_name": "1band_sr44100_hl256",
|
||||
"model_params": "lib/lib_v5/modelparams/1band_sr44100_hl256.json",
|
||||
"param_name": "1band_sr44100_hl256"
|
||||
},
|
||||
{
|
||||
"hash_name": "1band_sr44100_hl512",
|
||||
"model_params": "lib/lib_v5/modelparams/1band_sr44100_hl512.json",
|
||||
"param_name": "1band_sr44100_hl512"
|
||||
},
|
||||
{
|
||||
"hash_name": "1band_sr44100_hl1024",
|
||||
"model_params": "lib/lib_v5/modelparams/1band_sr44100_hl1024.json",
|
||||
"param_name": "1band_sr44100_hl1024"
|
||||
}
|
||||
],
|
||||
"2 Band": [
|
||||
{
|
||||
"hash_name": "2band_44100_lofi",
|
||||
"model_params": "lib/lib_v5/modelparams/2band_44100_lofi.json",
|
||||
"param_name": "2band_44100_lofi"
|
||||
},
|
||||
{
|
||||
"hash_name": "2band_32000",
|
||||
"model_params": "lib/lib_v5/modelparams/2band_32000.json",
|
||||
"param_name": "2band_32000"
|
||||
},
|
||||
{
|
||||
"hash_name": "2band_48000",
|
||||
"model_params": "lib/lib_v5/modelparams/2band_48000.json",
|
||||
"param_name": "2band_48000"
|
||||
}
|
||||
],
|
||||
"3 Band": [
|
||||
{
|
||||
"hash_name": "3band_44100",
|
||||
"model_params": "lib/lib_v5/modelparams/3band_44100.json",
|
||||
"param_name": "3band_44100"
|
||||
},
|
||||
{
|
||||
"hash_name": "3band_44100_mid",
|
||||
"model_params": "lib/lib_v5/modelparams/3band_44100_mid.json",
|
||||
"param_name": "3band_44100_mid"
|
||||
},
|
||||
{
|
||||
"hash_name": "3band_44100_msb2",
|
||||
"model_params": "lib/lib_v5/modelparams/3band_44100_msb2.json",
|
||||
"param_name": "3band_44100_msb2"
|
||||
}
|
||||
],
|
||||
"4 Band": [
|
||||
{
|
||||
"hash_name": "4band_44100",
|
||||
"model_params": "lib/lib_v5/modelparams/4band_44100.json",
|
||||
"param_name": "4band_44100"
|
||||
},
|
||||
{
|
||||
"hash_name": "4band_44100_mid",
|
||||
"model_params": "lib/lib_v5/modelparams/4band_44100_mid.json",
|
||||
"param_name": "4band_44100_mid"
|
||||
},
|
||||
{
|
||||
"hash_name": "4band_44100_msb",
|
||||
"model_params": "lib/lib_v5/modelparams/4band_44100_msb.json",
|
||||
"param_name": "4band_44100_msb"
|
||||
},
|
||||
{
|
||||
"hash_name": "4band_44100_msb2",
|
||||
"model_params": "lib/lib_v5/modelparams/4band_44100_msb2.json",
|
||||
"param_name": "4band_44100_msb2"
|
||||
},
|
||||
{
|
||||
"hash_name": "4band_44100_reverse",
|
||||
"model_params": "lib/lib_v5/modelparams/4band_44100_reverse.json",
|
||||
"param_name": "4band_44100_reverse"
|
||||
},
|
||||
{
|
||||
"hash_name": "4band_44100_sw",
|
||||
"model_params": "lib/lib_v5/modelparams/4band_44100_sw.json",
|
||||
"param_name": "4band_44100_sw"
|
||||
},
|
||||
{
|
||||
"hash_name": "4band_v2",
|
||||
"model_params": "lib/lib_v5/modelparams/4band_v2.json",
|
||||
"param_name": "4band_v2"
|
||||
},
|
||||
{
|
||||
"hash_name": "4band_v2_sn",
|
||||
"model_params": "lib/lib_v5/modelparams/4band_v2_sn.json",
|
||||
"param_name": "4band_v2_sn"
|
||||
},
|
||||
{
|
||||
"hash_name": "tmodelparam",
|
||||
"model_params": "lib/lib_v5/modelparams/tmodelparam.json",
|
||||
"param_name": "User Model Param Set"
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
}
|
||||
@ -1,17 +1,8 @@
|
||||
import json
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
def load_data(file_name: str = "./lib/name_params.json") -> dict:
|
||||
with open(file_name, "r") as f:
|
||||
data = json.load(f)
|
||||
|
||||
return data
|
||||
|
||||
|
||||
def make_padding(width, cropsize, offset):
|
||||
left = offset
|
||||
roi_size = cropsize - left * 2
|
||||
@ -89,25 +80,3 @@ def inference(X_spec, device, model, aggressiveness, data):
|
||||
return (pred + pred_tta) * 0.5 * coef, X_mag, np.exp(1.0j * X_phase)
|
||||
else:
|
||||
return pred * coef, X_mag, np.exp(1.0j * X_phase)
|
||||
|
||||
|
||||
def _get_name_params(model_path, model_hash):
|
||||
data = load_data()
|
||||
flag = False
|
||||
ModelName = model_path
|
||||
for type in list(data):
|
||||
for model in list(data[type][0]):
|
||||
for i in range(len(data[type][0][model])):
|
||||
if str(data[type][0][model][i]["hash_name"]) == model_hash:
|
||||
flag = True
|
||||
elif str(data[type][0][model][i]["hash_name"]) in ModelName:
|
||||
flag = True
|
||||
|
||||
if flag:
|
||||
model_params_auto = data[type][0][model][i]["model_params"]
|
||||
param_name_auto = data[type][0][model][i]["param_name"]
|
||||
if type == "equivalent":
|
||||
return param_name_auto, model_params_auto
|
||||
else:
|
||||
flag = False
|
||||
return param_name_auto, model_params_auto
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user