mirror of
https://github.com/RVC-Boss/GPT-SoVITS.git
synced 2026-08-13 01:53:43 +08:00
Remove redundant code
This commit is contained in:
RVC-Boss
@@ -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,37 +227,20 @@ 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)
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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")
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user