from fastapi import FastAPI, UploadFile
from fastapi.responses import FileResponse
from demucs.pretrained import get_model
from demucs.apply import apply_model
import torchaudio
import torch
import tempfile
批处理处理
app = FastAPI()
model = get_model('htdemucs')
model.cuda()
model.eval()
@app.post("/separate")
async def separate(file: UploadFile, stem: str = "vocals"):
# 保存上传的文件
with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp:
content = await file.read()
tmp.write(content)
tmp_path = tmp.name
# 加载并分离
wav, sr = torchaudio.load(tmp_path)
wav = wav.cuda()
with torch.no_grad():
sources = apply_model(model, wav.unsqueeze(0), split=True)[0]
stems = {'drums': 0, 'bass': 1, 'other': 2, 'vocals': 3}
output = sources[stems[stem]].cpu()
# 保存输出
with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as out:
torchaudio.save(out.name, output, sr)
return FileResponse(out.name, media_type="audio/wav")
@app.post("/instrumental")
async def get_instrumental(file: UploadFile):
with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp:
content = await file.read()
tmp.write(content)
tmp_path = tmp.name
wav, sr = torchaudio.load(tmp_path)
wav = wav.cuda()
with torch.no_grad():
sources = apply_model(model, wav.unsqueeze(0), split=True)[0]
# 合并非人声声部
instrumental = sources[0] + sources[1] + sources[2]
with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as out:
torchaudio.save(out.name, instrumental.cpu(), sr)
return FileResponse(out.name, media_type="audio/wav")
# 运行：uvicorn server:app --host 0.0.0.0 --port 8000
