update

.env.webui

@@ -9,8 +9,9 @@ AUTH=
 
 # Model inference options
 HALF=True
-OFF_TQDM=True
+OFF_TQDM=
 DEBUG_GENERATE=True
+PRELOAD_MODELS=True
 
 # Text-to-Speech (TTS) configuration
 TTS_MAX_LEN=1000

modules/ChatTTS/ChatTTS/core.py

@@ -112,6 +112,8 @@ class Chat:
         dtype_gpt = dtype_gpt or dtype
         dtype_decoder = dtype_decoder or dtype
 
+        map_location = torch.device("cpu")
+
         if vocos_config_path:
             vocos = (
                 Vocos.from_hparams(vocos_config_path)
@@ -119,7 +121,9 @@ class Chat:
                 .eval()
             )
             assert vocos_ckpt_path, "vocos_ckpt_path should not be None"
-            vocos.load_state_dict(torch.load(vocos_ckpt_path))
+            vocos.load_state_dict(
+                torch.load(vocos_ckpt_path, map_location=map_location)
+            )
             self.pretrain_models["vocos"] = vocos
             self.logger.log(logging.INFO, "vocos loaded.")
 
@@ -127,7 +131,7 @@ class Chat:
             cfg = OmegaConf.load(dvae_config_path)
             dvae = DVAE(**cfg).to(device=device, dtype=dtype_dvae).eval()
             assert dvae_ckpt_path, "dvae_ckpt_path should not be None"
-            dvae.load_state_dict(torch.load(dvae_ckpt_path, map_location=device))
+            dvae.load_state_dict(torch.load(dvae_ckpt_path, map_location=map_location))
             self.pretrain_models["dvae"] = dvae
             self.logger.log(logging.INFO, "dvae loaded.")
 
@@ -135,7 +139,7 @@ class Chat:
             cfg = OmegaConf.load(gpt_config_path)
             gpt = GPT_warpper(**cfg).to(device=device, dtype=dtype_gpt).eval()
             assert gpt_ckpt_path, "gpt_ckpt_path should not be None"
-            gpt.load_state_dict(torch.load(gpt_ckpt_path, map_location=device))
+            gpt.load_state_dict(torch.load(gpt_ckpt_path, map_location=map_location))
             if compile and "cuda" in str(device):
                 self.logger.info("compile gpt model")
                 gpt.gpt.forward = torch.compile(
@@ -146,21 +150,23 @@ class Chat:
             assert os.path.exists(
                 spk_stat_path
             ), f"Missing spk_stat.pt: {spk_stat_path}"
-            self.pretrain_models["spk_stat"] = torch.load(spk_stat_path).to(
-                device=device, dtype=dtype
-            )
+            self.pretrain_models["spk_stat"] = torch.load(
+                spk_stat_path, map_location=map_location
+            ).to(device=device, dtype=dtype)
             self.logger.log(logging.INFO, "gpt loaded.")
 
         if decoder_config_path:
             cfg = OmegaConf.load(decoder_config_path)
             decoder = DVAE(**cfg).to(device=device, dtype=dtype_decoder).eval()
             assert decoder_ckpt_path, "decoder_ckpt_path should not be None"
-            decoder.load_state_dict(torch.load(decoder_ckpt_path, map_location=device))
+            decoder.load_state_dict(
+                torch.load(decoder_ckpt_path, map_location=map_location)
+            )
             self.pretrain_models["decoder"] = decoder
             self.logger.log(logging.INFO, "decoder loaded.")
 
         if tokenizer_path:
-            tokenizer = torch.load(tokenizer_path, map_location=device)
+            tokenizer = torch.load(tokenizer_path, map_location=map_location)
             tokenizer.padding_side = "left"
             self.pretrain_models["tokenizer"] = tokenizer
             self.logger.log(logging.INFO, "tokenizer loaded.")

modules/Enhancer/ResembleEnhance.py

@@ -1,6 +1,5 @@
 import gc
-import os
-from typing import List, Literal
+from typing import Literal
 
 import numpy as np
 from modules.devices import devices
@@ -14,7 +13,6 @@ from modules.utils.constants import MODELS_DIR
 from pathlib import Path
 
 from threading import Lock
-from modules import config
 
 import logging
 
@@ -34,12 +32,13 @@ class ResembleEnhance:
 
     def load_model(self):
         hparams = HParams.load(Path(MODELS_DIR) / "resemble-enhance")
-        enhancer = Enhancer(hparams).to(device=self.device, dtype=self.dtype).eval()
+        enhancer = Enhancer(hparams)
         state_dict = torch.load(
             Path(MODELS_DIR) / "resemble-enhance" / "mp_rank_00_model_states.pt",
-            map_location=self.device,
+            map_location="cpu",
         )["module"]
         enhancer.load_state_dict(state_dict)
+        enhancer.to(device=self.device, dtype=self.dtype).eval()
 
         self.hparams = hparams
         self.enhancer = enhancer

modules/api/api_setup.py

@@ -1,9 +1,10 @@
 import logging
+from modules.Enhancer.ResembleEnhance import load_enhancer
 from modules.devices import devices
 import argparse
 
-import torch
 from modules import config
+from modules.models import load_chat_tts
 from modules.utils import env
 from modules import generate_audio
 from modules.api.Api import APIManager
@@ -77,6 +78,11 @@ def setup_model_args(parser: argparse.ArgumentParser):
         action="store_true",
         help="Enable debug mode for audio generation",
     )
