From a3e0bc1a8232899f339257c797c3f6f2ca66d721 Mon Sep 17 00:00:00 2001
From: Azalea <22280294+hykilpikonna@users.noreply.github.com>
Date: Sat, 13 Jul 2024 02:51:16 +0800
Subject: [PATCH] [+] Inference API

---
 VC_inference.py |  7 +++-
 api.py          | 96 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 102 insertions(+), 1 deletion(-)
 create mode 100644 api.py

diff --git a/VC_inference.py b/VC_inference.py
index 8b5db1f..86a27d0 100644
--- a/VC_inference.py
+++ b/VC_inference.py
@@ -9,7 +9,6 @@ import utils
 from models import SynthesizerTrn
 import gradio as gr
 import librosa
-import webbrowser
 
 from text import text_to_sequence, _clean_text
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
@@ -28,6 +27,8 @@ language_marks = {
     "Mix": "",
 }
 lang = ['日本語', '简体中文', 'English', 'Mix']
+
+
 def get_text(text, hps, is_symbol):
     text_norm = text_to_sequence(text, hps.symbols, [] if is_symbol else hps.data.text_cleaners)
     if hps.data.add_blank:
@@ -35,6 +36,7 @@ def get_text(text, hps, is_symbol):
     text_norm = LongTensor(text_norm)
     return text_norm
 
+
 def create_tts_fn(model, hps, speaker_ids):
     def tts_fn(text, speaker, language, speed):
         if language is not None:
@@ -52,6 +54,7 @@ def create_tts_fn(model, hps, speaker_ids):
 
     return tts_fn
 
+
 def create_vc_fn(model, hps, speaker_ids):
     def vc_fn(original_speaker, target_speaker, record_audio, upload_audio):
         input_audio = record_audio if record_audio is not None else upload_audio
@@ -83,6 +86,8 @@ def create_vc_fn(model, hps, speaker_ids):
         return "Success", (hps.data.sampling_rate, audio)
 
     return vc_fn
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument("--model_dir", default="./G_latest.pth", help="directory to your fine-tuned model")
diff --git a/api.py b/api.py
new file mode 100644
index 0000000..ad8656c
--- /dev/null
+++ b/api.py
@@ -0,0 +1,96 @@
+import argparse
+import io
+from pathlib import Path
+
+import soundfile as sf
+import torch
+import uvicorn
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import StreamingResponse
+from torch import no_grad, LongTensor
+
+import commons
+import utils
+from models import SynthesizerTrn
+from text import text_to_sequence
+
+app = FastAPI()
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+language_marks = {
+    "日本語": "[JA]",
+    "简体中文": "[ZH]",
+    "English": "[EN]",
+    "Mix": "",
+}
+
+
+def get_text(text: str, is_symbol: bool):
+    text_norm = text_to_sequence(text, hps.symbols, [] if is_symbol else hps.data.text_cleaners)
+    if hps.data.add_blank:
+        text_norm = commons.intersperse(text_norm, 0)
+    text_norm = LongTensor(text_norm)
+    return text_norm
+
+
+def tts_fn(text: str, speaker: str, language: str, speed: float):
+    if language is not None:
+        text = language_marks[language] + text + language_marks[language]
+    speaker_id = speaker_ids[speaker]
+    stn_tst = get_text(text, False)
+    with no_grad():
+        x_tst = stn_tst.unsqueeze(0).to(device)
+        x_tst_lengths = LongTensor([stn_tst.size(0)]).to(device)
+        sid = LongTensor([speaker_id]).to(device)
+        audio = model.infer(x_tst, x_tst_lengths, sid=sid, noise_scale=.667, noise_scale_w=0.8,
+                            length_scale=1.0 / speed)[0][0, 0].data.cpu().float().numpy()
+    del stn_tst, x_tst, x_tst_lengths, sid
+    return audio
+
+
+@app.get("/tts/options")
+async def get_options():
+    return {"speakers": list(speaker_ids.keys()), "languages": list(language_marks.keys())}
+
+
+@app.post("/tts")
+async def generate(request: Request):
+    data = await request.json()
+    text = data.get('text')
+    speaker = data.get('speaker')
+    language = data.get('language', '日本語')
+    speed = data.get('speed', 1.0)
+
+    if not text or not speaker or language not in language_marks:
+        raise HTTPException(status_code=400, detail="Invalid input parameters")
+
+    audio = tts_fn(text, speaker, language, speed)
+    audio_io = io.BytesIO()
+    sf.write(audio_io, audio, hps.data.sampling_rate, format='OGG')
+    audio_io.seek(0)
+
+    return StreamingResponse(audio_io, media_type='audio/ogg',
+                             headers={'Content-Disposition': 'attachment; filename="output.ogg"'})
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-d", default="./OUTPUT_MODEL",
+                        help="directory to your fine-tuned model (contains G_latest.pth and config.json)")
+    args = parser.parse_args()
+    d_config = Path(args.d) / "config.json"
+    d_model = Path(args.d) / "G_latest.pth"
+    hps = utils.get_hparams_from_file(d_config)
+
+    model = SynthesizerTrn(
+        len(hps.symbols),
+        hps.data.filter_length // 2 + 1,
+        hps.train.segment_size // hps.data.hop_length,
+        n_speakers=hps.data.n_speakers,
+        **hps.model).to(device)
+    _ = model.eval()
+
+    utils.load_checkpoint(d_model, model, None)
+    speaker_ids = hps.speakers
+
+    uvicorn.run(app, host='0.0.0.0', port=27519)