skeskinen / resemble-denoise-onnx-inference

Inference of resemble denoiser
MIT License
15 stars 3 forks source link

Nice! #1

Open matbee-eth opened 7 months ago

matbee-eth commented 7 months ago

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?

skeskinen commented 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

matbee-eth commented 3 months ago

Someone has released this https://github.com/mush42/istft-onnx to give us i/stft functions in onnx, going to give this a try

Choiyounsou commented 1 month ago

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!

imomin commented 1 month ago

@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.

skeskinen commented 1 month ago

@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

StephensCheng commented 1 week ago

Do we plan to release a C/C++ version resample-enhance model? Thanks