Update tts_engine.py

tts_engine.py

@@ -1,53 +1,71 @@
-# tts_engine.py - TTS engine wrapper for Nari DIA
+# tts_engine.py - TTS engine wrapper for CPU-friendly SpeechT5
 import logging
 import os
 from typing import Optional
 import tempfile
 import numpy as np
 import soundfile as sf
-import torch # Import torch for model operations
+import torch
 
-# Import the actual Nari DIA model
-try:
-    from dia.model import Dia
-except ImportError:
-    logging.error("Nari DIA library not found. Please ensure 'git+https://github.com/nari-labs/dia.git' is in your requirements.txt and installed.")
-    Dia = None # Set to None to prevent further errors
+from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
+from datasets import load_dataset # To get speaker embeddings from VCTK
 
 logger = logging.getLogger(__name__)
 
-class NariDIAEngine:
+class CPUMultiSpeakerTTS:
     def __init__(self):
+        self.processor = None
         self.model = None
-        # No separate processor object for Dia, it handles internal processing
+        self.vocoder = None
+        self.speaker_embeddings = {} # Will store speaker embeddings for S1, S2 etc.
         self._initialize_model()
     
     def _initialize_model(self):
-        """Initialize the Nari DIA 1.6B model."""
-        if Dia is None:
-            logger.error("Nari DIA library is not available. Cannot initialize model.")
-            return
-
+        """Initialize the SpeechT5 model and vocoder on CPU."""
         try:
-            logger.info("Initializing Nari DIA 1.6B model from nari-labs/Dia-1.6B...")
+            logger.info("Initializing SpeechT5 model for CPU...")
             
-            # Load the Nari DIA model
-            # Use compute_dtype="float16" for potentially better performance/memory on GPU
-            # Ensure you have a GPU with ~10GB VRAM for this.
-            self.model = Dia.from_pretrained("nari-labs/Dia-1.6B", compute_dtype="float16")
+            self.processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
+            self.model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
+            self.vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
             
-            # Move model to GPU if available
-            if torch.cuda.is_available():
-                self.model.to("cuda")
-                logger.info("Nari DIA model moved to GPU (CUDA).")
-            else:
-                logger.warning("CUDA not available. Nari DIA model will run on CPU, which is not officially supported and will be very slow.")
+            # Ensure all components are on CPU explicitly
+            self.model.to("cpu")
+            self.vocoder.to("cpu")
+            logger.info("SpeechT5 model and vocoder initialized successfully on CPU.")
 
-            logger.info("Nari DIA model initialized successfully.")
+            # Load speaker embeddings for multiple voices
+            logger.info("Loading VCTK dataset for speaker embeddings...")
+            # VCTK is a multi-speaker dataset used with SpeechT5
+            # We'll pick a few representative speaker embeddings for S1, S2, etc.
+            
+            # This loads the 'xvector' split of the vctk dataset which contains pre-computed embeddings
+            embeddings_dataset = load_dataset("Matthijs/cmu-arctic-xvectors", split="validation")
+            
+            # Map 'S1' and 'S2' to specific speaker embeddings from the dataset
+            # You can pick any speaker IDs from the dataset.
+            # Common ones from VCTK for examples are 'p280', 'p272', 'p232', 'p249' etc.
+            # Let's map S1 to a male voice and S2 to a female voice from common VCTK examples.
+            
+            # You can get a list of available speakers from the dataset:
+            # print(embeddings_dataset.features['speaker_id'].names)
+            
+            # Let's use two distinct speakers for S1 and S2
+            # These are common speaker IDs from VCTK used in SpeechT5 examples
+            self.speaker_embeddings["S1"] = torch.tensor(embeddings_dataset[0]["xvector"]).unsqueeze(0) # Speaker p280
+            self.speaker_embeddings["S2"] = torch.tensor(embeddings_dataset[1]["xvector"]).unsqueeze(0) # Speaker p272
+
+            # Ensure embeddings are also on CPU
+            self.speaker_embeddings["S1"] = self.speaker_embeddings["S1"].to("cpu")
+            self.speaker_embeddings["S2"] = self.speaker_embeddings["S2"].to("cpu")
+
+            logger.info("Speaker embeddings loaded for S1 and S2.")
             
         except Exception as e:
