voice-changer/server/MMVCServerSIO.py

import sys, os, struct, argparse, shutil, traceback, base64, struct
import numpy as np
import misc.log_control

from dataclasses import dataclass
from datetime import datetime
from distutils.util import strtobool

from scipy.io.wavfile import write, read

sys.path.append("MMVC_Trainer")
sys.path.append("MMVC_Trainer/text")

from fastapi.routing import APIRoute
from fastapi import HTTPException,  FastAPI, UploadFile, File, Form

from fastapi.encoders import jsonable_encoder
from fastapi.responses import JSONResponse

import uvicorn
import socketio
from pydantic import BaseModel

from typing import Callable

from mods.Trainer_Speakers import mod_get_speakers
from mods.Trainer_Training import mod_post_pre_training, mod_post_start_training, mod_post_stop_training, mod_get_related_files, mod_get_tail_training_log
from mods.Trainer_Model import mod_get_model, mod_delete_model

from mods.Trainer_Models import mod_get_models
from mods.Trainer_MultiSpeakerSetting import mod_get_multi_speaker_setting, mod_post_multi_speaker_setting
from mods.Trainer_Speaker_Voice import mod_get_speaker_voice
from mods.Trainer_Speaker_Voices import mod_get_speaker_voices

from mods.Trainer_Speaker import mod_delete_speaker
from mods.FileUploader import upload_file, concat_file_chunks

from mods.VoiceChanger import VoiceChanger

from mods.ssl import create_self_signed_cert

from voice_changer.VoiceChangerManager import VoiceChangerManager
from sio.MMVC_SocketIOApp import MMVC_SocketIOApp


from restapi.MMVC_Rest import MMVC_Rest
from pydantic import BaseModel

class VoiceModel(BaseModel):
    gpu: int
    srcId: int
    dstId: int
    timestamp: int
    prefixChunkSize: int
    buffer: str

@dataclass
class ExApplicationInfo():
    external_tensorboard_port: int


exApplitionInfo = ExApplicationInfo(external_tensorboard_port=0)




def setupArgParser():
    parser = argparse.ArgumentParser()
    parser.add_argument("-t", type=str, default="MMVC",
                        help="Server type. MMVC|TRAIN")
    parser.add_argument("-p", type=int, default=8080, help="port")
    parser.add_argument("-c", type=str, help="path for the config.json")
    parser.add_argument("-m", type=str, help="path for the model file")
    parser.add_argument("--https", type=strtobool,
                        default=False, help="use https")
    parser.add_argument("--httpsKey", type=str,
                        default="ssl.key", help="path for the key of https")
    parser.add_argument("--httpsCert", type=str,
                        default="ssl.cert", help="path for the cert of https")
    parser.add_argument("--httpsSelfSigned", type=strtobool,
                        default=True, help="generate self-signed certificate")
    parser.add_argument("--colab", type=strtobool,
                        default=False, help="run on colab")
    return parser


def printMessage(message, level=0):
    if level == 0:
        print(f"\033[17m{message}\033[0m")
    elif level == 1:
        print(f"\033[34m    {message}\033[0m")
    elif level == 2:
        print(f"\033[32m    {message}\033[0m")
    else:
        print(f"\033[47m    {message}\033[0m")



global app_socketio
global app_fastapi

parser = setupArgParser()
args = parser.parse_args()

printMessage(f"Phase name:{__name__}", level=2)
thisFilename = os.path.basename(__file__)[:-3]





