from __future__ import annotations import argparse import copy import random import tempfile from dataclasses import dataclass from pathlib import Path from typing import Any, Dict, Optional, Tuple import librosa import numpy as np import soundfile as sf import torch import torch.nn.functional as F from chichat.chatterbox.models.s3tokenizer import S3_SR, drop_invalid_tokens from chichat.chatterbox.models.s3gen import S3GEN_SR, S3Gen from chichat.chatterbox.models.t3 import T3 from chichat.chatterbox.models.t3.modules.cond_enc import T3Cond from chichat.chatterbox.models.tokenizers import EnTokenizer from chichat.chatterbox.models.voice_encoder import VoiceEncoder # ---------------------------------------------------------------------------- # CONFIG # ---------------------------------------------------------------------------- DEFAULT_BUNDLE_PATH = Path("tts.pt") DEFAULT_OUTPUT_PATH = Path("output.wav") DEVICE = "cuda" if torch.cuda.is_available() else "cpu" MAX_REF_SECONDS = 10.0 S3GEN_SR = 24000 S3_SR = 16000 # ---------------------------------------------------------------------------- # UTILITIES # ---------------------------------------------------------------------------- def set_seed(seed: int): if seed is None or int(seed) == 0: return seed = int(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) random.seed(seed) np.random.seed(seed) def clone_tensor(x: Optional[torch.Tensor], device=None) -> Optional[torch.Tensor]: if x is None: return None if not torch.is_tensor(x): return x out = x.detach().clone() if device is not None: out = out.to(device) return out def clone_ref_dict(ref_dict: Dict[str, Any], device=None) -> Dict[str, Any]: out: Dict[str, Any] = {} for k, v in ref_dict.items(): if torch.is_tensor(v): t = v.detach().clone() if device is not None: t = t.to(device) out[k] = t else: out[k] = copy.deepcopy(v) return out def normalize_name(name: str) -> str: import re return re.sub(r"[^a-z0-9]+", "", name.strip().lower()) # ---------------------------------------------------------------------------- # CONDITIONALS # ---------------------------------------------------------------------------- @dataclass class Conditionals: t3: T3Cond gen: dict def to(self, device): self.t3 = self.t3.to(device) self.t3.speaker_emb = clone_tensor(self.t3.speaker_emb, device) if getattr(self.t3, "cond_prompt_speech_tokens", None) is not None: self.t3.cond_prompt_speech_tokens = clone_tensor(self.t3.cond_prompt_speech_tokens, device) if getattr(self.t3, "emotion_adv", None) is not None: self.t3.emotion_adv = clone_tensor(self.t3.emotion_adv, device) for k, v in self.gen.items(): if torch.is_tensor(v): self.gen[k] = clone_tensor(v, device) return self # ---------------------------------------------------------------------------- # PACKED TTS # ---------------------------------------------------------------------------- class PackedTTS: def __init__(self, bundle: Dict[str, Any], device: str = DEVICE): self.bundle = bundle self.device = device self.t3: Optional[T3] = None self.s3gen: Optional[S3Gen] = None self.ve: Optional[VoiceEncoder] = None self.tokenizer: Optional[EnTokenizer] = None self.conds: Optional[Conditionals] = None self._tmpdir = tempfile.TemporaryDirectory(prefix="packed_tts_tokenizer_") self._load_models_from_bundle() @classmethod def load(cls, bundle_path: Path, device: str = DEVICE) -> "PackedTTS": bundle = torch.load(bundle_path, map_location="cpu") if not isinstance(bundle, dict): raise ValueError("Packed bundle did not contain a dictionary.") bundle.setdefault("voices", {}) bundle.setdefault("emotions", {}) bundle.setdefault("models", {}) bundle.setdefault("defaults", {}) bundle.setdefault("indexes", {}) return cls(bundle=bundle, device=device) def close(self): try: self._tmpdir.cleanup() except Exception: pass def __del__(self): self.close() # ------------------------------------------------------------------ # Model restore # ------------------------------------------------------------------ def _load_models_from_bundle(self): models = self.bundle.get("models", {}) if not models: raise ValueError("Bundle is missing packed model weights.") t3 = T3() t3.load_state_dict(models["t3_state"]) t3.to(self.device).eval() self.t3 = t3 s3gen = S3Gen() s3gen.load_state_dict(models["s3gen_state"], strict=False) s3gen.to(self.device).eval() self.s3gen = s3gen ve = VoiceEncoder() ve.load_state_dict(models["ve_state"]) ve.to(self.device).eval() self.ve = ve tokenizer_json = models.get("tokenizer_json") if not tokenizer_json: raise ValueError("Bundle is missing tokenizer_json.") tok_path = Path(self._tmpdir.name) / "tokenizer.json" tok_path.write_text(tokenizer_json, encoding="utf-8") self.tokenizer = EnTokenizer(str(tok_path)) # ------------------------------------------------------------------ # Audio extraction helpers # ------------------------------------------------------------------ def _load_reference_audio(self, ref_audio_path: str): wav, _ = librosa.load( ref_audio_path, sr=S3GEN_SR, mono=True, duration=MAX_REF_SECONDS, ) max_len = int(MAX_REF_SECONDS * S3GEN_SR) if len(wav) > max_len: wav = wav[:max_len] return wav def extract_conditionals_from_audio(self, ref_audio_path: str, exaggeration: float = 0.5) -> Dict[str, Any]: wav = self._load_reference_audio(ref_audio_path) with torch.inference_mode(): ref_dict_raw = self.s3gen.embed_ref(wav, S3GEN_SR, device=self.device) wav16k = librosa.resample(wav, orig_sr=S3GEN_SR, target_sr=S3_SR) wav16k = np.asarray(wav16k, dtype=np.float32) embed = self.ve.embeds_from_wavs([wav16k], sample_rate=S3_SR) if isinstance(embed, torch.Tensor): speaker_emb = clone_tensor(embed.mean(dim=0, keepdim=True), self.device) else: speaker_emb = torch.from_numpy(np.asarray(embed)).mean(dim=0, keepdim=True).to(self.device) plen = self.t3.hp.speech_cond_prompt_len tok = None if plen: tokens, _ = self.s3gen.tokenizer.forward([wav16k], max_len=plen) tok = torch.atleast_2d(tokens).clone().to(self.device) ref_dict = clone_ref_dict(ref_dict_raw, device=self.device) emotion_adv = torch.full((1, 1, 1), float(exaggeration), device=self.device) return { "speaker_emb": speaker_emb, "cond_prompt_speech_tokens": tok, "emotion_adv": emotion_adv, "gen": ref_dict, } # ------------------------------------------------------------------ # Resolution helpers # ------------------------------------------------------------------ def list_voices(self): return list(self.bundle.get("voices", {}).keys()) def list_emotions(self): return {k: len(v.get("variations", [])) for k, v in self.bundle.get("emotions", {}).items()} def resolve_voice(self, requested: Optional[str]) -> Tuple[str, Dict[str, Any]]: voices = self.bundle.get("voices", {}) if not voices: raise ValueError("No voices are packed in this bundle.") if not requested: default_voice = self.bundle.get("defaults", {}).get("default_voice") if default_voice and default_voice in voices: return default_voice, voices[default_voice] picked = random.choice(list(voices.keys())) return picked, voices[picked] norm = normalize_name(requested) idx = self.bundle.get("indexes", {}).get("voice_norm", {}) if norm in idx and idx[norm] in voices: name = idx[norm] return name, voices[name] from difflib import get_close_matches matches = get_close_matches(requested, list(voices.keys()), n=1, cutoff=self.bundle.get("defaults", {}).get("fuzzy_cutoff", 0.72)) if matches: name = matches[0] return name, voices[name] picked = random.choice(list(voices.keys())) return picked, voices[picked] def resolve_emotion(self, requested: Optional[str]) -> Tuple[str, Dict[str, Any]]: emotions = self.bundle.get("emotions", {}) if not emotions: raise ValueError("No emotions are packed in this bundle.") if not requested: default_emotion = self.bundle.get("defaults", {}).get("default_emotion") if default_emotion and default_emotion in emotions: emotion_name = default_emotion else: emotion_name = random.choice(list(emotions.keys())) else: norm = normalize_name(requested) idx = self.bundle.get("indexes", {}).get("emotion_norm", {}) if norm in idx and idx[norm] in emotions: emotion_name = idx[norm] else: from difflib import get_close_matches matches = get_close_matches(requested, list(emotions.keys()), n=1, cutoff=self.bundle.get("defaults", {}).get("fuzzy_cutoff", 0.72)) emotion_name = matches[0] if matches else random.choice(list(emotions.keys())) variations = emotions[emotion_name].get("variations", []) if not variations: raise ValueError(f"Emotion '{emotion_name}' has no variations.") return emotion_name, random.choice(variations) # ------------------------------------------------------------------ # Voice/emotion selection logic # ------------------------------------------------------------------ def _resolve_voice_source( self, voice: Optional[str], voice_ref: Optional[str], exaggeration: float, ) -> Tuple[str, Dict[str, Any], Dict[str, Any]]: """Return (voice_name, voice_entry_or_extracted, extracted_conditionals_if_any).""" if voice_ref: extracted = self.extract_conditionals_from_audio(voice_ref, exaggeration=exaggeration) return voice_ref, {"complete": True, **extracted}, extracted voice_name, entry = self.resolve_voice(voice) if entry.get("complete") and entry.get("speaker_emb") is not None: return voice_name, entry, entry raise ValueError( f"Voice '{voice_name}' does not have packed generation conditionals. Provide voice_ref or repack the voice with a sample.wav." ) def _resolve_emotion_source( self, emotion: Optional[str], emo_ref: Optional[str], voice_source_entry: Dict[str, Any], voice_extracted: Dict[str, Any], exaggeration: float, ) -> Tuple[str, Dict[str, Any]]: if emo_ref: extracted = self.extract_conditionals_from_audio(emo_ref, exaggeration=exaggeration) return emo_ref, extracted if emotion: emotion_name, variation = self.resolve_emotion(emotion) return emotion_name, variation # No explicit emotion: prefer the voice's stored emotion if available. if voice_source_entry.get("emotion_adv") is not None: return "voice_default", {"emotion_adv": clone_tensor(voice_source_entry["emotion_adv"], self.device)} # If the voice came from a ref audio, reuse its extracted emotion. if voice_extracted.get("emotion_adv") is not None: return "voice_ref", {"emotion_adv": clone_tensor(voice_extracted["emotion_adv"], self.device)} # Final fallback. return "fallback", {"emotion_adv": torch.full((1, 1, 1), float(exaggeration), device=self.device)} # ------------------------------------------------------------------ # Inference helpers # ------------------------------------------------------------------ def infer_t3(self, text: str, cfg_weight: float, temperature: float): assert self.conds is not None, "Conditionals not prepared." text = text.strip() sot, eot = self.t3.hp.start_text_token, self.t3.hp.stop_text_token tokens = self.tokenizer.text_to_tokens(text).to(self.device) if cfg_weight > 0: tokens = torch.cat([tokens, tokens], dim=0) tokens = F.pad(tokens, (1, 0), value=sot) tokens = F.pad(tokens, (0, 1), value=eot) with torch.inference_mode(): out = self.t3.inference( t3_cond=self.conds.t3, text_tokens=tokens, max_new_tokens=1000, temperature=temperature, cfg_weight=cfg_weight, ) return drop_invalid_tokens(out[0]).to(self.device) def infer_s3gen(self, speech_tokens: torch.Tensor): with torch.inference_mode(): wav, _ = self.s3gen.inference( speech_tokens=speech_tokens, ref_dict=self.conds.gen, ) return wav.squeeze(0).detach().cpu().numpy() # ------------------------------------------------------------------ # Public API # ------------------------------------------------------------------ def generate( self, text: str, voice: Optional[str] = None, emotion: Optional[str] = None, voice_ref: Optional[str] = None, emo_ref: Optional[str] = None, cfg_weight: float = 0.5, temperature: float = 0.8, exaggeration: float = 0.5, seed: int = 0, ): if seed: set_seed(seed) voice_name, voice_entry, voice_extracted = self._resolve_voice_source(voice, voice_ref, exaggeration) emotion_name, emotion_source = self._resolve_emotion_source( emotion=emotion, emo_ref=emo_ref, voice_source_entry=voice_entry, voice_extracted=voice_extracted, exaggeration=exaggeration, ) speaker_emb = voice_entry.get("speaker_emb") if speaker_emb is None: speaker_emb = voice_extracted.get("speaker_emb") speaker_emb = clone_tensor(speaker_emb, self.device) cond_prompt = voice_entry.get("cond_prompt_speech_tokens") if cond_prompt is None: cond_prompt = voice_extracted.get("cond_prompt_speech_tokens") cond_prompt = clone_tensor(cond_prompt, self.device) emotion_adv = emotion_source.get("emotion_adv") emotion_adv = clone_tensor(emotion_adv, self.device) gen = voice_entry.get("gen") if gen is None: gen = voice_extracted.get("gen") if gen is None: gen = {} gen = clone_ref_dict(gen, device=self.device) self.conds = Conditionals( t3=T3Cond( speaker_emb=speaker_emb, cond_prompt_speech_tokens=cond_prompt, emotion_adv=emotion_adv, ), gen=gen, ) tokens = self.infer_t3(text, cfg_weight, temperature) wav = self.infer_s3gen(tokens) return S3GEN_SR, wav, {"voice": voice_name, "emotion": emotion_name} forward = generate # ---------------------------------------------------------------------------- # CLI # ---------------------------------------------------------------------------- def build_parser() -> argparse.ArgumentParser: p = argparse.ArgumentParser(description="Use a packed TTS bundle to generate speech.") p.add_argument("--bundle", type=Path, default=DEFAULT_BUNDLE_PATH) p.add_argument("--text", type=str, default="Hello world, this is a test.") p.add_argument("--voice", type=str, default=None) p.add_argument("--emotion", type=str, default=None) p.add_argument("--voice-ref", type=Path, default=None) p.add_argument("--emo-ref", type=Path, default=None) p.add_argument("--cfg-weight", type=float, default=0.5) p.add_argument("--temperature", type=float, default=0.8) p.add_argument("--exaggeration", type=float, default=0.5) p.add_argument("--seed", type=int, default=42) p.add_argument("--output", type=Path, default=DEFAULT_OUTPUT_PATH) p.add_argument("--list", action="store_true", help="List packed voices and emotions, then exit") return p def main() -> None: args = build_parser().parse_args() tts = PackedTTS.load(args.bundle, device=DEVICE) if args.list: print("Voices:") for name in tts.list_voices(): print(f" - {name}") print("\nEmotions:") for name, count in tts.list_emotions().items(): print(f" - {name} ({count} variations)") return voice_ref = str(args.voice_ref) if args.voice_ref else None emo_ref = str(args.emo_ref) if args.emo_ref else None sr, audio, meta = tts.generate( text=args.text, voice=args.voice, emotion=args.emotion, voice_ref=voice_ref, emo_ref=emo_ref, cfg_weight=args.cfg_weight, temperature=args.temperature, exaggeration=args.exaggeration, seed=args.seed, ) sf.write(str(args.output), audio, sr) print(f"Saved {args.output}") print(f"Resolved voice={meta['voice']} emotion={meta['emotion']}") if __name__ == "__main__": bundle_path = DEFAULT_BUNDLE_PATH output_path = Path("sarah_happy_test.wav") tts = PackedTTS.load(bundle_path, device=DEVICE) sr, audio, meta = tts.generate( text="Hi, this is Sarah speaking with a angry emotion.", voice="Sarah", emotion="Disgust", cfg_weight=0.5, temperature=0.8, exaggeration=0.5, seed=42, ) sf.write(str(output_path), audio, sr) print(f"Saved {output_path}") print(f"Resolved voice={meta['voice']} emotion={meta['emotion']}")