add new vocoder model and denoiser

infer_onnx.py

@@ -31,7 +31,7 @@ def process_text(i: int, text: str, device: torch.device):
 
 MODEL_PATH_MATCHA_MEL="matcha_multispeaker_cat_opset_15_10_steps.onnx"
 MODEL_PATH_MATCHA="matcha_hifigan_multispeaker_cat.onnx"
-MODEL_PATH_VOCOS="mel_spec_22khz.onnx"
+MODEL_PATH_VOCOS="mel_spec_22khz_v2.onnx"
 CONFIG_PATH="config_22khz.yaml"
 
 sess_options = onnxruntime.SessionOptions()
@@ -40,7 +40,7 @@ model_vocos = onnxruntime.InferenceSession(str(MODEL_PATH_VOCOS), sess_options=s
 model_matcha = onnxruntime.InferenceSession(str(MODEL_PATH_MATCHA), sess_options=sess_options, providers=["CPUExecutionProvider"])
 
 
-def vocos_inference(mel):
+def vocos_inference(mel,denoise):
 
     with open(CONFIG_PATH, "r") as f:
         config = yaml.safe_load(f)
@@ -63,6 +63,37 @@ def vocos_inference(mel):
     spectrogram = mag * (x + 1j * y)
     window = torch.hann_window(win_length)
 
+    if denoise:
+        # Vocoder bias
+        mel_rand = torch.zeros_like(torch.tensor(mel))
+        mag_bias, x_bias, y_bias = model_vocos.run(
+            None,
+            {
+                "mels": mel_rand.float().numpy()
+            },
+        )
+
+        # complex spectrogram from vocos output
+        spectrogram_bias = mag_bias * (x_bias + 1j * y_bias)
+
+        # Denoising
+        spec = torch.view_as_real(torch.tensor(spectrogram))
+        # get magnitude of vocos spectrogram
+        mag_spec = torch.sqrt(spec.pow(2).sum(-1))
+
+        # get magnitude of bias spectrogram
+        spec_bias = torch.view_as_real(torch.tensor(spectrogram_bias))
+        mag_spec_bias = torch.sqrt(spec_bias.pow(2).sum(-1))
+
+        # substract 
+        strength = 0.0005
+        mag_spec_denoised = mag_spec - mag_spec_bias * strength
+        mag_spec_denoised = torch.clamp(mag_spec_denoised, 0.0)
+
+        # return to complex spectrogram from magnitude
+        angle = torch.atan2(spec[..., -1], spec[..., 0] )
+        spectrogram = torch.complex(mag_spec_denoised * torch.cos(angle), mag_spec_denoised * torch.sin(angle))
+
     # Inverse stft
     pad = (win_length - hop_length) // 2
     spectrogram = torch.tensor(spectrogram)
@@ -92,7 +123,7 @@ def vocos_inference(mel):
     return y
 
 
-def tts(text:str, spk_id:int, temperature:float, length_scale:float):
+def tts(text:str, spk_id:int, temperature:float, length_scale:float, denoise:bool):
     sid = np.array([int(spk_id)]) if spk_id is not None else None
     text_matcha , text_lengths = process_text(0,text,"cpu") 
 
@@ -111,7 +142,7 @@ def tts(text:str, spk_id:int, temperature:float, length_scale:float):
 
     vocos_t0 = perf_counter()
     # vocos inference
-    wavs_vocos = vocos_inference(mel)
+    wavs_vocos = vocos_inference(mel,denoise)
     vocos_infer_secs = perf_counter() - vocos_t0
     print("Vocos inference time", vocos_infer_secs)
 
@@ -193,7 +224,8 @@ vits2_inference = gr.Interface(
             step=0.01,
             label="Length scale",
             info=f"Controls speech pace, larger values for slower pace and smaller values for faster pace",
-        )
+        ),
+        gr.Checkbox(label="Denoise", info="Removes model bias from vocos"),
     ],
     outputs=[gr.Audio(label="Matcha vocos", interactive=False, type="filepath"),
              gr.Audio(label="Matcha hifigan", interactive=False, type="filepath")]

mel_spec_22khz_v2.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2b02c479881f89a8320024436e986f64b11e82b1fd48046d4b695c5fd9fb84e7
+size 53883652