Upload vizipvoice.py

vizipvoice.py

@@ -0,0 +1,556 @@
+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-<step>.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),
+        }