-            logger.error(f"Failed to initialize Nari DIA model: {e}", exc_info=True)
+            logger.error(f"Failed to initialize TTS model (SpeechT5): {e}", exc_info=True)
+            self.processor = None
             self.model = None
+            self.vocoder = None
     
     def synthesize_segment(
         self, 
@@ -56,7 +74,7 @@ class NariDIAEngine:
         output_path: str
     ) -> Optional[str]:
         """
-        Synthesize speech for a text segment using Nari DIA.
+        Synthesize speech for a text segment using SpeechT5.
         
         Args:
             text: Text to synthesize
@@ -66,45 +84,45 @@ class NariDIAEngine:
         Returns:
             Path to the generated audio file, or None if failed
         """
-        if not self.model:
-            logger.error("Nari DIA model not initialized. Cannot synthesize speech.")
+        if not self.model or not self.processor or not self.vocoder:
+            logger.error("SpeechT5 model, processor, or vocoder not initialized. Cannot synthesize speech.")
             return None
         
         try:
-            # Nari DIA expects [S1] or [S2] tags.
-            # The segmenter is directly outputting "S1" or "S2".
-            # We just need to wrap it in brackets.
-            if speaker in ["S1", "S2"]:
-                dia_speaker_tag = f"[{speaker}]"
-            else:
-                # Fallback in case segmenter outputs something unexpected
-                logger.warning(f"Unexpected speaker tag '{speaker}' from segmenter. Defaulting to [S1].")
-                dia_speaker_tag = "[S1]"
-
-            # Nari DIA expects the speaker tag at the beginning of the segment
-            full_text_input = f"{dia_speaker_tag} {text}"
+            # Get the correct speaker embedding
+            speaker_embedding = self.speaker_embeddings.get(speaker)
+            if speaker_embedding is None:
+                logger.warning(f"Speaker '{speaker}' not found in pre-loaded embeddings. Defaulting to S1.")
+                speaker_embedding = self.speaker_embeddings["S1"] # Fallback to S1
+            
+            logger.info(f"Synthesizing text for speaker {speaker}: {text[:100]}...")
             
-            # Generate audio using the Nari DIA model
-            logger.info(f"Synthesizing with Nari DIA: {full_text_input[:100]}...") # Log beginning of text
+            # Prepare inputs
+            inputs = self.processor(text=text, return_tensors="pt")
             
-            # Pass the text directly to the model's generate method
-            # Nari DIA's Dia class handles internal processing/tokenization
+            # Ensure inputs are on CPU
+            inputs = {k: v.to("cpu") for k, v in inputs.items()}
+
             with torch.no_grad():
-                # The .generate method should return audio waveform as a PyTorch tensor
-                audio_waveform_tensor = self.model.generate(full_text_input)
-                audio_waveform = audio_waveform_tensor.cpu().numpy().squeeze()
+                # Generate speech
+                # SpeechT5 returns logits/features, which then need to be passed to the vocoder
+                speech = self.model.generate_speech(
+                    inputs["input_ids"],
+                    speaker_embedding, # Pass the speaker embedding here
+                    vocoder=self.vocoder
+                )
+                
+                audio_waveform = speech.cpu().numpy().squeeze()
             
-            # Nari DIA's sampling rate is typically 22050 Hz.
-            # If the Dia model object itself exposes a sampling_rate attribute, use it.
-            # Otherwise, default to 22050 as it's common for TTS models.
-            sampling_rate = getattr(self.model, 'sampling_rate', 22050) 
+            # Sampling rate from the vocoder or model config (typically 16000 for SpeechT5)
+            sampling_rate = self.vocoder.config.sampling_rate if hasattr(self.vocoder.config, 'sampling_rate') else 16000
             
-            # Save as WAV file
             sf.write(output_path, audio_waveform, sampling_rate)
             
-            logger.info(f"Generated audio for {speaker} ({dia_speaker_tag}): {len(text)} characters to {output_path}")
+            logger.info(f"Generated audio for {speaker}: {len(text)} characters to {output_path}")
             return output_path
             
         except Exception as e:
-            logger.error(f"Failed to synthesize segment with Nari DIA: {e}", exc_info=True) # exc_info to print full traceback
+            logger.error(f"Failed to synthesize segment with SpeechT5: {e}", exc_info=True)
             return None
+