update

language/zh-CN.json

@@ -31,8 +31,8 @@
   "🔊Generate": "🔊生成",
   "Disable Normalize": "禁用文本预处理",
   "💪🏼Enhance": "💪🏼增强",
-  "Enable Enhance": "启用增强",
-  "Enable De-noise": "启用降噪",
+  "Enable Enhance": "启用人声增强",
+  "Enable De-noise": "启用背景降噪",
   "🔊Generate Audio": "🔊生成音频",
   "SSML": "SSML",
   "Editor": "编辑器",

launch.py

@@ -1,201 +1,69 @@
 import os
 import logging
 
+from modules.api.api_setup import setup_api_args, setup_model_args, setup_uvicon_args
+from modules.ffmpeg_env import setup_ffmpeg_path
+
+setup_ffmpeg_path()
 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,
-)
+from fastapi import FastAPI
 
 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)
+    setup_uvicon_args(parser=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",
+    host = env.get_and_update_env(args, "host", "0.0.0.0", str)
+    port = env.get_and_update_env(args, "port", 7870, int)
+    reload = env.get_and_update_env(args, "reload", False, bool)
+    workers = env.get_and_update_env(args, "workers", 1, int)
+    log_level = env.get_and_update_env(args, "log_level", "info", str)
+    access_log = env.get_and_update_env(args, "access_log", True, bool)
+    proxy_headers = env.get_and_update_env(args, "proxy_headers", True, bool)
+    timeout_keep_alive = env.get_and_update_env(args, "timeout_keep_alive", 5, int)
+    timeout_graceful_shutdown = env.get_and_update_env(
+        args, "timeout_graceful_shutdown", 0, int
+    )
+    ssl_keyfile = env.get_and_update_env(args, "ssl_keyfile", None, str)
+    ssl_certfile = env.get_and_update_env(args, "ssl_certfile", None, str)
+    ssl_keyfile_password = env.get_and_update_env(
+        args, "ssl_keyfile_password", None, str
     )
 
-    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)
+    uvicorn.run(
+        "modules.api.worker:app",
+        host=host,
+        port=port,
+        reload=reload,
+        workers=workers,
+        log_level=log_level,
+        access_log=access_log,
+        proxy_headers=proxy_headers,
+        timeout_keep_alive=timeout_keep_alive,
+        timeout_graceful_shutdown=timeout_graceful_shutdown,
+        ssl_keyfile=ssl_keyfile,
+        ssl_certfile=ssl_certfile,
+        ssl_keyfile_password=ssl_keyfile_password,
+    )

modules/Enhancer/ResembleEnhance.py

@@ -1,5 +1,8 @@
+import gc
 import os
 from typing import List, Literal
+
+import numpy as np
 from modules.devices import devices
 from modules.repos_static.resemble_enhance.enhancer.enhancer import Enhancer
 from modules.repos_static.resemble_enhance.enhancer.hparams import HParams
@@ -11,53 +14,54 @@ from modules.utils.constants import MODELS_DIR
 from pathlib import Path
 
 from threading import Lock
+from modules import config
 
-resemble_enhance = None
-lock = Lock()
+import logging
 
+logger = logging.getLogger(__name__)
 
-def load_enhancer(device: torch.device):
-    global resemble_enhance
-    with lock:
-        if resemble_enhance is None:
-            resemble_enhance = ResembleEnhance(device)
-            resemble_enhance.load_model()
-    return resemble_enhance
+resemble_enhance = None
+lock = Lock()
 
 
 class ResembleEnhance:
-    def __init__(self, device: torch.device):
+    def __init__(self, device: torch.device, dtype=torch.float32):
         self.device = device
+        self.dtype = dtype
 
         self.enhancer: HParams = None
         self.hparams: Enhancer = None
 
     def load_model(self):
         hparams = HParams.load(Path(MODELS_DIR) / "resemble-enhance")
-        enhancer = Enhancer(hparams)
+        enhancer = Enhancer(hparams).to(device=self.device, dtype=self.dtype).eval()
         state_dict = torch.load(
             Path(MODELS_DIR) / "resemble-enhance" / "mp_rank_00_model_states.pt",
             map_location=self.device,
         )["module"]
         enhancer.load_state_dict(state_dict)
-        enhancer.to(self.device).eval()
 
         self.hparams = hparams
         self.enhancer = enhancer
 
     @torch.inference_mode()
-    def denoise(self, dwav, sr, device) -> tuple[torch.Tensor, int]:
+    def denoise(self, dwav, sr) -> tuple[torch.Tensor, int]:
         assert self.enhancer is not None, "Model not loaded"
         assert self.enhancer.denoiser is not None, "Denoiser not loaded"
         enhancer = self.enhancer
