chattts / launch.py
zhzluke96
update
8c22399
raw
history blame
6.21 kB
import os
import logging
logging.basicConfig(
level=os.getenv("LOG_LEVEL", "INFO"),
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
from modules.devices import devices
import argparse
import uvicorn
import torch
from modules import config
from modules.utils import env
from modules import generate_audio as generate
from modules.api.Api import APIManager
from modules.api.impl import (
style_api,
tts_api,
ssml_api,
google_api,
openai_api,
refiner_api,
speaker_api,
ping_api,
models_api,
)
logger = logging.getLogger(__name__)
torch._dynamo.config.cache_size_limit = 64
torch._dynamo.config.suppress_errors = True
torch.set_float32_matmul_precision("high")
def create_api(app, no_docs=False, exclude=[]):
app_mgr = APIManager(app=app, no_docs=no_docs, exclude_patterns=exclude)
ping_api.setup(app_mgr)
models_api.setup(app_mgr)
style_api.setup(app_mgr)
speaker_api.setup(app_mgr)
tts_api.setup(app_mgr)
ssml_api.setup(app_mgr)
google_api.setup(app_mgr)
openai_api.setup(app_mgr)
refiner_api.setup(app_mgr)
return app_mgr
def get_and_update_env(*args):
val = env.get_env_or_arg(*args)
key = args[1]
config.runtime_env_vars[key] = val
return val
def setup_model_args(parser: argparse.ArgumentParser):
parser.add_argument("--compile", action="store_true", help="Enable model compile")
parser.add_argument(
"--half",
action="store_true",
help="Enable half precision for model inference",
)
parser.add_argument(
"--off_tqdm",
action="store_true",
help="Disable tqdm progress bar",
)
parser.add_argument(
"--device_id",
type=str,
help="Select the default CUDA device to use (export CUDA_VISIBLE_DEVICES=0,1,etc might be needed before)",
default=None,
)
parser.add_argument(
"--use_cpu",
nargs="+",
help="use CPU as torch device for specified modules",
default=[],
type=str.lower,
)
parser.add_argument(
"--lru_size",
type=int,
default=64,
help="Set the size of the request cache pool, set it to 0 will disable lru_cache",
)
def setup_api_args(parser: argparse.ArgumentParser):
parser.add_argument("--api_host", type=str, help="Host to run the server on")
parser.add_argument("--api_port", type=int, help="Port to run the server on")
parser.add_argument(
"--reload", action="store_true", help="Enable auto-reload for development"
)
parser.add_argument(
"--cors_origin",
type=str,
help="Allowed CORS origins. Use '*' to allow all origins.",
)
parser.add_argument(
"--no_playground",
action="store_true",
help="Disable the playground entry",
)
parser.add_argument(
"--no_docs",
action="store_true",
help="Disable the documentation entry",
)
# 配置哪些api要跳过 比如 exclude="/v1/speakers/*,/v1/tts/*"
parser.add_argument(
"--exclude",
type=str,
help="Exclude the specified API from the server",
)
def process_model_args(args):
lru_size = get_and_update_env(args, "lru_size", 64, int)
compile = get_and_update_env(args, "compile", False, bool)
device_id = get_and_update_env(args, "device_id", None, str)
use_cpu = get_and_update_env(args, "use_cpu", [], list)
half = get_and_update_env(args, "half", False, bool)
off_tqdm = get_and_update_env(args, "off_tqdm", False, bool)
generate.setup_lru_cache()
devices.reset_device()
devices.first_time_calculation()
def process_api_args(args, app):
cors_origin = get_and_update_env(args, "cors_origin", "*", str)
no_playground = get_and_update_env(args, "no_playground", False, bool)
no_docs = get_and_update_env(args, "no_docs", False, bool)
exclude = get_and_update_env(args, "exclude", "", str)
api = create_api(app=app, no_docs=no_docs, exclude=exclude.split(","))
config.api = api
if cors_origin:
api.set_cors(allow_origins=[cors_origin])
if not no_playground:
api.setup_playground()
if compile:
logger.info("Model compile is enabled")
app_description = """
ChatTTS-Forge 是一个功能强大的文本转语音生成工具,支持通过类 SSML 语法生成丰富的音频长文本,并提供全面的 API 服务,适用于各种场景。<br/>
ChatTTS-Forge is a powerful text-to-speech generation tool that supports generating rich audio long texts through class SSML syntax
项目地址: [https://github.com/lenML/ChatTTS-Forge](https://github.com/lenML/ChatTTS-Forge)
> 所有生成音频的 POST api都无法在此页面调试,调试建议使用 playground <br/>
> All audio generation POST APIs cannot be debugged on this page, it is recommended to use playground for debugging
> 如果你不熟悉本系统,建议从这个一键脚本开始,在colab中尝试一下:<br/>
> [https://colab.research.google.com/github/lenML/ChatTTS-Forge/blob/main/colab.ipynb](https://colab.research.google.com/github/lenML/ChatTTS-Forge/blob/main/colab.ipynb)
"""
app_title = "ChatTTS Forge API"
app_version = "0.1.0"
if __name__ == "__main__":
import dotenv
from fastapi import FastAPI
dotenv.load_dotenv(
dotenv_path=os.getenv("ENV_FILE", ".env.api"),
)
parser = argparse.ArgumentParser(
description="Start the FastAPI server with command line arguments"
)
setup_api_args(parser)
setup_model_args(parser)
args = parser.parse_args()
app = FastAPI(
title=app_title,
description=app_description,
version=app_version,
redoc_url=None if config.runtime_env_vars.no_docs else "/redoc",
docs_url=None if config.runtime_env_vars.no_docs else "/docs",
)
process_model_args(args)
process_api_args(args, app)
host = get_and_update_env(args, "api_host", "0.0.0.0", str)
port = get_and_update_env(args, "api_port", 7870, int)
reload = get_and_update_env(args, "reload", False, bool)
uvicorn.run(app, host=host, port=port, reload=reload)