diff --git a/ChatTTS/model/dvae.py b/ChatTTS/model/dvae.py index c767033..01802b6 100644 --- a/ChatTTS/model/dvae.py +++ b/ChatTTS/model/dvae.py @@ -251,7 +251,7 @@ class DVAE(nn.Module): self, inp: torch.Tensor, mode: Literal["encode", "decode"] = "decode" ) -> torch.Tensor: return super().__call__(inp, mode) - + @torch.inference_mode() def load_pretrained(self, filename: str, device: torch.device): state_dict_tensors = load_safetensors(filename) diff --git a/ChatTTS/utils/io.py b/ChatTTS/utils/io.py index fcb58f9..f3faa8f 100644 --- a/ChatTTS/utils/io.py +++ b/ChatTTS/utils/io.py @@ -17,6 +17,7 @@ def load_safetensors(filename: str): state_dict_tensors[k] = f.get_tensor(k) return state_dict_tensors + def get_latest_modified_file(directory): files = [os.path.join(directory, f) for f in os.listdir(directory)] diff --git a/examples/web/webui.py b/examples/web/webui.py index 77e6f2c..a0702c7 100644 --- a/examples/web/webui.py +++ b/examples/web/webui.py @@ -249,12 +249,8 @@ def main(): ) parser.add_argument("--server_port", type=int, default=8080, help="server port") parser.add_argument("--root_path", type=str, help="root path") - parser.add_argument( - "--custom_path", type=str, help="custom model path" - ) - parser.add_argument( - "--coef", type=str, help="custom dvae coefficient" - ) + parser.add_argument("--custom_path", type=str, help="custom model path") + parser.add_argument("--coef", type=str, help="custom dvae coefficient") args = parser.parse_args() logger.info("loading ChatTTS model...")