update

launch.py

@@ -1,109 +1,201 @@
+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 functools import lru_cache
-from typing import Callable
-
 from modules.api.Api import APIManager
 
 from modules.api.impl import (
-    base_api,
+    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():
-    api = APIManager()
+def create_api(app, no_docs=False, exclude=[]):
+    app_mgr = APIManager(app=app, no_docs=no_docs, exclude_patterns=exclude)
 
-    base_api.setup(api)
-    tts_api.setup(api)
-    ssml_api.setup(api)
-    google_api.setup(api)
-    openai_api.setup(api)
-    refiner_api.setup(api)
+    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 api
+    return app_mgr
 
 
-def conditional_cache(condition: Callable):
-    def decorator(func):
-        @lru_cache(None)
-        def cached_func(*args, **kwargs):
-            return func(*args, **kwargs)
+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 wrapper(*args, **kwargs):
-            if condition(*args, **kwargs):
-                return cached_func(*args, **kwargs)
-            else:
-                return func(*args, **kwargs)
 
-        return wrapper
-
-    return decorator
-
-
-if __name__ == "__main__":
-    import argparse
-    import uvicorn
-
-    parser = argparse.ArgumentParser(
-        description="Start the FastAPI server with command line arguments"
+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(
-        "--host", type=str, default="0.0.0.0", help="Host to run the server on"
+        "--off_tqdm",
+        action="store_true",
+        help="Disable tqdm progress bar",
     )
     parser.add_argument(
-        "--port", type=int, default=8000, help="Port to run the server on"
+        "--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(
-        "--reload", action="store_true", help="Enable auto-reload for development"
+        "--use_cpu",
+        nargs="+",
+        help="use CPU as torch device for specified modules",
+        default=[],
+        type=str.lower,
     )
-    parser.add_argument("--compile", action="store_true", help="Enable model compile")
     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,
-        default="*",
         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",
+    )
 
-    args = parser.parse_args()
 
-    config.args = args
+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)
 
-    if args.compile:
-        print("Model compile is enabled")
-        config.enable_model_compile = True
+    generate.setup_lru_cache()
+    devices.reset_device()
+    devices.first_time_calculation()
 
-    def should_cache(*args, **kwargs):
-        spk_seed = kwargs.get("spk_seed", -1)
-        infer_seed = kwargs.get("infer_seed", -1)
-        return spk_seed != -1 and infer_seed != -1
 
-    if args.lru_size > 0:
-        config.lru_size = args.lru_size
-        generate.generate_audio = conditional_cache(should_cache)(
-            generate.generate_audio
-        )
+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()
+    api = create_api(app=app, no_docs=no_docs, exclude=exclude.split(","))
     config.api = api
 
-    if args.cors_origin:
-        api.set_cors(allow_origins=[args.cors_origin])
+    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(api.app, host=args.host, port=args.port, reload=args.reload)
+    uvicorn.run(app, host=host, port=port, reload=reload)

modules/api/Api.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, APIRouter
+from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 
 import logging
@@ -24,25 +24,8 @@ def is_excluded(path, exclude_patterns):
 
 
 class APIManager:
-    def __init__(self, no_docs=False, exclude_patterns=[]):
-        self.app = FastAPI(
-            title="ChatTTS Forge API",
-            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)
-            """,
-            version="0.1.0",
-            redoc_url=None if no_docs else "/redoc",
-            docs_url=None if no_docs else "/docs",
-        )
+    def __init__(self, app: FastAPI, no_docs=False, exclude_patterns=[]):
+        self.app = app
         self.registered_apis = {}
         self.logger = logging.getLogger(__name__)
         self.exclude = exclude_patterns
@@ -57,6 +40,8 @@ ChatTTS-Forge is a powerful text-to-speech generation tool that supports generat
         allow_methods: list = ["*"],
         allow_headers: list = ["*"],
     ):
+        # reset middleware stack
+        self.app.middleware_stack = None
         self.app.add_middleware(
             CORSMiddleware,
             allow_origins=allow_origins,
@@ -64,6 +49,7 @@ ChatTTS-Forge is a powerful text-to-speech generation tool that supports generat
             allow_methods=allow_methods,
             allow_headers=allow_headers,
         )
+        self.app.build_middleware_stack()
 
     def setup_playground(self):
         app = self.app

modules/api/impl/refiner_api.py

@@ -7,6 +7,7 @@ from modules import refiner
 
 from modules.api import utils as api_utils
 from modules.api.Api import APIManager
+from modules.normalization import text_normalize
 
 
 class RefineTextRequest(BaseModel):
@@ -18,6 +19,7 @@ class RefineTextRequest(BaseModel):
     temperature: float = 0.7
     repetition_penalty: float = 1.0
     max_new_token: int = 384
+    normalize: bool = True
 
 
 async def refiner_prompt_post(request: RefineTextRequest):
@@ -26,8 +28,11 @@ async def refiner_prompt_post(request: RefineTextRequest):
     """
 
     try:
