Update server.py

server.py

@@ -1,25 +1,72 @@
-from flask import Flask, request, Response
-from io import BytesIO
+import argparse
+import os
+from pathlib import Path
+
+import logging
+import re_matching
+
+from flask import Flask, request, jsonify
+from flask_cors import CORS
+
+logging.getLogger("numba").setLevel(logging.WARNING)
+logging.getLogger("markdown_it").setLevel(logging.WARNING)
+logging.getLogger("urllib3").setLevel(logging.WARNING)
+logging.getLogger("matplotlib").setLevel(logging.WARNING)
+
+logging.basicConfig(
+    level=logging.INFO, format="| %(name)s | %(levelname)s | %(message)s"
+)
+
+logger = logging.getLogger(__name__)
+import librosa
+import numpy as np
 import torch
-from av import open as avopen
+import torch.nn as nn
+from torch.utils.data import Dataset
+from torch.utils.data import DataLoader, Dataset
+from tqdm import tqdm
+from transformers import Wav2Vec2Processor
+from transformers.models.wav2vec2.modeling_wav2vec2 import (
+    Wav2Vec2Model,
+    Wav2Vec2PreTrainedModel,
+)
+
+import utils
+from config import config
 
+import torch
 import commons
+from text import cleaned_text_to_sequence, get_bert
+from emo_gen import process_func, EmotionModel, Wav2Vec2Processor, Wav2Vec2Model, Wav2Vec2PreTrainedModel, RegressionHead
+from text.cleaner import clean_text
 import utils
+
 from models import SynthesizerTrn
 from text.symbols import symbols
-from text import cleaned_text_to_sequence, get_bert
-from text.cleaner import clean_text
-from scipy.io import wavfile
+import sys
 
-# Flask Init
-app = Flask(__name__)
-app.config["JSON_AS_ASCII"] = False
+from scipy.io.wavfile import write
+
+net_g = None
 
+device = 'cpu'
 
