from spleeter import Splitter
import torchaudio
from torchaudio.transforms import Resample
import torch
import gradio as gr
def separate(audio_path:str,inst_no:int,progress=gr.Progress(True)):
    """
    Separate audio into instrument tracks.

    Args:
        audio_path (str): Path to input audio.
        inst_no (int): Number of instruments to separate.

    Returns:
        tuple: Up to 5 MP3 file paths for separated tracks.
    """
    
    model = Splitter(inst_no)
    wav, sr = torchaudio.load(audio_path)
    target_sr = 44100
    if sr != target_sr:
        resampler = Resample(sr, target_sr)
        wav = resampler(wav)
        sr = target_sr
    with torch.no_grad():
        results = model.forward(wav)
    for i in results:
        torchaudio.save(f"{i}.mp3", results[i], sr)
    return tuple([i+".mp3" for i in results] + [None for _ in range(5-len(results))])

gr.Interface(separate, [gr.Audio(type="filepath"),gr.Dropdown([2,4,5],2)], [gr.Audio(type="filepath"), gr.Audio(type="filepath"),gr.Audio(type="filepath"),gr.Audio(type="filepath"),gr.Audio(type="filepath")]).launch(mcp_server=True)