Update voice_cloning_engine.py

voice_cloning_engine.py

@@ -6,22 +6,62 @@ from pathlib import Path
 import tempfile
 import os
 import logging
-from typing import Optional, Union, List
+from typing import Optional, Union, List, Tuple
 import warnings
-warnings.filterwarnings('ignore')
-
-# Handle numpy 2.0 compatibility
-if np.__version__.startswith('2.'):
-    import numpy as np
-    # Patch numpy functions if needed
-    np.float_ = np.float64
+import sys
+import importlib.util
 
+# Check if chatterbox is installed
+CHATTERBOX_AVAILABLE = False
 try:
-    from chatterbox import ChatterboxTTS
-    from chatterbox.utils import prepare_text
-except ImportError as e:
-    logging.error(f"Failed to import chatterbox: {e}")
-    logging.error("Please ensure chatterbox is properly installed")
+    # Try different import paths
+    try:
+        from chatterbox import ChatterboxTTS
+        from chatterbox.utils import prepare_text
+        CHATTERBOX_AVAILABLE = True
+        print("✅ Chatterbox TTS imported successfully")
+    except ImportError:
+        try:
+            from chatterbox_tts import ChatterboxTTS
+            from chatterbox_tts.utils import prepare_text
+            CHATTERBOX_AVAILABLE = True
+            print("✅ Chatterbox TTS (chatterbox_tts) imported successfully")
+        except ImportError:
+            print("⚠️ Chatterbox TTS not found, using mock engine for testing")
+except Exception as e:
+    print(f"⚠️ Error importing Chatterbox: {e}")
+
+warnings.filterwarnings('ignore')
+
+# Mock TTS engine for testing when Chatterbox is not available
+class MockChatterboxTTS:
+    """Mock TTS engine for development and testing"""
+    
+    def __init__(self, model_name=None, device=None, use_cache=None):
+        self.model_name = model_name
+        self.device = device or "cpu"
+        self.sample_rate = 24000
+        
+    def synthesize(self, text, reference_audio=None, **kwargs):
+        """Generate mock audio (sine wave) for testing"""
+        import numpy as np
+        
+        # Generate a simple sine wave based on text length
+        duration = min(10, len(text) * 0.1)  # Max 10 seconds
+        t = np.linspace(0, duration, int(self.sample_rate * duration))
+        
+        # Create a simple tone that varies with text
+        frequency = 220 + (hash(text) % 100)  # Base frequency around 220-320 Hz
+        audio = 0.5 * np.sin(2 * np.pi * frequency * t)
+        
+        # Add some harmonics to make it sound more "speech-like"
+        audio += 0.25 * np.sin(4 * np.pi * frequency * t)
+        audio += 0.125 * np.sin(6 * np.pi * frequency * t)
+        
+        # Normalize
+        audio = audio / np.max(np.abs(audio))
+        
+        return audio.astype(np.float32)
 
 class VoiceCloningEngine:
     """
@@ -38,15 +78,38 @@ class VoiceCloningEngine:
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.model_name = model_name
         self.model = None
-        self.sample_rate = 24000  # Default sample rate for chatterbox
-        self.max_chunk_length = 300  # Maximum characters per chunk
-        self._load_model()
+        self.sample_rate = 24000
+        self.max_chunk_length = 300
+        self.is_mock = not CHATTERBOX_AVAILABLE
         
+        try:
+            self._load_model()
+        except Exception as e:
+            logging.error(f"Failed to load model: {e}")
+            if not self.is_mock:
+                print(f"\n❌ Error: Could not load Chatterbox TTS model.")
+                print(f"   Details: {e}")
+                print(f"   Using mock TTS engine for testing.\n")
+                self.is_mock = True
+                self.model = MockChatterboxTTS()
+    
     def _load_model(self):
-        """Load the Chatterbox model"""
+        """Load the Chatterbox model or create mock"""
+        if self.is_mock:
+            logging.info("Using mock TTS engine")
+            self.model = MockChatterboxTTS(
+                self.model_name,
+                device=self.device
+            )
+            return
+            
         try:
             logging.info(f"Loading model {self.model_name} on {self.device}")
-            self.model = ChatterboxTTS(self.model_name, device=self.device)
+            self.model = ChatterboxTTS(
+                self.model_name, 
+                device=self.device,
+                use_cache=True
+            )
             logging.info("Model loaded successfully")
         except Exception as e:
             logging.error(f"Failed to load model: {e}")
@@ -54,159 +117,232 @@ class VoiceCloningEngine:
     
     def _split_text_into_chunks(self, text: str, max_chunk_size: int = 300) -> List[str]:
         """
