import torch

import gradio as gr
import yt_dlp as youtube_dl
from transformers import pipeline
from transformers.pipelines.audio_utils import ffmpeg_read
from speechbrain.pretrained import SepformerSeparation as separator
import torchaudio

import tempfile
import os




model = separator.from_hparams(source="speechbrain/sepformer-libri2mix", savedir='pretrained_models/sepformer-libri2mix')





demo = gr.Blocks()


def separateaudio(filepath):
    est_sources = model.separate_file(path=filepath)
    output_path = "file.wav"
    torchaudio.save(output_path, est_sources[:, :, 0].detach().cpu(), 8000)
    return output_path
    

separation = gr.Interface(
    fn=separateaudio,
    inputs=gr.Audio( type="filepath"),
    outputs=gr.Audio(type="filepath"),
)


with demo:
    gr.TabbedInterface(
        [separation],
        ["Separate audio file"],
    )

demo.launch()