diodiogod
f3168e940e
Implemented full support for ChatterBox v2 with enhanced tokenization featuring 30+ special tokens for emotions, sounds, and vocal effects. Key Features: - Model version selection (v1/v2) in engine configuration node - Auto-download of v2 model files (t3_mtl23ls_v2.safetensors, grapheme mappings) - Special tags system using <emotion> syntax to avoid character switching conflicts - Automatic conversion of <giggle>, <sigh>, <whisper> etc. to v2 format - Cache invalidation between v1/v2 to prevent conflicts - Full integration with character switching and pause tags Technical Changes: - Added model_version parameter throughout TTS and SRT processing chain - Created chatterbox_v2_special_tags.py for tag conversion - Updated cache keys to include model_version - Version-aware model loading and validation - Tag conversion happens after character/pause parsing, before engine User can now write: [Alice] Hello! <giggle> Nice to meet you. [pause:0.5] <whisper> Secret! Documentation added in docs/CHATTERBOX_V2_SPECIAL_TOKENS.md
353 lines
16 KiB
Python
353 lines
16 KiB
Python
"""
|
|
Audio Cache Module - Unified caching system for TTS engines
|
|
Provides centralized cache management with engine-specific cache key generation
|
|
"""
|
|
|
|
import hashlib
|
|
import torch
|
|
from typing import Dict, Any, Optional, Tuple, Callable
|
|
from abc import ABC, abstractmethod
|
|
|
|
|
|
# Global audio cache shared across all engines
|
|
GLOBAL_AUDIO_CACHE = {}
|
|
|
|
|
|
class CacheKeyGenerator(ABC):
|
|
"""Abstract base class for engine-specific cache key generation."""
|
|
|
|
@abstractmethod
|
|
def generate_cache_key(self, **params) -> str:
|
|
"""Generate cache key from engine-specific parameters."""
|
|
pass
|
|
|
|
|
|
class F5TTSCacheKeyGenerator(CacheKeyGenerator):
|
|
"""Cache key generator for F5-TTS engine."""
|
|
|
|
def generate_cache_key(self, **params) -> str:
|
|
"""Generate F5-TTS cache key from parameters."""
|
|
cache_data = {
|
|
'text': params.get('text', ''),
|
|
'model_name': params.get('model_name', ''),
|
|
'device': params.get('device', ''),
|
|
'audio_component': params.get('audio_component', ''),
|
|
'ref_text': params.get('ref_text', ''),
|
|
'temperature': params.get('temperature', 0.8),
|
|
'speed': params.get('speed', 1.0),
|
|
'target_rms': params.get('target_rms', 0.1),
|
|
'cross_fade_duration': params.get('cross_fade_duration', 0.15),
|
|
'nfe_step': params.get('nfe_step', 32),
|
|
'cfg_strength': params.get('cfg_strength', 2.0),
|
|
'seed': params.get('seed', 0),
|
|
'character': params.get('character', 'narrator'),
|
|
'engine': 'f5tts'
|
|
}
|
|
cache_string = str(sorted(cache_data.items()))
|
|
return hashlib.md5(cache_string.encode()).hexdigest()
|
|
|
|
|
|
class ChatterBoxCacheKeyGenerator(CacheKeyGenerator):
|
|
"""Cache key generator for ChatterBox engine."""
|
|
|
|
def generate_cache_key(self, **params) -> str:
|
|
"""Generate ChatterBox cache key from parameters."""
|
|
cache_data = {
|
|
'text': params.get('text', ''),
|
|
'exaggeration': params.get('exaggeration', 1.0),
|
|
'temperature': params.get('temperature', 0.8),
|
|
'cfg_weight': params.get('cfg_weight', 1.0),
|
|
'seed': params.get('seed', 0),
|
|
'audio_component': params.get('audio_component', ''),
|
|
'model_source': params.get('model_source', ''),
|
|
'device': params.get('device', ''),
|
|
'language': params.get('language', 'English'),
|
|
'character': params.get('character', 'narrator'),
|
|
'engine': 'chatterbox',
|
|
# ChatterBox Official 23-Lang specific parameters
|
|
'repetition_penalty': params.get('repetition_penalty', 1.0),
|
|
'min_p': params.get('min_p', 0.0),
|
|
'top_p': params.get('top_p', 1.0)
|
|
}
|
|
cache_string = str(sorted(cache_data.items()))
|
|
return hashlib.md5(cache_string.encode()).hexdigest()
|
|
|
|
|
|
class ChatterBoxOfficial23LangCacheKeyGenerator(CacheKeyGenerator):
|
|
"""Cache key generator for ChatterBox Official 23-Lang engine."""
|
|
|
|
def generate_cache_key(self, **params) -> str:
|
|
"""Generate ChatterBox Official 23-Lang cache key from parameters."""
|
|
cache_data = {
|
|
'text': params.get('text', ''),
|
|
'exaggeration': params.get('exaggeration', 0.5),
|
|
'temperature': params.get('temperature', 0.8),
|
|
'cfg_weight': params.get('cfg_weight', 0.5),
|
|
'seed': params.get('seed', 0),
|
|
'audio_component': params.get('audio_component', ''),
|
|
'model_source': params.get('model_source', ''),
|
|
'model_version': params.get('model_version', 'v1'), # v1 or v2 - ensures cache invalidation between versions
|
|
'device': params.get('device', ''),
|
|
'language': params.get('language', 'English'),
|
|
'character': params.get('character', 'narrator'),
|
|
'engine': 'chatterbox_official_23lang',
|
|
# ChatterBox Official 23-Lang specific parameters
|
|
'repetition_penalty': params.get('repetition_penalty', 1.0),
|
|
'min_p': params.get('min_p', 0.0),
|
|
'top_p': params.get('top_p', 1.0)
|
|
}
|
|
cache_string = str(sorted(cache_data.items()))
|
|
return hashlib.md5(cache_string.encode()).hexdigest()
|
|
|
|
|
|
class HiggsAudioCacheKeyGenerator(CacheKeyGenerator):
|
|
"""Cache key generator for Higgs Audio engine."""
|
|
|
|
def generate_cache_key(self, **params) -> str:
|
|
"""Generate Higgs Audio cache key from parameters."""
|
|
cache_data = {
|
|
'text': params.get('text', ''),
|
|
'model_path': params.get('model_path', ''),
|
|
'tokenizer_path': params.get('tokenizer_path', ''),
|
|
'device': params.get('device', ''),
|
|
'reference_text': params.get('reference_text', ''),
|
|
'system_prompt': params.get('system_prompt', ''),
|
|
'temperature': params.get('temperature', 1.0),
|
|
'top_p': params.get('top_p', 0.95),
|
|
'top_k': params.get('top_k', 50),
|
|
'max_new_tokens': params.get('max_new_tokens', 2048),
|
|
'force_audio_gen': params.get('force_audio_gen', False),
|
|
'ras_win_len': params.get('ras_win_len', 7),
|
|
'ras_max_num_repeat': params.get('ras_max_num_repeat', 2),
|
|
'seed': params.get('seed', 0),
|
|
'character': params.get('character', 'narrator'),
|
|
'max_chars_per_chunk': params.get('max_chars_per_chunk', 400),
|
|
# Note: silence_between_chunks_ms excluded - it's post-processing, doesn't affect generation
|
|
'engine': 'higgs_audio'
|
|
}
|
|
|
|
# Add reference audio hash if present
|
|
if 'reference_audio' in params and params['reference_audio'] is not None:
|
|
audio_dict = params['reference_audio']
|
|
if isinstance(audio_dict, dict) and 'waveform' in audio_dict:
|
|
# Create hash of audio waveform
|
|
import hashlib
|
|
waveform = audio_dict['waveform']
|
|
if hasattr(waveform, 'cpu'): # torch tensor
|
|
audio_hash = hashlib.md5(waveform.cpu().numpy().tobytes()).hexdigest()[:8]
|
|
cache_data['audio_hash'] = audio_hash
|
|
|
|
cache_string = str(sorted(cache_data.items()))
|
|
return hashlib.md5(cache_string.encode()).hexdigest()
|
|
|
|
|
|
class VibeVoiceCacheKeyGenerator(CacheKeyGenerator):
|
|
"""Cache key generator for VibeVoice engine."""
|
|
|
|
def generate_cache_key(self, **params) -> str:
|
|
"""Generate VibeVoice cache key from parameters."""
|
|
# Fix floating point precision issues by rounding to 3 decimal places
|
|
cfg_scale = params.get('cfg_scale', 1.3)
|
|
temperature = params.get('temperature', 0.95)
|
|
top_p = params.get('top_p', 0.95)
|
|
|
|
if isinstance(cfg_scale, (int, float)):
|
|
cfg_scale = round(float(cfg_scale), 3)
|
|
if isinstance(temperature, (int, float)):
|
|
temperature = round(float(temperature), 3)
|
|
if isinstance(top_p, (int, float)):
|
|
top_p = round(float(top_p), 3)
|
|
|
|
cache_data = {
|
|
'text': params.get('text', ''),
|
|
'cfg_scale': cfg_scale,
|
|
'temperature': temperature,
|
|
'top_p': top_p,
|
|
'use_sampling': params.get('use_sampling', False),
|
|
'seed': params.get('seed', 42),
|
|
'model_source': params.get('model_source', 'vibevoice-1.5B'),
|
|
'device': params.get('device', 'auto'),
|
|
'max_new_tokens': params.get('max_new_tokens'),
|
|
'multi_speaker_mode': params.get('multi_speaker_mode', 'Custom Character Switching'),
|
|
'audio_component': params.get('audio_component', ''),
|
|
'character': params.get('character', 'narrator'),
|
|
'inference_steps': params.get('inference_steps', 20), # Include diffusion inference steps in cache key
|
|
'attention_mode': params.get('attention_mode', 'eager'), # Include attention mode in cache key
|
|
'quantize': params.get('quantize', False), # Include quantization in cache key
|
|
'engine': 'vibevoice'
|
|
}
|
|
|
|
cache_string = str(sorted(cache_data.items()))
|
|
return hashlib.md5(cache_string.encode()).hexdigest()
|
|
|
|
|
|
class IndexTTSCacheKeyGenerator(CacheKeyGenerator):
|
|
"""Cache key generator for IndexTTS-2 engine."""
|
|
|
|
def generate_cache_key(self, **params) -> str:
|
|
"""Generate IndexTTS-2 cache key from parameters."""
|
|
# Round floating point values to avoid precision issues
|
|
temperature = params.get('temperature', 0.8)
|
|
top_p = params.get('top_p', 0.8)
|
|
emotion_alpha = params.get('emotion_alpha', 1.0)
|
|
|
|
if isinstance(temperature, (int, float)):
|
|
temperature = round(float(temperature), 3)
|
|
if isinstance(top_p, (int, float)):
|
|
top_p = round(float(top_p), 3)
|
|
if isinstance(emotion_alpha, (int, float)):
|
|
emotion_alpha = round(float(emotion_alpha), 3)
|
|
|
|
cache_data = {
|
|
'text': params.get('text', ''),
|
|
'speaker_audio': params.get('speaker_audio', ''),
|
|
'emotion_audio': params.get('emotion_audio', ''),
|
|
'emotion_alpha': emotion_alpha,
|
|
'emotion_vector': params.get('emotion_vector'), # List or None
|
|
'use_emotion_text': params.get('use_emotion_text', False),
|
|
'emotion_text': params.get('emotion_text', ''),
|
|
'use_random': params.get('use_random', False),
|
|
'seed': params.get('seed', 0), # Seed for reproducible generation
|
|
'temperature': temperature,
|
|
'top_p': top_p,
|
|
'top_k': params.get('top_k', 30),
|
|
'length_penalty': params.get('length_penalty', 0.0),
|
|
'repetition_penalty': params.get('repetition_penalty', 10.0),
|
|
'max_mel_tokens': params.get('max_mel_tokens', 1500),
|
|
'max_text_tokens_per_segment': params.get('max_text_tokens_per_segment', 120),
|
|
'interval_silence': params.get('interval_silence', 200),
|
|
'model_name': params.get('model_name', 'IndexTTS-2'),
|
|
'device': params.get('device', 'auto'),
|
|
'character': params.get('character', 'narrator'),
|
|
'engine': 'index_tts'
|
|
}
|
|
|
|
cache_string = str(sorted(cache_data.items()))
|
|
return hashlib.md5(cache_string.encode()).hexdigest()
|
|
|
|
|
|
class AudioCache:
|
|
"""Unified audio cache manager for all TTS engines."""
|
|
|
|
def __init__(self):
|
|
self.cache_key_generators = {
|
|
'f5tts': F5TTSCacheKeyGenerator(),
|
|
'chatterbox': ChatterBoxCacheKeyGenerator(),
|
|
'chatterbox_official_23lang': ChatterBoxOfficial23LangCacheKeyGenerator(), # Uses specialized generator with advanced params
|
|
'higgs_audio': HiggsAudioCacheKeyGenerator(),
|
|
'vibevoice': VibeVoiceCacheKeyGenerator(),
|
|
'index_tts': IndexTTSCacheKeyGenerator()
|
|
}
|
|
|
|
def register_cache_key_generator(self, engine_type: str, generator: CacheKeyGenerator):
|
|
"""Register a cache key generator for a specific engine."""
|
|
self.cache_key_generators[engine_type] = generator
|
|
|
|
def generate_cache_key(self, engine_type: str, **params) -> str:
|
|
"""Generate cache key for specified engine type."""
|
|
if engine_type not in self.cache_key_generators:
|
|
raise ValueError(f"Unknown engine type: {engine_type}")
|
|
|
|
generator = self.cache_key_generators[engine_type]
|
|
return generator.generate_cache_key(**params)
|
|
|
|
def get_cached_audio(self, cache_key: str) -> Optional[Tuple[torch.Tensor, float]]:
|
|
"""Retrieve cached audio by cache key."""
|
|
return GLOBAL_AUDIO_CACHE.get(cache_key)
|
|
|
|
def cache_audio(self, cache_key: str, audio_tensor: torch.Tensor, duration: float):
|
|
"""Cache audio tensor with duration."""
|
|
GLOBAL_AUDIO_CACHE[cache_key] = (audio_tensor.clone(), duration)
|
|
|
|
def create_cache_function(self, engine_type: str, **static_params) -> Callable:
|
|
"""
|
|
Create a cache function for use with TTS generation.
|
|
|
|
Args:
|
|
engine_type: "f5tts" or "chatterbox"
|
|
**static_params: Parameters that don't change between calls
|
|
|
|
Returns:
|
|
Cache function that can be called with (text, audio_result=None)
|
|
"""
|
|
def cache_fn(text_content: str, audio_result=None):
|
|
# Combine static params with dynamic text
|
|
cache_params = static_params.copy()
|
|
cache_params['text'] = f"{cache_params.get('character', 'narrator')}:{text_content}"
|
|
|
|
# Generate cache key
|
|
cache_key = self.generate_cache_key(engine_type, **cache_params)
|
|
|
|
if audio_result is None:
|
|
# Get from cache
|
|
cached_data = self.get_cached_audio(cache_key)
|
|
if cached_data:
|
|
character = cache_params.get('character', 'narrator')
|
|
language = cache_params.get('language', cache_params.get('model_name', ''))
|
|
if language and language != 'English':
|
|
print(f"💾 Using cached audio for '{character}' ({language}): '{text_content[:30]}...'")
|
|
else:
|
|
print(f"💾 Using cached audio for '{character}': '{text_content[:30]}...'")
|
|
return cached_data[0]
|
|
return None
|
|
else:
|
|
# Store in cache
|
|
duration = self._calculate_duration(audio_result, engine_type)
|
|
self.cache_audio(cache_key, audio_result, duration)
|
|
|
|
return cache_fn
|
|
|
|
def _calculate_duration(self, audio_tensor: torch.Tensor, engine_type: str) -> float:
|
|
"""Calculate audio duration based on engine type."""
|
|
if audio_tensor.dim() == 1:
|
|
num_samples = audio_tensor.shape[0]
|
|
elif audio_tensor.dim() == 2:
|
|
num_samples = audio_tensor.shape[1]
|
|
else:
|
|
num_samples = audio_tensor.numel()
|
|
|
|
# Use engine-specific sample rates
|
|
sample_rate = 24000 if engine_type == 'f5tts' else 44100
|
|
return num_samples / sample_rate
|
|
|
|
def clear_cache(self):
|
|
"""Clear all cached audio."""
|
|
global GLOBAL_AUDIO_CACHE
|
|
GLOBAL_AUDIO_CACHE.clear()
|
|
|
|
def get_cache_stats(self) -> Dict[str, Any]:
|
|
"""Get cache statistics."""
|
|
total_items = len(GLOBAL_AUDIO_CACHE)
|
|
total_memory = sum(
|
|
audio.numel() * audio.element_size() + 8 # 8 bytes for duration float
|
|
for audio, _ in GLOBAL_AUDIO_CACHE.values()
|
|
)
|
|
|
|
return {
|
|
'total_items': total_items,
|
|
'total_memory_bytes': total_memory,
|
|
'total_memory_mb': total_memory / (1024 * 1024)
|
|
}
|
|
|
|
|
|
# Global cache instance
|
|
audio_cache = AudioCache()
|
|
|
|
|
|
def get_audio_cache() -> AudioCache:
|
|
"""Get the global audio cache instance."""
|
|
return audio_cache
|
|
|
|
|
|
def create_cache_function(engine_type: str, **static_params) -> Callable:
|
|
"""
|
|
Convenience function to create a cache function.
|
|
|
|
Args:
|
|
engine_type: "f5tts" or "chatterbox"
|
|
**static_params: Parameters that don't change between calls
|
|
|
|
Returns:
|
|
Cache function for use with TTS generation
|
|
"""
|
|
return audio_cache.create_cache_function(engine_type, **static_params) |