Shokoufehhh's picture
Update app.py
7cfdfe4 verified
raw
history blame
1.11 kB
import torch
import gradio as gr
from sgmse.model import ScoreModel
# Load your model
model_path = "https://huggingface.co/sp-uhh/speech-enhancement-sgmse/resolve/main/pretrained_checkpoints/speech_enhancement/train_vb_29nqe0uh_epoch%3D115.ckpt"
model = SGMSE() # Initialize your model class
model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
model.eval() # Set the model to evaluation mode
def enhance_audio(input_audio):
import torchaudio
# Load the input audio file
waveform, sample_rate = torchaudio.load(input_audio)
with torch.no_grad():
enhanced_waveform = model(waveform)
output_path = "enhanced_audio.wav"
torchaudio.save(output_path, enhanced_waveform.cpu(), sample_rate)
return output_path
# Create the Gradio interface
iface = gr.Interface(
fn=enhance_audio,
inputs=gr.Audio(source="upload", type="filepath"),
outputs=gr.Audio(type="file"),
title="Speech Enhancement Model",
description="Upload a noisy audio file to enhance it using the SGMSE model."
)
if __name__ == "__main__":
iface.launch()