+        text = request.text
+        if request.normalize:
+            text = text_normalize(request.text)
         refined_text = refiner.refine_text(
-            text=request.text,
+            text=text,
             prompt=request.prompt,
             seed=request.seed,
             top_P=request.top_P,

modules/api/impl/speaker_api.py

@@ -35,10 +35,14 @@ def setup(app: APIManager):
 
     @app.get("/v1/speakers/list", response_model=api_utils.BaseResponse)
     async def list_speakers():
-        return {
-            "message": "ok",
-            "data": [spk.to_json() for spk in speaker_mgr.list_speakers()],
-        }
+        return api_utils.success_response(
+            [spk.to_json() for spk in speaker_mgr.list_speakers()]
+        )
+
+    @app.post("/v1/speakers/refresh", response_model=api_utils.BaseResponse)
+    async def refresh_speakers():
+        speaker_mgr.refresh_speakers()
+        return api_utils.success_response(None)
 
     @app.post("/v1/speakers/update", response_model=api_utils.BaseResponse)
     async def update_speakers(request: SpeakersUpdate):
@@ -59,7 +63,8 @@ def setup(app: APIManager):
                 # number array => Tensor
                 speaker.emb = torch.tensor(spk["tensor"])
         speaker_mgr.save_all()
-        return {"message": "ok", "data": None}
+
+        return api_utils.success_response(None)
 
     @app.post("/v1/speaker/create", response_model=api_utils.BaseResponse)
     async def create_speaker(request: CreateSpeaker):
@@ -88,12 +93,7 @@ def setup(app: APIManager):
             raise HTTPException(
                 status_code=400, detail="Missing tensor or seed in request"
             )
-        return {"message": "ok", "data": speaker.to_json()}
-
-    @app.post("/v1/speaker/refresh", response_model=api_utils.BaseResponse)
-    async def refresh_speakers():
-        speaker_mgr.refresh_speakers()
-        return {"message": "ok"}
+        return api_utils.success_response(speaker.to_json())
 
     @app.post("/v1/speaker/update", response_model=api_utils.BaseResponse)
     async def update_speaker(request: UpdateSpeaker):
@@ -113,11 +113,11 @@ def setup(app: APIManager):
             # number array => Tensor
             speaker.emb = torch.tensor(request.tensor)
         speaker_mgr.update_speaker(speaker)
-        return {"message": "ok"}
+        return api_utils.success_response(None)
 
     @app.post("/v1/speaker/detail", response_model=api_utils.BaseResponse)
     async def speaker_detail(request: SpeakerDetail):
         speaker = speaker_mgr.get_speaker_by_id(request.id)
         if speaker is None:
             raise HTTPException(status_code=404, detail="Speaker not found")
-        return {"message": "ok", "data": speaker.to_json(with_emb=request.with_emb)}
+        return api_utils.success_response(speaker.to_json(with_emb=request.with_emb))

modules/api/impl/tts_api.py

@@ -9,8 +9,6 @@ from fastapi.responses import FileResponse
 
 from modules.normalization import text_normalize
 
-from modules import generate_audio as generate
-
 from modules.api import utils as api_utils
 from modules.api.Api import APIManager
 from modules.synthesize_audio import synthesize_audio

modules/gradio_dcls_fix.py

@@ -0,0 +1,6 @@
+def dcls_patch():
+    from gradio import data_classes
+
+    data_classes.PredictBody.__get_pydantic_json_schema__ = lambda x, y: {
+        "type": "object",
+    }

modules/webui/app.py

@@ -46,11 +46,19 @@ def create_app_footer():
 
     config.versions.gradio_version = gradio_version
 
+    footer_items = ["🍦 [ChatTTS-Forge](https://github.com/lenML/ChatTTS-Forge)"]
+    footer_items.append(
+        f"version: [{git_tag}](https://github.com/lenML/ChatTTS-Forge/commit/{git_commit})"
+    )
+    footer_items.append(f"branch: `{git_branch}`")
+    footer_items.append(f"python: `{python_version}`")
+    footer_items.append(f"torch: `{torch_version}`")
+
+    if config.runtime_env_vars.api and not config.runtime_env_vars.no_docs:
+        footer_items.append(f"[API](/docs)")
+
     gr.Markdown(
-        f"""
-🍦 [ChatTTS-Forge](https://github.com/lenML/ChatTTS-Forge)
-version: [{git_tag}](https://github.com/lenML/ChatTTS-Forge/commit/{git_commit}) | branch: `{git_branch}` | python: `{python_version}` | torch: `{torch_version}`
-        """,
+        " | ".join(footer_items),
         elem_classes=["no-translate"],
     )
 

modules/webui/js/localization.js

@@ -163,6 +163,23 @@ function localizeWholePage() {
   }
 }
 