if __name__ == thisFilename or args.colab == True:
    printMessage(f"PHASE3:{__name__}", level=2)
    TYPE = args.t
    PORT = args.p
    CONFIG = args.c
    MODEL = args.m

    if os.getenv("EX_TB_PORT"):
        EX_TB_PORT = os.environ["EX_TB_PORT"]
        exApplitionInfo.external_tensorboard_port = int(EX_TB_PORT)

    voiceChangerManager = VoiceChangerManager.get_instance()    
    if CONFIG and MODEL:
        voiceChangerManager.loadModel(CONFIG, MODEL)
    app_fastapi = MMVC_Rest.get_instance(voiceChangerManager)
    app_socketio = MMVC_SocketIOApp.get_instance(app_fastapi, voiceChangerManager)

    ############
    # File Uploder
    # ##########
    # @app_fastapi.post("/load_model_for_train")
    # async def post_load_model_for_train(
    #     modelGFilename: str = Form(...),
    #     modelGFilenameChunkNum: int = Form(...),
    #     modelDFilename: str = Form(...),
    #     modelDFilenameChunkNum: int = Form(...),
    # ):



    # @app_fastapi.post("/extract_voices")
    # async def post_load_model(
    #     zipFilename: str = Form(...),
    #     zipFileChunkNum: int = Form(...),
    # ):
    #     zipFilePath = concat_file_chunks(
    #         UPLOAD_DIR, zipFilename, zipFileChunkNum, UPLOAD_DIR)
    #     shutil.unpack_archive(zipFilePath, "MMVC_Trainer/dataset/textful/")
    #     return {"Zip file unpacked": f"{zipFilePath}"}

    ############
    # Voice Changer
    # ##########

    # Trainer REST API ※ ColabがTop直下のパスにしかPOSTを投げれないようなので"REST風"

    @app_fastapi.get("/get_speakers")
    async def get_speakers():
        return mod_get_speakers()

    @app_fastapi.delete("/delete_speaker")
    async def delete_speaker(speaker: str = Form(...)):
        return mod_delete_speaker(speaker)

    @app_fastapi.get("/get_speaker_voices")
    async def get_speaker_voices(speaker: str):
        return mod_get_speaker_voices(speaker)

    @app_fastapi.get("/get_speaker_voice")
    async def get_speaker_voices(speaker: str, voice: str):
        return mod_get_speaker_voice(speaker, voice)

    @app_fastapi.get("/get_multi_speaker_setting")
    async def get_multi_speaker_setting():
        return mod_get_multi_speaker_setting()

    @app_fastapi.post("/post_multi_speaker_setting")
    async def post_multi_speaker_setting(setting: str = Form(...)):
        return mod_post_multi_speaker_setting(setting)

    @app_fastapi.get("/get_models")
    async def get_models():
        return mod_get_models()

    @app_fastapi.get("/get_model")
    async def get_model(model: str):
        return mod_get_model(model)

    @app_fastapi.delete("/delete_model")
    async def delete_model(model: str = Form(...)):
        return mod_delete_model(model)

    @app_fastapi.post("/post_pre_training")
    async def post_pre_training(batch: int = Form(...)):
        return mod_post_pre_training(batch)

    @app_fastapi.post("/post_start_training")
    async def post_start_training(enable_finetuning: bool = Form(...),GModel: str = Form(...),DModel: str = Form(...)):
        print("POST START TRAINING..")
        return mod_post_start_training(enable_finetuning, GModel, DModel)

    @app_fastapi.post("/post_stop_training")
    async def post_stop_training():
        print("POST STOP TRAINING..")
        return mod_post_stop_training()

    @app_fastapi.get("/get_related_files")
    async def get_related_files():
        return mod_get_related_files()

    @app_fastapi.get("/get_tail_training_log")
    async def get_tail_training_log(num: int):
        return mod_get_tail_training_log(num)

    @app_fastapi.get("/get_ex_application_info")
    async def get_ex_application_info():
        json_compatible_item_data = jsonable_encoder(exApplitionInfo)
        return JSONResponse(content=json_compatible_item_data)

if __name__ == '__mp_main__':
    printMessage(f"PHASE2:{__name__}", level=2)

