diff --git a/StepAudioTTS.py b/StepAudioTTS.py index 8c1c6ef..d4a082a 100644 --- a/StepAudioTTS.py +++ b/StepAudioTTS.py @@ -526,15 +526,17 @@ def gen_tags(*args): formatted_args.append(f"({arg})") return formatted_args + +FUNASR_MODEL = None +KMS = None +ORT_COSY_TOKENIZER = None +LLM = None +AUTOTOKENIZER = None +COMMON_COSY_MODEL = None +MUSIC_COSY_MODEL = None + class StepAudioRun: def __init__(self): - self.funasr_model = None - self.kms = None - self.ort_cosy_tokenizer = None - self.llm = None - self.autotokenizer = None - self.common_cosy_model = None - self.music_cosy_model = None self.device = "cuda" if torch.cuda.is_available() else "cpu" @classmethod @@ -576,14 +578,14 @@ class StepAudioRun: custom_mark="", unload_model=False, ): - - if self.funasr_model is None: - self.funasr_model, self.kms, self.ort_cosy_tokenizer, self.llm, self.autotokenizer, self.common_cosy_model, self.music_cosy_model = load_models(self.device) + global FUNASR_MODEL, KMS, ORT_COSY_TOKENIZER, LLM, AUTOTOKENIZER, COMMON_COSY_MODEL, MUSIC_COSY_MODEL + if FUNASR_MODEL is None: + FUNASR_MODEL, KMS, ORT_COSY_TOKENIZER, LLM, AUTOTOKENIZER, COMMON_COSY_MODEL, MUSIC_COSY_MODEL = load_models(self.device) encoder = StepAudioTokenizer( - self.funasr_model, - self.kms, - self.ort_cosy_tokenizer, + FUNASR_MODEL, + KMS, + ORT_COSY_TOKENIZER, self.device, ) custom_mark = custom_mark.strip() if custom_mark.strip() else None @@ -596,14 +598,14 @@ class StepAudioRun: marks = gen_tags(emotion, language, speed, custom_mark) if "(RAP)" in marks or "(哼唱)" in marks: - cosy_model = self.music_cosy_model + cosy_model = MUSIC_COSY_MODEL else: - cosy_model = self.common_cosy_model + cosy_model = COMMON_COSY_MODEL tts_engine = StepAudioTTS( encoder, - self.llm, - self.autotokenizer, + LLM, + AUTOTOKENIZER, cosy_model, self.device, ) @@ -625,13 +627,13 @@ class StepAudioRun: if unload_model: tts_engine.cleanup() encoder.cleanup() - self.funasr_model = None - self.kms = None - self.ort_cosy_tokenizer = None - self.llm = None - self.autotokenizer = None - self.common_cosy_model = None - self.music_cosy_model = None + FUNASR_MODEL = None + KMS = None + ORT_COSY_TOKENIZER = None + LLM = None + AUTOTOKENIZER = None + COMMON_COSY_MODEL = None + MUSIC_COSY_MODEL = None del cosy_model gc.collect() torch.cuda.empty_cache() @@ -641,13 +643,6 @@ class StepAudioRun: class StepAudioClone: def __init__(self): - self.funasr_model = None - self.kms = None - self.ort_cosy_tokenizer = None - self.llm = None - self.autotokenizer = None - self.common_cosy_model = None - self.music_cosy_model = None self.device = "cuda" if torch.cuda.is_available() else "cpu" @classmethod @@ -690,14 +685,14 @@ class StepAudioClone: custom_mark="", unload_model=False, ): - - if self.funasr_model is None: - self.funasr_model, self.kms, self.ort_cosy_tokenizer, self.llm, self.autotokenizer, self.common_cosy_model, self.music_cosy_model = load_models(self.device) + global FUNASR_MODEL, KMS, ORT_COSY_TOKENIZER, LLM, AUTOTOKENIZER, COMMON_COSY_MODEL, MUSIC_COSY_MODEL + if FUNASR_MODEL is None: + FUNASR_MODEL, KMS, ORT_COSY_TOKENIZER, LLM, AUTOTOKENIZER, COMMON_COSY_MODEL, MUSIC_COSY_MODEL = load_models(self.device) encoder = StepAudioTokenizer( - self.funasr_model, - self.kms, - self.ort_cosy_tokenizer, + FUNASR_MODEL, + KMS, + ORT_COSY_TOKENIZER, self.device, ) @@ -711,14 +706,14 @@ class StepAudioClone: marks = gen_tags(emotion, language, speed, custom_mark) if "(RAP)" in marks or "(哼唱)" in marks: - cosy_model = self.music_cosy_model + cosy_model = MUSIC_COSY_MODEL else: - cosy_model = self.common_cosy_model + cosy_model = COMMON_COSY_MODEL tts_engine = StepAudioTTS( encoder, - self.llm, - self.autotokenizer, + LLM, + AUTOTOKENIZER, cosy_model, self.device, ) @@ -741,13 +736,13 @@ class StepAudioClone: if unload_model: tts_engine.cleanup() encoder.cleanup() - self.funasr_model = None - self.kms = None - self.ort_cosy_tokenizer = None - self.llm = None - self.autotokenizer = None - self.common_cosy_model = None - self.music_cosy_model = None + FUNASR_MODEL = None + KMS = None + ORT_COSY_TOKENIZER = None + LLM = None + AUTOTOKENIZER = None + COMMON_COSY_MODEL = None + MUSIC_COSY_MODEL = None del cosy_model gc.collect() torch.cuda.empty_cache() diff --git a/pyproject.toml b/pyproject.toml index cc99195..e10d30b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "stepaudiotts_mw" description = "A Text To Speech node using Step-Audio-TTS in ComfyUI. Can speak, rap, sing, or clone voice." -version = "2.0.3" +version = "2.0.4" license = {file = "LICENSE"} dependencies = ["transformers>=4.48.3", "openai-whisper>=20231117", "sox>=1.5.0", "hyperpyyaml", "conformer>=0.3.2", "funasr>=1.1.3"]