File size: 5,191 Bytes
1f90181 034b456 68fd114 034b456 1f90181 01cab82 1f90181 682df44 1f90181 01cab82 1f90181 68fd114 1f90181 034b456 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | """Hugging Face remote-code entry point for the ONNX Runtime TeraTTS release."""
from __future__ import annotations
from pathlib import Path
from typing import Iterator
import numpy as np
from transformers import PreTrainedModel
from .configuration_teratts import TeraTTSConfig
# Keep the bundled stress-runtime module in Transformers' remote-code module
# dependency set. ``teratts.py`` imports it at runtime after this entry point
# has been loaded.
from .teratts_ruaccent import RUAccent as _BundledRUAccent
from .teratts import (
LoadedTTS,
generate_speech,
generate_speech_stream,
load_model,
normalize_text,
write_wav,
)
class TeraTTSModel(PreTrainedModel):
"""A reusable ONNX Runtime speech generator loaded through ``AutoModel``."""
config_class = TeraTTSConfig
def __init__(self, config: TeraTTSConfig) -> None:
super().__init__(config)
self.runtime: LoadedTTS | None = None
@classmethod
def from_pretrained(
cls,
pretrained_model_name_or_path: str | Path,
*model_args: object,
config: TeraTTSConfig | None = None,
provider: str = "CPUExecutionProvider",
threads: int | None = None,
diffusion_model: str | None = None,
russian_stress: bool = True,
ruaccent_model_size: str = "turbo3.1",
ruaccent_device: str = "CPU",
ruaccent_mode: str = "full",
**kwargs: object,
) -> "TeraTTSModel":
"""Download/load the release and initialize reusable ONNX sessions.
``provider``, ``threads``, and ``diffusion_model`` are TeraTTS-specific
arguments. Standard Hub arguments such as ``revision``, ``token``,
``cache_dir``, and ``local_files_only`` are forwarded to the snapshot
download when a Hub model ID is supplied.
"""
if model_args:
raise TypeError("TeraTTSModel.from_pretrained accepts no positional model arguments")
# ``AutoModel`` forwards these framework-only values to a custom
# class after it has already used them. This ONNX runtime has no
# adapters or PyTorch state dict to load, so safely discard them.
for key in (
"trust_remote_code",
"adapter_kwargs",
"_from_auto",
"use_safetensors",
"weights_only",
"low_cpu_mem_usage",
"device_map",
"torch_dtype",
"dtype",
):
kwargs.pop(key, None)
source = Path(pretrained_model_name_or_path)
if source.is_dir():
release = source
else:
from huggingface_hub import snapshot_download
download_keys = {"revision", "token", "cache_dir", "local_files_only", "force_download"}
download_kwargs = {key: kwargs.pop(key) for key in list(kwargs) if key in download_keys}
if kwargs:
unexpected = ", ".join(sorted(kwargs))
raise TypeError(f"unexpected TeraTTS loading arguments: {unexpected}")
release = Path(
snapshot_download(repo_id=str(pretrained_model_name_or_path), **download_kwargs)
)
if config is None:
config = TeraTTSConfig.from_pretrained(release)
instance = cls(config)
instance.runtime = load_model(
release,
model=diffusion_model or config.default_diffusion_model,
provider=provider,
threads=threads,
russian_stress=russian_stress,
ruaccent_model_size=ruaccent_model_size,
ruaccent_device=ruaccent_device,
ruaccent_mode=ruaccent_mode,
)
return instance
def _runtime(self) -> LoadedTTS:
if self.runtime is None:
raise RuntimeError("load this model with from_pretrained() before generating speech")
return self.runtime
def generate_speech(
self,
text: str,
voice: str,
*,
duration_scale: float = 1.0,
guidance: float = 3.0,
seed: int = 1234,
) -> np.ndarray:
return generate_speech(
self._runtime(),
text,
voice,
duration_scale=duration_scale,
guidance=guidance,
seed=seed,
)
def normalize_text(self, text: str) -> str:
"""Show the final tagged text passed to the TTS text encoder."""
return normalize_text(self._runtime(), text)
def generate_speech_stream(
self,
text: str,
voice: str,
*,
duration_scale: float = 1.0,
guidance: float = 3.0,
seed: int = 1234,
chunk_frames: int = 16,
) -> Iterator[np.ndarray]:
yield from generate_speech_stream(
self._runtime(),
text,
voice,
duration_scale=duration_scale,
guidance=guidance,
seed=seed,
chunk_frames=chunk_frames,
)
def save_wav(self, path: str | Path, waveform: np.ndarray) -> None:
"""Write a generated mono float32 waveform as a 44.1 kHz WAV file."""
write_wav(Path(path), waveform)
|