if __name__ == '__main__':
    printMessage(f"PHASE1:{__name__}", level=2)
    TYPE = args.t
    PORT = args.p
    CONFIG = args.c
    MODEL = args.m

    if TYPE != "MMVC" and TYPE != "TRAIN":
        print("Type(-t) should be MMVC or TRAIN")
        exit(1)

    printMessage(f"Start MMVC SocketIO Server", level=0)
    printMessage(f"CONFIG:{CONFIG}, MODEL:{MODEL}", level=1)

    if args.colab == False:
        if os.getenv("EX_PORT"):
            EX_PORT = os.environ["EX_PORT"]
            printMessage(
                f"External_Port:{EX_PORT} Internal_Port:{PORT}", level=1)
        else:
            printMessage(f"Internal_Port:{PORT}", level=1)

        if os.getenv("EX_TB_PORT"):
            EX_TB_PORT = os.environ["EX_TB_PORT"]
            printMessage(f"External_TeonsorBord_Port:{EX_TB_PORT}", level=1)

        if os.getenv("EX_IP"):
            EX_IP = os.environ["EX_IP"]
            printMessage(f"External_IP:{EX_IP}", level=1)

        # HTTPS key/cert作成
        if args.https and args.httpsSelfSigned == 1:
            # HTTPS(おれおれ証明書生成)
            os.makedirs("./key", exist_ok=True)
            key_base_name = f"{datetime.now().strftime('%Y%m%d_%H%M%S')}"
            keyname = f"{key_base_name}.key"
            certname = f"{key_base_name}.cert"
            create_self_signed_cert(certname, keyname, certargs={"Country": "JP",
                                                                 "State": "Tokyo",
                                                                 "City": "Chuo-ku",
                                                                 "Organization": "F",
                                                                 "Org. Unit": "F"}, cert_dir="./key")
            key_path = os.path.join("./key", keyname)
            cert_path = os.path.join("./key", certname)
            printMessage(
                f"protocol: HTTPS(self-signed), key:{key_path}, cert:{cert_path}", level=1)
        elif args.https and args.httpsSelfSigned == 0:
            # HTTPS
            key_path = args.httpsKey
            cert_path = args.httpsCert
            printMessage(
                f"protocol: HTTPS, key:{key_path}, cert:{cert_path}", level=1)
        else:
            # HTTP
            printMessage(f"protocol: HTTP", level=1)

        # アドレス表示
        if args.https == 1:
            printMessage(
                f"open https://<IP>:<PORT>/ with your browser.", level=0)
        else:
            printMessage(
                f"open http://<IP>:<PORT>/ with your browser.", level=0)

        if TYPE == "MMVC":
            path = ""
        else:
            path = "trainer"
        if "EX_PORT" in locals() and "EX_IP" in locals() and args.https == 1:
            printMessage(f"In many cases it is one of the following", level=1)
            printMessage(f"https://localhost:{EX_PORT}/{path}", level=1)
            for ip in EX_IP.strip().split(" "):
                printMessage(f"https://{ip}:{EX_PORT}/{path}", level=1)
        elif "EX_PORT" in locals() and "EX_IP" in locals() and args.https == 0:
            printMessage(f"In many cases it is one of the following", level=1)
            printMessage(f"http://localhost:{EX_PORT}/{path}", level=1)

    # サーバ起動
    if args.https:
        # HTTPS サーバ起動
        uvicorn.run(
            f"{os.path.basename(__file__)[:-3]}:app_socketio",
            host="0.0.0.0",
            port=int(PORT),
            reload=True,
            ssl_keyfile=key_path,
            ssl_certfile=cert_path,
            log_level="critical"
        )
    else:
        # HTTP サーバ起動
        if args.colab == True:
            uvicorn.run(
                f"{os.path.basename(__file__)[:-3]}:app_fastapi",
                host="0.0.0.0",
                port=int(PORT),
                log_level="critical"
            )
        else:
            uvicorn.run(
                f"{os.path.basename(__file__)[:-3]}:app_socketio",
                host="0.0.0.0",
                port=int(PORT),
                reload=True,
                log_level="critical"
            )