+    parser.add_argument(
+        "--preload_models",
+        action="store_true",
+        help="Preload all models at startup",
+    )
 
 
 def process_model_args(args):
@@ -87,6 +93,7 @@ def process_model_args(args):
     no_half = env.get_and_update_env(args, "no_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)
+    preload_models = env.get_and_update_env(args, "preload_models", False, bool)
 
     generate_audio.setup_lru_cache()
     devices.reset_device()
@@ -95,6 +102,10 @@ def process_model_args(args):
     if debug_generate:
         generate_audio.logger.setLevel(logging.DEBUG)
 
+    if preload_models:
+        load_chat_tts()
+        load_enhancer()
+
 
 def setup_uvicon_args(parser: argparse.ArgumentParser):
     parser.add_argument("--host", type=str, help="Host to run the server on")

modules/webui/css/style.css

@@ -66,9 +66,9 @@
     display: none !important;
   }
   
-  .progress-bar{
+  /* .progress-bar{
     height: 30px !important;
-  }
+  } */
   
   .progress-bar span {
       text-align: right;

modules/webui/webui_utils.py

@@ -1,14 +1,13 @@
-import io
 from typing import Union
 import numpy as np
 
 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.utils.hf import spaces
 from modules.webui import webui_config
 
 import torch
+import gradio as gr
 
 from modules.ssml_parser.SSMLParser import create_ssml_parser, SSMLBreak, SSMLSegment
 from modules.SynthesizeSegments import SynthesizeSegments, combine_audio_segments
@@ -24,7 +23,6 @@ from modules import refiner
 from modules.utils import audio
 from modules.SentenceSplitter import SentenceSplitter
 
-from pydub import AudioSegment
 import torch.profiler
 
 
@@ -97,6 +95,7 @@ def synthesize_ssml(
     enable_denoise=False,
     eos: str = "[uv_break]",
     spliter_thr: int = 100,
+    progress=gr.Progress(track_tqdm=True),
 ):
     try:
         batch_size = int(batch_size)
@@ -157,6 +156,7 @@ def tts_generate(
     spk_file=None,
     spliter_thr: int = 100,
     eos: str = "[uv_break]",
+    progress=gr.Progress(track_tqdm=True),
 ):
     try:
         batch_size = int(batch_size)
@@ -219,7 +219,11 @@ def tts_generate(
 
 @torch.inference_mode()
 @spaces.GPU(duration=120)
-def refine_text(text: str, prompt: str):
+def refine_text(
+    text: str,
+    prompt: str,
+    progress=gr.Progress(track_tqdm=True),
+):
     text = text_normalize(text)
     return refiner.refine_text(text, prompt=prompt)