From adcc0a2d862d5728d762213639928b48cc869fe3 Mon Sep 17 00:00:00 2001 From: frodo821 Date: Mon, 10 Apr 2023 09:05:14 +0900 Subject: [PATCH] Add ModelType for valid model type names --- server/const.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/server/const.py b/server/const.py index 9a5ac820..e96feb8c 100644 --- a/server/const.py +++ b/server/const.py @@ -1,6 +1,10 @@ import os import sys import tempfile +from typing import Literal, TypeAlias + + +ModelType: TypeAlias = Literal['MMVCv15', 'MMVCv13', 'so-vits-svc-40v2', 'so-vits-svc-40', 'so-vits-svc-40_c', 'DDSP-SVC', 'RVC'] ERROR_NO_ONNX_SESSION = "ERROR_NO_ONNX_SESSION" @@ -22,14 +26,14 @@ TMP_DIR = os.path.join(tmpdir.name, "tmp_dir") if hasattr(sys, "_MEIPASS") else os.makedirs(TMP_DIR, exist_ok=True) -modelType = "MMVCv15" +modelType: ModelType = "MMVCv15" -def getModelType(): +def getModelType() -> ModelType: return modelType -def setModelType(_modelType: str): +def setModelType(_modelType: ModelType): global modelType modelType = _modelType