Clearer messages commit correction

system/tts_engines/f5tts/model_engine.py

@@ -79,7 +79,7 @@ def install_and_restart():
         raise ImportError("Could not install required packages")
 
 try:
-    from f5_tts.model import CFM, DiT
+    from f5_tts.model import CFM, DiT, UNetT
     from f5_tts.model.utils import (
         get_tokenizer,
         convert_char_to_pinyin,
@@ -190,7 +190,7 @@ class tts_class:
         self.sway_sampling_coef = -1.0
         self.speed = 1.0
         # F5-TTS model configuration
-        self.model_cfg = dict(
+        self.f5_model_cfg = dict(
             dim=1024,
             depth=22,
             heads=16,
@@ -198,6 +198,13 @@ class tts_class:
             text_dim=512,
             conv_layers=4
         )
+        # Add E2-TTS model configuration
+        self.e2_model_cfg = dict(
+            dim=1024,
+            depth=24,
+            heads=16,
+            ff_mult=4
+        )       
     ################################################################
     # DONT CHANGE #  Print out Python, CUDA, DeepSpeed versions ####
     ################################################################
@@ -387,7 +394,7 @@ class tts_class:
             return self.available_models
             
         # Look for model directories that match the pattern f5tts_v*
-        for model_dir in models_dir.glob("f5tts_v*"):
+        for model_dir in models_dir.glob("*tts_v*"):
             if model_dir.is_dir():
                 # First try to find model_*.safetensors files
                 model_files = list(model_dir.glob("model_*.safetensors"))
@@ -509,7 +516,10 @@ class tts_class:
             self.vocoder.load_state_dict(vocoder_state)
             self.vocoder = self.vocoder.eval()
             
-            # Initialize F5-TTS model with same mel_spec_kwargs as original
+            # Determine if this is an E2-TTS model from the folder name
+            is_e2_model = "e2tts" in model_folder.lower()
+            
+            # Initialize model with appropriate architecture
             vocab_char_map, vocab_size = get_tokenizer(str(vocab_path), "custom")
             
             mel_spec_kwargs = dict(
@@ -521,14 +531,34 @@ class tts_class:
                 mel_spec_type=self.mel_spec_type,
             )
             
-            self.model = CFM(
+            if is_e2_model:
-                transformer=DiT(**self.model_cfg, text_num_embeds=vocab_size, mel_dim=self.n_mel_channels),
+                print(f"[{self.branding}ENG] Initializing E2-TTS model") if self.debug_tts else None
-                mel_spec_kwargs=mel_spec_kwargs,
+                self.model = CFM(
-                odeint_kwargs=dict(
+                    transformer=UNetT(
-                    method=self.ode_method,
+                        **self.e2_model_cfg,
-                ),
+                        text_num_embeds=vocab_size,
-                vocab_char_map=vocab_char_map
+                        mel_dim=self.n_mel_channels
-            ).float()  # Ensure float32
+                    ),
+                    mel_spec_kwargs=mel_spec_kwargs,
+                    odeint_kwargs=dict(
+                        method=self.ode_method,
+                    ),
+                    vocab_char_map=vocab_char_map
+                ).float()  # Ensure float32
+            else:
+                print(f"[{self.branding}ENG] Initializing F5-TTS model") if self.debug_tts else None
+                self.model = CFM(
+                    transformer=DiT(
+                        **self.f5_model_cfg,
+                        text_num_embeds=vocab_size,
+                        mel_dim=self.n_mel_channels
+                    ),
+                    mel_spec_kwargs=mel_spec_kwargs,
+                    odeint_kwargs=dict(
+                        method=self.ode_method,
+                    ),
+                    vocab_char_map=vocab_char_map
+                ).float()  # Ensure float32
             
             # Load model weights
             if str(model_path).endswith('.safetensors'):
@@ -574,7 +604,8 @@ class tts_class:
                 self.vocoder = self.vocoder.to(self.device)
                 
             self.is_tts_model_loaded = True
-            print(f"[{self.branding}ENG] F5-TTS model loaded successfully") if self.debug_tts else None
+            model_type = "E2-TTS" if is_e2_model else "F5-TTS"
+            print(f"[{self.branding}ENG] {model_type} model loaded successfully") if self.debug_tts else None
             
         except Exception as e:
             print(f"[{self.branding}ENG] \033[91mError loading model: {str(e)}\033[0m")
@@ -897,6 +928,38 @@ class tts_class:
 
         return chunks
 
+    #############################################################################################################################
+    #############################################################################################################################
+    # DONT CHANGE ME # De-Capitalise words (Not used currently) #################################################################
+    #############################################################################################################################
+    #############################################################################################################################
+    def process_text_for_tts(self, text):
+        """
+        Process text by converting fully capitalized words to lowercase,
+        while preserving other words' capitalization.
+        
+        Args:
+            self: The class instance
+            text (str): Input text string
+            
+        Returns:
+            str: Processed text with fully capitalized words converted to lowercase
+        """
+        # Split text into words
+        words = text.split()
+        
+        # Process each word
+        processed_words = []
+        for word in words:
+            # Check if word is fully uppercase and longer than 1 character
+            if word.isupper() and len(word) > 1:
+                processed_words.append(word.lower())
+            else:
+                processed_words.append(word)
+        
+        # Join words back together
+        return ' '.join(processed_words)
+
     ##########################################################################################################################################
     ##########################################################################################################################################
     # CHANGE ME # Model changing. Unload out old model and load in a new one #################################################################
@@ -983,6 +1046,8 @@ class tts_class:
                 str(ref_audio_path), ref_text
             )
             
+            text = self.process_text_for_tts(text)
+            
             # Generate the audio using infer_process
             final_wave, final_sample_rate, _ = await self.infer_process(
                 ref_audio,