-def get_text(text, language_str, hps):
+def get_net_g(model_path: str, version: str, device: str, hps):
+    net_g = SynthesizerTrn(
+        len(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)
+    _ = net_g.eval()
+    _ = utils.load_checkpoint(model_path, net_g, None, skip_optimizer=True)
+    return net_g
+
+def get_text(text, language_str, hps, device):
     norm_text, phone, tone, word2ph = clean_text(text, language_str)
     phone, tone, language = cleaned_text_to_sequence(phone, tone, language_str)
-
+    #print(text)
     if hps.data.add_blank:
         phone = commons.intersperse(phone, 0)
         tone = commons.intersperse(tone, 0)
@@ -27,38 +74,83 @@ def get_text(text, language_str, hps):
         for i in range(len(word2ph)):
             word2ph[i] = word2ph[i] * 2
         word2ph[0] += 1
-    bert = get_bert(norm_text, word2ph, language_str)
+    bert_ori = get_bert(norm_text, word2ph, language_str, device)
     del word2ph
-    assert bert.shape[-1] == len(phone), phone
+    assert bert_ori.shape[-1] == len(phone), phone
 
     if language_str == "ZH":
-        bert = bert
-        ja_bert = torch.zeros(768, len(phone))
-    elif language_str == "JA":
-        ja_bert = bert
+        bert = bert_ori
+        ja_bert = torch.zeros(1024, len(phone))
+        en_bert = torch.zeros(1024, len(phone))
+    elif language_str == "JP":
         bert = torch.zeros(1024, len(phone))
-    else:
+        ja_bert = bert_ori
+        en_bert = torch.zeros(1024, len(phone))
+    elif language_str == "EN":
         bert = torch.zeros(1024, len(phone))
-        ja_bert = torch.zeros(768, len(phone))
+        ja_bert = torch.zeros(1024, len(phone))
+        en_bert = bert_ori
+    else:
+        raise ValueError("language_str should be ZH, JP or EN")
+
     assert bert.shape[-1] == len(
         phone
     ), f"Bert seq len {bert.shape[-1]} != {len(phone)}"
+
     phone = torch.LongTensor(phone)
     tone = torch.LongTensor(tone)
     language = torch.LongTensor(language)
-    return bert, ja_bert, phone, tone, language
+    return bert, ja_bert, en_bert, phone, tone, language
 
+def get_emo_(reference_audio, emotion):
 
-def infer(text, sdp_ratio, noise_scale, noise_scale_w, length_scale, sid, language):
-    bert, ja_bert, phones, tones, lang_ids = get_text(text, language, hps)
+    if (emotion == 10 and reference_audio):
+        emo = torch.from_numpy(get_emo(reference_audio))
+    else:
+        emo = torch.Tensor([emotion])
+
+    return emo
+
+def get_emo(path):
+    wav, sr = librosa.load(path, 16000)
+    device = config.bert_gen_config.device
+    return process_func(
+        np.expand_dims(wav, 0).astype(np.float64),
+        sr,
+        emotional_model,
+        emotional_processor,
+        device,
+        embeddings=True,
+    ).squeeze(0)
+
+def infer(
+    text,
+    sdp_ratio,
+    noise_scale,
+    noise_scale_w,
+    length_scale,
+    sid,
+    reference_audio=None,
+    emotion=0,
+):
+
+    language= 'JP' if is_japanese(text) else 'ZH'
+    bert, ja_bert, en_bert, phones, tones, lang_ids = get_text(
+        text, language, hps, device
+    )
+    emo = get_emo_(reference_audio, emotion)
     with torch.no_grad():
-        x_tst = phones.to(dev).unsqueeze(0)
-        tones = tones.to(dev).unsqueeze(0)
-        lang_ids = lang_ids.to(dev).unsqueeze(0)
-        bert = bert.to(dev).unsqueeze(0)
+        x_tst = phones.to(device).unsqueeze(0)
+        tones = tones.to(device).unsqueeze(0)
+        lang_ids = lang_ids.to(device).unsqueeze(0)
+        bert = bert.to(device).unsqueeze(0)
         ja_bert = ja_bert.to(device).unsqueeze(0)
-        x_tst_lengths = torch.LongTensor([phones.size(0)]).to(dev)
-        speakers = torch.LongTensor([hps.data.spk2id[sid]]).to(dev)
+        en_bert = en_bert.to(device).unsqueeze(0)
+        x_tst_lengths = torch.LongTensor([phones.size(0)]).to(device)
+        emo = emo.to(device).unsqueeze(0)
+        print(emo)
+        del phones
+        speakers = torch.LongTensor([hps.data.spk2id[sid]]).to(device)
         audio = (
             net_g.infer(
                 x_tst,
@@ -68,6 +160,8 @@ def infer(text, sdp_ratio, noise_scale, noise_scale_w, length_scale, sid, langua
                 lang_ids,
                 bert,
                 ja_bert,
+                en_bert,
+                emo,
                 sdp_ratio=sdp_ratio,
                 noise_scale=noise_scale,
                 noise_scale_w=noise_scale_w,
@@ -77,94 +171,59 @@ def infer(text, sdp_ratio, noise_scale, noise_scale_w, length_scale, sid, langua
             .float()
             .numpy()
         )
-        return audio
-
-
-def replace_punctuation(text, i=2):
-    punctuation = "，。？！"
-    for char in punctuation:
-        text = text.replace(char, char * i)
-    return text
-
-
-def wav2(i, o, format):
-    inp = avopen(i, "rb")
-    out = avopen(o, "wb", format=format)
-    if format == "ogg":
-        format = "libvorbis"
-
-    ostream = out.add_stream(format)
-
-    for frame in inp.decode(audio=0):
-        for p in ostream.encode(frame):
-            out.mux(p)
-
-    for p in ostream.encode(None):
-        out.mux(p)
-
-    out.close()
-    inp.close()
-
-
-# Load Generator
-hps = utils.get_hparams_from_file("./configs/config.json")
-
-dev = "cuda"
-net_g = SynthesizerTrn(
-    len(symbols),
-    hps.data.filter_length // 2 + 1,
-    hps.train.segment_size // hps.data.hop_length,
-    n_speakers=hps.data.n_speakers,
-    **hps.model,
-).to(dev)
-_ = net_g.eval()
-
-_ = utils.load_checkpoint("logs/G_649000.pth", net_g, None, skip_optimizer=True)
-
-
-@app.route("/")
-def main():
-    try:
-        speaker = request.args.get("speaker")
-        text = request.args.get("text").replace("/n", "")
-        sdp_ratio = float(request.args.get("sdp_ratio", 0.2))
-        noise = float(request.args.get("noise", 0.5))
-        noisew = float(request.args.get("noisew", 0.6))
-        length = float(request.args.get("length", 1.2))
-        language = request.args.get("language")
-        if length >= 2:
-            return "Too big length"
-        if len(text) >= 250:
-            return "Too long text"
-        fmt = request.args.get("format", "wav")
-        if None in (speaker, text):
-            return "Missing Parameter"
-        if fmt not in ("mp3", "wav", "ogg"):
-            return "Invalid Format"
-        if language not in ("JA", "ZH"):
-            return "Invalid language"
-    except:
-        return "Invalid Parameter"
-
-    with torch.no_grad():
-        audio = infer(
-            text,
-            sdp_ratio=sdp_ratio,
-            noise_scale=noise,
-            noise_scale_w=noisew,
-            length_scale=length,
-            sid=speaker,
-            language=language,
-        )
+        del x_tst, tones, lang_ids, bert, x_tst_lengths, speakers, ja_bert, en_bert, emo
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        write("temp.wav", 44100, audio)
+        return 'success'
+
+def is_japanese(string):
+        for ch in string:
+            if ord(ch) > 0x3040 and ord(ch) < 0x30FF:
+                return True
+        return False
+
+def loadmodel(model):
+    _ = net_g.eval()
+    _ = utils.load_checkpoint(model, net_g, None, skip_optimizer=True)
+    return "success"
 
-    with BytesIO() as wav:
-        wavfile.write(wav, hps.data.sampling_rate, audio)
-        torch.cuda.empty_cache()
-        if fmt == "wav":
-            return Response(wav.getvalue(), mimetype="audio/wav")
-        wav.seek(0, 0)
-        with BytesIO() as ofp:
-            wav2(wav, ofp, fmt)
-            return Response(
-                ofp.getvalue(), mimetype="audio/mpeg" if fmt == "mp3" else "audio/ogg"
-            )
+app = Flask(__name__)
+CORS(app)
+@app.route('/tts')
+
+def tts():
+    # 这些没必要改
+    speaker = request.args.get('speaker')
+    sdp_ratio = float(request.args.get('sdp_ratio', 0.2))
+    noise_scale = float(request.args.get('noise_scale', 0.6))
+    noise_scale_w = float(request.args.get('noise_scale_w', 0.8))
+    length_scale = float(request.args.get('length_scale', 1))
+    text = request.args.get('text')
+    status = infer(text, sdp_ratio=sdp_ratio, noise_scale=noise_scale, noise_scale_w=noise_scale_w, length_scale=length_scale,sid = speaker, reference_audio=None, emotion=0)
+    with open('temp.wav','rb') as bit:
+        wav_bytes = bit.read()
+    
+    headers = {
+            'Content-Type': 'audio/wav',
+            'Text': status.encode('utf-8')}
+    return wav_bytes, 200, headers
+
+
+if __name__ == "__main__":
+    emotional_model_name = "./emotional/wav2vec2-large-robust-12-ft-emotion-msp-dim"
+    REPO_ID = "audeering/wav2vec2-large-robust-12-ft-emotion-msp-dim"
+    emotional_processor = Wav2Vec2Processor.from_pretrained(emotional_model_name)
+    emotional_model = EmotionModel.from_pretrained(emotional_model_name).to(device)
+    languages = [ "Auto", "ZH", "JP"]
+    modelPaths = []
+    for dirpath, dirnames, filenames in os.walk("Data/Bushiroad/models/"):
+        for filename in filenames:
+            modelPaths.append(os.path.join(dirpath, filename))
+    hps = utils.get_hparams_from_file('Data/Bushiroad/configs/config.json')
+    net_g = get_net_g(
+        model_path=modelPaths[-1], version="2.1", device=device, hps=hps
+    )
+    speaker_ids = hps.data.spk2id
+    speakers = list(speaker_ids.keys())
+    app.run(host="0.0.0.0", port=5000)