-        Split text into chunks of approximately max_chunk_size characters
-        
-        Args:
-            text: Input text to split
-            max_chunk_size: Maximum characters per chunk
-            
-        Returns:
-            List of text chunks
+        Split text into chunks intelligently
         """
-        # Split by sentences (simple approach)
-        sentences = text.replace('!', '.').replace('?', '.').split('.')
-        sentences = [s.strip() + '.' for s in sentences if s.strip()]
+        import re
+        
+        # Clean up text
+        text = re.sub(r'\s+', ' ', text).strip()
+        
+        # Try to split by sentences first
+        sentences = re.split(r'(?<=[.!?])\s+', text)
         
         chunks = []
         current_chunk = ""
         
         for sentence in sentences:
-            if len(current_chunk) + len(sentence) <= max_chunk_size:
-                current_chunk += sentence + " "
-            else:
-                if current_chunk:
-                    chunks.append(current_chunk.strip())
-                # If a single sentence is longer than max_chunk_size, split it
-                if len(sentence) > max_chunk_size:
-                    words = sentence.split()
-                    temp_chunk = ""
-                    for word in words:
-                        if len(temp_chunk) + len(word) <= max_chunk_size:
-                            temp_chunk += word + " "
+            # If a single sentence is too long, split it further
+            if len(sentence) > max_chunk_size:
+                # Split by phrases (commas, semicolons, etc.)
+                phrases = re.split(r'(?<=[,;:])\s+', sentence)
+                for phrase in phrases:
+                    if len(current_chunk) + len(phrase) <= max_chunk_size:
+                        if current_chunk:
+                            current_chunk += " " + phrase
+                        else:
+                            current_chunk = phrase
+                    else:
+                        if current_chunk:
+                            chunks.append(current_chunk)
+                        # If phrase itself is too long, split by words
+                        if len(phrase) > max_chunk_size:
+                            words = phrase.split()
+                            temp_chunk = ""
+                            for word in words:
+                                if len(temp_chunk) + len(word) + 1 <= max_chunk_size:
+                                    if temp_chunk:
+                                        temp_chunk += " " + word
+                                    else:
+                                        temp_chunk = word
+                                else:
+                                    if temp_chunk:
+                                        chunks.append(temp_chunk)
+                                    temp_chunk = word
+                            if temp_chunk:
+                                current_chunk = temp_chunk
+                            else:
+                                current_chunk = ""
                         else:
-                            chunks.append(temp_chunk.strip())
-                            temp_chunk = word + " "
-                    if temp_chunk:
-                        chunks.append(temp_chunk.strip())
+                            current_chunk = phrase
+            else:
+                if len(current_chunk) + len(sentence) <= max_chunk_size:
+                    if current_chunk:
+                        current_chunk += " " + sentence
+                    else:
+                        current_chunk = sentence
                 else:
-                    current_chunk = sentence + " "
+                    if current_chunk:
+                        chunks.append(current_chunk)
+                    current_chunk = sentence
         
         if current_chunk:
-            chunks.append(current_chunk.strip())
+            chunks.append(current_chunk)
+        
+        # If no chunks were created (shouldn't happen), split by words
+        if not chunks:
+            words = text.split()
+            chunks = []
+            current_chunk = ""
+            for word in words:
+                if len(current_chunk) + len(word) + 1 <= max_chunk_size:
+                    if current_chunk:
+                        current_chunk += " " + word
+                    else:
+                        current_chunk = word
+                else:
+                    if current_chunk:
+                        chunks.append(current_chunk)
+                    current_chunk = word
+            if current_chunk:
+                chunks.append(current_chunk)
         
         return chunks
     
-    def _merge_audio_chunks(self, audio_chunks: List[np.ndarray]) -> np.ndarray:
-        """
-        Merge multiple audio chunks into a single audio array
-        
-        Args:
-            audio_chunks: List of audio arrays
-            
-        Returns:
-            Merged audio array
-        """
+    def _merge_audio_chunks(self, audio_chunks: List[np.ndarray], crossfade: int = 50) -> np.ndarray:
+        """Merge audio chunks with crossfade"""
         if not audio_chunks:
             return np.array([])
-        return np.concatenate(audio_chunks)
+        
+        if len(audio_chunks) == 1:
+            return audio_chunks[0]
+        
+        merged = audio_chunks[0]
+        
+        for chunk in audio_chunks[1:]:
+            if crossfade > 0 and len(merged) > crossfade and len(chunk) > crossfade:
+                fade_out = np.linspace(1, 0, crossfade)
+                fade_in = np.linspace(0, 1, crossfade)
+                
+                merged[-crossfade:] *= fade_out
+                chunk[:crossfade] *= fade_in
+                
+                merged = np.concatenate([merged, chunk])
+            else:
+                merged = np.concatenate([merged, chunk])
+        
+        # Normalize to prevent clipping
+        max_val = np.max(np.abs(merged))
+        if max_val > 0:
+            merged = merged / max_val * 0.95
+        
+        return merged
     
     def clone_voice(
         self,
         text: str,
-        reference_audio: Union[str, np.ndarray],
+        reference_audio: Union[str, np.ndarray, Tuple[np.ndarray, int]],
         temperature: float = 0.7,
         seed: Optional[int] = None,
         top_p: float = 0.9,
         top_k: int = 50,
         max_chunk_size: int = 300,
-        auto_chunk: bool = True
-    ) -> tuple:
+        auto_chunk: bool = True,
+        crossfade_chunks: bool = True
+    ) -> Tuple[np.ndarray, int]:
         """
         Clone voice and generate speech
