Chatterbox_tts_long_handling / chatterbox_engine.py
codewithjarair's picture
Upload 6 files
332cb9e verified
Raw
History Blame Contribute Delete
3.72 kB
import torch
from chatterbox.tts import ChatterboxTTS
import numpy as np
import os
import tempfile
import sys
import types
# --- Perth Watermarker Monkeypatch ---
# Chatterbox requires resemble-perth, but sometimes 'perth' package is installed instead.
# This patch prevents the AttributeError if the correct attribute is missing.
try:
import perth
if not hasattr(perth, 'PerthImplicitWatermarker'):
class MockWatermarker:
def apply_watermark(self, wav, sample_rate): return wav
def get_watermark(self, wav, sample_rate): return 0.0
perth.PerthImplicitWatermarker = MockWatermarker
except ImportError:
perth = types.ModuleType("perth")
class MockWatermarker:
def apply_watermark(self, wav, sample_rate): return wav
def get_watermark(self, wav, sample_rate): return 0.0
perth.PerthImplicitWatermarker = MockWatermarker
sys.modules["perth"] = perth
# -------------------------------------
class ChatterboxEngine:
def __init__(self, model_type="voice_clone", device=None):
self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
self.model_type = model_type
print(f"Initializing ChatterboxEngine on {self.device} with model type: {model_type}...")
# Load the Chatterbox model (English only for voice cloning)
self.model = ChatterboxTTS.from_pretrained(device=self.device)
# Available voices (English only for voice cloning)
self.voices = {
"English": ["default", "male_1", "male_2", "female_1", "female_2", "child_1", "narrator", "announcer"]
}
# Custom voices storage (for cloned voices)
self.custom_voices = {}
def get_voice_list(self):
"""Get list of all available voices including custom ones."""
all_voices = []
for category in self.voices.values():
all_voices.extend(category)
all_voices.extend(self.custom_voices.keys())
return all_voices
def clone_voice(self, audio_path, voice_name):
"""
Clone a voice from an audio file.
For ChatterboxTTS, we just store the path for use with audio_prompt_path.
"""
try:
# Store the audio path for voice cloning
self.custom_voices[voice_name] = audio_path
print(f"Voice '{voice_name}' registered for cloning")
return voice_name
except Exception as e:
print(f"Error cloning voice: {str(e)}")
return None
def generate(self, text, voice="default", speed=1.0, lang='en', custom_voice_path=None, exaggeration=0.5, cfg_weight=0.5, seed=None, temperature=1.0):
"""
Generates audio from text using a specified voice.
"""
try:
# Check if this is a custom voice
if voice in self.custom_voices:
# Use stored audio path for voice cloning
audio = self.model.generate(text, audio_prompt_path=self.custom_voices[voice])
elif custom_voice_path:
# Use provided audio path for one-shot cloning
audio = self.model.generate(text, audio_prompt_path=custom_voice_path)
# Ensure audio is a flat numpy array for soundfile compatibility
if torch.is_tensor(audio):
audio = audio.detach().cpu().numpy().flatten()
elif isinstance(audio, np.ndarray):
audio = audio.flatten()
sample_rate = self.model.sr
return audio, sample_rate
except Exception as e:
print(f"Error generating audio: {str(e)}")
return None, 22050