MikuAI/api.py

from fastapi import FastAPI, File, Query, Response
from fastapi.responses import FileResponse
from config import EDGETTS_VOICE, TOKEN
import edge_tts
import hmac
import model
from typing import Annotated, List, Optional
import tempfile

from rvc.main import song_cover_pipeline

app = FastAPI()


@app.post("/")
async def root(token: str,
    messages: List[model.DiscordMessage],
    response: Response,
    max_new_tokens: Optional[int] = 128,
    temperature: Optional[float] = 0.9):
    if not hmac.compare_digest(token, TOKEN):
        response.status_code = 401
        return {"error": "Bad token"}

    return model.inference(messages, max_new_tokens=max_new_tokens, temperature=temperature)


@app.post("/rvc")
async def rvc(token: str,
    file: Annotated[bytes, File()],
    response: Response,
    pitch_change_oct: Annotated[int, Query()] = 1,
    pitch_change_sem: Annotated[int, Query()] = 0):
    if not hmac.compare_digest(token, TOKEN):
        response.status_code = 401
        return {"error": "Bad token"}

    with tempfile.NamedTemporaryFile() as tmp:
        tmp.write(file)
        ai_vocals_path = song_cover_pipeline(tmp.name, pitch_change_oct, voice_model='miku', pitch_change_sem=pitch_change_sem)

    return FileResponse(ai_vocals_path)


@app.post("/tts")
async def tts(token: str,
    text: str,
    response: Response,
    pitch_change_oct: Annotated[int, Query()] = 1,
    pitch_change_sem: Annotated[int, Query()] = 0):
    if not hmac.compare_digest(token, TOKEN):
        response.status_code = 401
        return {"error": "Bad token"}

    with tempfile.NamedTemporaryFile() as tmp:
        communicate = edge_tts.Communicate(text, EDGETTS_VOICE)
        await communicate.save(tmp.name)
        ai_vocals_path = song_cover_pipeline(tmp.name, pitch_change_oct, voice_model='miku', pitch_change_sem=pitch_change_sem)

    return FileResponse(ai_vocals_path)


@app.get("/ping")
def ping():
    return {"message": "pong"}