|
import gradio as gr |
|
from miipher.dataset.preprocess_for_infer import PreprocessForInfer |
|
from miipher.lightning_module import MiipherLightningModule |
|
from lightning_vocoders.models.hifigan.xvector_lightning_module import HiFiGANXvectorLightningModule |
|
import torch |
|
import torchaudio |
|
import hydra |
|
import tempfile |
|
|
|
miipher_path = "miipher.ckpt" |
|
miipher = MiipherLightningModule.load_from_checkpoint(miipher_path,map_location='cpu') |
|
vocoder = HiFiGANXvectorLightningModule.load_from_checkpoint("vocoder_finetuned.ckpt",map_location='cpu') |
|
xvector_model = hydra.utils.instantiate(vocoder.cfg.data.xvector.model) |
|
xvector_model = xvector_model.to('cpu') |
|
preprocessor = PreprocessForInfer(miipher.cfg) |
|
|
|
@torch.inference_mode() |
|
def main(wav_path,transcript,lang_code): |
|
wav,sr =torchaudio.load(wav_path) |
|
wav = wav[0].unsqueeze(0) |
|
batch = preprocessor.process( |
|
'test', |
|
(torch.tensor(wav),sr), |
|
word_segmented_text=transcript, |
|
lang_code=lang_code |
|
) |
|
|
|
miipher.feature_extractor(batch) |
|
( |
|
phone_feature, |
|
speaker_feature, |
|
degraded_ssl_feature, |
|
_, |
|
) = miipher.feature_extractor(batch) |
|
cleaned_ssl_feature, _ = miipher(phone_feature,speaker_feature,degraded_ssl_feature) |
|
vocoder_xvector = xvector_model.encode_batch(batch['degraded_wav_16k'].view(1,-1).cpu()).squeeze(1) |
|
cleaned_wav = vocoder.generator_forward({"input_feature": cleaned_ssl_feature, "xvector": vocoder_xvector})[0].T |
|
with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as fp: |
|
torchaudio.save(fp,cleaned_wav.view(1,-1), sample_rate=22050,format='wav') |
|
return fp.name |
|
|
|
inputs = [gr.Audio(label="noisy audio",type='filepath'),gr.Textbox(label="Transcript", value="Your transcript here", max_lines=1), |
|
gr.Radio(label="Language", choices=["eng-us", "jpn"], value="eng-us")] |
|
outputs = gr.Audio(label="Output") |
|
|
|
demo = gr.Interface(fn=main, inputs=inputs, outputs=outputs) |
|
|
|
demo.launch() |
|
|
