""" Multilingual Engine - Central orchestrator for multilingual TTS generation Handles language switching, character management, and cache optimization for any TTS engine """ import torch from typing import Dict, Any, Optional, List, Tuple, Callable from dataclasses import dataclass from utils.text.character_parser import character_parser from utils.voice.discovery import get_available_characters, get_character_mapping from utils.voice.character_logging import resolved_character_label from utils.text.pause_processor import PauseTagProcessor @dataclass class AudioSegmentResult: """Result of audio generation for a single segment""" audio: torch.Tensor duration: float character: str text: str language: str original_index: int edit_tags: list = None # Edit tags for this specific segment @dataclass class MultilingualResult: """Complete result of multilingual processing""" audio: torch.Tensor total_duration: float segments: List[AudioSegmentResult] languages_used: List[str] characters_used: List[str] info_message: str class MultilingualEngine: """ Central orchestrator for multilingual TTS generation. Handles: - Language-aware character parsing - Smart model loading with cache optimization - Language grouping for efficient processing - Proper segment ordering and audio assembly """ def __init__(self, engine_type: str): """ Initialize multilingual engine. Args: engine_type: "f5tts" or "chatterbox" """ self.engine_type = engine_type self.sample_rate = 24000 if engine_type == "f5tts" else 44100 self.loaded_models = set() # Track loaded models across calls def process_multilingual_text(self, text: str, engine_adapter, **params) -> MultilingualResult: """ Main entry point for multilingual processing. Args: text: Input text with character/language tags engine_adapter: Engine-specific adapter (F5TTSEngineAdapter or ChatterBoxEngineAdapter) **params: Engine-specific parameters Returns: MultilingualResult with generated audio and metadata """ # 1. Parse character segments with languages from original text character_segments_with_lang = character_parser.split_by_character_with_language(text) # Get detailed segments to access original character information detailed_segments = character_parser.parse_text_segments(text) # Extract edit tags per segment and handle pause-splitting if needed segment_edit_tags = [] if self.engine_type == "chatterbox": # Check if we need edit tag extraction (only for classic ChatterBox, not official 23-lang) extract_edit_for_chatterbox = params.get('extract_edit_tags', False) if extract_edit_for_chatterbox: from utils.text.step_audio_editx_special_tags import parse_edit_tags_with_iterations from dataclasses import replace # First pass: extract edit tags and check for pause tags new_detailed_segments = [] for seg_idx, seg in enumerate(detailed_segments): # Extract edit tags from this segment's original text seg_clean_text, seg_edits = parse_edit_tags_with_iterations(seg.text) # Check if this segment has BOTH pause tags and edit tags if seg_edits and PauseTagProcessor.has_pause_tags(seg_clean_text): # Split by pause tags FIRST to preserve pauses during edit processing pause_segments, _ = PauseTagProcessor.parse_pause_tags(seg_clean_text) # Create new segments for each pause-split part for pause_seg_type, pause_content in pause_segments: if pause_seg_type == 'text': # Text segment - create new detailed segment new_seg = replace(seg, text=pause_content) new_detailed_segments.append(new_seg) segment_edit_tags.append(seg_edits) # This text segment can be edited elif pause_seg_type == 'pause': # Pause segment - create silence marker pause_marker_seg = replace(seg, text=f"__PAUSE_{pause_content}__") new_detailed_segments.append(pause_marker_seg) segment_edit_tags.append([]) # Silence has no edit tags else: # No pause or no edit tags - keep segment as-is seg.text = seg_clean_text new_detailed_segments.append(seg) segment_edit_tags.append(seg_edits) # Replace detailed_segments with pause-split version detailed_segments = new_detailed_segments else: segment_edit_tags = [[] for _ in detailed_segments] else: segment_edit_tags = [[] for _ in detailed_segments] # 2. Analyze segments characters = list(set(char for char, _, _ in character_segments_with_lang)) languages = list(set(lang for _, _, lang in character_segments_with_lang)) has_multiple_characters = len(characters) > 1 or (len(characters) == 1 and characters[0] != "narrator") has_multiple_languages = len(languages) > 1 # Print analysis if has_multiple_languages: print(f"🌍 {self.engine_type.title()}: Language switching mode - found languages: {', '.join(languages)}") if has_multiple_characters: print(f"🎭 {self.engine_type.title()}: Character switching mode - found characters: {', '.join(characters)}") # 3. Group segments by language to optimize model loading language_groups = self._group_segments_by_language_with_original(detailed_segments) # 4. Get character mapping for all characters character_mapping = get_character_mapping(characters, engine_type=self.engine_type) # 5. Check cache optimization opportunities cache_info = self._analyze_cache_coverage(language_groups, character_mapping, engine_adapter, **params) # 6. Process each language group with smart model loading all_audio_segments = [] base_model_loaded = False # Initialize loaded_models tracking with currently loaded model (if any) current_model = getattr(engine_adapter.node, 'current_model_name', None) if current_model and current_model not in self.loaded_models: self.loaded_models.add(current_model) print(f"💾 Multilingual engine: Detected already loaded model '{current_model}'") # Use persistent loaded_models tracking across calls for lang_code, lang_segments in language_groups.items(): # Get required model for this language required_model = engine_adapter.get_model_for_language(lang_code, params.get("model", "default")) # Check if all segments in this language group are cached if cache_info.get(lang_code, {}).get("all_cached", False): print(f"💾 Skipping model load for language '{lang_code}' - all {len(lang_segments)} segments cached") else: # Load base model if this is the first time we need to generate anything if not base_model_loaded: # CRITICAL FIX: Load the required model for this language, not the default model # This ensures the correct tokenizer is loaded print(f"🔄 Loading base {self.engine_type.title()} model for generation ({required_model})") engine_adapter.load_base_model(required_model, params.get("device", "auto")) base_model_loaded = True # Only load language model if we haven't loaded this specific model yet if required_model not in self.loaded_models: print(f"🌍 Loading {self.engine_type.title()} model '{required_model}' for language '{lang_code}' ({len(lang_segments)} segments)") try: engine_adapter.load_language_model(required_model, params.get("device", "auto")) self.loaded_models.add(required_model) # Track that this model is now loaded # Update node's current language state engine_adapter.node.current_language = required_model engine_adapter.node.current_model_name = required_model print(f"🔄 Updated current_language to '{required_model}'") except Exception as e: print(f"⚠️ Failed to load model '{required_model}' for language '{lang_code}': {e}") print(f"🔄 Falling back to English model") engine_adapter.load_base_model("English", params.get("device", "auto")) else: print(f"💾 Model '{required_model}' already loaded - reusing for language '{lang_code}' ({len(lang_segments)} segments)") # CRITICAL FIX: Force model manager to switch to the correct cached model instance # The model is cached but self.tts_model needs to point to the right instance engine_adapter.node.load_tts_model(params.get("device", "auto"), required_model) engine_adapter.node.current_language = required_model engine_adapter.node.current_model_name = required_model print(f"🔄 Updated current_language to '{required_model}' (cached)") # Process each segment in this language group for segment_data in lang_segments: original_idx, character, segment_text, segment_lang, original_character = segment_data segment_display_idx = original_idx + 1 # For display (1-based) # Check if this is a pause marker segment if segment_text.startswith("__PAUSE_") and segment_text.endswith("__"): # Extract pause duration and create silence pause_duration = float(segment_text[8:-2]) # Extract from __PAUSE_X__ silence = PauseTagProcessor.create_silence_segment(pause_duration, self.sample_rate) # Get edit tags for this segment (should be empty for pause) seg_edit_tags = segment_edit_tags[original_idx] if original_idx < len(segment_edit_tags) else [] # Store silence segment all_audio_segments.append(AudioSegmentResult( audio=silence, duration=pause_duration, character=character, text="", # Empty text for silence language=segment_lang, original_index=original_idx, edit_tags=seg_edit_tags )) continue # Skip TTS generation for pause segments # Get character voice or fallback to main if self.engine_type == "f5tts": char_audio, char_text = character_mapping.get(character, (None, None)) if not char_audio or not char_text: char_audio = params.get("main_audio_reference") char_text = params.get("main_text_reference") else: # chatterbox # CRITICAL FIX: Check if original character (before alias resolution) was "narrator" # to detect language-only tags like [de:] that should use main voice main_ref = params.get("main_audio_reference") # If the original tag was language-only (defaulted to "narrator") and we have main ref, # prioritize main reference over any character alias mapping should_use_main_ref = (original_character == "narrator" and main_ref) if should_use_main_ref: # Language-only tag like [de:] - use main narrator voice (user's selected voice) print(f"✅ Using main narrator voice (user-selected) for language-only tag in {segment_lang}") char_audio = main_ref else: # Explicit character tag like [de:Alice] - use character voice char_audio_tuple = character_mapping.get(character, (None, None)) if char_audio_tuple[0]: char_audio = char_audio_tuple[0] # Only get the audio path else: char_audio = main_ref # Show generation message with character and language info # Check if we're using main voice for narrator (language-only tags) is_using_main_voice = (self.engine_type == "chatterbox" and original_character == "narrator" and params.get("main_audio_reference")) if is_using_main_voice: # Language-only tag using main voice if segment_lang != 'en': print(f"🎤 Generating {self.engine_type.title()} segment {segment_display_idx} using main voice in {segment_lang}...") else: print(f"🎤 Generating {self.engine_type.title()} segment {segment_display_idx} using main voice...") elif character == "narrator": if segment_lang != 'en': print(f"🎤 Generating {self.engine_type.title()} segment {segment_display_idx} in {segment_lang}...") else: print(f"🎤 Generating {self.engine_type.title()} segment {segment_display_idx}...") else: if segment_lang != 'en': print(f"🎭 Generating {self.engine_type.title()} segment {segment_display_idx} using '{resolved_character_label(character, char_audio)}' in {segment_lang}") else: print(f"🎭 Generating {self.engine_type.title()} segment {segment_display_idx} using '{resolved_character_label(character, char_audio)}'") # Show what model is actually being used for generation (for verification) current_model = getattr(engine_adapter.node, 'current_language', 'unknown') print(f"🔧 ACTUAL MODEL: Generating segment {segment_display_idx} using '{current_model}' model") # Show the final text that will go to the TTS model print(f"🔤 Final text to {self.engine_type.upper()} via multilingual engine ({resolved_character_label(character, char_audio)}): '{segment_text}'") # CRITICAL FIX: For language-only tags, use "narrator" as character for cache consistency cache_character = character if self.engine_type == "chatterbox" and original_character == "narrator": cache_character = "narrator" # CRITICAL FIX: Update current_language in params to match loaded model # The multilingual engine loads models but the adapter needs the updated language updated_params = params.copy() updated_params['current_language'] = getattr(engine_adapter.node, 'current_language', segment_lang) # CRITICAL FIX: Handle pause tags within character segments # This ensures pause changes don't invalidate cache for text content # (PauseTagProcessor already imported at top of file) if PauseTagProcessor.has_pause_tags(segment_text): # Process pause tags for this character segment if self.engine_type == "f5tts": segment_audio = engine_adapter.generate_segment_audio( text=segment_text, # Let adapter handle pause tags internally char_audio=char_audio, char_text=char_text, character=cache_character, enable_pause_tags=True, # Enable pause processing in adapter **updated_params ) else: # chatterbox segment_audio = engine_adapter.generate_segment_audio( text=segment_text, # Let adapter handle pause tags internally char_audio=char_audio, character=cache_character, enable_pause_tags=True, # Enable pause processing in adapter **updated_params ) else: # No pause tags, standard generation if self.engine_type == "f5tts": segment_audio = engine_adapter.generate_segment_audio( text=segment_text, char_audio=char_audio, char_text=char_text, character=cache_character, **updated_params ) else: # chatterbox segment_audio = engine_adapter.generate_segment_audio( text=segment_text, char_audio=char_audio, character=cache_character, **updated_params ) # Calculate duration duration = self._get_audio_duration(segment_audio) # Get edit tags for this segment seg_edit_tags = segment_edit_tags[original_idx] if original_idx < len(segment_edit_tags) else [] # Store result with original index for proper ordering all_audio_segments.append(AudioSegmentResult( audio=segment_audio, duration=duration, character=character, text=segment_text, language=segment_lang, original_index=original_idx, edit_tags=seg_edit_tags )) # 7. Reorder segments back to original order and combine all_audio_segments.sort(key=lambda x: x.original_index) ordered_audio = [seg.audio for seg in all_audio_segments] combined_audio = torch.cat(ordered_audio, dim=1) if ordered_audio else torch.zeros(1, 0) # 8. Calculate total duration and create info message total_duration = sum(seg.duration for seg in all_audio_segments) info_message = self._generate_info_message( total_duration, len(all_audio_segments), characters, languages, has_multiple_characters, has_multiple_languages ) return MultilingualResult( audio=combined_audio, total_duration=total_duration, segments=all_audio_segments, languages_used=languages, characters_used=characters, info_message=info_message ) def _group_segments_by_language(self, character_segments_with_lang: List[Tuple[str, str, str]]) -> Dict[str, List[Tuple[int, str, str, str]]]: """Group character segments by language for efficient processing.""" language_groups = {} for original_idx, (character, segment_text, segment_lang) in enumerate(character_segments_with_lang): if segment_lang not in language_groups: language_groups[segment_lang] = [] language_groups[segment_lang].append((original_idx, character, segment_text, segment_lang)) return language_groups def _group_segments_by_language_with_original(self, segments: List) -> Dict[str, List[Tuple[int, str, str, str, str]]]: """Group segments by language for efficient processing, preserving original character info.""" language_groups = {} for original_idx, segment in enumerate(segments): segment_lang = segment.language if segment_lang not in language_groups: language_groups[segment_lang] = [] language_groups[segment_lang].append(( original_idx, segment.character, segment.text, segment_lang, segment.original_character or segment.character )) return language_groups def _analyze_cache_coverage(self, language_groups: Dict, character_mapping: Dict, engine_adapter, **params) -> Dict[str, Dict[str, Any]]: """Analyze cache coverage for each language group to optimize model loading.""" cache_info = {} # Only check cache if caching is enabled enable_cache = params.get("enable_audio_cache", True) if not enable_cache: for lang_code, lang_segments in language_groups.items(): cache_info[lang_code] = { "all_cached": False, "segments_count": len(lang_segments) } return cache_info for lang_code, lang_segments in language_groups.items(): cached_segments = 0 total_segments = len(lang_segments) for segment_data in lang_segments: original_idx, character, segment_text, segment_lang = segment_data[:4] # Take only first 4 elements # Check if this specific segment is cached if self._is_segment_cached(character, segment_text, segment_lang, character_mapping, **params): cached_segments += 1 all_cached = cached_segments == total_segments cache_info[lang_code] = { "all_cached": all_cached, "cached_segments": cached_segments, "segments_count": total_segments } if all_cached: print(f"💾 Cache optimization: All {total_segments} segments in '{lang_code}' are cached") return cache_info def _is_segment_cached(self, character: str, text: str, language: str, character_mapping: Dict, **params) -> bool: """Check if a specific segment is cached.""" try: # Import cache function from utils.audio.cache import create_cache_function # Get character voice information if self.engine_type == "f5tts": char_audio, char_text = character_mapping.get(character, (None, None)) if not char_audio or not char_text: char_audio = params.get("main_audio_reference") char_text = params.get("main_text_reference", "") # CRITICAL FIX: Use main narrator's stable component instead of generic fallback main_stable_component = params.get("stable_audio_component", "main_reference") audio_component = char_audio or main_stable_component ref_text_component = char_text or "" else: # chatterbox char_audio_tuple = character_mapping.get(character, (None, None)) char_audio = char_audio_tuple[0] if char_audio_tuple[0] else params.get("main_audio_reference") # CRITICAL FIX: Use main narrator's stable component instead of generic fallback main_stable_component = params.get("stable_audio_component", "main_reference") audio_component = char_audio or main_stable_component ref_text_component = "" # CRITICAL FIX: For language-only tags using main narrator voice, cache as "narrator" # This prevents cache pollution when narrator voice changes cache_character = character if audio_component == main_stable_component: # This character is using main narrator voice (language-only tag) cache_character = "narrator" # Create cache function to check if segment exists cache_fn = create_cache_function( text_content=text, audio_component=str(audio_component), ref_text_component=ref_text_component, character=cache_character, language=language, model_name=params.get("model", "default"), temperature=params.get("temperature", 0.8), speed=params.get("speed", 1.0) if self.engine_type == "f5tts" else None, nfe_step=params.get("nfe_step", 32) if self.engine_type == "f5tts" else None, cfg_strength=params.get("cfg_strength", 2.0) if self.engine_type == "f5tts" else None, exaggeration=params.get("exaggeration", 0.5) if self.engine_type == "chatterbox" else None, cfg_weight=params.get("cfg_weight", 0.5) if self.engine_type == "chatterbox" else None ) # Check cache - if it returns data, segment is cached cached_data = cache_fn(text, audio_result=None) return cached_data is not None except Exception as e: # If cache checking fails, assume not cached return False def _get_audio_duration(self, audio_tensor: torch.Tensor) -> float: """Calculate audio duration in seconds.""" if audio_tensor.dim() == 1: num_samples = audio_tensor.shape[0] elif audio_tensor.dim() == 2: num_samples = audio_tensor.shape[1] # Assume shape (channels, samples) else: num_samples = audio_tensor.numel() return num_samples / self.sample_rate def _generate_info_message(self, total_duration: float, num_segments: int, characters: List[str], languages: List[str], has_multiple_characters: bool, has_multiple_languages: bool) -> str: """Generate descriptive info message about the generation.""" character_info = f"characters: {', '.join(characters)}" if has_multiple_characters else "narrator" language_info = f" across {len(languages)} languages ({', '.join(languages)})" if has_multiple_languages else "" return f"Generated {total_duration:.1f}s audio from {num_segments} segments using {character_info}{language_info} ({self.engine_type.title()} models)" def is_multilingual_or_multicharacter(self, text: str) -> bool: """Quick check if text needs multilingual processing.""" character_segments_with_lang = character_parser.split_by_character_with_language(text) characters = list(set(char for char, _, _ in character_segments_with_lang)) languages = list(set(lang for _, _, lang in character_segments_with_lang)) has_multiple_characters = len(characters) > 1 or (len(characters) == 1 and characters[0] != "narrator") has_multiple_languages = len(languages) > 1 return has_multiple_characters or has_multiple_languages