""" Voice cloning module — wraps Qwen3-TTS for zero-shot voice cloning. Supports two backends: - Base 1.7B: zero-shot voice cloning from reference audio - CustomVoice 0.6B: fast predefined speakers (no cloning, lower latency) Voice profiles are: - Cached in-memory for fast access during a session - Persisted to Voice_Profile/ folder as .pt files for reuse across restarts - Auto-loaded on startup if saved profiles exist Latency optimizations applied: - bfloat16 / float16 precision (auto-detected per GPU arch) - FlashAttention-2 when available - Streaming mode (non_streaming_mode=False) - Reduced sampling params (top_k=20, temperature=0.7) - Capped max_new_tokens=1024 - torch.set_float32_matmul_precision('high') - Reference audio trimmed to 3-5s for faster embedding extraction Usage: profile_id = create_voice_profile(ref_audio_path, voice_name="Mom") wav, sr = synthesize_cloned(text, profile_id) """ from __future__ import annotations import json import logging import os import uuid import threading from pathlib import Path import numpy as np import soundfile as sf import torch from runtime_config import GPU_INFERENCE_LOCK, VOICE_PROFILE_DIR logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Global PyTorch optimizations # --------------------------------------------------------------------------- torch.set_float32_matmul_precision("high") # --------------------------------------------------------------------------- # Model configuration # --------------------------------------------------------------------------- # Base model for zero-shot voice cloning (1.7B) — used for BOTH cloned and stock voice BASE_MODEL_ID = "Qwen/Qwen3-TTS-12Hz-1.7B-Base" # Stock voice reference audio (pre-generated "vivian" sample) VIVIAN_REF_PATH = Path(__file__).parent / "assets" / "vivian_reference.wav" # Profile ID for the built-in stock voice STOCK_VOICE_PROFILE_ID = "__stock_vivian__" # Optimized generation parameters (reduced from defaults: top_k=50, temp=0.9, max=2048) GENERATION_PARAMS = dict( top_k=20, temperature=0.7, subtalker_top_k=20, subtalker_temperature=0.7, max_new_tokens=1024, ) # Reference audio limits (seconds) — 3-5s is optimal for Qwen3-TTS REF_AUDIO_MIN_SEC = 3.0 REF_AUDIO_MAX_SEC = 10.0 REF_AUDIO_TARGET_SR = 24000 # --------------------------------------------------------------------------- # Voice profile persistence # --------------------------------------------------------------------------- VOICE_PROFILE_DIR.mkdir(exist_ok=True) # Default profile ID — used when no cloned voice exists DEFAULT_PROFILE_ID = "__default__" # --------------------------------------------------------------------------- # Server-side cache: { profile_id -> VoiceClonePromptItem list } # --------------------------------------------------------------------------- _PROFILE_CACHE: dict[str, list] = {} _cache_lock = threading.Lock() _qwen_tts_model = None _model_lock = threading.Lock() def _select_dtype() -> torch.dtype: """Pick optimal dtype based on GPU architecture.""" if not torch.cuda.is_available(): return torch.float32 cap = torch.cuda.get_device_capability() # bfloat16 requires compute capability >= 8.0 (Ampere+) if cap[0] >= 8: return torch.bfloat16 return torch.float16 def _select_attn_impl() -> str: """Use FlashAttention-2 if available, else SDPA (PyTorch native).""" try: import flash_attn # noqa: F401 return "flash_attention_2" except ImportError: logger.info("flash-attn not installed — using SDPA attention.") return "sdpa" def get_qwen_tts(): """Lazy-load Qwen3-TTS Base 1.7B model for cloning. Thread-safe.""" global _qwen_tts_model if _qwen_tts_model is None: with _model_lock: if _qwen_tts_model is None: from qwen_tts import Qwen3TTSModel model_id = BASE_MODEL_ID device = "cuda" if torch.cuda.is_available() else "cpu" logger.info("Loading %s on %s (cloning model)...", model_id, device) attn_impl = _select_attn_impl() _qwen_tts_model = Qwen3TTSModel.from_pretrained( model_id, device_map=device, attn_implementation=attn_impl, ) logger.info("Qwen3-TTS Base loaded on %s (attn=%s).", device, attn_impl) return _qwen_tts_model def _try_torch_compile(wrapper): """Best-effort torch.compile on model submodules. Disabled for stability.""" # torch.compile can cause CUDA asserts on some GPU architectures (T4/Turing) # Disable for now in favor of stability return def _trim_reference_audio(audio_path: str) -> str: """ Trim reference audio to REF_AUDIO_MAX_SEC seconds if longer. Returns path to trimmed file (or original if already short enough). """ try: info = sf.info(audio_path) duration = info.duration if duration <= REF_AUDIO_MAX_SEC: return audio_path logger.info( "Reference audio %.1fs exceeds %.0fs limit — trimming.", duration, REF_AUDIO_MAX_SEC, ) data, sr = sf.read(audio_path) max_samples = int(REF_AUDIO_MAX_SEC * sr) trimmed = data[:max_samples] trimmed_path = audio_path + ".trimmed.wav" sf.write(trimmed_path, trimmed, sr) return trimmed_path except Exception as exc: logger.warning("Could not trim reference audio: %s", exc) return audio_path # --------------------------------------------------------------------------- # Profile persistence (disk ↔ memory) # --------------------------------------------------------------------------- def save_profile_to_disk(profile_id: str, voice_name: str = "Cloned Voice") -> Path: """ Save a cached voice profile to Voice_Profile/ as a .pt file + metadata JSON. Returns the path to the saved .pt file. """ with _cache_lock: prompt_items = _PROFILE_CACHE.get(profile_id) if prompt_items is None: raise ValueError(f"Profile '{profile_id}' not found in cache.") profile_dir = VOICE_PROFILE_DIR / profile_id profile_dir.mkdir(exist_ok=True) # Serialize VoiceClonePromptItem fields serializable = [] for item in prompt_items: serializable.append({ "ref_code": item.ref_code.cpu() if item.ref_code is not None else None, "ref_spk_embedding": item.ref_spk_embedding.cpu(), "x_vector_only_mode": item.x_vector_only_mode, "icl_mode": item.icl_mode, "ref_text": item.ref_text, }) pt_path = profile_dir / "profile.pt" torch.save(serializable, pt_path) # Save metadata meta = {"profile_id": profile_id, "voice_name": voice_name} meta_path = profile_dir / "metadata.json" with open(meta_path, "w", encoding="utf-8") as f: json.dump(meta, f, indent=2) logger.info("Voice profile '%s' (%s) saved to %s", profile_id, voice_name, profile_dir) return pt_path def load_profile_from_disk(profile_id: str) -> bool: """ Load a voice profile from Voice_Profile//profile.pt into memory cache. Returns True if loaded successfully, False otherwise. """ profile_dir = VOICE_PROFILE_DIR / profile_id pt_path = profile_dir / "profile.pt" if not pt_path.exists(): logger.warning("Profile file not found: %s", pt_path) return False try: from qwen_tts.inference.qwen3_tts_model import VoiceClonePromptItem raw_items = torch.load(pt_path, map_location="cpu", weights_only=False) prompt_items = [] for item_dict in raw_items: prompt_items.append(VoiceClonePromptItem( ref_code=item_dict["ref_code"], ref_spk_embedding=item_dict["ref_spk_embedding"], x_vector_only_mode=item_dict["x_vector_only_mode"], icl_mode=item_dict["icl_mode"], ref_text=item_dict.get("ref_text"), )) with _cache_lock: _PROFILE_CACHE[profile_id] = prompt_items logger.info("Voice profile '%s' loaded from disk.", profile_id) return True except Exception as exc: logger.exception("Failed to load profile '%s': %s", profile_id, exc) return False def list_saved_profiles() -> list[dict]: """ List all saved voice profiles from Voice_Profile/ directory. Returns list of {profile_id, voice_name, path} dicts, newest first. """ profiles = [] if not VOICE_PROFILE_DIR.exists(): return profiles for entry in VOICE_PROFILE_DIR.iterdir(): if not entry.is_dir(): continue pt_path = entry / "profile.pt" meta_path = entry / "metadata.json" if not pt_path.exists(): continue voice_name = "Cloned Voice" if meta_path.exists(): try: with open(meta_path, encoding="utf-8") as f: meta = json.load(f) voice_name = meta.get("voice_name", voice_name) except Exception: pass profiles.append({ "profile_id": entry.name, "voice_name": voice_name, "path": str(entry), "mtime": pt_path.stat().st_mtime, }) # Newest first profiles.sort(key=lambda p: p["mtime"], reverse=True) return profiles def load_default_profile() -> str | None: """ Load the most recently saved voice profile from Voice_Profile/ into memory. Returns the profile_id if found, None if no saved profiles exist. This is called on app startup to restore the last cloned voice. """ saved = list_saved_profiles() if not saved: logger.info("No saved voice profiles found — using stock voice as default.") return None newest = saved[0] profile_id = newest["profile_id"] if load_profile_from_disk(profile_id): logger.info( "Default voice profile loaded: '%s' (%s)", profile_id, newest["voice_name"], ) return profile_id return None # --------------------------------------------------------------------------- # Core API # --------------------------------------------------------------------------- def create_voice_profile(ref_audio_path: str, voice_name: str = "Cloned Voice", profile_id_override: str | None = None) -> str: """ Extract speaker embedding from reference audio, cache it, and save to disk. Returns a profile_id string for later synthesis. Reference audio is trimmed to 3-10s for optimal latency. """ trimmed_path = _trim_reference_audio(ref_audio_path) model = get_qwen_tts() logger.info("Creating voice profile from %s...", ref_audio_path) with GPU_INFERENCE_LOCK: prompt_items = model.create_voice_clone_prompt( ref_audio=trimmed_path, x_vector_only_mode=True, ) profile_id = profile_id_override or uuid.uuid4().hex[:12] with _cache_lock: _PROFILE_CACHE[profile_id] = prompt_items # Persist to disk try: save_profile_to_disk(profile_id, voice_name=voice_name) except Exception as exc: logger.warning("Failed to save profile to disk: %s", exc) logger.info("Voice profile %s created and saved.", profile_id) return profile_id def synthesize_cloned(text: str, profile_id: str) -> tuple[np.ndarray, int]: """ Synthesize text using a cached voice profile. Returns (wav_array, sample_rate). """ with _cache_lock: prompt_items = _PROFILE_CACHE.get(profile_id) if prompt_items is None: # Try loading from disk if not in memory if load_profile_from_disk(profile_id): with _cache_lock: prompt_items = _PROFILE_CACHE.get(profile_id) if prompt_items is None: raise ValueError(f"Voice profile '{profile_id}' not found. Record voice first.") model = get_qwen_tts() with GPU_INFERENCE_LOCK: audio_list, sample_rate = model.generate_voice_clone( text=text, language="english", voice_clone_prompt=prompt_items, non_streaming_mode=False, **GENERATION_PARAMS, ) wav = np.concatenate(audio_list) if audio_list else np.zeros(0, dtype=np.float32) return wav, sample_rate def _ensure_stock_voice_profile(): """Ensure the stock vivian voice profile is loaded (uses Base 1.7B with reference audio).""" with _cache_lock: if STOCK_VOICE_PROFILE_ID in _PROFILE_CACHE: return # Create profile from vivian reference audio if not VIVIAN_REF_PATH.exists(): logger.error("Stock voice reference not found at %s", VIVIAN_REF_PATH) raise FileNotFoundError(f"Stock voice reference not found: {VIVIAN_REF_PATH}") create_voice_profile(str(VIVIAN_REF_PATH), voice_name="Vivian (Stock)", profile_id_override=STOCK_VOICE_PROFILE_ID) logger.info("Stock vivian voice profile created from reference audio.") def synthesize_custom_voice_streaming( text: str, speaker: str = "vivian", language: str = "english" ): """ Synthesize text with the Base 1.7B model using stock vivian reference. Yields (wav_segment, sample_rate) tuples. """ _ensure_stock_voice_profile() with _cache_lock: prompt_items = _PROFILE_CACHE.get(STOCK_VOICE_PROFILE_ID) if prompt_items is None: raise RuntimeError( "Stock voice profile could not be created. " "Check that assets/vivian_reference.wav exists and the TTS model loaded correctly." ) model = get_qwen_tts() with GPU_INFERENCE_LOCK: audio_list, sample_rate = model.generate_voice_clone( text=text, language=language, voice_clone_prompt=prompt_items, non_streaming_mode=False, **GENERATION_PARAMS, ) for segment in audio_list: if segment is not None and len(segment) > 0: yield segment, sample_rate def synthesize_custom_voice( text: str, speaker: str = "vivian", language: str = "english" ) -> tuple[np.ndarray, int]: """ Synthesize text with the Base 1.7B model using stock vivian reference. Returns (wav_array, sample_rate). """ segments = list(synthesize_custom_voice_streaming(text, speaker, language)) if segments: wav = np.concatenate([s for s, _ in segments]) sr = segments[0][1] return wav, sr # Fallback: empty audio return np.zeros(0, dtype=np.float32), 24000 def synthesize_cloned_preview(profile_id: str) -> tuple[np.ndarray, int]: """Short preview sentence to verify the clone sounds right.""" return synthesize_cloned( "Hello! I'm ready to read a bedtime story for you tonight.", profile_id, ) def has_profile(profile_id: str | None) -> bool: """Check if a voice profile exists in cache or on disk.""" if not profile_id: return False with _cache_lock: if profile_id in _PROFILE_CACHE: return True # Check disk pt_path = VOICE_PROFILE_DIR / profile_id / "profile.pt" return pt_path.exists()