Open matbee-eth opened 7 months ago
Hi, what a nice coincidence :)
Conversion script goes like this:
from resemble_enhance.denoiser import denoiser_stub
from resemble_enhance.enhancer.inference import denoise, enhance, load_enhancer
import torch
import onnx
from onnxconverter_common import float16
import onnxruntime
import numpy as np
enhancer = load_enhancer(None, 'cpu')
denoiser = denoiser_stub.Denoiser(enhancer.hp)
denoiser.load_state_dict(enhancer.denoiser.state_dict())
torch_input = torch.randn(1, 841, 1600)
torch.onnx.export(denoiser, (torch_input, torch_input, torch_input), 'denoicer_torchjit.onnx', verbose=False,
input_names=['mag', 'cos', 'sin'], output_names=['out_mag', 'out_cos', 'out_sin'],
dynamic_axes={'mag': [0, 2], 'cos': [0, 2], 'sin': [0, 2], 'out_mag': [0, 2], 'out_cos': [0, 2], 'out_sin': [0, 2]})
Where denoiser_stub
is the denoiser nn.Module where forward has been replaced with this:
def forward(self, mag: Tensor, cos: Tensor, sin: Tensor):
mag_mask, sin_res, cos_res = self._predict(mag, cos, sin)
sep_mag, sep_cos, sep_sin = self._separate(mag, cos, sin, mag_mask, cos_res, sin_res)
return sep_mag, sep_cos, sep_sin
This is because the onnx doesn't like the fourier transform operations (stft, istft)... Altho there should be support for stft in new versions of onnx and istft is probably coming sometime in the future https://github.com/onnx/onnx/issues/4777
Someone has released this https://github.com/mush42/istft-onnx to give us i/stft functions in onnx, going to give this a try
That's really cool I also want to find the denoiser model and refer to it, I'd like to ask how you can refer to only denoiser separately when the model currently distributed by resemble-enhance is a training model including enhancer!
@skeskinen just wanted to thank you for the project. I have used the code from your project and onnx file in my personal project here. It's a work in progress but wanted to share and let you know.
@Choiyounsou You look at what the original code is doing and get the model weights and code only for the denoiser parts.
@imomin Nice, cool
Do we plan to release a C/C++ version resample-enhance model? Thanks
Was just checking to see if an ONNX model had been made for this. Saw a post on audacity forums. Released yesterday. How cool. Thank you!
Do you have the conversion scripts by chance?