from fastapi import FastAPI, File, UploadFile, Form
from starlette.responses import JSONResponse
import torch
from pipeline import build_audiosep, separate_audio

app = FastAPI()

def process_audio(audio_file_path, text):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    model = build_audiosep(
        config_yaml='config/audiosep_base.yaml',
        checkpoint_path='checkpoint/audiosep_base_4M_steps.ckpt',
        device=device)

    output_file = 'SeparatedAudio.wav'
    separate_audio(model, audio_file_path, text, output_file, device)
    return output_file

@app.post("/upload")
async def upload_file(file: UploadFile = File(...), text: str = Form(...)):
    contents = await file.read()
    # Save the audio file
    with open(file.filename, "wb") as f:
        f.write(contents)
    
    # Process the audio file
    processed_audio_file_path = process_audio(file.filename, text)

    # Response with processed audio file path
    return JSONResponse(status_code=200, content={"message": "Audio file received and processed.", "processed_audio_file": processed_audio_file_path})