+/**
+ *
+ * @param {HTMLElement} node
+ */
+function isNeedTranslate(node) {
+  if (!node) return false;
+  if (!(node instanceof HTMLElement)) return true;
+  while (node.parentElement !== document.body) {
+    if (node.classList.contains("no-translate")) {
+      return false;
+    }
+    node = node.parentElement;
+    if (!node) break;
+  }
+  return true;
+}
+
 document.addEventListener("DOMContentLoaded", function () {
   if (!hasLocalization()) {
     return;
@@ -170,9 +187,11 @@ document.addEventListener("DOMContentLoaded", function () {
 
   onUiUpdate(function (m) {
     m.forEach(function (mutation) {
-      mutation.addedNodes.forEach(function (node) {
-        processNode(node);
-      });
+      Array.from(mutation.addedNodes)
+        .filter(isNeedTranslate)
+        .forEach(function (node) {
+          processNode(node);
+        });
     });
   });
 

modules/webui/tts_tab.py

@@ -96,7 +96,7 @@ def create_tts_interface():
                             )
 
                             gr.Markdown("📝Speaker info")
-                            infos = gr.Markdown("empty")
+                            infos = gr.Markdown("empty", elem_classes=["no-translate"])
 
                             spk_file_upload.change(
                                 fn=load_spk_info,

webui.py

@@ -1,27 +1,30 @@
 import os
 import logging
 
-# logging.basicConfig(
-#     level=os.getenv("LOG_LEVEL", "INFO"),
-#     format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-# )
+logging.basicConfig(
+    level=os.getenv("LOG_LEVEL", "INFO"),
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
 
-from modules.devices import devices
-from modules.utils import env
+from launch import (
+    get_and_update_env,
+    setup_api_args,
+    setup_model_args,
+    process_api_args,
+    process_model_args,
+    app_description,
+    app_title,
+    app_version,
+)
 from modules.webui import webui_config
 from modules.webui.app import webui_init, create_interface
-from modules import generate_audio
-from modules import config
+import argparse
+from modules.gradio_dcls_fix import dcls_patch
 
-if __name__ == "__main__":
-    import argparse
-    import dotenv
+dcls_patch()
 
-    dotenv.load_dotenv(
-        dotenv_path=os.getenv("ENV_FILE", ".env.webui"),
-    )
 
-    parser = argparse.ArgumentParser(description="Gradio App")
+def setup_webui_args(parser: argparse.ArgumentParser):
     parser.add_argument("--server_name", type=str, help="server name")
     parser.add_argument("--server_port", type=int, help="server port")
     parser.add_argument(
@@ -29,16 +32,6 @@ if __name__ == "__main__":
     )
     parser.add_argument("--debug", action="store_true", help="enable debug mode")
     parser.add_argument("--auth", type=str, help="username:password for authentication")
-    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(
         "--tts_max_len",
         type=int,
@@ -54,58 +47,39 @@ if __name__ == "__main__":
         type=int,
         help="Max batch size for TTS",
     )
-    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",
-    )
-    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("--compile", action="store_true", help="Enable model compile")
     # webui_Experimental
     parser.add_argument(
         "--webui_experimental",
         action="store_true",
         help="Enable webui_experimental features",
     )
-
     parser.add_argument(
         "--language",
         type=str,
         help="Set the default language for the webui",
     )
-    args = parser.parse_args()
+    parser.add_argument(
+        "--api",
+        action="store_true",
+        help="use api=True to launch the API together with the webui (run launch.py for only API server)",
+    )
 
-    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 process_webui_args(args):
     server_name = get_and_update_env(args, "server_name", "0.0.0.0", str)
     server_port = get_and_update_env(args, "server_port", 7860, int)
     share = get_and_update_env(args, "share", False, bool)
     debug = get_and_update_env(args, "debug", False, bool)
     auth = get_and_update_env(args, "auth", None, str)
-    half = get_and_update_env(args, "half", False, bool)
-    off_tqdm = get_and_update_env(args, "off_tqdm", False, bool)
-    lru_size = get_and_update_env(args, "lru_size", 64, int)
-    device_id = get_and_update_env(args, "device_id", None, str)
-    use_cpu = get_and_update_env(args, "use_cpu", [], list)
-    compile = get_and_update_env(args, "compile", False, bool)
     language = get_and_update_env(args, "language", "zh-CN", str)
+    api = get_and_update_env(args, "api", "zh-CN", str)
+
+    webui_config.experimental = get_and_update_env(
+        args, "webui_experimental", False, bool
+    )
+    webui_config.tts_max = get_and_update_env(args, "tts_max_len", 1000, int)
+    webui_config.ssml_max = get_and_update_env(args, "ssml_max_len", 5000, int)
+    webui_config.max_batch_size = get_and_update_env(args, "max_batch_size", 8, int)
 
     webui_config.experimental = get_and_update_env(
         args, "webui_experimental", False, bool
@@ -120,15 +94,57 @@ if __name__ == "__main__":
     if auth:
         auth = tuple(auth.split(":"))
 
-    generate_audio.setup_lru_cache()
-    devices.reset_device()
-    devices.first_time_calculation()
-
-    demo.queue().launch(
+    app, local_url, share_url = demo.queue().launch(
         server_name=server_name,
         server_port=server_port,
         share=share,
         debug=debug,
         auth=auth,
         show_api=False,
+        prevent_thread_lock=True,
+        app_kwargs={
+            "title": app_title,
+            "description": app_description,
+            "version": app_version,
+            # "redoc_url": (
+            #     None
+            #     if api is False
+            #     else None if config.runtime_env_vars.no_docs else "/redoc"
+            # ),
+            # "docs_url": (
+            #     None
+            #     if api is False
+            #     else None if config.runtime_env_vars.no_docs else "/docs"
+            # ),
+            "docs_url": "/docs",
+        },
+    )
+    # gradio uses a very open CORS policy via app.user_middleware, which makes it possible for
+    # an attacker to trick the user into opening a malicious HTML page, which makes a request to the
+    # running web ui and do whatever the attacker wants, including installing an extension and
+    # running its code. We disable this here. Suggested by RyotaK.
+    app.user_middleware = [
+        x for x in app.user_middleware if x.cls.__name__ != "CustomCORSMiddleware"
+    ]
+
+    if api:
+        process_api_args(args, app)
+
+
+if __name__ == "__main__":
+    import dotenv
+
+    dotenv.load_dotenv(
+        dotenv_path=os.getenv("ENV_FILE", ".env.webui"),
     )
+
+    parser = argparse.ArgumentParser(description="Gradio App")
+
+    setup_webui_args(parser)
+    setup_model_args(parser)
+    setup_api_args(parser)
+
+    args = parser.parse_args()
+
+    process_model_args(args)
+    process_webui_args(args)