import json import logging import re import tempfile import time from pathlib import Path from typing import Optional, Union import safetensors.torch import torch import torchaudio from huggingface_hub import hf_hub_download, list_repo_files from lhotse.utils import fix_random_seed from zipvoice.bin.infer_zipvoice import generate_sentence, get_vocoder from zipvoice.models.zipvoice import ZipVoice from zipvoice.tokenizer.tokenizer import SimpleTokenizer from zipvoice.utils.checkpoint import load_checkpoint from zipvoice.utils.feature import VocosFbank DEFAULT_REPO_ID = "contextboxai/ViZipvoice" DEFAULT_CHECKPOINT_NAME = "latest" CHECKPOINT_RE = re.compile(r"^checkpoint-(\d+)\.pt$") SENTENCE_SPLIT_PATTERN = re.compile(r"[^.?？。]+(?:[.?？。]+|$)", re.S) PUNCTUATION_NO_SPACE_BEFORE = r",.;:!?…%" OPENING_QUOTES_AND_BRACKETS = r"\(\[\{«“‘" CLOSING_QUOTES_AND_BRACKETS = r"\)\]\}»”’" def _resolve_device(device: Optional[Union[str, torch.device]] = None) -> torch.device: if device is not None: return torch.device(device) if torch.cuda.is_available(): return torch.device("cuda", 0) if torch.backends.mps.is_available(): return torch.device("mps") return torch.device("cpu") def _download_model_files( repo_id: str, revision: Optional[str], checkpoint_name: str, ) -> tuple[Path, Path, Path]: checkpoint_name = _resolve_hf_checkpoint_name( repo_id=repo_id, revision=revision, checkpoint_name=checkpoint_name, ) checkpoint_path = Path( hf_hub_download( repo_id=repo_id, filename=checkpoint_name, revision=revision, ) ) model_config_path = _download_config_file( repo_id=repo_id, revision=revision, ) token_file = Path( hf_hub_download( repo_id=repo_id, filename="tokens.txt", revision=revision, ) ) return checkpoint_path, model_config_path, token_file def _download_config_file(repo_id: str, revision: Optional[str]) -> Path: last_error = None for filename in ("config.json", "model.json"): try: return Path( hf_hub_download( repo_id=repo_id, filename=filename, revision=revision, ) ) except Exception as exc: last_error = exc raise FileNotFoundError("No config.json or model.json file found.") from last_error def _checkpoint_step(filename: str) -> int: match = CHECKPOINT_RE.match(Path(filename).name) return int(match.group(1)) if match else -1 def _select_latest_checkpoint(filenames: list[str]) -> str: checkpoints = [ filename for filename in filenames if _checkpoint_step(filename) >= 0 ] if checkpoints: return max(checkpoints, key=lambda filename: _checkpoint_step(filename)) raise FileNotFoundError("No checkpoint-.pt file found.") def _resolve_hf_checkpoint_name( repo_id: str, revision: Optional[str], checkpoint_name: str, ) -> str: if checkpoint_name != "latest": return checkpoint_name filenames = list_repo_files(repo_id=repo_id, revision=revision) return _select_latest_checkpoint(filenames) def _resolve_local_checkpoint_path( model_dir: Path, checkpoint_name: str, ) -> Path: if checkpoint_name != "latest": return model_dir / checkpoint_name filenames = [path.name for path in model_dir.iterdir() if path.is_file()] return model_dir / _select_latest_checkpoint(filenames) def _resolve_local_config_path(model_dir: Path) -> Path: for filename in ("config.json", "model.json"): config_path = model_dir / filename if config_path.is_file(): return config_path raise FileNotFoundError(f"No config.json or model.json file found in {model_dir}") def cleanup_vietnamese_spacing(text: str) -> str: text = re.sub(r"\s+", " ", text.strip()) text = re.sub( rf"\s+([{re.escape(PUNCTUATION_NO_SPACE_BEFORE)}])", r"\1", text, ) text = re.sub( rf"\s+([{CLOSING_QUOTES_AND_BRACKETS}])", r"\1", text, ) text = re.sub( rf"([{OPENING_QUOTES_AND_BRACKETS}])\s+", r"\1", text, ) text = re.sub( rf"([{re.escape(PUNCTUATION_NO_SPACE_BEFORE)}])" rf"([^\s{CLOSING_QUOTES_AND_BRACKETS}])", r"\1 \2", text, ) return text.strip() def normalize_vietnamese_text(text: str, enabled: bool = True) -> str: if not enabled: return text.strip() try: from soe_vinorm import normalize_text except ImportError as exc: raise RuntimeError( "Vietnamese normalization requires soe-vinorm. " "Install it with `pip install soe-vinorm`." ) from exc return cleanup_vietnamese_spacing(normalize_text(text)) def split_text_into_sentences(text: str) -> list[str]: text = text.strip() if not text: return [] sentences = [ match.group(0).strip() for match in SENTENCE_SPLIT_PATTERN.finditer(text) if match.group(0).strip() ] return sentences or [text] def count_sentence_words(sentence: str) -> int: return len(re.findall(r"\w+", sentence, flags=re.UNICODE)) def get_sentence_inference_params( sentence: str, base_num_step: int, base_speed: float, ) -> tuple[int, float, int]: word_count = count_sentence_words(sentence) if word_count == 1: return max(base_num_step, 24), 0.6, word_count if 2 <= word_count <= 4: return base_num_step, 0.8, word_count return base_num_step, base_speed, word_count def match_audio_channels( first: torch.Tensor, second: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: if first.shape[0] == second.shape[0]: return first, second if first.shape[0] == 1: return first.expand(second.shape[0], -1), second if second.shape[0] == 1: return first, second.expand(first.shape[0], -1) channels = min(first.shape[0], second.shape[0]) return first[:channels], second[:channels] def append_with_crossfade( first: torch.Tensor, second: torch.Tensor, crossfade_samples: int, ) -> torch.Tensor: first, second = match_audio_channels(first, second) fade_len = min(crossfade_samples, first.shape[1], second.shape[1]) if fade_len <= 0: return torch.cat([first, second], dim=1) fade_out = torch.linspace( 1.0, 0.0, fade_len, dtype=first.dtype, device=first.device, ).unsqueeze(0) fade_in = torch.linspace( 0.0, 1.0, fade_len, dtype=second.dtype, device=second.device, ).unsqueeze(0) overlap = first[:, -fade_len:] * fade_out + second[:, :fade_len] * fade_in return torch.cat([first[:, :-fade_len], overlap, second[:, fade_len:]], dim=1) def apply_fade(audio: torch.Tensor, fade_in_samples: int, fade_out_samples: int) -> torch.Tensor: if audio.numel() == 0: return audio audio = audio.clone() if fade_in_samples > 0: fade_len = min(fade_in_samples, audio.shape[1]) fade = torch.linspace( 0.0, 1.0, fade_len, dtype=audio.dtype, device=audio.device, ).unsqueeze(0) audio[:, :fade_len] *= fade if fade_out_samples > 0: fade_len = min(fade_out_samples, audio.shape[1]) fade = torch.linspace( 1.0, 0.0, fade_len, dtype=audio.dtype, device=audio.device, ).unsqueeze(0) audio[:, -fade_len:] *= fade return audio def postprocess_audio_segments( segment_paths: list[Path], output_path: Path, sampling_rate: int, crossfade_ms: int, silence_ms: int, fade_in_ms: int, fade_out_ms: int, ) -> None: if not segment_paths: raise RuntimeError("No generated audio segments to postprocess.") crossfade_samples = int(sampling_rate * max(crossfade_ms, 0) / 1000) silence_samples = int(sampling_rate * max(silence_ms, 0) / 1000) fade_in_samples = int(sampling_rate * max(fade_in_ms, 0) / 1000) fade_out_samples = int(sampling_rate * max(fade_out_ms, 0) / 1000) combined = None for index, segment_path in enumerate(segment_paths): audio, sr = torchaudio.load(str(segment_path)) if sr != sampling_rate: audio = torchaudio.functional.resample(audio, sr, sampling_rate) if index < len(segment_paths) - 1 and silence_samples > 0: silence = torch.zeros( audio.shape[0], silence_samples, dtype=audio.dtype, device=audio.device, ) audio = torch.cat([audio, silence], dim=1) if combined is None: combined = audio else: combined = append_with_crossfade(combined, audio, crossfade_samples) combined = apply_fade( combined, fade_in_samples=fade_in_samples, fade_out_samples=fade_out_samples, ) combined = combined.clamp(min=-1.0, max=1.0).cpu() torchaudio.save(str(output_path), combined, sampling_rate) def wav_seconds(path: Union[str, Path]) -> float: try: import soundfile as sf info = sf.info(str(path)) return float(info.frames) / float(info.samplerate) except Exception: audio, sr = torchaudio.load(str(path)) return float(audio.shape[-1]) / float(sr) class ViZipVoiceTTS: """Small wrapper for Vietnamese ZipVoice inference. The wrapper downloads model files from Hugging Face by default, builds the ZipVoice model with the Vietnamese character tokenizer, and exposes a single synthesize method. """ def __init__( self, repo_id: str = DEFAULT_REPO_ID, revision: Optional[str] = None, model_dir: Optional[Union[str, Path]] = None, checkpoint_name: str = DEFAULT_CHECKPOINT_NAME, vocoder_path: Optional[Union[str, Path]] = None, device: Optional[Union[str, torch.device]] = None, use_fp16: bool = True, num_threads: int = 1, ) -> None: try: torch.set_num_threads(num_threads) torch.set_num_interop_threads(num_threads) except RuntimeError: logging.debug("PyTorch thread settings were already initialized.") self.repo_id = repo_id self.revision = revision self.device = _resolve_device(device) self.use_fp16 = bool(use_fp16 and self.device.type == "cuda") if model_dir is None: checkpoint_path, model_config_path, token_file = _download_model_files( repo_id=repo_id, revision=revision, checkpoint_name=checkpoint_name, ) else: model_dir = Path(model_dir) checkpoint_path = _resolve_local_checkpoint_path( model_dir=model_dir, checkpoint_name=checkpoint_name, ) model_config_path = _resolve_local_config_path(model_dir) token_file = model_dir / "tokens.txt" self.checkpoint_path = Path(checkpoint_path) self.model_config_path = Path(model_config_path) self.token_file = Path(token_file) self._validate_model_files() with self.model_config_path.open("r", encoding="utf-8") as f: self.model_config = json.load(f) self.tokenizer = SimpleTokenizer(token_file=str(self.token_file)) self.model = ZipVoice( **self.model_config["model"], vocab_size=self.tokenizer.vocab_size, pad_id=self.tokenizer.pad_id, ) self._load_checkpoint() self.model.to(self.device) self.model.eval() self.feature_extractor = VocosFbank() self.vocoder = get_vocoder(str(vocoder_path) if vocoder_path else None) self.vocoder.to(self.device) self.vocoder.eval() self.sampling_rate = int(self.model_config["feature"]["sampling_rate"]) logging.info( "Loaded ViZipVoice from %s on %s | fp16 autocast: %s", self.checkpoint_path, self.device, self.use_fp16, ) def _validate_model_files(self) -> None: missing = [ path for path in [self.checkpoint_path, self.model_config_path, self.token_file] if not path.is_file() ] if missing: missing_text = ", ".join(str(path) for path in missing) raise FileNotFoundError(f"Missing ViZipVoice model file(s): {missing_text}") def _load_checkpoint(self) -> None: suffix = self.checkpoint_path.suffix.lower() if suffix == ".safetensors": safetensors.torch.load_model(self.model, str(self.checkpoint_path)) elif suffix == ".pt": load_checkpoint( filename=self.checkpoint_path, model=self.model, strict=True, ) else: raise ValueError(f"Unsupported checkpoint format: {self.checkpoint_path}") @torch.inference_mode() def synthesize( self, prompt_wav: Union[str, Path], prompt_text: str, text: str, output_path: Union[str, Path] = "output.wav", num_step: int = 16, guidance_scale: float = 1.0, speed: float = 1.0, t_shift: float = 0.5, target_rms: float = 0.1, feat_scale: float = 0.1, max_duration: float = 100, remove_long_sil: bool = False, seed: Optional[int] = 666, normalize_vietnamese: bool = True, split_sentences: bool = True, crossfade_ms: int = 80, silence_ms: int = 180, fade_in_ms: int = 20, fade_out_ms: int = 80, ) -> dict: if seed is not None and seed >= 0: fix_random_seed(int(seed)) prompt_text = normalize_vietnamese_text( prompt_text, enabled=normalize_vietnamese, ) text = normalize_vietnamese_text( text, enabled=normalize_vietnamese, ) target_sentences = split_text_into_sentences(text) if split_sentences else [text] if not target_sentences: raise ValueError("No valid text to synthesize.") output_path = Path(output_path) output_path.parent.mkdir(parents=True, exist_ok=True) segment_paths = [] segment_metrics = [] segment_settings = [] start_time = time.time() with tempfile.TemporaryDirectory( prefix=f"{output_path.stem}_segments_", dir=str(output_path.parent), ) as segment_dir_name: segment_dir = Path(segment_dir_name) with torch.autocast( device_type="cuda", dtype=torch.float16, enabled=self.use_fp16, ): for index, sentence in enumerate(target_sentences, start=1): sentence_num_step, sentence_speed, word_count = ( get_sentence_inference_params( sentence=sentence, base_num_step=int(num_step), base_speed=float(speed), ) ) segment_path = segment_dir / f"segment_{index:03d}.wav" metrics = generate_sentence( save_path=str(segment_path), prompt_text=prompt_text, prompt_wav=str(prompt_wav), text=sentence, model=self.model, vocoder=self.vocoder, tokenizer=self.tokenizer, feature_extractor=self.feature_extractor, device=self.device, num_step=sentence_num_step, guidance_scale=float(guidance_scale), speed=sentence_speed, t_shift=float(t_shift), target_rms=float(target_rms), feat_scale=float(feat_scale), sampling_rate=self.sampling_rate, max_duration=float(max_duration), remove_long_sil=bool(remove_long_sil), ) segment_paths.append(segment_path) segment_metrics.append(metrics) segment_settings.append( { "index": index, "word_count": word_count, "speed": sentence_speed, "num_step": sentence_num_step, "text": sentence, } ) postprocess_audio_segments( segment_paths=segment_paths, output_path=output_path, sampling_rate=self.sampling_rate, crossfade_ms=int(crossfade_ms), silence_ms=int(silence_ms), fade_in_ms=int(fade_in_ms), fade_out_ms=int(fade_out_ms), ) elapsed = time.time() - start_time audio_seconds = wav_seconds(output_path) t_no_vocoder = sum(item.get("t_no_vocoder", 0.0) for item in segment_metrics) t_vocoder = sum(item.get("t_vocoder", 0.0) for item in segment_metrics) rtf = elapsed / audio_seconds if audio_seconds else 0.0 return { "t": elapsed, "t_no_vocoder": t_no_vocoder, "t_vocoder": t_vocoder, "wav_seconds": audio_seconds, "rtf": rtf, "rtf_no_vocoder": t_no_vocoder / audio_seconds if audio_seconds else 0.0, "rtf_vocoder": t_vocoder / audio_seconds if audio_seconds else 0.0, "segments": len(segment_paths), "segment_settings": segment_settings, "segment_metrics": segment_metrics, "crossfade_ms": int(crossfade_ms), "silence_ms": int(silence_ms), "fade_in_ms": int(fade_in_ms), "fade_out_ms": int(fade_out_ms), }