from fastapi import FastAPI, File, UploadFile
from model_handler import MySeparationModel
import torchaudio

app = FastAPI()
model = MySeparationModel()


@app.post("/separate/")
async def separate_audio(file: UploadFile = File(...)):
    # Save input
    input_path = "input.wav"
    with open(input_path, "wb") as f:
        f.write(await file.read())

    # Process
    separated = model.predict(input_path)

    # Save outputs
    output_paths = []
    for i in range(separated.shape[-1]):
        path = f"output_{i}.wav"
        torchaudio.save(path, separated[..., i].detach().cpu(), 8000)
        output_paths.append(path)

    return {"outputs": output_paths}