RustyMark/dots.tts
0
1from __future__ import annotations2 3import hashlib4import json5import os6import time7from pathlib import Path8from typing import Any, Iterator, TypedDict9 10import librosa11import torch12from huggingface_hub import snapshot_download13from loguru import logger14 15from dots_tts.data.pipelines.tokenizing import build_generation_schedule16from dots_tts.data.pipelines.tts_pipeline import (17 DEFAULT_INSTRUCTION_TTS_TEMPLATE,18 DEFAULT_INTERLEAVE_TRAIN_TEMPLATE,19 DEFAULT_TEXT_TO_AUDIO_TEMPLATE,20 DEFAULT_TRAIN_TEMPLATE,21)22from dots_tts.models.dots_tts.model import DotsTtsModel23from dots_tts.utils.audio import high_quality_resample24from dots_tts.utils.profiling import (25 InferenceProfiler,26 activate_inference_profiler,27 inference_profiling,28 log_inference_profile,29)30from dots_tts.utils.text import (31 attach_language_tag,32 detect,33 normalize_language_code,34 normalize_text,35)36from dots_tts.utils.util import get_dtype37 38RUNTIME_TEMPLATE_BY_NAME = {39 "tts": DEFAULT_TRAIN_TEMPLATE,40 "instruction_tts": DEFAULT_INSTRUCTION_TTS_TEMPLATE,41 "text_to_audio": DEFAULT_TEXT_TO_AUDIO_TEMPLATE,42 "tts_interleave": DEFAULT_INTERLEAVE_TRAIN_TEMPLATE,43}44 45 46class RuntimeInputs(TypedDict, total=False):47 fid: str48 language: str49 text: str50 prompt_text: str51 template_name: str52 generation_schedule: torch.Tensor53 prompt_audio: torch.Tensor54 55 56class DotsTtsRuntime:57 # region Lifecycle and pretrained loading58 def __init__(59 self,60 model: DotsTtsModel,61 pretrained_path: Path,62 *,63 precision: str = "bfloat16",64 optimize: bool = False,65 max_generate_length: int = 500,66 ):67 self.model = model68 self.pretrained_path = pretrained_path69 self.precision = precision70 if torch.cuda.is_available():71 self.device = torch.device("cuda")72 else:73 self.device = torch.device("cpu")74 torch.set_num_threads(1)75 if self.device.type == "cuda" and self.precision.lower() in {76 "fp32",77 "torch.float32",78 "float32",79 }:80 torch.set_float32_matmul_precision("high")81 target_dtype = get_dtype(self.precision)82 self.model.core.to(dtype=target_dtype)83 self.model = self.model.to(self.device).eval()84 self.optimize = bool(optimize)85 self.max_generate_length = int(max_generate_length)86 self.model.set_optimize(self.optimize)87 self.sample_rate = int(self.model.config.vocoder.sample_rate)88 skip_init_warmup = os.environ.get("DOTS_TTS_SKIP_INIT_WARMUP", "0") == "1"89 if self.optimize and hasattr(self.model, "run_warmup") and not skip_init_warmup:90 self.model.run_warmup(91 max_generate_length=self.max_generate_length,92 precision=self.precision,93 )94 logger.info(95 "Runtime initialized: pretrained_path={} device={} sample_rate={} "96 "precision={} "97 "optimize={} max_audio_patch_count={}",98 self.pretrained_path,99 self.device,100 self.sample_rate,101 self.precision,102 self.optimize,103 self.max_generate_length,104 )105 106 @classmethod107 def from_pretrained(108 cls,109 model_name_or_path: str,110 *,111 revision: str | None = None,112 cache_dir: str | None = None,113 precision: str = "bfloat16",114 optimize: bool = False,115 max_generate_length: int = 500,116 ) -> DotsTtsRuntime:117 logger.info(118 "Runtime load started: model={} revision={} cache_dir={} precision={}",119 model_name_or_path,120 revision,121 cache_dir,122 precision,123 )124 pretrained_path = cls._resolve_pretrained_path(125 model_name_or_path,126 revision=revision,127 cache_dir=cache_dir,128 )129 loaded_model = DotsTtsModel.from_pretrained(pretrained_path)130 logger.info("Runtime load completed: pretrained_path={}", pretrained_path)131 return cls(132 model=loaded_model,133 pretrained_path=pretrained_path,134 precision=precision,135 optimize=optimize,136 max_generate_length=max_generate_length,137 )138 139 @classmethod140 def _resolve_pretrained_path(141 cls,142 model_name_or_path: str,143 revision: str | None = None,144 cache_dir: str | None = None,145 ) -> Path:146 logger.info(147 "Resolving pretrained path: model={} revision={} cache_dir={}",148 model_name_or_path,149 revision,150 cache_dir,151 )152 resolved_path = Path(model_name_or_path).expanduser().resolve()153 if resolved_path.exists():154 logger.info("Using local pretrained directory: path={}", resolved_path)155 return resolved_path156 157 logger.info(158 "Downloading pretrained snapshot: repo_id={} revision={} cache_dir={}",159 model_name_or_path,160 revision,161 cache_dir,162 )163 snapshot_dir = snapshot_download(164 repo_id=model_name_or_path,165 revision=revision,166 cache_dir=cache_dir,167 )168 resolved_path = Path(snapshot_dir).expanduser().resolve()169 logger.info("Pretrained snapshot ready: path={}", resolved_path)170 return resolved_path171 # endregion Lifecycle and pretrained loading172 173 # region Request normalization and metadata174 @staticmethod175 def _build_request_id(176 *,177 text: str,178 prompt_audio_path: str | None,179 prompt_text: str | None,180 template_name: str,181 language: str | None = None,182 ) -> str:183 payload = {184 "text": text,185 "prompt_audio_path": prompt_audio_path,186 "prompt_text": prompt_text,187 "template_name": template_name,188 }189 if language is not None:190 payload["language"] = language191 digest = hashlib.sha1(192 json.dumps(payload, ensure_ascii=False, sort_keys=True).encode("utf-8")193 ).hexdigest()194 return digest[:16]195 196 def _load_prompt_audio(197 self,198 prompt_audio_path: str,199 ) -> torch.Tensor:200 logger.info("Loading prompt audio: path={}", prompt_audio_path)201 prompt_audio, sample_rate = librosa.load(prompt_audio_path, sr=None, mono=True)202 prompt_audio = librosa.effects.trim(prompt_audio, top_db=30)[0]203 prompt_audio = torch.from_numpy(prompt_audio).unsqueeze(0)204 prompt_audio = high_quality_resample(205 prompt_audio,206 orig_sr=sample_rate,207 target_sr=self.sample_rate,208 )209 if prompt_audio.ndim == 1:210 prompt_audio = prompt_audio.unsqueeze(0)211 logger.info(212 "Prompt audio loaded: path={} original_sample_rate={} resampled_sample_rate={} "213 "samples={}",214 prompt_audio_path,215 sample_rate,216 self.sample_rate,217 prompt_audio.shape[-1],218 )219 return prompt_audio220 221 def _resolve_language(222 self,223 language: str | None,224 *,225 text: str,226 ) -> str | None:227 if language is None:228 return None229 230 stripped = language.strip()231 if not stripped or stripped.lower() == "none":232 return None233 if stripped.lower() == "auto_detect":234 return normalize_language_code(detect(text))235 236 normalized_language = normalize_language_code(stripped)237 if normalized_language is None:238 raise ValueError(239 f"Unsupported language={language!r}. "240 "Expected 'none', 'auto_detect', or a valid language code/name."241 )242 return normalized_language243 244 def _process_prompt_text(245 self,246 prompt_text: str | None,247 *,248 language: str | None = None,249 ) -> str:250 if prompt_text is None:251 return ""252 prompt_text = prompt_text.strip()253 if not prompt_text:254 return ""255 256 prompt_language = language257 if prompt_language is None:258 prompt_language = normalize_language_code(detect(prompt_text))259 260 if prompt_language not in {"ZH", "YUE", "JA", "口音:粤语"}:261 prompt_text += " "262 if language is not None:263 prompt_text = attach_language_tag(prompt_text, language)264 return prompt_text265 266 def _process_text(267 self,268 text: str,269 *,270 language: str | None = None,271 normalize: bool = False,272 ) -> tuple[str, str | None]:273 stripped = text.strip()274 if normalize:275 stripped = normalize_text(stripped)276 resolved_language = self._resolve_language(language, text=stripped)277 return stripped, resolved_language278 279 def _estimate_prompt_audio_patch_count(280 self,281 *,282 prompt_audio: torch.Tensor | None,283 prompt_text: str,284 ) -> int:285 if prompt_audio is None or not prompt_text:286 return 0287 samples_per_patch = int(self.model.config.patch_size * self.model.hop_size)288 prompt_samples = int(prompt_audio.shape[-1])289 return (prompt_samples + samples_per_patch - 1) // samples_per_patch290 # endregion Request normalization and metadata291 292 # region Generation schedule assembly293 def _normalize_template_name(self, template_name: str | None) -> str:294 if template_name is None:295 return "tts"296 if template_name not in RUNTIME_TEMPLATE_BY_NAME:297 raise ValueError(298 f"Unknown template_name={template_name!r}. "299 f"Expected one of {sorted(RUNTIME_TEMPLATE_BY_NAME)}."300 )301 return template_name302 303 def _prepare_inputs(304 self,305 *,306 text: str,307 prompt_audio_path: str | None,308 prompt_text: str | None,309 template_name: str | None,310 language: str | None = None,311 normalize_text: bool = False,312 ) -> RuntimeInputs:313 normalized_template_name = self._normalize_template_name(template_name)314 template = RUNTIME_TEMPLATE_BY_NAME[normalized_template_name]315 if prompt_text and not prompt_audio_path:316 raise ValueError("prompt_text requires prompt_audio_path.")317 318 normalized_text, normalized_language = self._process_text(319 text,320 language=language,321 normalize=normalize_text,322 )323 normalized_prompt_text = self._process_prompt_text(324 prompt_text,325 language=normalized_language,326 )327 if normalized_language is not None and not normalized_prompt_text:328 normalized_text = attach_language_tag(normalized_text, normalized_language)329 inputs: RuntimeInputs = {330 "fid": self._build_request_id(331 text=normalized_text,332 prompt_audio_path=prompt_audio_path,333 prompt_text=normalized_prompt_text,334 template_name=normalized_template_name,335 language=normalized_language,336 ),337 "language": normalized_language or "",338 "text": normalized_text,339 "prompt_text": normalized_prompt_text,340 "template_name": normalized_template_name,341 }342 343 if prompt_audio_path:344 inputs["prompt_audio"] = self._load_prompt_audio(prompt_audio_path)345 prompt_audio_patch_count = self._estimate_prompt_audio_patch_count(346 prompt_audio=inputs.get("prompt_audio"),347 prompt_text=normalized_prompt_text,348 )349 if (350 prompt_audio_patch_count > 0351 and self.max_generate_length <= prompt_audio_patch_count352 ):353 raise ValueError(354 "max_generate_length must exceed prompt audio patch count when prompt_text is provided: "355 f"max_generate_length={self.max_generate_length} "356 f"prompt_audio_patch_count={prompt_audio_patch_count}."357 )358 359 schedule_spec = build_generation_schedule(360 text=f"{normalized_prompt_text}{normalized_text}",361 tokenizer=self.model.tokenizer,362 template=template,363 max_audio_tokens=self.max_generate_length,364 )365 schedule = torch.tensor(366 schedule_spec["schedule_ids"],367 dtype=torch.long,368 device=self.device,369 )370 inputs["generation_schedule"] = schedule.unsqueeze(0)371 logger.info(372 "Inputs prepared: request_id={} template_name={} "373 "language={} text_len={} prompt_text_len={} schedule_length={} "374 "prompt_audio_patch_count={} max_audio_patch_count={} has_prompt_audio={}",375 inputs["fid"],376 normalized_template_name,377 normalized_language,378 len(normalized_text),379 len(normalized_prompt_text),380 schedule.numel(),381 prompt_audio_patch_count,382 self.max_generate_length,383 bool(prompt_audio_path),384 )385 return inputs386 # endregion Generation schedule assembly387 388 # region Public generation APIs389 def generate_stream(390 self,391 *,392 text: str,393 prompt_audio_path: str | None = None,394 prompt_text: str | None = None,395 template_name: str | None = None,396 language: str | None = None,397 speaker_scale: float = 1.5,398 ode_method: str = "euler",399 num_steps: int = 10,400 guidance_scale: float = 1.2,401 normalize_text: bool = False,402 profile_inference: bool = False,403 ) -> Iterator[torch.Tensor]:404 inputs = self._prepare_inputs(405 text=text,406 prompt_audio_path=prompt_audio_path,407 prompt_text=prompt_text,408 template_name=template_name,409 language=language,410 normalize_text=normalize_text,411 )412 logger.info(413 "Streaming generation started: request_id={} text_len={} has_prompt_audio={} "414 "has_prompt_text={} template_name={} language={} precision={} ode_method={} num_steps={} "415 "guidance_scale={} speaker_scale={} max_audio_patch_count={} normalize_text={}",416 inputs["fid"],417 len(inputs["text"]),418 bool(prompt_audio_path),419 bool(inputs["prompt_text"]),420 inputs["template_name"],421 inputs["language"] or None,422 self.precision,423 ode_method,424 num_steps,425 guidance_scale,426 speaker_scale,427 self.max_generate_length,428 normalize_text,429 )430 start_time = time.time()431 emitted_samples = 0432 chunk_count = 0433 profiler: InferenceProfiler | None = None434 try:435 profiler = (436 InferenceProfiler(self.device) if profile_inference else None437 )438 stream = self.model.generate_audio_stream(439 inputs,440 precision=self.precision,441 ode_method=ode_method,442 num_steps=num_steps,443 guidance_scale=guidance_scale,444 speaker_scale=speaker_scale,445 )446 while True:447 try:448 with activate_inference_profiler(profiler):449 chunk = next(stream)450 except StopIteration:451 break452 emitted_samples += int(chunk.shape[-1])453 chunk_count += 1454 yield chunk455 except Exception:456 logger.exception(457 "Streaming generation failed: request_id={}",458 inputs["fid"],459 )460 raise461 time_used = time.time() - start_time462 duration_seconds = emitted_samples / self.sample_rate463 rtf = time_used / duration_seconds if duration_seconds > 0 else float("inf")464 if profile_inference and profiler is not None:465 log_inference_profile(466 request_id=inputs["fid"],467 profiling=profiler.summary(duration_seconds=duration_seconds),468 duration_seconds=duration_seconds,469 )470 logger.info(471 "Streaming generation finished: request_id={} chunk_count={} elapsed_seconds={:.3f} "472 "audio_seconds={:.3f} rtf={:.4f} sample_rate={}",473 inputs["fid"],474 chunk_count,475 time_used,476 duration_seconds,477 rtf,478 self.sample_rate,479 )480 481 def generate(482 self,483 *,484 text: str,485 prompt_audio_path: str | None = None,486 prompt_text: str | None = None,487 template_name: str | None = None,488 language: str | None = None,489 speaker_scale: float = 1.5,490 ode_method: str = "euler",491 num_steps: int = 10,492 guidance_scale: float = 1.2,493 normalize_text: bool = False,494 profile_inference: bool = False,495 ) -> dict[str, Any]:496 inputs = self._prepare_inputs(497 text=text,498 prompt_audio_path=prompt_audio_path,499 prompt_text=prompt_text,500 template_name=template_name,501 language=language,502 normalize_text=normalize_text,503 )504 logger.info(505 "Generation started: request_id={} text_len={} has_prompt_audio={} "506 "has_prompt_text={} template_name={} language={} precision={} ode_method={} num_steps={} "507 "guidance_scale={} speaker_scale={} max_audio_patch_count={} normalize_text={}",508 inputs["fid"],509 len(inputs["text"]),510 bool(prompt_audio_path),511 bool(inputs["prompt_text"]),512 inputs["template_name"],513 inputs["language"] or None,514 self.precision,515 ode_method,516 num_steps,517 guidance_scale,518 speaker_scale,519 self.max_generate_length,520 normalize_text,521 )522 start_time = time.time()523 profiling = None524 try:525 with inference_profiling(526 enabled=profile_inference,527 device=self.device,528 ) as profiler:529 audio = self.model.generate_audio(530 inputs,531 precision=self.precision,532 ode_method=ode_method,533 num_steps=num_steps,534 guidance_scale=guidance_scale,535 speaker_scale=speaker_scale,536 )537 except Exception:538 logger.exception("Generation failed: request_id={}", inputs["fid"])539 raise540 time_used = time.time() - start_time541 duration_seconds = audio.shape[-1] / self.sample_rate542 rtf = time_used / duration_seconds if duration_seconds > 0 else float("inf")543 if profiler is not None:544 profiling = profiler.summary(duration_seconds=duration_seconds)545 log_inference_profile(546 request_id=inputs["fid"],547 profiling=profiling,548 duration_seconds=duration_seconds,549 )550 logger.info(551 "Generation completed: request_id={} elapsed_seconds={:.3f} audio_seconds={:.3f} "552 "rtf={:.4f} sample_rate={}",553 inputs["fid"],554 time_used,555 duration_seconds,556 rtf,557 self.sample_rate,558 )559 return {560 "fid": inputs["fid"],561 "audio": audio,562 "sample_rate": self.sample_rate,563 "time_used": time_used,564 "rtf": rtf,565 "profiling": profiling,566 }567 # endregion Public generation APIs568