-        
-        Args:
-            text: Text to synthesize
-            reference_audio: Path to reference audio or audio array
-            temperature: Sampling temperature
-            seed: Random seed for reproducibility
-            top_p: Top-p sampling parameter
-            top_k: Top-k sampling parameter
-            max_chunk_size: Maximum characters per chunk
-            auto_chunk: Whether to automatically chunk long text
-            
-        Returns:
-            Tuple of (audio_array, sample_rate)
         """
         if self.model is None:
             raise RuntimeError("Model not loaded")
         
-        # Set seed for reproducibility
+        # Set seed
         if seed is not None:
+            np.random.seed(seed)
             torch.manual_seed(seed)
             if torch.cuda.is_available():
                 torch.cuda.manual_seed(seed)
         
-        # Load reference audio if path provided
-        if isinstance(reference_audio, str):
-            audio_array, sr = torchaudio.load(reference_audio)
-            # Resample if necessary
-            if sr != self.sample_rate:
-                resampler = torchaudio.transforms.Resample(sr, self.sample_rate)
-                audio_array = resampler(audio_array)
-            reference_audio = audio_array.squeeze().numpy()
+        # Process reference audio
+        reference_array = self._process_reference_audio(reference_audio)
         
-        # Determine if chunking is needed
+        # Generate speech
         if auto_chunk and len(text) > max_chunk_size:
             chunks = self._split_text_into_chunks(text, max_chunk_size)
             audio_chunks = []
             
-            for chunk in chunks:
-                # Process each chunk
-                audio = self.model.synthesize(
-                    text=chunk,
-                    reference_audio=reference_audio,
-                    temperature=temperature,
-                    top_p=top_p,
-                    top_k=top_k
-                )
-                audio_chunks.append(audio)
+            for i, chunk in enumerate(chunks):
+                logging.info(f"Processing chunk {i+1}/{len(chunks)}: {len(chunk)} chars")
+                
+                try:
+                    audio = self.model.synthesize(
+                        text=chunk,
+                        reference_audio=reference_array,
+                        temperature=temperature,
+                        top_p=top_p,
+                        top_k=top_k
+                    )
+                    
+                    # Ensure audio is numpy array
+                    if torch.is_tensor(audio):
+                        audio = audio.cpu().numpy()
+                    
+                    audio_chunks.append(audio)
+                    
+                except Exception as e:
+                    logging.error(f"Error processing chunk {i+1}: {e}")
+                    continue
+            
+            if not audio_chunks:
+                raise RuntimeError("No audio chunks were generated")
             
-            # Merge chunks
-            final_audio = self._merge_audio_chunks(audio_chunks)
+            final_audio = self._merge_audio_chunks(
+                audio_chunks,
+                crossfade=50 if crossfade_chunks else 0
+            )
         else:
-            # Process as single chunk
-            final_audio = self.model.synthesize(
+            audio = self.model.synthesize(
                 text=text,
-                reference_audio=reference_audio,
+                reference_audio=reference_array,
                 temperature=temperature,
                 top_p=top_p,
                 top_k=top_k
             )
+            
+            if torch.is_tensor(audio):
+                audio = audio.cpu().numpy()
+            
+            final_audio = audio
         
-        return final_audio, self.sample_rate
+        return final_audio.astype(np.float32), self.sample_rate
+    
+    def _process_reference_audio(self, reference_audio):
+        """Process reference audio input"""
+        if isinstance(reference_audio, str):
+            # Load audio file
+            audio_array, sr = torchaudio.load(reference_audio)
+            audio_array = audio_array.squeeze().numpy()
+            
+            # Resample if needed
+            if sr != self.sample_rate:
+                import scipy.signal
+                resample_ratio = self.sample_rate / sr
+                new_length = int(len(audio_array) * resample_ratio)
+                audio_array = scipy.signal.resample(audio_array, new_length)
+            
+            return audio_array
+            
+        elif isinstance(reference_audio, tuple) and len(reference_audio) == 2:
+            audio_array, sr = reference_audio
+            
+            if sr != self.sample_rate:
+                import scipy.signal
+                resample_ratio = self.sample_rate / sr
+                new_length = int(len(audio_array) * resample_ratio)
+                audio_array = scipy.signal.resample(audio_array, new_length)
+            
+            return audio_array
+            
+        else:
+            # Assume it's already a numpy array at correct sample rate
+            return reference_audio
     
     def save_audio(self, audio_array: np.ndarray, sample_rate: int, output_path: str):
-        """
-        Save audio array to file
-        
-        Args:
-            audio_array: Audio data
-            sample_rate: Sample rate
-            output_path: Output file path
-        """
+        """Save audio to file"""
+        os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
         sf.write(output_path, audio_array, sample_rate)
     
     def get_temp_audio_path(self, audio_array: np.ndarray, sample_rate: int) -> str:
-        """
-        Save audio to temporary file and return path
-        
-        Args:
-            audio_array: Audio data
-            sample_rate: Sample rate
-            
-        Returns:
-            Path to temporary audio file
-        """
+        """Save to temp file and return path"""
         with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as f:
             temp_path = f.name
         self.save_audio(audio_array, sample_rate, temp_path)