-        return inference(model=enhancer.denoiser, dwav=dwav, sr=sr, device=device)
+        return inference(
+            model=enhancer.denoiser,
+            dwav=dwav,
+            sr=sr,
+            device=self.devicem,
+            dtype=self.dtype,
+        )
 
     @torch.inference_mode()
     def enhance(
         self,
         dwav,
         sr,
-        device,
         nfe=32,
         solver: Literal["midpoint", "rk4", "euler"] = "midpoint",
         lambd=0.5,
@@ -74,7 +78,81 @@ class ResembleEnhance:
         assert self.enhancer is not None, "Model not loaded"
         enhancer = self.enhancer
         enhancer.configurate_(nfe=nfe, solver=solver, lambd=lambd, tau=tau)
-        return inference(model=enhancer, dwav=dwav, sr=sr, device=device)
+        return inference(
+            model=enhancer, dwav=dwav, sr=sr, device=self.device, dtype=self.dtype
+        )
+
+
+def load_enhancer() -> ResembleEnhance:
+    global resemble_enhance
+    with lock:
+        if resemble_enhance is None:
+            logger.info("Loading ResembleEnhance model")
+            resemble_enhance = ResembleEnhance(
+                device=devices.device, dtype=devices.dtype
+            )
+            resemble_enhance.load_model()
+            logger.info("ResembleEnhance model loaded")
+    return resemble_enhance
+
+
+def unload_enhancer():
+    global resemble_enhance
+    with lock:
+        if resemble_enhance is not None:
+            logger.info("Unloading ResembleEnhance model")
+            del resemble_enhance
+            resemble_enhance = None
+            devices.torch_gc()
+            gc.collect()
+            logger.info("ResembleEnhance model unloaded")
+
+
+def reload_enhancer():
+    logger.info("Reloading ResembleEnhance model")
+    unload_enhancer()
+    load_enhancer()
+    logger.info("ResembleEnhance model reloaded")
+
+
+def apply_audio_enhance_full(
+    audio_data: np.ndarray,
+    sr: int,
+    nfe=32,
+    solver: Literal["midpoint", "rk4", "euler"] = "midpoint",
+    lambd=0.5,
+    tau=0.5,
+):
+    # FIXME: 这里可能改成 to(device) 会优化一点？
+    tensor = torch.from_numpy(audio_data).float().squeeze().cpu()
+    enhancer = load_enhancer()
+
+    tensor, sr = enhancer.enhance(
+        tensor, sr, tau=tau, nfe=nfe, solver=solver, lambd=lambd
+    )
+
+    audio_data = tensor.cpu().numpy()
+    return audio_data, int(sr)
+
+
+def apply_audio_enhance(
+    audio_data: np.ndarray, sr: int, enable_denoise: bool, enable_enhance: bool
+):
+    if not enable_denoise and not enable_enhance:
+        return audio_data, sr
+
+    # FIXME: 这里可能改成 to(device) 会优化一点？
+    tensor = torch.from_numpy(audio_data).float().squeeze().cpu()
+    enhancer = load_enhancer()
+
+    if enable_enhance or enable_denoise:
+        lambd = 0.9 if enable_denoise else 0.1
+        tensor, sr = enhancer.enhance(
+            tensor, sr, tau=0.5, nfe=64, solver="rk4", lambd=lambd
+        )
+
+    audio_data = tensor.cpu().numpy()
+    return audio_data, int(sr)
 
 
 if __name__ == "__main__":

modules/api/Api.py

@@ -24,7 +24,7 @@ def is_excluded(path, exclude_patterns):
 
 
 class APIManager:
-    def __init__(self, app: FastAPI, no_docs=False, exclude_patterns=[]):
+    def __init__(self, app: FastAPI, exclude_patterns=[]):
         self.app = app
         self.registered_apis = {}
         self.logger = logging.getLogger(__name__)

modules/api/api_setup.py

@@ -0,0 +1,164 @@
+import logging
+from modules.devices import devices
+import argparse
+
+import torch
+from modules import config
+from modules.utils import env
+from modules import generate_audio
+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__)
+
+
+def create_api(app, exclude=[]):
+    app_mgr = APIManager(app=app, 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 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",
+    )
+    parser.add_argument(
+        "--debug_generate",
+        action="store_true",
+        help="Enable debug mode for audio generation",
+    )
+
+
+def process_model_args(args):
+    lru_size = env.get_and_update_env(args, "lru_size", 64, int)
+    compile = env.get_and_update_env(args, "compile", False, bool)
+    device_id = env.get_and_update_env(args, "device_id", None, str)
+    use_cpu = env.get_and_update_env(args, "use_cpu", [], list)
+    half = env.get_and_update_env(args, "half", False, bool)
+    off_tqdm = env.get_and_update_env(args, "off_tqdm", False, bool)
+    debug_generate = env.get_and_update_env(args, "debug_generate", False, bool)
+
+    generate_audio.setup_lru_cache()
+    devices.reset_device()
+    devices.first_time_calculation()
+
+    if debug_generate:
+        generate_audio.logger.setLevel(logging.DEBUG)
+
+
+def setup_uvicon_args(parser: argparse.ArgumentParser):
+    parser.add_argument("--host", type=str, help="Host to run the server on")
+    parser.add_argument("--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("--workers", type=int, help="Number of worker processes")
+    parser.add_argument("--log_level", type=str, help="Log level")
+    parser.add_argument("--access_log", action="store_true", help="Enable access log")
+    parser.add_argument(
+        "--proxy_headers", action="store_true", help="Enable proxy headers"
+    )
+    parser.add_argument(
+        "--timeout_keep_alive", type=int, help="Keep-alive timeout duration"
+    )
+    parser.add_argument(
+        "--timeout_graceful_shutdown",
+        type=int,
+        help="Graceful shutdown timeout duration",
+    )
+    parser.add_argument("--ssl_keyfile", type=str, help="SSL key file path")
+    parser.add_argument("--ssl_certfile", type=str, help="SSL certificate file path")
+    parser.add_argument(
+        "--ssl_keyfile_password", type=str, help="SSL key file password"
+    )
+
+
+def setup_api_args(parser: argparse.ArgumentParser):
+    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_api_args(args, app):
+    cors_origin = env.get_and_update_env(args, "cors_origin", "*", str)
+    no_playground = env.get_and_update_env(args, "no_playground", False, bool)
+    no_docs = env.get_and_update_env(args, "no_docs", False, bool)
+    exclude = env.get_and_update_env(args, "exclude", "", str)
+
+    api = create_api(app=app, 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")

modules/api/app_config.py

@@ -0,0 +1,14 @@
+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"

modules/api/impl/google_api.py

@@ -1,4 +1,5 @@
 import base64
+from typing import Literal
 from fastapi import HTTPException
 
 import io
@@ -6,7 +7,12 @@ import soundfile as sf
 from pydantic import BaseModel
 
 
+from modules.Enhancer.ResembleEnhance import (
+    apply_audio_enhance,
+    apply_audio_enhance_full,
+)
 from modules.api.Api import APIManager
+from modules.synthesize_audio import synthesize_audio
 from modules.utils.audio import apply_prosody_to_audio_data
 from modules.normalization import text_normalize
 
@@ -44,15 +50,25 @@ class AudioConfig(BaseModel):
     speakingRate: float = 1
     pitch: float = 0
     volumeGainDb: float = 0
-    sampleRateHertz: int
+    sampleRateHertz: int = 24000
     batchSize: int = 1
     spliterThreshold: int = 100
 
 
+class EnhancerConfig(BaseModel):
+    enabled: bool = False
+    model: str = "resemble-enhance"
+    nfe: int = 32
+    solver: Literal["midpoint", "rk4", "euler"] = "midpoint"
+    lambd: float = 0.5
+    tau: float = 0.5
+
+
 class GoogleTextSynthesizeRequest(BaseModel):
     input: SynthesisInput
     voice: VoiceSelectionParams
-    audioConfig: dict
+    audioConfig: AudioConfig
+    enhancerConfig: EnhancerConfig = None
 
 
 class GoogleTextSynthesizeResponse(BaseModel):
@@ -63,6 +79,7 @@ async def google_text_synthesize(request: GoogleTextSynthesizeRequest):
     input = request.input
     voice = request.voice
     audioConfig = request.audioConfig
+    enhancerConfig = request.enhancerConfig
 
     # 提取参数
 
@@ -70,40 +87,41 @@ async def google_text_synthesize(request: GoogleTextSynthesizeRequest):
     language_code = voice.languageCode
     voice_name = voice.name
     infer_seed = voice.seed or 42
-    audio_format = audioConfig.get("audioEncoding", "mp3")
-    speaking_rate = audioConfig.get("speakingRate", 1)
-    pitch = audioConfig.get("pitch", 0)
-    volume_gain_db = audioConfig.get("volumeGainDb", 0)
+    audio_format = audioConfig.audioEncoding or "mp3"
+    speaking_rate = audioConfig.speakingRate or 1
+    pitch = audioConfig.pitch or 0
+    volume_gain_db = audioConfig.volumeGainDb or 0
 
-    batch_size = audioConfig.get("batchSize", 1)
+    batch_size = audioConfig.batchSize or 1
 
     # TODO spliter_threshold
-    spliter_threshold = audioConfig.get("spliterThreshold", 100)
+    spliter_threshold = audioConfig.spliterThreshold or 100
 
     # TODO sample_rate
-    sample_rate_hertz = audioConfig.get("sampleRateHertz", 24000)
+    sample_rate_hertz = audioConfig.sampleRateHertz or 24000
 
     params = api_utils.calc_spk_style(spk=voice.name, style=voice.style)
 
-    # TODO maybe need to change the sample rate
-    sample_rate = 24000
-
     # 虽然 calc_spk_style 可以解析 seed 形式，但是这个接口只准备支持 speakers list 中存在的 speaker
     if speaker_mgr.get_speaker(voice_name) is None:
         raise HTTPException(
-            status_code=400, detail="The specified voice name is not supported."
+            status_code=422, detail="The specified voice name is not supported."
         )
 
     if audio_format != "mp3" and audio_format != "wav":
         raise HTTPException(
-            status_code=400, detail="Invalid audio encoding format specified."
+            status_code=422, detail="Invalid audio encoding format specified."
         )
 
+    if enhancerConfig.enabled:
+        # TODO enhancer params checker
+        pass
+
     try:
         if input.text:
             # 处理文本合成逻辑
             text = text_normalize(input.text, is_end=True)
-            sample_rate, audio_data = generate.generate_audio(
+            sample_rate, audio_data = synthesize_audio(
                 text,
                 temperature=(
                     voice.temperature
@@ -117,6 +135,8 @@ async def google_text_synthesize(request: GoogleTextSynthesizeRequest):
                 prompt1=params.get("prompt1", ""),
                 prompt2=params.get("prompt2", ""),
                 prefix=params.get("prefix", ""),
+                batch_size=batch_size,
+                spliter_threshold=spliter_threshold,
             )
 
         elif input.ssml:
@@ -128,7 +148,7 @@ async def google_text_synthesize(request: GoogleTextSynthesizeRequest):
 
             if len(segments) == 0:
                 raise HTTPException(
-                    status_code=400, detail="The SSML text is empty or parsing failed."
+                    status_code=422, detail="The SSML text is empty or parsing failed."
                 )
 
             synthesize = SynthesizeSegments(batch_size=batch_size)
@@ -144,7 +164,17 @@ async def google_text_synthesize(request: GoogleTextSynthesizeRequest):
 
         else:
             raise HTTPException(
-                status_code=400, detail="Either text or SSML input must be provided."
+                status_code=422, detail="Either text or SSML input must be provided."
+            )
+
+        if enhancerConfig.enabled:
+            audio_data, sample_rate = apply_audio_enhance_full(
+                audio_data=audio_data,
+                sr=sample_rate,
+                nfe=enhancerConfig.nfe,
+                solver=enhancerConfig.solver,
+                lambd=enhancerConfig.lambd,
+                tau=enhancerConfig.tau,
             )
 
         audio_data = apply_prosody_to_audio_data(

modules/api/impl/models_api.py

@@ -1,11 +1,18 @@
+from modules.Enhancer.ResembleEnhance import reload_enhancer, unload_enhancer
 from modules.api import utils as api_utils
 from modules.api.Api import APIManager
-from modules.models import reload_chat_tts
+from modules.models import reload_chat_tts, unload_chat_tts
 
 
 def setup(app: APIManager):
     @app.get("/v1/models/reload", response_model=api_utils.BaseResponse)
     async def reload_models():
-        # Reload models
         reload_chat_tts()
+        reload_enhancer()
         return api_utils.success_response("Models reloaded")
+
+    @app.get("/v1/models/unload", response_model=api_utils.BaseResponse)
+    async def reload_models():
+        unload_chat_tts()
+        unload_enhancer()
+        return api_utils.success_response("Models unloaded")

modules/api/worker.py

@@ -0,0 +1,49 @@
+import argparse
+import logging
+import os
+import dotenv
+from fastapi import FastAPI
+
+from modules.ffmpeg_env import setup_ffmpeg_path
+
+setup_ffmpeg_path()
+logging.basicConfig(
+    level=os.getenv("LOG_LEVEL", "INFO"),
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+
+from modules.api.api_setup import (
+    process_api_args,
+    process_model_args,
+    setup_api_args,
+    setup_model_args,
+    setup_uvicon_args,
+)
+from modules.api.app_config import app_description, app_title, app_version
+from modules import config
+from modules.utils.torch_opt import configure_torch_optimizations
+
+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)
+setup_uvicon_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)
+
+configure_torch_optimizations()

modules/config.py

@@ -3,7 +3,7 @@ import sys
 import torch
 from modules.utils.JsonObject import JsonObject
 
-from modules.utils import git
+from modules.utils import git, ffmpeg
 
 # TODO impl RuntimeEnvVars() class
 runtime_env_vars = JsonObject({})
@@ -20,5 +20,6 @@ versions = JsonObject(
         "git_tag": git.git_tag(),
         "git_branch": git.branch_name(),
         "git_commit": git.commit_hash(),
+        "ffmpeg_version": ffmpeg.ffmpeg_version(),
     }
 )

modules/ffmpeg_env.py

@@ -0,0 +1,16 @@
+import os
+from modules.utils.constants import ROOT_DIR
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+def setup_ffmpeg_path():
+    ffmpeg_path = os.path.join(ROOT_DIR, "ffmpeg")
+    os.environ["PATH"] = ffmpeg_path + os.pathsep + os.environ["PATH"]
+
+    import pydub.utils
+
+    if pydub.utils.which("ffmpeg") is None:
+        logger.error("ffmpeg not found in PATH")
+        raise Exception("ffmpeg not found in PATH")

modules/generate_audio.py

@@ -74,10 +74,10 @@ def generate_audio_batch(
     if isinstance(spk, int):
         with SeedContext(spk, True):
             params_infer_code["spk_emb"] = chat_tts.sample_random_speaker()
-        logger.info(("spk", spk))
+        logger.debug(("spk", spk))
     elif isinstance(spk, Speaker):
         params_infer_code["spk_emb"] = spk.emb
-        logger.info(("spk", spk.name))
+        logger.debug(("spk", spk.name))
     else:
         logger.warn(
             f"spk must be int or Speaker, but: <{type(spk)}> {spk}, wiil set to default voice"
@@ -85,7 +85,7 @@ def generate_audio_batch(
         with SeedContext(2, True):
             params_infer_code["spk_emb"] = chat_tts.sample_random_speaker()
 
-    logger.info(
+    logger.debug(
         {
             "text": texts,
             "infer_seed": infer_seed,

modules/gradio_dcls_fix.py

@@ -1,6 +1,7 @@
 def dcls_patch():
     from gradio import data_classes
 
+    # https://github.com/gradio-app/gradio/pull/8530
     data_classes.PredictBody.__get_pydantic_json_schema__ = lambda x, y: {
         "type": "object",
     }

modules/models.py

@@ -55,10 +55,9 @@ def unload_chat_tts():
             if isinstance(model, torch.nn.Module):
                 model.cpu()
                 del model
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-    gc.collect()
     chat_tts = None
+    devices.torch_gc()
+    gc.collect()
     logger.info("ChatTTS models unloaded")
 
 

modules/repos_static/resemble_enhance/denoiser/denoiser.py

@@ -65,7 +65,9 @@ class Denoiser(nn.Module):
             x = x.cpu()
 
         window = torch.hann_window(self.stft_cfg["win_length"], device=x.device)
-        s = torch.stft(x.float(), **self.stft_cfg, window=window, return_complex=True)  # (b f t+1)
+        s = torch.stft(
+            x.float(), **self.stft_cfg, window=window, return_complex=True
+        )  # (b f t+1)
 
         s = s[..., :-1]  # (b f t)
 
@@ -106,6 +108,7 @@ class Denoiser(nn.Module):
         if s.isnan().any():
             logger.warning("NaN detected in ISTFT input.")
 
+        s = s.to(torch.complex64)
         s = F.pad(s, (0, 1), "replicate")  # (b f t+1)
 
         window = torch.hann_window(self.stft_cfg["win_length"], device=s.device)
@@ -168,7 +171,9 @@ class Denoiser(nn.Module):
 
         mag, cos, sin = self._stft(x)  # (b 2f t)
         mag_mask, sin_res, cos_res = self._predict(mag, cos, sin)
-        sep_mag, sep_cos, sep_sin = self._separate(mag, cos, sin, mag_mask, cos_res, sin_res)
+        sep_mag, sep_cos, sep_sin = self._separate(
+            mag, cos, sin, mag_mask, cos_res, sin_res
+        )
 
         o = self._istft(sep_mag, sep_cos, sep_sin)
 

modules/repos_static/resemble_enhance/enhancer/lcfm/irmae.py

@@ -64,7 +64,12 @@ class IRMAE(nn.Module):
             nn.Conv1d(input_dim, hidden_dim, 3, padding="same"),
             *[ResBlock(hidden_dim) for _ in range(4)],
             # Try to obtain compact representation (https://proceedings.neurips.cc/paper/2020/file/a9078e8653368c9c291ae2f8b74012e7-Paper.pdf)
-            *[nn.Conv1d(hidden_dim if i == 0 else latent_dim, latent_dim, 1, bias=False) for i in range(num_irms)],
+            *[
+                nn.Conv1d(
+                    hidden_dim if i == 0 else latent_dim, latent_dim, 1, bias=False
+                )
+                for i in range(num_irms)
+            ],
             nn.Tanh(),
         )
 
@@ -92,9 +97,10 @@ class IRMAE(nn.Module):
         self.stats = {}
         self.stats["z_mean"] = z.mean().item()
         self.stats["z_std"] = z.std().item()
-        self.stats["z_abs_68"] = z.abs().quantile(0.6827).item()
-        self.stats["z_abs_95"] = z.abs().quantile(0.9545).item()
-        self.stats["z_abs_99"] = z.abs().quantile(0.9973).item()
+        z_float = z.float()
+        self.stats["z_abs_68"] = z_float.abs().quantile(0.6827).item()
+        self.stats["z_abs_95"] = z_float.abs().quantile(0.9545).item()
+        self.stats["z_abs_99"] = z_float.abs().quantile(0.9973).item()
         return z
 
     def decode(self, z):

modules/repos_static/resemble_enhance/inference.py

@@ -8,6 +8,8 @@ from torchaudio.functional import resample
 from torchaudio.transforms import MelSpectrogram
 from tqdm import trange
 
+from modules.devices import devices
+
 from .hparams import HParams
 
 from modules import config
@@ -16,7 +18,14 @@ logger = logging.getLogger(__name__)
 
 
 @torch.inference_mode()
-def inference_chunk(model, dwav, sr, device, npad=441):
+def inference_chunk(
+    model,
+    dwav: torch.Tensor,
+    sr: int,
+    device: torch.device,
+    dtype: torch.dtype,
+    npad=441,
+) -> torch.Tensor:
     assert model.hp.wav_rate == sr, f"Expected {model.hp.wav_rate} Hz, got {sr} Hz"
     del sr
 
@@ -24,10 +33,10 @@ def inference_chunk(model, dwav, sr, device, npad=441):
     abs_max = dwav.abs().max().clamp(min=1e-7)
 
     assert dwav.dim() == 1, f"Expected 1D waveform, got {dwav.dim()}D"
-    dwav = dwav.to(device)
+    dwav = dwav.to(device=device, dtype=dtype)
     dwav = dwav / abs_max  # Normalize
     dwav = F.pad(dwav, (0, npad))
-    hwav = model(dwav[None])[0].cpu()  # (T,)
+    hwav: torch.Tensor = model(dwav[None])[0].cpu()  # (T,)
     hwav = hwav[:length]  # Trim padding
     hwav = hwav * abs_max  # Unnormalize
 
@@ -60,6 +69,9 @@ def compute_offset(chunk1, chunk2, sr=44100):
         f_max=sr // 2,
     )
 
+    chunk1 = chunk1.float()
+    chunk2 = chunk2.float()
+
     spec1 = mel_fn(chunk1).log1p()
     spec2 = mel_fn(chunk2).log1p()
 
@@ -123,7 +135,13 @@ def remove_weight_norm_recursively(module):
 
 
 def inference(
-    model, dwav, sr, device, chunk_seconds: float = 30.0, overlap_seconds: float = 1.0
+    model,
+    dwav,
+    sr,
+    device,
+    dtype,
+    chunk_seconds: float = 30.0,
+    overlap_seconds: float = 1.0,
 ):
     if config.runtime_env_vars.off_tqdm:
         trange = range
@@ -159,9 +177,11 @@ def inference(
 
     chunks = []
     for start in trange(0, dwav.shape[-1], hop_length):
-        chunks.append(
-            inference_chunk(model, dwav[start : start + chunk_length], sr, device)
+        chunk_dwav = inference_chunk(
+            model, dwav[start : start + chunk_length], sr, device, dtype
         )
+        chunks.append(chunk_dwav.cpu())
+        devices.torch_gc()
 
     hwav = merge_chunks(chunks, chunk_length, hop_length, sr=sr, length=dwav.shape[-1])
 
@@ -172,5 +192,6 @@ def inference(
     logger.info(
         f"Elapsed time: {elapsed_time:.3f} s, {hwav.shape[-1] / elapsed_time / 1000:.3f} kHz"
     )
+    devices.torch_gc()
 
     return hwav, sr

modules/speaker.py

@@ -104,7 +104,7 @@ class SpeakerManager:
             if not os.path.exists(self.speaker_dir + fname):
                 del self.speakers[fname]
 
-    def list_speakers(self):
+    def list_speakers(self) -> list[Speaker]:
         return list(self.speakers.values())
 
     def create_speaker_from_seed(self, seed, name="", gender="", describe=""):

modules/utils/env.py

@@ -1,5 +1,7 @@
 import os
 
+from modules import config
+
 
 def get_env_val(key, val_type):
     env_val = os.getenv(key.upper())
@@ -27,3 +29,10 @@ def get_env_or_arg(args, arg_name, default, arg_type):
         return env_val
 
     return default
+
+
+def get_and_update_env(*args):
+    val = get_env_or_arg(*args)
+    key = args[1]
+    config.runtime_env_vars[key] = val
+    return val

modules/utils/ffmpeg.py

@@ -0,0 +1,20 @@
+import subprocess
+from functools import lru_cache
+
+
+@lru_cache()
+def ffmpeg_version():
+    try:
+        result = subprocess.check_output(
+            ["ffmpeg", "-version"], shell=False, encoding="utf8"
+        )
+        version_info = result.split("\n")[0]
+        version_info = version_info.split("ffmpeg version")[1].strip()
+        version_info = version_info.split("Copyright")[0].strip()
+        return version_info
+    except Exception:
+        return "<none>"
+
+
+if __name__ == "__main__":
+    print(ffmpeg_version())

modules/utils/git.py

@@ -3,23 +3,15 @@ import os
 import subprocess
 
 
-from modules.utils import constants
-
-# 用于判断是否在hf spaces
-try:
-    import spaces
-except:
-    spaces = None
+from modules.utils import constants, hf
 
 git = os.environ.get("GIT", "git")
 
-in_hf_spaces = spaces is not None
-
 
 @lru_cache()
 def commit_hash():
     try:
-        if in_hf_spaces:
+        if hf.is_spaces_env:
             return "<hf>"
         return subprocess.check_output(
             [git, "-C", constants.ROOT_DIR, "rev-parse", "HEAD"],
@@ -33,7 +25,7 @@ def commit_hash():
 @lru_cache()
 def git_tag():
     try:
-        if in_hf_spaces:
+        if hf.is_spaces_env:
             return "<hf>"
         return subprocess.check_output(
             [git, "-C", constants.ROOT_DIR, "describe", "--tags"],
@@ -57,7 +49,7 @@ def git_tag():
 @lru_cache()
 def branch_name():
     try:
-        if in_hf_spaces:
+        if hf.is_spaces_env:
             return "<hf>"
         return subprocess.check_output(
             [git, "-C", constants.ROOT_DIR, "rev-parse", "--abbrev-ref", "HEAD"],

modules/utils/hf.py

@@ -0,0 +1,17 @@
+# 给huggingface space写的兼容代码
+
+try:
+    import spaces
+
+    is_spaces_env = True
+except:
+
+    class NoneSpaces:
+        def __init__(self):
+            pass
+
+        def GPU(self, fn):
+            return fn
+
+    spaces = NoneSpaces()
+    is_spaces_env = False

modules/utils/torch_opt.py

@@ -0,0 +1,7 @@
+import torch
+
+
+def configure_torch_optimizations():
+    torch._dynamo.config.cache_size_limit = 64
+    torch._dynamo.config.suppress_errors = True
+    torch.set_float32_matmul_precision("high")

modules/webui/app.py

@@ -1,11 +1,10 @@
 import logging
 import os
 
-import torch
 import gradio as gr
 
 from modules import config
-from modules.webui import gradio_extensions, localization, webui_config, gradio_hijack
+from modules.webui import gradio_extensions, webui_config
 
 from modules.webui.changelog_tab import create_changelog_tab
 from modules.webui.localization_runtime import ENLocalizationVars, ZHLocalizationVars
@@ -24,10 +23,6 @@ def webui_init():
     # fix: If the system proxy is enabled in the Windows system, you need to skip these
     os.environ["NO_PROXY"] = "localhost,127.0.0.1,0.0.0.0"
 
-    torch._dynamo.config.cache_size_limit = 64
-    torch._dynamo.config.suppress_errors = True
-    torch.set_float32_matmul_precision("high")
-
     if config.runtime_env_vars.language == "en":
         webui_config.localization = ENLocalizationVars()
     else:
@@ -43,6 +38,7 @@ def create_app_footer():
     git_branch = os.environ.get("V_GIT_BRANCH") or config.versions.git_branch
     python_version = config.versions.python_version
     torch_version = config.versions.torch_version
+    ffmpeg_version = config.versions.ffmpeg_version
 
     config.versions.gradio_version = gradio_version
 
@@ -53,9 +49,10 @@ def create_app_footer():
     footer_items.append(f"branch: `{git_branch}`")
     footer_items.append(f"python: `{python_version}`")
     footer_items.append(f"torch: `{torch_version}`")
+    footer_items.append(f"ffmpeg: `{ffmpeg_version}`")
 
     if config.runtime_env_vars.api and not config.runtime_env_vars.no_docs:
-        footer_items.append(f"[API](/docs)")
+        footer_items.append(f"[api](/docs)")
 
     gr.Markdown(
         " | ".join(footer_items),

modules/webui/gradio_extensions.py

@@ -14,7 +14,7 @@ WEBUI_DIR_PATH = Path(os.path.dirname(os.path.realpath(__file__)))
 
 
 def read_file(fp):
-    with open(WEBUI_DIR_PATH / fp, "r") as f:
+    with open(WEBUI_DIR_PATH / fp, "r", encoding="utf-8") as f:
         return f.read()
 
 

modules/webui/js/index.js

@@ -154,6 +154,7 @@ addObserverIfDesiredNodeAvailable(".toast-wrap", function (added) {
   added.forEach(function (element) {
     if (element.innerText.includes("Connection errored out.")) {
       window.setTimeout(function () {
+        // FIXME: 这几个button好像是没有...a1111里面的gradio版本和我们的不一样
         document.getElementById("reset_button")?.classList.remove("hidden");
         document.getElementById("generate_button")?.classList.add("hidden");
         document.getElementById("skip_button")?.classList.add("hidden");

modules/webui/speaker/speaker_creator.py

@@ -2,7 +2,7 @@ import gradio as gr
 import torch
 from modules.speaker import Speaker
 from modules.utils.SeedContext import SeedContext
-from modules.hf import spaces
+from modules.utils.hf import spaces
 from modules.models import load_chat_tts
 from modules.utils.rng import np_rng
 from modules.webui import webui_config

modules/webui/speaker/speaker_editor.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import torch
 from modules.speaker import Speaker
-from modules.hf import spaces
+from modules.utils.hf import spaces
 from modules.webui import webui_config
 from modules.webui.webui_utils import tts_generate
 

modules/webui/speaker/speaker_merger.py

@@ -2,7 +2,7 @@ import io
 import gradio as gr
 import torch
 
-from modules.hf import spaces
+from modules.utils.hf import spaces
 from modules.webui import webui_config, webui_utils
 from modules.webui.webui_utils import get_speakers, tts_generate
 from modules.speaker import speaker_mgr, Speaker

modules/webui/ssml/podcast_tab.py

@@ -4,68 +4,68 @@ import torch
 
 from modules.normalization import text_normalize
 from modules.webui import webui_utils
-from modules.hf import spaces
+from modules.utils.hf import spaces
 
 podcast_default_case = [
     [
         1,
         "female2",
         "你好，欢迎收听今天的播客内容。今天我们要聊的是中华料理。 [lbreak]",
-        "podcast_p",
+        "podcast",
     ],
     [
         2,
         "Alice",
         "嗨，我特别期待这个话题！中华料理真的是博大精深。 [lbreak]",
-        "podcast_p",
+        "podcast",
     ],
     [
         3,
         "Bob",
         "没错，中华料理有着几千年的历史，而且每个地区都有自己的特色菜。 [lbreak]",
-        "podcast_p",
+        "podcast",
     ],
     [
         4,
         "female2",
         "那我们先从最有名的川菜开始吧。川菜以其麻辣著称，是很多人的最爱。 [lbreak]",
-        "podcast_p",
+        "podcast",
     ],
     [
         5,
         "Alice",
         "对，我特别喜欢吃麻婆豆腐和辣子鸡。那种麻辣的感觉真是让人难以忘怀。 [lbreak]",
-        "podcast_p",
+        "podcast",
     ],
     [
         6,
         "Bob",
         "除了川菜，粤菜也是很受欢迎的。粤菜讲究鲜美，像是白切鸡和蒸鱼都是经典。 [lbreak]",
-        "podcast_p",
+        "podcast",
     ],
     [
         7,
         "female2",
         "对啊，粤菜的烹饪方式比较清淡，更注重食材本身的味道。 [lbreak]",
-        "podcast_p",
+        "podcast",
     ],
     [
         8,
         "Alice",
         "还有北京的京菜，像北京烤鸭，那可是来北京必吃的美食。 [lbreak]",
-        "podcast_p",
+        "podcast",
     ],
     [
         9,
         "Bob",
         "不仅如此，还有淮扬菜、湘菜、鲁菜等等，每个菜系都有其独特的风味。 [lbreak]",
-        "podcast_p",
+        "podcast",
     ],
     [
         10,
         "female2",
         "对对对，像淮扬菜的狮子头，湘菜的剁椒鱼头，都是让人垂涎三尺的美味。 [lbreak]",
-        "podcast_p",
+        "podcast",
     ],
 ]
 

modules/webui/ssml/spliter_tab.py

@@ -7,7 +7,7 @@ from modules.webui.webui_utils import (
     get_styles,
     split_long_text,
 )
-from modules.hf import spaces
+from modules.utils.hf import spaces
 
 
 # NOTE: 因为 text_normalize 需要使用 tokenizer

modules/webui/webui_utils.py

@@ -2,10 +2,10 @@ import io
 from typing import Union
 import numpy as np
 
-from modules.Enhancer.ResembleEnhance import load_enhancer
+from modules.Enhancer.ResembleEnhance import apply_audio_enhance as _apply_audio_enhance
 from modules.devices import devices
 from modules.synthesize_audio import synthesize_audio
-from modules.hf import spaces
+from modules.utils.hf import spaces
 from modules.webui import webui_config
 
 import torch
@@ -85,22 +85,7 @@ def segments_length_limit(
 @torch.inference_mode()
 @spaces.GPU
 def apply_audio_enhance(audio_data, sr, enable_denoise, enable_enhance):
-    if not enable_denoise and not enable_enhance:
-        return audio_data, sr
-
-    device = devices.device
-    # NOTE: 这里很奇怪按道理得放到 device 上，但是 enhancer 做 chunk 的时候会报错...所以得 cpu()
-    tensor = torch.from_numpy(audio_data).float().squeeze().cpu()
-    enhancer = load_enhancer(device)
-
-    if enable_enhance or enable_denoise:
-        lambd = 0.9 if enable_denoise else 0.1
-        tensor, sr = enhancer.enhance(
-            tensor, sr, tau=0.5, nfe=64, solver="rk4", lambd=lambd, device=device
-        )
-
-    audio_data = tensor.cpu().numpy()
-    return audio_data, int(sr)
+    return _apply_audio_enhance(audio_data, sr, enable_denoise, enable_enhance)
 
 
 @torch.inference_mode()

webui.py

@@ -1,21 +1,23 @@
 import os
 import logging
 
+from modules.api.api_setup import (
+    process_api_args,
+    process_model_args,
+    setup_api_args,
+    setup_model_args,
+)
+from modules.ffmpeg_env import setup_ffmpeg_path
+from modules.utils.env import get_and_update_env
+from modules.api.app_config import app_description, app_title, app_version
+from modules.utils.torch_opt import configure_torch_optimizations
+
+setup_ffmpeg_path()
 logging.basicConfig(
     level=os.getenv("LOG_LEVEL", "INFO"),
     format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
 )
 
-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 import config
 from modules.webui.app import webui_init, create_interface
@@ -89,6 +91,7 @@ def process_webui_args(args):
     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)
 
+    configure_torch_optimizations()
     webui_init()
     demo = create_interface()
 
@@ -102,7 +105,7 @@ def process_webui_args(args):
         debug=debug,
         auth=auth,
         show_api=False,
-        prevent_thread_lock=False,
+        prevent_thread_lock=True,
         app_kwargs={
             "title": app_title,
             "description": app_description,
@@ -119,6 +122,18 @@ def process_webui_args(args):
             ),
         },
     )
+    # 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)
+
+    demo.block_thread()
 
 
 if __name__ == "__main__":